@@ -85,15 +85,33 @@ def generate_secret_bytes() -> bytes:
8585 from cryptography .fernet import Fernet
8686 return Fernet .generate_key ()
8787
88+ def apply_pipeline_options (self , options : Any ) -> None :
89+ """Configures this secret from pipeline options.
90+
91+ Called by :meth:`from_json` and :meth:`parse_secret_option` when pipeline
92+ options are given. The default implementation does nothing; secrets that
93+ need settings from pipeline options (e.g. :class:`AwsSecret`) override it.
94+
95+ Args:
96+ options: PipelineOptions, or a dict of option names to values.
97+ """
98+ pass
99+
88100 @classmethod
89- def parse_secret_option (cls , secret : str ) -> 'Secret' :
101+ def parse_secret_option (
102+ cls , secret : str , options : Optional [Any ] = None ) -> 'Secret' :
90103 """Parses a secret string and returns the appropriate secret type.
91104
92105 The secret string should be formatted like:
93106 'type:<secret_type>;<secret_param>:<value>'
94107
95108 For example, 'type:GcpSecret;version_name:my_secret/versions/latest'
96109 would return a GcpSecret initialized with 'my_secret/versions/latest'.
110+
111+ Args:
112+ secret: The secret string.
113+ options: Optional pipeline options, passed to
114+ :meth:`apply_pipeline_options` of the returned secret.
97115 """
98116 param_map = {}
99117 for param in secret .split (';' ):
@@ -108,24 +126,39 @@ def parse_secret_option(cls, secret: str) -> 'Secret':
108126 secret_type = raw_type .lower ()
109127 secret_manager = _SECRET_TYPE_TO_SECRET_MANAGER .get (secret_type )
110128 if not secret_manager :
129+ supported = ', ' .join (
130+ getattr (cls , '__name__' , cls ) for cls in _SECRET_CLASSES .values ())
111131 raise ValueError (
112- f'Invalid secret type { secret_type } , currently only '
113- 'GcpSecret and GcpHsmGeneratedSecret are supported' )
132+ f'Invalid secret type { secret_type } , currently supported types: '
133+ f' { supported } ' )
114134
115- return cls .from_json (json .dumps (param_map ), secret_manager )
135+ return cls .from_json (json .dumps (param_map ), secret_manager , options )
116136
117137 @classmethod
118138 def from_json (
119- cls , spec : str , secret_manager : Optional [str ] = None ) -> 'Secret' :
139+ cls ,
140+ spec : str ,
141+ secret_manager : Optional [str ] = None ,
142+ options : Optional [Any ] = None ) -> 'Secret' :
120143 """Return a Secret instance based on secret_manager provider and secret specification.
121144
122145 Args:
123146 spec: Secret string (raw secret or JSON specification string).
124147 secret_manager: Secret manager string (e.g. 'GoogleCloudSecretManager').
148+ options: Optional pipeline options (PipelineOptions or a dict), passed to
149+ :meth:`apply_pipeline_options` of the returned secret. For example,
150+ AwsSecret uses them to get AWS credentials from AwsOptions.
125151
126152 Returns:
127153 An instance of Secret.
128154 """
155+ secret = cls ._from_json (spec , secret_manager )
156+ if options is not None :
157+ secret .apply_pipeline_options (options )
158+ return secret
159+
160+ @classmethod
161+ def _from_json (cls , spec : str , secret_manager : Optional [str ]) -> 'Secret' :
129162 if not isinstance (spec , str ):
130163 raise TypeError (
131164 f"Secret 'spec' must be a string, got { type (spec ).__name__ } " )
@@ -162,9 +195,11 @@ def from_json(
162195 else :
163196 return secret_cls (spec )
164197 else :
198+ supported = ', ' .join (
199+ f"'{ name } '" for name in _SECRET_TYPE_TO_SECRET_MANAGER .values ())
165200 raise ValueError (
166- f"Unsupported secret manager: '{ secret_manager_name } '. Currently supported options: 'GoogleCloudSecretManager', 'GoogleCloudHsmGeneratedSecretManager'. "
167- )
201+ f"Unsupported secret manager: '{ secret_manager_name } '. "
202+ f"Currently supported options: { supported } ." )
168203
169204 # If secret_manager is not set or empty, check if spec is a JSON specification dict
170205 if spec_dict is not None :
@@ -271,6 +306,155 @@ def __eq__(self, secret):
271306 return self ._version_name == getattr (secret , '_version_name' , None )
272307
273308
309+ class AwsSecret (Secret ):
310+ """A secret manager implementation that retrieves secrets from AWS Secrets
311+ Manager.
312+
313+ AWS credentials and connection settings (region, endpoint, SSL) come from the
314+ ``--aws_*`` defaults in :class:`AwsOptions` when pipeline options are given
315+ (see :meth:`apply_pipeline_options`), and otherwise from boto3's defaults
316+ (e.g. environment variables, ~/.aws, or an IAM role). Connection settings
317+ can be overridden per secret with keyword arguments, e.g.
318+ ``AwsSecret('my-secret', region_name='us-west-2')``, or in the secret spec,
319+ e.g. ``type:awssecret;name:my-secret;region_name:us-west-2``.
320+ """
321+ # Spec keys that identify the secret. Other spec keys are client settings.
322+ _SECRET_SPEC_KEYS = frozenset ({
323+ 'secret_id' ,
324+ 'name' ,
325+ 'arn' ,
326+ 'version_id' ,
327+ 'version' ,
328+ 'version_stage' ,
329+ 'stage' ,
330+ })
331+
332+ def __init__ (
333+ self ,
334+ secret_id : str ,
335+ version_id : Optional [str ] = None ,
336+ version_stage : Optional [str ] = None ,
337+ ** client_settings : Any ):
338+ """Initializes an AwsSecret object.
339+
340+ Args:
341+ secret_id: The ARN or name of the secret in AWS Secrets Manager.
342+ version_id: The unique identifier of the version of the secret.
343+ version_stage: The staging label of the version of the secret (e.g. 'AWSCURRENT').
344+ **client_settings: AWS client connection settings that override the
345+ ``--aws_*`` pipeline options for this secret. Names are AwsOptions
346+ setting names without the ``aws_`` prefix (e.g. ``region_name``,
347+ ``endpoint_url``, ``verify``, ``disable_ssl``). Credentials cannot be
348+ set here; use the ``--aws_*`` pipeline options instead.
349+ """
350+ super ().__init__ ()
351+ self ._secret_id = secret_id
352+ self ._version_id = version_id
353+ self ._version_stage = version_stage
354+ # Set by apply_pipeline_options.
355+ self ._aws_options : Optional [dict [str , Any ]] = None
356+ self ._client_settings : Optional [dict [str , Any ]] = None
357+ if client_settings :
358+ from apache_beam .io .aws .clients import client_options
359+ self ._client_settings = client_options .parse_connection_settings (
360+ client_settings )
361+
362+ @classmethod
363+ def from_dict (cls , spec_dict : dict [str , str ]) -> 'AwsSecret' :
364+ """Initialize AwsSecret from a dictionary specification.
365+
366+ Keys other than the secret identifiers are AWS client connection settings
367+ (see the ``**client_settings`` keyword arguments of :meth:`__init__`).
368+ """
369+ secret_id = (
370+ spec_dict .get ("secret_id" ) or spec_dict .get ("name" ) or
371+ spec_dict .get ("arn" ))
372+ if not secret_id :
373+ raise ValueError ("Secret ID or name must be specified in secret spec." )
374+
375+ version_id = spec_dict .get ("version_id" ) or spec_dict .get ("version" )
376+ version_stage = spec_dict .get ("version_stage" ) or spec_dict .get ("stage" )
377+ client_settings = {
378+ k : v
379+ for k , v in spec_dict .items () if k not in cls ._SECRET_SPEC_KEYS
380+ }
381+
382+ return cls (
383+ secret_id = secret_id ,
384+ version_id = version_id ,
385+ version_stage = version_stage ,
386+ ** client_settings ,
387+ )
388+
389+ def apply_pipeline_options (self , options : Any ) -> None :
390+ """Uses AWS credentials and connection settings from AwsOptions.
391+
392+ Only the (unresolved) option values are stored; credential secret specs
393+ are resolved when the secret is fetched, e.g. on workers. Settings in
394+ ``client_settings`` still take precedence.
395+
396+ Args:
397+ options: PipelineOptions, or a dict of option names to values.
398+ """
399+ from apache_beam .io .aws .clients import client_options
400+ from apache_beam .options .pipeline_options import AwsOptions
401+ aws_options = client_options .view_options (options , AwsOptions )
402+ self ._aws_options = aws_options .get_all_options (
403+ drop_default = True , current_only = True ) or None
404+
405+ def _get_client_kwargs (self ) -> dict [str , Any ]:
406+ if not self ._aws_options and not self ._client_settings :
407+ return {}
408+ from apache_beam .io .aws .clients import client_options
409+ # Unset settings are dropped so boto3 falls back to its own defaults.
410+ return {
411+ k : v
412+ for k , v in client_options .get_client_kwargs (
413+ self ._aws_options or {}, overrides = self ._client_settings ).items ()
414+ if v is not None
415+ }
416+
417+ def get_secret_bytes (self ) -> bytes :
418+ try :
419+ import base64
420+ import boto3
421+
422+ client = boto3 .client ('secretsmanager' , ** self ._get_client_kwargs ())
423+
424+ request_kwargs = {'SecretId' : self ._secret_id }
425+ if self ._version_id :
426+ request_kwargs ['VersionId' ] = self ._version_id
427+ if self ._version_stage :
428+ request_kwargs ['VersionStage' ] = self ._version_stage
429+
430+ response = client .get_secret_value (** request_kwargs )
431+ if 'SecretString' in response :
432+ return response ['SecretString' ].encode ('utf-8' )
433+ elif 'SecretBinary' in response :
434+ binary_data = response ['SecretBinary' ]
435+ if isinstance (binary_data , str ):
436+ return base64 .b64decode (binary_data )
437+ return binary_data
438+ else :
439+ raise ValueError (
440+ f"Secret '{ self ._secret_id } ' contains neither SecretString nor SecretBinary."
441+ )
442+ except Exception as e :
443+ raise RuntimeError (
444+ 'Failed to retrieve secret bytes for secret '
445+ f'{ self ._secret_id } with exception { e } ' )
446+
447+ def __eq__ (self , other : Any ) -> bool :
448+ if not isinstance (other , AwsSecret ):
449+ return False
450+ return (
451+ self ._secret_id == other ._secret_id and
452+ self ._version_id == other ._version_id and
453+ self ._version_stage == other ._version_stage and
454+ self ._aws_options == other ._aws_options and
455+ self ._client_settings == other ._client_settings )
456+
457+
274458class GcpHsmGeneratedSecret (Secret ):
275459 """A secret manager implementation that generates a secret using a GCP HSM key
276460 and stores it in Google Cloud Secret Manager. If the secret already exists,
@@ -458,9 +642,11 @@ def generate_dek(self, dek_size: int = 32) -> bytes:
458642_SECRET_TYPE_TO_SECRET_MANAGER : dict [str , str ] = {
459643 "gcpsecret" : "GoogleCloudSecretManager" ,
460644 "gcphsmgeneratedsecret" : "GoogleCloudHsmGeneratedSecretManager" ,
645+ "awssecret" : "AwsSecretsManager" ,
461646}
462647
463648_SECRET_CLASSES : dict [str , Any ] = {
464649 "googlecloudsecretmanager" : "GcpSecret" ,
465650 "googlecloudhsmgeneratedsecretmanager" : "GcpHsmGeneratedSecret" ,
651+ "awssecretsmanager" : "AwsSecret" ,
466652}
0 commit comments