From c3c06b0b9d47452cfae8d15da089b27779690fbf Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Mon, 12 Oct 2026 21:52:46 -0400 Subject: [PATCH 1/4] Add AwsSecret to the Secret utility Read secrets from AWS Secrets Manager with Secret.from_json(spec, 'AwsSecretsManager') or 'type:awssecret;name:...'. Credentials and connection settings come from the --aws_* options, passed through a new Secret.apply_pipeline_options() hook (also used for --gbek). A spec can override non-credential settings such as region_name or endpoint_url. --- .../io/aws/clients/client_options.py | 13 + sdks/python/apache_beam/transforms/core.py | 3 +- sdks/python/apache_beam/transforms/util.py | 2 + sdks/python/apache_beam/utils/secret.py | 192 ++++++++++++- sdks/python/apache_beam/utils/secret_test.py | 267 ++++++++++++++++++ 5 files changed, 471 insertions(+), 6 deletions(-) diff --git a/sdks/python/apache_beam/io/aws/clients/client_options.py b/sdks/python/apache_beam/io/aws/clients/client_options.py index 4057886a1528..6a484fd220e7 100644 --- a/sdks/python/apache_beam/io/aws/clients/client_options.py +++ b/sdks/python/apache_beam/io/aws/clients/client_options.py @@ -32,8 +32,15 @@ from apache_beam.options.pipeline_options import AwsOptions from apache_beam.options.pipeline_options import PipelineOptions +from apache_beam.utils.secret import _SECRET_CLASSES from apache_beam.utils.secret import Secret +# Secret manager names that resolve to AwsSecret. AWS credentials cannot be +# stored there, because the credentials are needed to access it. +_AWS_SECRET_MANAGERS = frozenset( + name for name, secret_cls in _SECRET_CLASSES.items() + if secret_cls == 'AwsSecret') + # Client settings (see AwsOptions.get_client_settings) that hold credentials. CREDENTIAL_SETTINGS = frozenset( {'access_key_id', 'secret_access_key', 'session_token'}) @@ -127,6 +134,12 @@ def _parse_credential(value: Optional[str], """ if not value: return value + if (secret_manager and + secret_manager.strip().lower() in _AWS_SECRET_MANAGERS): + raise ValueError( + '--aws_credentials_secret_manager cannot be AWS Secrets Manager, ' + 'because the AWS credentials are needed to access it. Store the ' + 'credentials in another secret manager, e.g. GoogleCloudSecretManager.') return _fetch_secret(value, secret_manager) diff --git a/sdks/python/apache_beam/transforms/core.py b/sdks/python/apache_beam/transforms/core.py index 07c6955af7e4..16af71b5d0ba 100644 --- a/sdks/python/apache_beam/transforms/core.py +++ b/sdks/python/apache_beam/transforms/core.py @@ -3497,7 +3497,8 @@ def expand(self, pcoll): from apache_beam.transforms.util import GroupByEncryptedKey from apache_beam.transforms.util import Secret - secret = Secret.parse_secret_option(replace_with_gbek_secret) + secret = Secret.parse_secret_option( + replace_with_gbek_secret, pcoll.pipeline._options) return (pcoll | "Group by encrypted key" >> GroupByEncryptedKey(secret)) from apache_beam.transforms.trigger import DataLossReason diff --git a/sdks/python/apache_beam/transforms/util.py b/sdks/python/apache_beam/transforms/util.py index b7885544dbfb..60b590447387 100644 --- a/sdks/python/apache_beam/transforms/util.py +++ b/sdks/python/apache_beam/transforms/util.py @@ -82,6 +82,7 @@ from apache_beam.utils import shared from apache_beam.utils import windowed_value from apache_beam.utils.annotations import deprecated +from apache_beam.utils.secret import AwsSecret from apache_beam.utils.secret import GcpHsmGeneratedSecret from apache_beam.utils.secret import GcpSecret from apache_beam.utils.secret import Secret @@ -94,6 +95,7 @@ _LOGGER = logging.getLogger(__name__) __all__ = [ + 'AwsSecret', 'BatchElements', 'CoGroupByKey', 'Distinct', diff --git a/sdks/python/apache_beam/utils/secret.py b/sdks/python/apache_beam/utils/secret.py index 9e2772f70e96..36702522b06a 100644 --- a/sdks/python/apache_beam/utils/secret.py +++ b/sdks/python/apache_beam/utils/secret.py @@ -85,8 +85,21 @@ def generate_secret_bytes() -> bytes: from cryptography.fernet import Fernet return Fernet.generate_key() + def apply_pipeline_options(self, options: Any) -> None: + """Configures this secret from pipeline options. + + Called by :meth:`from_json` and :meth:`parse_secret_option` when pipeline + options are given. The default implementation does nothing; secrets that + need settings from pipeline options (e.g. :class:`AwsSecret`) override it. + + Args: + options: PipelineOptions, or a dict of option names to values. + """ + pass + @classmethod - def parse_secret_option(cls, secret: str) -> 'Secret': + def parse_secret_option( + cls, secret: str, options: Optional[Any] = None) -> 'Secret': """Parses a secret string and returns the appropriate secret type. The secret string should be formatted like: @@ -94,6 +107,11 @@ def parse_secret_option(cls, secret: str) -> 'Secret': For example, 'type:GcpSecret;version_name:my_secret/versions/latest' would return a GcpSecret initialized with 'my_secret/versions/latest'. + + Args: + secret: The secret string. + options: Optional pipeline options, passed to + :meth:`apply_pipeline_options` of the returned secret. """ param_map = {} for param in secret.split(';'): @@ -110,22 +128,35 @@ def parse_secret_option(cls, secret: str) -> 'Secret': if not secret_manager: raise ValueError( f'Invalid secret type {secret_type}, currently only ' - 'GcpSecret and GcpHsmGeneratedSecret are supported') + 'GcpSecret, GcpHsmGeneratedSecret, and AwsSecret are supported') - return cls.from_json(json.dumps(param_map), secret_manager) + return cls.from_json(json.dumps(param_map), secret_manager, options) @classmethod def from_json( - cls, spec: str, secret_manager: Optional[str] = None) -> 'Secret': + cls, + spec: str, + secret_manager: Optional[str] = None, + options: Optional[Any] = None) -> 'Secret': """Return a Secret instance based on secret_manager provider and secret specification. Args: spec: Secret string (raw secret or JSON specification string). secret_manager: Secret manager string (e.g. 'GoogleCloudSecretManager'). + options: Optional pipeline options (PipelineOptions or a dict), passed to + :meth:`apply_pipeline_options` of the returned secret. For example, + AwsSecret uses them to get AWS credentials from AwsOptions. Returns: An instance of Secret. """ + secret = cls._from_json(spec, secret_manager) + if options is not None: + secret.apply_pipeline_options(options) + return secret + + @classmethod + def _from_json(cls, spec: str, secret_manager: Optional[str]) -> 'Secret': if not isinstance(spec, str): raise TypeError( f"Secret 'spec' must be a string, got {type(spec).__name__}") @@ -163,7 +194,7 @@ def from_json( return secret_cls(spec) else: raise ValueError( - f"Unsupported secret manager: '{secret_manager_name}'. Currently supported options: 'GoogleCloudSecretManager', 'GoogleCloudHsmGeneratedSecretManager'." + f"Unsupported secret manager: '{secret_manager_name}'. Currently supported options: 'GoogleCloudSecretManager', 'GoogleCloudHsmGeneratedSecretManager', 'AwsSecretsManager'." ) # If secret_manager is not set or empty, check if spec is a JSON specification dict @@ -271,6 +302,155 @@ def __eq__(self, secret): return self._version_name == getattr(secret, '_version_name', None) +class AwsSecret(Secret): + """A secret manager implementation that retrieves secrets from AWS Secrets + Manager. + + AWS credentials and connection settings (region, endpoint, SSL) come from the + ``--aws_*`` defaults in :class:`AwsOptions` when pipeline options are given + (see :meth:`apply_pipeline_options`), and otherwise from boto3's defaults + (e.g. environment variables, ~/.aws, or an IAM role). Connection settings + can be overridden per secret with keyword arguments, e.g. + ``AwsSecret('my-secret', region_name='us-west-2')``, or in the secret spec, + e.g. ``type:awssecret;name:my-secret;region_name:us-west-2``. + """ + # Spec keys that identify the secret. Other spec keys are client settings. + _SECRET_SPEC_KEYS = frozenset({ + 'secret_id', + 'name', + 'arn', + 'version_id', + 'version', + 'version_stage', + 'stage', + }) + + def __init__( + self, + secret_id: str, + version_id: Optional[str] = None, + version_stage: Optional[str] = None, + **client_settings: Any): + """Initializes an AwsSecret object. + + Args: + secret_id: The ARN or name of the secret in AWS Secrets Manager. + version_id: The unique identifier of the version of the secret. + version_stage: The staging label of the version of the secret (e.g. 'AWSCURRENT'). + **client_settings: AWS client connection settings that override the + ``--aws_*`` pipeline options for this secret. Names are AwsOptions + setting names without the ``aws_`` prefix (e.g. ``region_name``, + ``endpoint_url``, ``verify``, ``disable_ssl``). Credentials cannot be + set here; use the ``--aws_*`` pipeline options instead. + """ + super().__init__() + self._secret_id = secret_id + self._version_id = version_id + self._version_stage = version_stage + # Set by apply_pipeline_options. + self._aws_options: Optional[dict[str, Any]] = None + self._client_settings: Optional[dict[str, Any]] = None + if client_settings: + from apache_beam.io.aws.clients import client_options + self._client_settings = client_options.parse_connection_settings( + client_settings) + + @classmethod + def from_dict(cls, spec_dict: dict[str, str]) -> 'AwsSecret': + """Initialize AwsSecret from a dictionary specification. + + Keys other than the secret identifiers are AWS client connection settings + (see the ``**client_settings`` keyword arguments of :meth:`__init__`). + """ + secret_id = ( + spec_dict.get("secret_id") or spec_dict.get("name") or + spec_dict.get("arn")) + if not secret_id: + raise ValueError("Secret ID or name must be specified in secret spec.") + + version_id = spec_dict.get("version_id") or spec_dict.get("version") + version_stage = spec_dict.get("version_stage") or spec_dict.get("stage") + client_settings = { + k: v + for k, v in spec_dict.items() if k not in cls._SECRET_SPEC_KEYS + } + + return cls( + secret_id=secret_id, + version_id=version_id, + version_stage=version_stage, + **client_settings, + ) + + def apply_pipeline_options(self, options: Any) -> None: + """Uses AWS credentials and connection settings from AwsOptions. + + Only the (unresolved) option values are stored; credential secret specs + are resolved when the secret is fetched, e.g. on workers. Settings in + ``client_settings`` still take precedence. + + Args: + options: PipelineOptions, or a dict of option names to values. + """ + from apache_beam.io.aws.clients import client_options + from apache_beam.options.pipeline_options import AwsOptions + aws_options = client_options.view_options(options, AwsOptions) + self._aws_options = aws_options.get_all_options( + drop_default=True, current_only=True) or None + + def _get_client_kwargs(self) -> dict[str, Any]: + if not self._aws_options and not self._client_settings: + return {} + from apache_beam.io.aws.clients import client_options + # Unset settings are dropped so boto3 falls back to its own defaults. + return { + k: v + for k, v in client_options.get_client_kwargs( + self._aws_options or {}, overrides=self._client_settings).items() + if v is not None + } + + def get_secret_bytes(self) -> bytes: + try: + import base64 + import boto3 + + client = boto3.client('secretsmanager', **self._get_client_kwargs()) + + request_kwargs = {'SecretId': self._secret_id} + if self._version_id: + request_kwargs['VersionId'] = self._version_id + if self._version_stage: + request_kwargs['VersionStage'] = self._version_stage + + response = client.get_secret_value(**request_kwargs) + if 'SecretString' in response: + return response['SecretString'].encode('utf-8') + elif 'SecretBinary' in response: + binary_data = response['SecretBinary'] + if isinstance(binary_data, str): + return base64.b64decode(binary_data) + return binary_data + else: + raise ValueError( + f"Secret '{self._secret_id}' contains neither SecretString nor SecretBinary." + ) + except Exception as e: + raise RuntimeError( + 'Failed to retrieve secret bytes for secret ' + f'{self._secret_id} with exception {e}') + + def __eq__(self, other: Any) -> bool: + if not isinstance(other, AwsSecret): + return False + return ( + self._secret_id == other._secret_id and + self._version_id == other._version_id and + self._version_stage == other._version_stage and + self._aws_options == other._aws_options and + self._client_settings == other._client_settings) + + class GcpHsmGeneratedSecret(Secret): """A secret manager implementation that generates a secret using a GCP HSM key and stores it in Google Cloud Secret Manager. If the secret already exists, @@ -458,9 +638,11 @@ def generate_dek(self, dek_size: int = 32) -> bytes: _SECRET_TYPE_TO_SECRET_MANAGER: dict[str, str] = { "gcpsecret": "GoogleCloudSecretManager", "gcphsmgeneratedsecret": "GoogleCloudHsmGeneratedSecretManager", + "awssecret": "AwsSecretsManager", } _SECRET_CLASSES: dict[str, Any] = { "googlecloudsecretmanager": "GcpSecret", "googlecloudhsmgeneratedsecretmanager": "GcpHsmGeneratedSecret", + "awssecretsmanager": "AwsSecret", } \ No newline at end of file diff --git a/sdks/python/apache_beam/utils/secret_test.py b/sdks/python/apache_beam/utils/secret_test.py index 179b7ca3a5f7..d5867facf6f8 100644 --- a/sdks/python/apache_beam/utils/secret_test.py +++ b/sdks/python/apache_beam/utils/secret_test.py @@ -15,6 +15,7 @@ # limitations under the License. # +import base64 import json import unittest from unittest import mock @@ -23,6 +24,7 @@ from parameterized import parameterized from apache_beam.utils.annotations import BeamDeprecationWarning +from apache_beam.utils.secret import AwsSecret from apache_beam.utils.secret import GcpHsmGeneratedSecret from apache_beam.utils.secret import GcpSecret from apache_beam.utils.secret import RawSecret @@ -33,6 +35,11 @@ except ImportError: secretmanager = None # type: ignore[assignment] +try: + import boto3 +except ImportError: + boto3 = None # type: ignore[assignment] + class SecretTest(unittest.TestCase): @parameterized.expand([ @@ -45,6 +52,12 @@ class SecretTest(unittest.TestCase): param( secret_string='type:gcpsecreT;version_name:my_secret/versions/latest', secret=GcpSecret('my_secret/versions/latest')), + param( + secret_string='type:AwsSecret;secret_id:my_aws_secret', + secret=AwsSecret('my_aws_secret')), + param( + secret_string='type:awssecreT;name:my_aws_secret;version:v1', + secret=AwsSecret('my_aws_secret', version_id='v1')), ]) def test_secret_manager_parses_correctly(self, secret_string, secret): self.assertEqual(secret, Secret.parse_secret_option(secret_string)) @@ -59,6 +72,12 @@ def test_secret_manager_parses_correctly(self, secret_string, secret): param( secret_string='type:gcpsecreT;version_name:foo;extra:val', exception_str='Invalid secret parameter extra'), + param( + secret_string='type:awssecreT', + exception_str='Secret ID or name must be specified in secret spec'), + param( + secret_string='type:awssecreT;name:foo;invalid_key:bar', + exception_str='Invalid secret parameter invalid_key'), ]) def test_secret_manager_throws_on_invalid(self, secret_string, exception_str): with self.assertRaisesRegex(Exception, exception_str): @@ -348,6 +367,231 @@ def test_getstate_clears_cached_secret(self, mock_sm_client_cls): self.assertIsNone(state["_cached_secret_bytes"]) +@unittest.skipIf(boto3 is None, 'AWS dependencies are not installed') +class AwsSecretTest(unittest.TestCase): + def setUp(self): + from apache_beam.io.aws.clients import client_options + client_options._fetch_secret.cache_clear() + patcher = mock.patch('boto3.client') + self.mock_boto3_client = patcher.start() + self.addCleanup(patcher.stop) + self.mock_sm = self.mock_boto3_client.return_value + self.mock_sm.get_secret_value.return_value = {'SecretString': 'aws-value'} + + def _options(self, *flags): + from apache_beam.options.pipeline_options import PipelineOptions + return PipelineOptions(list(flags)) + + def _client_kwargs(self): + args, kwargs = self.mock_boto3_client.call_args + self.assertEqual(args, ('secretsmanager', )) + return kwargs + + def test_get_secret_string_cached(self): + secret = AwsSecret('my-secret') + self.assertEqual(secret.get_str(cacheSecret=True), 'aws-value') + self.assertEqual(secret.get_bytes(cacheSecret=True), b'aws-value') + # Fetched once; without options or settings boto3's defaults are used. + self.mock_boto3_client.assert_called_once_with('secretsmanager') + self.mock_sm.get_secret_value.assert_called_once_with(SecretId='my-secret') + # The cached value is not pickled. + self.assertIsNone(secret.__getstate__()['_cached_secret_bytes']) + + @parameterized.expand([ + ('bytes', b'binary'), + ('base64_string', base64.b64encode(b'binary').decode('utf-8')), + ]) + def test_get_secret_binary(self, _, payload): + self.mock_sm.get_secret_value.return_value = {'SecretBinary': payload} + self.assertEqual(AwsSecret('my-secret').get_bytes(), b'binary') + + def test_version_and_stage(self): + AwsSecret.from_dict({ + 'name': 'my-secret', 'version_id': 'v-12345', 'stage': 'AWSCURRENT' + }).get_str() + self.mock_sm.get_secret_value.assert_called_once_with( + SecretId='my-secret', VersionId='v-12345', VersionStage='AWSCURRENT') + + @parameterized.expand([ + ('api_error', { + 'side_effect': Exception('Secret not found') + }, 'Secret not found'), + ('empty_payload', { + 'return_value': {} + }, 'contains neither SecretString nor SecretBinary'), + ]) + def test_fetch_failure_raises_runtime_error(self, _, behavior, message): + self.mock_sm.get_secret_value.configure_mock(**behavior) + with self.assertRaisesRegex(RuntimeError, message): + AwsSecret('my-secret').get_str() + + def test_connection_settings_in_spec(self): + secret = AwsSecret.from_dict({ + 'name': 'my-secret', + 'region_name': 'us-west-2', + 'endpoint_url': 'http://localhost:4566', + 'verify': 'false', + 'disable_ssl': 'true', + }) + self.assertEqual( + secret._client_settings, + { + 'region_name': 'us-west-2', + 'endpoint_url': 'http://localhost:4566', + 'verify': 'false', + 'disable_ssl': True, + }) + # Settings apply without pipeline options too. + AwsSecret('my-secret', region_name='eu-west-1').get_str() + self.assertEqual( + self._client_kwargs(), { + 'region_name': 'eu-west-1', 'use_ssl': True + }) + + @parameterized.expand([ + ('credential', { + 'access_key_id': 'x' + }, 'Use the --aws_access_key_id pipeline option'), + ('alias', { + 'region': 'x' + }, 'Invalid secret parameter region'), + ('option_name', { + 'aws_region_name': 'x' + }, 'Invalid secret parameter aws_region_name'), + ('bad_bool', { + 'disable_ssl': 'maybe' + }, 'Invalid boolean value for disable_ssl'), + ]) + def test_invalid_spec_raises_value_error(self, _, extra, message): + with self.assertRaisesRegex(ValueError, message): + AwsSecret.from_dict({'name': 'my-secret', **extra}) + + @mock.patch.dict('os.environ', {'AWS_DEFAULT_REGION': 'eu-central-1'}) + def test_settings_from_aws_options(self): + # --aws_* options take precedence over the environment, and --s3_* + # overrides apply to the S3 client only. + secret = Secret.parse_secret_option( + 'type:awssecret;name:my-secret', + self._options( + '--aws_access_key_id=AKIDEXAMPLE', + '--aws_secret_access_key=SECRETEXAMPLE', + '--aws_session_token=TOKEN', + '--aws_region_name=us-east-2', + '--aws_endpoint_url=http://localhost:4566', + '--s3_access_key_id=S3KEY', + '--s3_endpoint_url=http://localhost:9000')) + self.assertEqual(secret.get_str(), 'aws-value') + self.assertEqual( + self._client_kwargs(), + { + 'aws_access_key_id': 'AKIDEXAMPLE', + 'aws_secret_access_key': 'SECRETEXAMPLE', + 'aws_session_token': 'TOKEN', + 'region_name': 'us-east-2', + 'endpoint_url': 'http://localhost:4566', + 'use_ssl': True, + }) + + def test_spec_settings_override_aws_options(self): + secret = Secret.from_json( + '{"name": "my-secret", "region_name": "us-west-2", ' + '"endpoint_url": "http://localhost:4566", "disable_ssl": "true"}', + 'AwsSecretsManager', + self._options( + '--aws_access_key_id=AKIDEXAMPLE', + '--aws_region_name=us-east-2', + '--aws_endpoint_url=http://other:4566')) + secret.get_str() + self.assertEqual( + self._client_kwargs(), + { + 'aws_access_key_id': 'AKIDEXAMPLE', + 'region_name': 'us-west-2', + 'endpoint_url': 'http://localhost:4566', + 'use_ssl': False, + }) + + # An explicit false in the spec overrides --aws_disable_ssl. + Secret.parse_secret_option( + 'type:awssecret;name:my-secret;disable_ssl:false', + self._options('--aws_disable_ssl')).get_str() + self.assertTrue(self._client_kwargs()['use_ssl']) + + @mock.patch('google.cloud.secretmanager.SecretManagerServiceClient') + def test_credentials_resolved_through_gcp_secret_manager( + self, mock_gcp_client): + payloads = { + 'projects/p/secrets/key-id/versions/latest': b'AKIDEXAMPLE', + 'projects/p/secrets/secret-key/versions/latest': b'SECRETEXAMPLE', + } + + def access(request): + response = mock.MagicMock() + response.payload.data = payloads[request['name']] + return response + + mock_gcp_client.return_value.access_secret_version.side_effect = access + + secret = Secret.from_json( + '{"name": "my-secret"}', + 'AwsSecretsManager', + self._options( + '--aws_credentials_secret_manager=GoogleCloudSecretManager', + '--aws_access_key_id={"name": "key-id", "project": "p"}', + '--aws_secret_access_key={"name": "secret-key", "project": "p"}')) + self.assertEqual(secret.get_str(), 'aws-value') + kwargs = self._client_kwargs() + self.assertEqual(kwargs['aws_access_key_id'], 'AKIDEXAMPLE') + self.assertEqual(kwargs['aws_secret_access_key'], 'SECRETEXAMPLE') + + def test_credentials_in_aws_secrets_manager_rejected(self): + secret = Secret.from_json( + '{"name": "my-secret"}', + 'AwsSecretsManager', + self._options( + '--aws_credentials_secret_manager=AwsSecretsManager', + '--aws_access_key_id={"name": "key-id"}')) + with self.assertRaisesRegex(RuntimeError, 'cannot be AWS Secrets Manager'): + secret.get_str() + self.mock_boto3_client.assert_not_called() + + def test_pickle_and_equality(self): + import pickle + with mock.patch('apache_beam.utils.secret.Secret.from_json', + wraps=Secret.from_json) as spy: + secret = Secret.from_json( + '{"name": "my-secret", "region_name": "us-west-2"}', + 'AwsSecretsManager', + self._options( + '--aws_credentials_secret_manager=GoogleCloudSecretManager', + '--aws_access_key_id=projects/p/secrets/key-id/versions/1')) + # Only the outer from_json call; the credential spec is not resolved yet. + self.assertEqual(spy.call_count, 1) + + restored = pickle.loads(pickle.dumps(secret)) + self.assertEqual(restored, secret) + self.assertEqual( + restored._aws_options['aws_access_key_id'], + 'projects/p/secrets/key-id/versions/1') + self.assertEqual(restored._client_settings, {'region_name': 'us-west-2'}) + + with_options_1 = AwsSecret('s') + with_options_1.apply_pipeline_options({'aws_region_name': 'us-east-1'}) + with_options_2 = AwsSecret('s') + with_options_2.apply_pipeline_options({'aws_region_name': 'us-east-2'}) + self.assertNotEqual(with_options_1, with_options_2) + self.assertNotEqual( + AwsSecret('s', region_name='us-east-1'), + AwsSecret('s', region_name='us-east-2')) + + def test_other_secrets_ignore_pipeline_options(self): + secret = Secret.from_json( + 'projects/p/secrets/s/versions/1', + 'GoogleCloudSecretManager', + self._options('--aws_region_name=us-east-2')) + self.assertEqual(secret, GcpSecret('projects/p/secrets/s/versions/1')) + + class RawSecretTest(unittest.TestCase): def test_raw_secret_str(self): secret = RawSecret("STATIC_SECRET_") @@ -403,6 +647,21 @@ def test_secret_factory(self): spec=spec_dict, # type: ignore[arg-type] secret_manager="GoogleCloudSecretManager") + def test_secret_factory_aws(self): + spec = json.dumps({"name": "test-aws-secret", "version": "v1"}) + + secret_aws = Secret.from_json( + spec=spec, secret_manager="AwsSecretsManager") + self.assertIsInstance(secret_aws, AwsSecret) + self.assertEqual(secret_aws._secret_id, "test-aws-secret") + self.assertEqual(secret_aws._version_id, "v1") + + # When spec is a plain string + secret_plain = Secret.from_json( + spec="my-secret-arn", secret_manager="AwsSecretsManager") + self.assertIsInstance(secret_plain, AwsSecret) + self.assertEqual(secret_plain._secret_id, "my-secret-arn") + def test_secret_factory_hsm(self): hsm_spec = json.dumps({ "project_id": "p", @@ -442,6 +701,14 @@ def test_equality(self): self.assertNotEqual(gcp1, gcp3) self.assertNotEqual(gcp1, raw1) + aws1 = AwsSecret.from_dict({"name": "sec", "version": "v1"}) + aws2 = AwsSecret.from_dict({"name": "sec", "version": "v1"}) + aws3 = AwsSecret.from_dict({"name": "other", "version": "v1"}) + self.assertEqual(aws1, aws2) + self.assertNotEqual(aws1, aws3) + self.assertNotEqual(aws1, gcp1) + self.assertNotEqual(aws1, raw1) + hsm1 = GcpHsmGeneratedSecret("p", "l", "r", "k", "j") hsm2 = GcpHsmGeneratedSecret("p", "l", "r", "k", "j") hsm3 = GcpHsmGeneratedSecret("p", "l", "r", "k", "other") From 82d7b0762c9341626f24bf21e92a170b1bbcbe42 Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Tue, 13 Oct 2026 13:39:21 -0400 Subject: [PATCH 2/4] Edited CHANGES.md --- CHANGES.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGES.md b/CHANGES.md index 7693e65d78a0..b5976ab01f61 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -75,6 +75,7 @@ * 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)). * (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)). * (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)). +* (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)). ## Breaking Changes From 2780090af1c067181f4b37e70f6f038ce74481c1 Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Tue, 13 Oct 2026 14:33:55 -0400 Subject: [PATCH 3/4] Reformat --- sdks/python/apache_beam/utils/secret_test.py | 43 ++++++++++++-------- 1 file changed, 26 insertions(+), 17 deletions(-) diff --git a/sdks/python/apache_beam/utils/secret_test.py b/sdks/python/apache_beam/utils/secret_test.py index d5867facf6f8..4b84be318ce5 100644 --- a/sdks/python/apache_beam/utils/secret_test.py +++ b/sdks/python/apache_beam/utils/secret_test.py @@ -413,12 +413,16 @@ def test_version_and_stage(self): SecretId='my-secret', VersionId='v-12345', VersionStage='AWSCURRENT') @parameterized.expand([ - ('api_error', { - 'side_effect': Exception('Secret not found') - }, 'Secret not found'), - ('empty_payload', { - 'return_value': {} - }, 'contains neither SecretString nor SecretBinary'), + ( + 'api_error', { + 'side_effect': Exception('Secret not found') + }, + 'Secret not found'), + ( + 'empty_payload', { + 'return_value': {} + }, + 'contains neither SecretString nor SecretBinary'), ]) def test_fetch_failure_raises_runtime_error(self, _, behavior, message): self.mock_sm.get_secret_value.configure_mock(**behavior) @@ -449,18 +453,24 @@ def test_connection_settings_in_spec(self): }) @parameterized.expand([ - ('credential', { - 'access_key_id': 'x' - }, 'Use the --aws_access_key_id pipeline option'), + ( + 'credential', { + 'access_key_id': 'x' + }, + 'Use the --aws_access_key_id pipeline option'), ('alias', { 'region': 'x' }, 'Invalid secret parameter region'), - ('option_name', { - 'aws_region_name': 'x' - }, 'Invalid secret parameter aws_region_name'), - ('bad_bool', { - 'disable_ssl': 'maybe' - }, 'Invalid boolean value for disable_ssl'), + ( + 'option_name', { + 'aws_region_name': 'x' + }, + 'Invalid secret parameter aws_region_name'), + ( + 'bad_bool', { + 'disable_ssl': 'maybe' + }, + 'Invalid boolean value for disable_ssl'), ]) def test_invalid_spec_raises_value_error(self, _, extra, message): with self.assertRaisesRegex(ValueError, message): @@ -650,8 +660,7 @@ def test_secret_factory(self): def test_secret_factory_aws(self): spec = json.dumps({"name": "test-aws-secret", "version": "v1"}) - secret_aws = Secret.from_json( - spec=spec, secret_manager="AwsSecretsManager") + secret_aws = Secret.from_json(spec=spec, secret_manager="AwsSecretsManager") self.assertIsInstance(secret_aws, AwsSecret) self.assertEqual(secret_aws._secret_id, "test-aws-secret") self.assertEqual(secret_aws._version_id, "v1") From bc846b3b0143f5e5f004de8975590cf1f27c2e43 Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Tue, 13 Oct 2026 20:29:02 -0400 Subject: [PATCH 4/4] Simplify tests. --- sdks/python/apache_beam/utils/secret_test.py | 221 +++++-------------- 1 file changed, 51 insertions(+), 170 deletions(-) diff --git a/sdks/python/apache_beam/utils/secret_test.py b/sdks/python/apache_beam/utils/secret_test.py index 4b84be318ce5..a3c0e3b7c90a 100644 --- a/sdks/python/apache_beam/utils/secret_test.py +++ b/sdks/python/apache_beam/utils/secret_test.py @@ -17,6 +17,7 @@ import base64 import json +import pickle import unittest from unittest import mock @@ -78,6 +79,9 @@ def test_secret_manager_parses_correctly(self, secret_string, secret): param( secret_string='type:awssecreT;name:foo;invalid_key:bar', exception_str='Invalid secret parameter invalid_key'), + param( + secret_string='type:awssecreT;name:foo;access_key_id:bar', + exception_str='Use the --aws_access_key_id pipeline option'), ]) def test_secret_manager_throws_on_invalid(self, secret_string, exception_str): with self.assertRaisesRegex(Exception, exception_str): @@ -391,170 +395,70 @@ def test_get_secret_string_cached(self): secret = AwsSecret('my-secret') self.assertEqual(secret.get_str(cacheSecret=True), 'aws-value') self.assertEqual(secret.get_bytes(cacheSecret=True), b'aws-value') - # Fetched once; without options or settings boto3's defaults are used. self.mock_boto3_client.assert_called_once_with('secretsmanager') self.mock_sm.get_secret_value.assert_called_once_with(SecretId='my-secret') - # The cached value is not pickled. self.assertIsNone(secret.__getstate__()['_cached_secret_bytes']) - @parameterized.expand([ - ('bytes', b'binary'), - ('base64_string', base64.b64encode(b'binary').decode('utf-8')), - ]) - def test_get_secret_binary(self, _, payload): - self.mock_sm.get_secret_value.return_value = {'SecretBinary': payload} - self.assertEqual(AwsSecret('my-secret').get_bytes(), b'binary') + # Keyword settings apply without pipeline options. + AwsSecret('my-secret', region_name='eu-west-1').get_str() + self.assertEqual( + self._client_kwargs(), { + 'region_name': 'eu-west-1', 'use_ssl': True + }) - def test_version_and_stage(self): AwsSecret.from_dict({ 'name': 'my-secret', 'version_id': 'v-12345', 'stage': 'AWSCURRENT' }).get_str() - self.mock_sm.get_secret_value.assert_called_once_with( + self.mock_sm.get_secret_value.assert_called_with( SecretId='my-secret', VersionId='v-12345', VersionStage='AWSCURRENT') - @parameterized.expand([ - ( - 'api_error', { - 'side_effect': Exception('Secret not found') - }, - 'Secret not found'), - ( - 'empty_payload', { - 'return_value': {} - }, - 'contains neither SecretString nor SecretBinary'), - ]) - def test_fetch_failure_raises_runtime_error(self, _, behavior, message): - self.mock_sm.get_secret_value.configure_mock(**behavior) - with self.assertRaisesRegex(RuntimeError, message): + def test_get_secret_binary(self): + self.mock_sm.get_secret_value.return_value = {'SecretBinary': b'binary'} + self.assertEqual(AwsSecret('my-secret').get_bytes(), b'binary') + + # A base64-encoded string payload is decoded. + self.mock_sm.get_secret_value.return_value = { + 'SecretBinary': base64.b64encode(b'binary').decode('utf-8') + } + self.assertEqual(AwsSecret('my-secret').get_bytes(), b'binary') + + def test_fetch_failure_raises_runtime_error(self): + self.mock_sm.get_secret_value.side_effect = Exception('Secret not found') + with self.assertRaisesRegex(RuntimeError, 'Secret not found'): AwsSecret('my-secret').get_str() - def test_connection_settings_in_spec(self): - secret = AwsSecret.from_dict({ - 'name': 'my-secret', - 'region_name': 'us-west-2', - 'endpoint_url': 'http://localhost:4566', - 'verify': 'false', - 'disable_ssl': 'true', - }) - self.assertEqual( - secret._client_settings, - { - 'region_name': 'us-west-2', - 'endpoint_url': 'http://localhost:4566', - 'verify': 'false', - 'disable_ssl': True, - }) - # Settings apply without pipeline options too. - AwsSecret('my-secret', region_name='eu-west-1').get_str() - self.assertEqual( - self._client_kwargs(), { - 'region_name': 'eu-west-1', 'use_ssl': True - }) + self.mock_sm.get_secret_value.side_effect = None + self.mock_sm.get_secret_value.return_value = {} + with self.assertRaisesRegex( + RuntimeError, 'contains neither SecretString nor SecretBinary'): + AwsSecret('my-secret').get_str() - @parameterized.expand([ - ( - 'credential', { - 'access_key_id': 'x' - }, - 'Use the --aws_access_key_id pipeline option'), - ('alias', { - 'region': 'x' - }, 'Invalid secret parameter region'), - ( - 'option_name', { - 'aws_region_name': 'x' - }, - 'Invalid secret parameter aws_region_name'), - ( - 'bad_bool', { - 'disable_ssl': 'maybe' - }, - 'Invalid boolean value for disable_ssl'), - ]) - def test_invalid_spec_raises_value_error(self, _, extra, message): - with self.assertRaisesRegex(ValueError, message): - AwsSecret.from_dict({'name': 'my-secret', **extra}) - - @mock.patch.dict('os.environ', {'AWS_DEFAULT_REGION': 'eu-central-1'}) - def test_settings_from_aws_options(self): - # --aws_* options take precedence over the environment, and --s3_* - # overrides apply to the S3 client only. + def test_aws_options_and_spec_overrides(self): secret = Secret.parse_secret_option( - 'type:awssecret;name:my-secret', + 'type:awssecret;name:my-secret;region_name:us-west-2;disable_ssl:true', self._options( '--aws_access_key_id=AKIDEXAMPLE', '--aws_secret_access_key=SECRETEXAMPLE', - '--aws_session_token=TOKEN', '--aws_region_name=us-east-2', '--aws_endpoint_url=http://localhost:4566', - '--s3_access_key_id=S3KEY', - '--s3_endpoint_url=http://localhost:9000')) - self.assertEqual(secret.get_str(), 'aws-value') + '--s3_region_name=eu-west-1')) + restored = pickle.loads(pickle.dumps(secret)) + self.assertEqual(restored, secret) + self.assertNotEqual(secret, AwsSecret('my-secret', region_name='us-west-2')) + self.assertEqual(restored.get_str(), 'aws-value') self.assertEqual( self._client_kwargs(), { 'aws_access_key_id': 'AKIDEXAMPLE', 'aws_secret_access_key': 'SECRETEXAMPLE', - 'aws_session_token': 'TOKEN', - 'region_name': 'us-east-2', - 'endpoint_url': 'http://localhost:4566', - 'use_ssl': True, - }) - - def test_spec_settings_override_aws_options(self): - secret = Secret.from_json( - '{"name": "my-secret", "region_name": "us-west-2", ' - '"endpoint_url": "http://localhost:4566", "disable_ssl": "true"}', - 'AwsSecretsManager', - self._options( - '--aws_access_key_id=AKIDEXAMPLE', - '--aws_region_name=us-east-2', - '--aws_endpoint_url=http://other:4566')) - secret.get_str() - self.assertEqual( - self._client_kwargs(), - { - 'aws_access_key_id': 'AKIDEXAMPLE', 'region_name': 'us-west-2', 'endpoint_url': 'http://localhost:4566', 'use_ssl': False, }) - # An explicit false in the spec overrides --aws_disable_ssl. - Secret.parse_secret_option( - 'type:awssecret;name:my-secret;disable_ssl:false', - self._options('--aws_disable_ssl')).get_str() - self.assertTrue(self._client_kwargs()['use_ssl']) - - @mock.patch('google.cloud.secretmanager.SecretManagerServiceClient') - def test_credentials_resolved_through_gcp_secret_manager( - self, mock_gcp_client): - payloads = { - 'projects/p/secrets/key-id/versions/latest': b'AKIDEXAMPLE', - 'projects/p/secrets/secret-key/versions/latest': b'SECRETEXAMPLE', - } - - def access(request): - response = mock.MagicMock() - response.payload.data = payloads[request['name']] - return response - - mock_gcp_client.return_value.access_secret_version.side_effect = access - - secret = Secret.from_json( - '{"name": "my-secret"}', - 'AwsSecretsManager', - self._options( - '--aws_credentials_secret_manager=GoogleCloudSecretManager', - '--aws_access_key_id={"name": "key-id", "project": "p"}', - '--aws_secret_access_key={"name": "secret-key", "project": "p"}')) - self.assertEqual(secret.get_str(), 'aws-value') - kwargs = self._client_kwargs() - self.assertEqual(kwargs['aws_access_key_id'], 'AKIDEXAMPLE') - self.assertEqual(kwargs['aws_secret_access_key'], 'SECRETEXAMPLE') - def test_credentials_in_aws_secrets_manager_rejected(self): + # Credential specs are resolved lazily, so construction succeeds and the + # error is only raised on fetch. secret = Secret.from_json( '{"name": "my-secret"}', 'AwsSecretsManager', @@ -565,42 +469,6 @@ def test_credentials_in_aws_secrets_manager_rejected(self): secret.get_str() self.mock_boto3_client.assert_not_called() - def test_pickle_and_equality(self): - import pickle - with mock.patch('apache_beam.utils.secret.Secret.from_json', - wraps=Secret.from_json) as spy: - secret = Secret.from_json( - '{"name": "my-secret", "region_name": "us-west-2"}', - 'AwsSecretsManager', - self._options( - '--aws_credentials_secret_manager=GoogleCloudSecretManager', - '--aws_access_key_id=projects/p/secrets/key-id/versions/1')) - # Only the outer from_json call; the credential spec is not resolved yet. - self.assertEqual(spy.call_count, 1) - - restored = pickle.loads(pickle.dumps(secret)) - self.assertEqual(restored, secret) - self.assertEqual( - restored._aws_options['aws_access_key_id'], - 'projects/p/secrets/key-id/versions/1') - self.assertEqual(restored._client_settings, {'region_name': 'us-west-2'}) - - with_options_1 = AwsSecret('s') - with_options_1.apply_pipeline_options({'aws_region_name': 'us-east-1'}) - with_options_2 = AwsSecret('s') - with_options_2.apply_pipeline_options({'aws_region_name': 'us-east-2'}) - self.assertNotEqual(with_options_1, with_options_2) - self.assertNotEqual( - AwsSecret('s', region_name='us-east-1'), - AwsSecret('s', region_name='us-east-2')) - - def test_other_secrets_ignore_pipeline_options(self): - secret = Secret.from_json( - 'projects/p/secrets/s/versions/1', - 'GoogleCloudSecretManager', - self._options('--aws_region_name=us-east-2')) - self.assertEqual(secret, GcpSecret('projects/p/secrets/s/versions/1')) - class RawSecretTest(unittest.TestCase): def test_raw_secret_str(self): @@ -671,6 +539,19 @@ def test_secret_factory_aws(self): self.assertIsInstance(secret_plain, AwsSecret) self.assertEqual(secret_plain._secret_id, "my-secret-arn") + # When pipeline options are provided + secret_with_options = Secret.from_json( + spec=json.dumps({ + "name": "test-aws-secret", "region_name": "us-west-2" + }), + secret_manager="AwsSecretsManager", + options={"aws_region_name": "us-east-2"}) + self.assertIsInstance(secret_with_options, AwsSecret) + self.assertEqual( + secret_with_options._aws_options, {"aws_region_name": "us-east-2"}) + self.assertEqual( + secret_with_options._client_settings, {"region_name": "us-west-2"}) + def test_secret_factory_hsm(self): hsm_spec = json.dumps({ "project_id": "p",