Skip to content

Commit 5faad8f

Browse files
authored
[Python] Add AwsSecret to the Secret utility (#40558)
1 parent b764d86 commit 5faad8f

6 files changed

Lines changed: 368 additions & 8 deletions

File tree

‎CHANGES.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,7 @@
7979
* ReadFromBigQuery now supports Lakehouse runtime catalog (BigLake metastore) tables with `method=DIRECT_READ`, using 4-part `project.catalog.namespace.table` identifiers. Previously `project:catalog.namespace.table` was silently mis-parsed and tables that report no `numBytes` failed to split (Python) ([#39597](https://github.com/apache/beam/issues/39597)).
8080
* (Python) Added `AwsOptions` with `--aws_*` pipeline options (e.g. `--aws_access_key_id`, `--aws_region_name`, `--aws_endpoint_url`) that set defaults for all AWS service clients. The existing `--s3_*` options in `S3Options` override these defaults for the S3 client only ([#40528](https://github.com/apache/beam/issues/40528)).
8181
* (Python) Added `--aws_credentials_secret_manager` so AWS credentials (`--aws_access_key_id`, `--aws_secret_access_key`, `--aws_session_token` and their `--s3_*` overrides) can be secret specs resolved through a secret manager such as `GoogleCloudSecretManager`, instead of plain-text values ([#40528](https://github.com/apache/beam/issues/40528)).
82+
* (Python) Added `AwsSecret` (`AwsSecretsManager` / `type:awssecret`) to retrieve secrets from AWS Secrets Manager, configured from `AwsOptions` (`--aws_*`) with optional per-secret connection overrides ([#40528](https://github.com/apache/beam/issues/40528)).
8283

8384
## Breaking Changes
8485

‎sdks/python/apache_beam/io/aws/clients/client_options.py‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,8 +32,15 @@
3232

3333
from apache_beam.options.pipeline_options import AwsOptions
3434
from apache_beam.options.pipeline_options import PipelineOptions
35+
from apache_beam.utils.secret import _SECRET_CLASSES
3536
from apache_beam.utils.secret import Secret
3637

38+
# Secret manager names that resolve to AwsSecret. AWS credentials cannot be
39+
# stored there, because the credentials are needed to access it.
40+
_AWS_SECRET_MANAGERS = frozenset(
41+
name for name, secret_cls in _SECRET_CLASSES.items()
42+
if secret_cls == 'AwsSecret')
43+
3744
# Client settings (see AwsOptions.get_client_settings) that hold credentials.
3845
CREDENTIAL_SETTINGS = frozenset(
3946
{'access_key_id', 'secret_access_key', 'session_token'})
@@ -127,6 +134,12 @@ def _parse_credential(value: Optional[str],
127134
"""
128135
if not value:
129136
return value
137+
if (secret_manager and
138+
secret_manager.strip().lower() in _AWS_SECRET_MANAGERS):
139+
raise ValueError(
140+
'--aws_credentials_secret_manager cannot be AWS Secrets Manager, '
141+
'because the AWS credentials are needed to access it. Store the '
142+
'credentials in another secret manager, e.g. GoogleCloudSecretManager.')
130143
return _fetch_secret(value, secret_manager)
131144

132145

‎sdks/python/apache_beam/transforms/core.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3497,7 +3497,8 @@ def expand(self, pcoll):
34973497
from apache_beam.transforms.util import GroupByEncryptedKey
34983498
from apache_beam.transforms.util import Secret
34993499

3500-
secret = Secret.parse_secret_option(replace_with_gbek_secret)
3500+
secret = Secret.parse_secret_option(
3501+
replace_with_gbek_secret, pcoll.pipeline._options)
35013502
return (pcoll | "Group by encrypted key" >> GroupByEncryptedKey(secret))
35023503

35033504
from apache_beam.transforms.trigger import DataLossReason

‎sdks/python/apache_beam/transforms/util.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,7 @@
8282
from apache_beam.utils import shared
8383
from apache_beam.utils import windowed_value
8484
from apache_beam.utils.annotations import deprecated
85+
from apache_beam.utils.secret import AwsSecret
8586
from apache_beam.utils.secret import GcpHsmGeneratedSecret
8687
from apache_beam.utils.secret import GcpSecret
8788
from apache_beam.utils.secret import Secret
@@ -94,6 +95,7 @@
9495
_LOGGER = logging.getLogger(__name__)
9596

9697
__all__ = [
98+
'AwsSecret',
9799
'BatchElements',
98100
'CoGroupByKey',
99101
'Distinct',

‎sdks/python/apache_beam/utils/secret.py‎

Lines changed: 193 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -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+
274458
class 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

Comments
 (0)