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 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..4fbf9d56e85f 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(';'): @@ -108,24 +126,39 @@ def parse_secret_option(cls, secret: str) -> 'Secret': secret_type = raw_type.lower() secret_manager = _SECRET_TYPE_TO_SECRET_MANAGER.get(secret_type) if not secret_manager: + supported = ', '.join( + getattr(cls, '__name__', cls) for cls in _SECRET_CLASSES.values()) raise ValueError( - f'Invalid secret type {secret_type}, currently only ' - 'GcpSecret and GcpHsmGeneratedSecret are supported') + f'Invalid secret type {secret_type}, currently supported types: ' + f'{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__}") @@ -162,9 +195,11 @@ def from_json( else: return secret_cls(spec) else: + supported = ', '.join( + f"'{name}'" for name in _SECRET_TYPE_TO_SECRET_MANAGER.values()) raise ValueError( - f"Unsupported secret manager: '{secret_manager_name}'. Currently supported options: 'GoogleCloudSecretManager', 'GoogleCloudHsmGeneratedSecretManager'." - ) + f"Unsupported secret manager: '{secret_manager_name}'. " + f"Currently supported options: {supported}.") # If secret_manager is not set or empty, check if spec is a JSON specification dict if spec_dict is not None: @@ -271,6 +306,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 +642,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..a3c0e3b7c90a 100644 --- a/sdks/python/apache_beam/utils/secret_test.py +++ b/sdks/python/apache_beam/utils/secret_test.py @@ -15,7 +15,9 @@ # limitations under the License. # +import base64 import json +import pickle import unittest from unittest import mock @@ -23,6 +25,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 +36,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 +53,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 +73,15 @@ 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'), + 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): @@ -348,6 +371,105 @@ 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') + self.mock_boto3_client.assert_called_once_with('secretsmanager') + self.mock_sm.get_secret_value.assert_called_once_with(SecretId='my-secret') + self.assertIsNone(secret.__getstate__()['_cached_secret_bytes']) + + # 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 + }) + + AwsSecret.from_dict({ + 'name': 'my-secret', 'version_id': 'v-12345', 'stage': 'AWSCURRENT' + }).get_str() + self.mock_sm.get_secret_value.assert_called_with( + SecretId='my-secret', VersionId='v-12345', VersionStage='AWSCURRENT') + + 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() + + 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() + + def test_aws_options_and_spec_overrides(self): + secret = Secret.parse_secret_option( + '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_region_name=us-east-2', + '--aws_endpoint_url=http://localhost:4566', + '--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', + 'region_name': 'us-west-2', + 'endpoint_url': 'http://localhost:4566', + 'use_ssl': False, + }) + + 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', + 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() + + class RawSecretTest(unittest.TestCase): def test_raw_secret_str(self): secret = RawSecret("STATIC_SECRET_") @@ -403,6 +525,33 @@ 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") + + # 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", @@ -442,6 +591,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")