From 1efda3065fa9460e70e32b640c34700c05f04e7f Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Wed, 5 Aug 2026 10:15:56 -0400 Subject: [PATCH 01/11] [python] Add Secret management interface and implementations in utils.secret * Add `Secret` base class with shared caching (`get`, `get_bytes`) and pickling protection. * Implement `RawSecret`, `GcpSecret`, and `GcpHsmGeneratedSecret` secret providers. * Add factory methods (`from_spec`, `from_option_string`) and unit test suite in `utils/secret_test.py`. --- sdks/python/apache_beam/utils/secret.py | 425 +++++++++++++++++++ sdks/python/apache_beam/utils/secret_test.py | 295 +++++++++++++ 2 files changed, 720 insertions(+) create mode 100644 sdks/python/apache_beam/utils/secret.py create mode 100644 sdks/python/apache_beam/utils/secret_test.py diff --git a/sdks/python/apache_beam/utils/secret.py b/sdks/python/apache_beam/utils/secret.py new file mode 100644 index 000000000000..b078db3618f4 --- /dev/null +++ b/sdks/python/apache_beam/utils/secret.py @@ -0,0 +1,425 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""Interface and implementations for Secret providers in Apache Beam.""" + +import abc +import json +import logging +import os +import warnings +from typing import Any, Dict, Optional, Union + +from google.cloud import secretmanager + +from apache_beam.utils.annotations import deprecated + +__all__ = [ + 'Secret', + 'RawSecret', + 'GcpSecret', + 'GcpHsmGeneratedSecret', +] + + +class Secret(abc.ABC): + """Abstract base interface for Secrets in Apache Beam.""" + def __init__(self): + self._cached_secret_bytes: Optional[bytes] = None + + def get(self, cacheSecret: bool = False) -> str: + """Retrieve secret value as string. + + Args: + cacheSecret: If True, caches secret value in memory after first fetch. + + Returns: + The retrieved secret value as string. + """ + return self.get_bytes(cacheSecret=cacheSecret).decode("utf-8") + + def get_bytes(self, cacheSecret: bool = False) -> bytes: + """Retrieve secret value as bytes. + + Args: + cacheSecret: If True, caches secret value in memory after first fetch. + + Returns: + The retrieved secret value as bytes. + """ + if cacheSecret and getattr(self, '_cached_secret_bytes', None) is not None: + return self._cached_secret_bytes + + secret_val_bytes = self.get_secret_bytes() + + if cacheSecret: + self._cached_secret_bytes = secret_val_bytes + + return secret_val_bytes + + @abc.abstractmethod + def get_secret_bytes(self) -> bytes: + """Retrieve secret value as bytes from the underlying secret provider. + + Returns: + The retrieved secret value as bytes. + """ + raise NotImplementedError + + def __getstate__(self): + """Strip cached secrets before pickling for pipeline submission/transmission.""" + state = self.__dict__.copy() + state['_cached_secret_bytes'] = None + return state + + @classmethod + def from_spec( + cls, + spec: Union[str, Dict[str, str]], + secret_manager: Optional[str] = None, + secret_type: Optional[str] = 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: Provider type string (e.g. 'GoogleCloudSecretManager'). + secret_type: Provider type string (e.g. 'gcpsecret'). + + Returns: + An instance of Secret. + """ + sm_manager = secret_manager.strip( + ) if secret_manager and secret_manager.strip() else None + sm_type = secret_type.strip( + ) if secret_type and secret_type.strip() else None + + if sm_manager and sm_type: + raise ValueError( + f"Cannot specify both 'secret_manager' ('{secret_manager}') and 'secret_type' ('{secret_type}'). " + "Please specify only one.") + + if isinstance(spec, str): + spec_dict = None + try: + spec_dict = json.loads(spec) + if not isinstance(spec_dict, dict): + spec_dict = None + except Exception: + pass + elif isinstance(spec, dict): + spec_dict = spec + else: + spec_dict = None + + provider_str = sm_manager or sm_type + if provider_str: + secret_cls = _SECRET_CLASSES.get(provider_str.lower()) + if secret_cls: + if isinstance(spec_dict, dict) and hasattr(secret_cls, 'from_dict'): + return secret_cls.from_dict(spec_dict) + else: + raise ValueError( + f"Unsupported secret provider: '{provider_str}'. Currently supported options: 'GoogleCloudSecretManager', 'gcpsecret'." + ) + + # If secret_manager is not set or empty, check if spec is a JSON specification dict + if spec_dict is not None: + msg = ( + "The 'spec' parameter appears to be a JSON specification, but " + "'secret_manager' is not set. Defaulting to Raw.") + logging.warning(msg) + warnings.warn(msg, UserWarning) + + return RawSecret(spec) + + @classmethod + def from_option_string(cls, option: str) -> 'Secret': + param_map = {} + for param in option.split(';'): + parts = param.split(':') + param_map[parts[0]] = parts[1] + + secret_manager = param_map.get('type', None) + del param_map['type'] + return cls.from_spec(json.dumps(param_map), secret_manager) + + +class RawSecret(Secret): + """Secret implementation wrapping a raw secret string or bytes directly.""" + def __init__(self, secret: Union[str, bytes]): + super().__init__() + if isinstance(secret, str): + self.secret = secret.encode("utf-8") + else: + self.secret = secret + + def get_secret_bytes(self) -> bytes: + return self.secret + + def __eq__(self, other: Any) -> bool: + if not isinstance(other, RawSecret): + return False + return self.secret == other.secret + + +class GcpSecret(Secret): + """Secret implementation using Google Cloud Secret Manager.""" + def __init__(self, version_name: str): + super().__init__() + self.version_name = version_name + + @classmethod + def from_dict(cls, spec_dict: Dict[str, str]) -> 'GcpSecret': + """Initialize GcpSecret from a dictionary specification.""" + version_name = cls._parse_version_name(spec_dict) + return cls(version_name) + + @classmethod + def _parse_version_name(cls, spec_dict: Dict[str, str]) -> str: + if "version_name" in spec_dict: + return spec_dict["version_name"] + + secret_id = spec_dict.get("name") + if not secret_id: + raise ValueError("Secret name ('name') must be specified in secret spec.") + + # Resolve project ID from spec, environment variables, or Application Default Credentials + project_id = ( + spec_dict.get("project") or os.environ.get("GOOGLE_CLOUD_PROJECT") or + os.environ.get("GCP_PROJECT")) + + if not project_id: + try: + import google.auth + _, project_id = google.auth.default() + except Exception: + pass + + version_id = spec_dict.get("version", "latest") + + if not project_id: + raise ValueError( + f"Could not resolve GCP project ID for secret '{secret_id}'. " + "Please specify 'project' in the secret spec, set GOOGLE_CLOUD_PROJECT environment variable, " + "or configure Application Default Credentials.") + + return f"projects/{project_id}/secrets/{secret_id}/versions/{version_id}" + + def __eq__(self, other: Any) -> bool: + if not isinstance(other, GcpSecret): + return False + return self.version_name == other.version_name + + def get_secret_bytes(self) -> bytes: + """Get the secret value as bytes from GCP Secret Manager. + + Returns: + The secret bytes. + """ + try: + client = secretmanager.SecretManagerServiceClient() + response = client.access_secret_version( + request={"name": self.version_name}) + secret_val_bytes = response.payload.data + logging.info( + "Successfully fetched secret from GCP Secret Manager (version_name '%s')", + self.version_name) + return secret_val_bytes + except Exception as e: + raise RuntimeError( + f'Failed to retrieve secret bytes for secret ' + f'{self.version_name} with exception {e}') + + +class GcpHsmGeneratedSecret(Secret): + """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, + it will be retrieved. + """ + def __init__( + self, + project_id: str, + location_id: str, + key_ring_id: str, + key_id: str, + job_name: str): + """Initializes a GcpHsmGeneratedSecret object. + + Args: + project_id: The GCP project ID. + location_id: The GCP location ID for the HSM key. + key_ring_id: The ID of the KMS key ring. + key_id: The ID of the KMS key. + job_name: The name of the job, used to generate a unique secret name. + """ + super().__init__() + self.project_id = project_id + self.location_id = location_id + self.key_ring_id = key_ring_id + self.key_id = key_id + self.job_name = job_name + self._secret_version_name = f'HsmGeneratedSecret_{job_name}' + + @classmethod + def from_dict(cls, spec_dict: Dict[str, str]) -> 'GcpHsmGeneratedSecret': + """Initialize GcpHsmGeneratedSecret from a dictionary specification.""" + required_keys = [ + 'project_id', 'location_id', 'key_ring_id', 'key_id', 'job_name' + ] + missing = [k for k in required_keys if k not in spec_dict] + if missing: + raise ValueError( + f"Missing required parameter(s) for GcpHsmGeneratedSecret: {missing}") + return cls( + project_id=spec_dict['project_id'], + location_id=spec_dict['location_id'], + key_ring_id=spec_dict['key_ring_id'], + key_id=spec_dict['key_id'], + job_name=spec_dict['job_name'], + ) + + def __eq__(self, other: Any) -> bool: + if not isinstance(other, GcpHsmGeneratedSecret): + return False + return ( + self.project_id == other.project_id and + self.location_id == other.location_id and + self.key_ring_id == other.key_ring_id and + self.key_id == other.key_id and self.job_name == other.job_name) + + def get_secret_bytes(self) -> bytes: + """Retrieves the secret bytes from GCP Secret Manager, creating it if needed. + + Returns: + The secret bytes. + """ + from google.api_core import exceptions as api_exceptions + from google.cloud import secretmanager + + client = secretmanager.SecretManagerServiceClient() + + project_path = f"projects/{self.project_id}" + secret_path = f"{project_path}/secrets/{self._secret_version_name}" + secret_version_path = f"{secret_path}/versions/1" + + try: + try: + response = client.access_secret_version( + request={"name": secret_version_path}) + return response.payload.data + except api_exceptions.NotFound: + pass + + try: + client.create_secret( + request={ + "parent": project_path, + "secret_id": self._secret_version_name, + "secret": { + "replication": { + "automatic": {} + } + }, + }) + except api_exceptions.AlreadyExists: + pass + + new_key = self.generate_dek() + try: + # Try one more time in case it was created while we were generating the DEK. + response = client.access_secret_version( + request={"name": secret_version_path}) + return response.payload.data + except api_exceptions.NotFound: + logging.info( + "Secret version %s not found. Creating new secret and version.", + secret_version_path) + client.add_secret_version( + request={ + "parent": secret_path, "payload": { + "data": new_key + } + }) + response = client.access_secret_version( + request={"name": secret_version_path}) + return response.payload.data + + except Exception as e: + raise RuntimeError( + f'Failed to retrieve or create secret bytes for secret ' + f'{self._secret_version_name} with exception {e}') + + def generate_dek(self, dek_size: int = 32) -> bytes: + """Generates a new Data Encryption Key (DEK) using an HSM-backed key. + + This function follows a key derivation process that incorporates entropy + from the HSM-backed key into the nonce used for key derivation. + + Args: + dek_size: The size of the DEK to generate. + + Returns: + A new DEK of the specified size, url-safe base64-encoded. + """ + try: + import base64 + import os + + from cryptography.hazmat.primitives import hashes + from cryptography.hazmat.primitives.kdf.hkdf import HKDF + from google.cloud import kms + + # 1. Generate a random nonce (nonce_one) + nonce_one = os.urandom(dek_size) + + # 2. Use the HSM-backed key to encrypt nonce_one to create nonce_two + kms_client = kms.KeyManagementServiceClient() + key_path = kms_client.crypto_key_path( + self.project_id, self.location_id, self.key_ring_id, self.key_id) + response = kms_client.encrypt( + request={ + 'name': key_path, 'plaintext': nonce_one + }) + nonce_two = response.ciphertext + + # 3. Generate a Derivation Key (DK) + dk = os.urandom(dek_size) + + # 4. Use a KDF to derive the DEK using DK and nonce_two + hkdf = HKDF( + algorithm=hashes.SHA256(), + length=dek_size, + salt=nonce_two, + info=None, + ) + dek = hkdf.derive(dk) + return base64.urlsafe_b64encode(dek) + except Exception as e: + raise RuntimeError(f'Failed to generate DEK with exception {e}') + + +_SECRET_CLASSES: Dict[str, type] = { + "googlecloudsecretmanager": GcpSecret, + "gcpsecret": GcpSecret, + "gcphsmgeneratedsecret": GcpHsmGeneratedSecret, +} + + +def generate_secret_bytes() -> bytes: + """Generates a new secret key using Fernet.""" + from cryptography.fernet import Fernet + return Fernet.generate_key() diff --git a/sdks/python/apache_beam/utils/secret_test.py b/sdks/python/apache_beam/utils/secret_test.py new file mode 100644 index 000000000000..1cc2bda412fd --- /dev/null +++ b/sdks/python/apache_beam/utils/secret_test.py @@ -0,0 +1,295 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import json +import unittest +from unittest.mock import MagicMock, patch + +from apache_beam.utils.annotations import BeamDeprecationWarning +from apache_beam.utils.secret import GcpHsmGeneratedSecret +from apache_beam.utils.secret import GcpSecret +from apache_beam.utils.secret import RawSecret +from apache_beam.utils.secret import Secret +from apache_beam.utils.secret import generate_secret_bytes + + +class GcpSecretTest(unittest.TestCase): + @patch("google.cloud.secretmanager.SecretManagerServiceClient") + def test_gcp_secret_success(self, mock_client_cls): + mock_client = MagicMock() + mock_client_cls.return_value = mock_client + mock_response = MagicMock() + mock_response.payload.data = b"secret-payload-value" + mock_client.access_secret_version.return_value = mock_response + + spec = {"name": "my-secret", "version": "1", "project": "my-project"} + secret = GcpSecret.from_dict(spec) + + secret_val = secret.get(cacheSecret=True) + self.assertEqual(secret_val, "secret-payload-value") + secret_bytes = secret.get_bytes(cacheSecret=True) + self.assertEqual(secret_bytes, b"secret-payload-value") + mock_client.access_secret_version.assert_called_once_with( + request={"name": "projects/my-project/secrets/my-secret/versions/1"}) + + # Second call with cacheSecret=True should return cached value without calling client again + mock_client.reset_mock() + secret_val_cached = secret.get(cacheSecret=True) + self.assertEqual(secret_val_cached, "secret-payload-value") + mock_client.access_secret_version.assert_not_called() + + @patch("google.cloud.secretmanager.SecretManagerServiceClient") + def test_gcp_secret_get_bytes_uncached(self, mock_client_cls): + mock_client = MagicMock() + mock_client_cls.return_value = mock_client + mock_response = MagicMock() + mock_response.payload.data = b"secret-payload-value" + mock_client.access_secret_version.return_value = mock_response + + spec = {"name": "my-secret", "project": "my-project"} + secret = GcpSecret.from_dict(spec) + + secret_bytes = secret.get_bytes() + self.assertEqual(secret_bytes, b"secret-payload-value") + self.assertIsNone(secret._cached_secret_bytes) + + @patch("google.cloud.secretmanager.SecretManagerServiceClient") + def test_gcp_secret_getstate_clears_cached_secret(self, mock_client_cls): + mock_client = MagicMock() + mock_client_cls.return_value = mock_client + mock_response = MagicMock() + mock_response.payload.data = b"secret-payload-value" + mock_client.access_secret_version.return_value = mock_response + + spec = {"name": "my-secret", "project": "my-project"} + secret = GcpSecret.from_dict(spec) + + # Cache the secret in memory + secret.get(cacheSecret=True) + self.assertEqual(secret._cached_secret_bytes, b"secret-payload-value") + + # When pickled / getstate is called during pipeline submission + state = secret.__getstate__() + self.assertIsNone(state["_cached_secret_bytes"]) + + @patch.dict("os.environ", {"GOOGLE_CLOUD_PROJECT": "env-project-123"}) + @patch("google.cloud.secretmanager.SecretManagerServiceClient") + def test_gcp_secret_env_project_fallback(self, mock_client_cls): + mock_client = MagicMock() + mock_client_cls.return_value = mock_client + mock_response = MagicMock() + mock_response.payload.data = b"env-secret-val" + mock_client.access_secret_version.return_value = mock_response + + # Project omitted from spec + spec = {"name": "env-secret", "version": "latest"} + secret = GcpSecret.from_dict(spec) + + secret_val = secret.get(cacheSecret=False) + self.assertEqual(secret_val, "env-secret-val") + self.assertEqual(secret.get_bytes(cacheSecret=False), b"env-secret-val") + mock_client.access_secret_version.assert_called_with( + request={ + "name": "projects/env-project-123/secrets/env-secret/versions/latest" + }) + + @patch("google.cloud.secretmanager.SecretManagerServiceClient") + def test_gcp_secret_failure_raises_exception(self, mock_client_cls): + mock_client = MagicMock() + mock_client_cls.return_value = mock_client + mock_client.access_secret_version.side_effect = RuntimeError( + "Permission denied or secret not found") + + spec = {"name": "non-existent-secret", "project": "my-project"} + secret = GcpSecret.from_dict(spec) + + with self.assertRaises(RuntimeError) as ctx: + secret.get(cacheSecret=False) + self.assertIn("Permission denied or secret not found", str(ctx.exception)) + + @patch.dict("os.environ", {}, clear=True) + @patch("google.auth.default", side_effect=Exception("No ADC")) + def test_ill_formed_missing_project_raises_value_error( + self, mock_auth_default): + spec = {"name": "my-secret"} + with self.assertRaises(ValueError) as ctx: + GcpSecret.from_dict(spec) + self.assertIn("Could not resolve GCP project ID", str(ctx.exception)) + + def test_ill_formed_missing_secret_name_raises_value_error(self): + spec = {"project": "my-project"} + with self.assertRaises(ValueError) as ctx: + GcpSecret.from_dict(spec) + self.assertIn("Secret name ('name') must be specified", str(ctx.exception)) + + +class GcpHsmGeneratedSecretTest(unittest.TestCase): + def test_from_dict_success(self): + spec = { + "project_id": "test-proj", + "location_id": "global", + "key_ring_id": "ring", + "key_id": "key", + "job_name": "my-job" + } + secret = GcpHsmGeneratedSecret.from_dict(spec) + self.assertEqual(secret.project_id, "test-proj") + self.assertEqual(secret.location_id, "global") + self.assertEqual(secret.key_ring_id, "ring") + self.assertEqual(secret.key_id, "key") + self.assertEqual(secret.job_name, "my-job") + self.assertEqual(secret._secret_version_name, "HsmGeneratedSecret_my-job") + + def test_from_dict_missing_params_raises_value_error(self): + spec = {"project_id": "test-proj", "location_id": "global"} + with self.assertRaises(ValueError) as ctx: + GcpHsmGeneratedSecret.from_dict(spec) + self.assertIn("Missing required parameter(s)", str(ctx.exception)) + + @patch("google.cloud.secretmanager.SecretManagerServiceClient") + def test_get_bytes_cached(self, mock_sm_client_cls): + mock_client = MagicMock() + mock_sm_client_cls.return_value = mock_client + mock_response = MagicMock() + mock_response.payload.data = b"hsm-derived-key" + mock_client.access_secret_version.return_value = mock_response + + secret = GcpHsmGeneratedSecret("p", "l", "r", "k", "j") + secret_bytes = secret.get_bytes(cacheSecret=True) + self.assertEqual(secret_bytes, b"hsm-derived-key") + + # Second call uses cache + mock_client.reset_mock() + self.assertEqual(secret.get_bytes(cacheSecret=True), b"hsm-derived-key") + mock_client.access_secret_version.assert_not_called() + + @patch("google.cloud.secretmanager.SecretManagerServiceClient") + def test_getstate_clears_cached_secret(self, mock_sm_client_cls): + mock_client = MagicMock() + mock_sm_client_cls.return_value = mock_client + mock_response = MagicMock() + mock_response.payload.data = b"hsm-derived-key" + mock_client.access_secret_version.return_value = mock_response + + secret = GcpHsmGeneratedSecret("p", "l", "r", "k", "j") + secret.get_bytes(cacheSecret=True) + self.assertEqual(secret._cached_secret_bytes, b"hsm-derived-key") + + state = secret.__getstate__() + self.assertIsNone(state["_cached_secret_bytes"]) + + +class RawSecretTest(unittest.TestCase): + def test_raw_secret_str(self): + secret = RawSecret("STATIC_SECRET_") + self.assertEqual(secret.get(cacheSecret=True), "STATIC_SECRET_") + self.assertEqual(secret.get_bytes(cacheSecret=True), b"STATIC_SECRET_") + + def test_raw_secret_bytes(self): + secret = RawSecret(b"STATIC_BYTES_") + self.assertEqual(secret.get(cacheSecret=True), "STATIC_BYTES_") + self.assertEqual(secret.get_bytes(cacheSecret=True), b"STATIC_BYTES_") + + +class SecretFactoryTest(unittest.TestCase): + def test_secret_factory(self): + spec = json.dumps({"name": "test-secret", "project": "proj"}) + + # When provider is set to 'GoogleCloudSecretManager' + secret_gcp = Secret.from_spec( + spec=spec, secret_manager="GoogleCloudSecretManager") + self.assertIsInstance(secret_gcp, GcpSecret) + + # When provider is set to 'gcpsecret' + secret_gcp_alias = Secret.from_spec(spec=spec, secret_manager="gcpsecret") + self.assertIsInstance(secret_gcp_alias, GcpSecret) + + # When provider is None or empty with plain string + secret_raw = Secret.from_spec(spec="STATIC_SECRET_", secret_manager=None) + self.assertIsInstance(secret_raw, RawSecret) + + # Unsupported provider raises ValueError + with self.assertRaises(ValueError): + Secret.from_spec(spec="spec", secret_manager="unsupported_provider") + + def test_secret_factory_secret_type(self): + spec = json.dumps({"name": "test-secret", "project": "proj"}) + + # Using secret_type alone + secret_type_obj = Secret.from_spec(spec=spec, secret_type="gcpsecret") + self.assertIsInstance(secret_type_obj, GcpSecret) + + def test_secret_factory_both_providers_set_raises_exception(self): + spec = json.dumps({"name": "test-secret", "project": "proj"}) + with self.assertRaises(ValueError) as ctx: + Secret.from_spec( + spec=spec, + secret_manager="GoogleCloudSecretManager", + secret_type="gcpsecret") + self.assertIn("Cannot specify both 'secret_manager'", str(ctx.exception)) + + def test_json_secret_without_secret_manager_warning(self): + json_spec = json.dumps({"name": "my-secret", "project": "my-proj"}) + with self.assertWarns(UserWarning): + secret = Secret.from_spec(spec=json_spec, secret_manager=None) + self.assertIsInstance(secret, RawSecret) + + def test_from_option_string(self): + option_str = "type:gcpsecreT;version_name:my_secret/versions/latest" + secret = Secret.from_option_string(option_str) + self.assertIsInstance(secret, GcpSecret) + self.assertEqual(secret.version_name, "my_secret/versions/latest") + + def test_from_option_string_hsm(self): + option_str = ( + "type:gcphsmgeneratedsecret;project_id:p;location_id:l;" + "key_ring_id:r;key_id:k;job_name:j") + secret = Secret.from_option_string(option_str) + self.assertIsInstance(secret, GcpHsmGeneratedSecret) + self.assertEqual(secret.project_id, "p") + self.assertEqual(secret.job_name, "j") + + def test_generate_secret_bytes(self): + key = generate_secret_bytes() + self.assertIsInstance(key, bytes) + self.assertTrue(len(key) > 0) + + def test_equality(self): + raw1 = RawSecret("secret_value") + raw2 = RawSecret("secret_value") + raw3 = RawSecret("other_value") + self.assertEqual(raw1, raw2) + self.assertNotEqual(raw1, raw3) + self.assertNotEqual(raw1, "secret_value") + + gcp1 = GcpSecret.from_dict({"name": "sec", "project": "proj"}) + gcp2 = GcpSecret.from_dict({"name": "sec", "project": "proj"}) + gcp3 = GcpSecret.from_dict({"name": "other", "project": "proj"}) + self.assertEqual(gcp1, gcp2) + self.assertNotEqual(gcp1, gcp3) + self.assertNotEqual(gcp1, raw1) + + hsm1 = GcpHsmGeneratedSecret("p", "l", "r", "k", "j") + hsm2 = GcpHsmGeneratedSecret("p", "l", "r", "k", "j") + hsm3 = GcpHsmGeneratedSecret("p", "l", "r", "k", "other") + self.assertEqual(hsm1, hsm2) + self.assertNotEqual(hsm1, hsm3) + self.assertNotEqual(hsm1, gcp1) + + +if __name__ == "__main__": + unittest.main() From 2d4396b2c62b78e419796c9a5473d3bd3624b47a Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Wed, 5 Aug 2026 10:40:12 -0400 Subject: [PATCH 02/11] Make it easy to migrate from the old secret classes. --- sdks/python/apache_beam/utils/secret.py | 57 +++++++++++++++----- sdks/python/apache_beam/utils/secret_test.py | 13 +++++ 2 files changed, 57 insertions(+), 13 deletions(-) diff --git a/sdks/python/apache_beam/utils/secret.py b/sdks/python/apache_beam/utils/secret.py index b078db3618f4..dfbc3a52daa9 100644 --- a/sdks/python/apache_beam/utils/secret.py +++ b/sdks/python/apache_beam/utils/secret.py @@ -33,9 +33,16 @@ 'RawSecret', 'GcpSecret', 'GcpHsmGeneratedSecret', + 'generate_secret_bytes', ] +def generate_secret_bytes() -> bytes: + """Generates a new secret key using Fernet.""" + from cryptography.fernet import Fernet + return Fernet.generate_key() + + class Secret(abc.ABC): """Abstract base interface for Secrets in Apache Beam.""" def __init__(self): @@ -80,6 +87,24 @@ def get_secret_bytes(self) -> bytes: """ raise NotImplementedError + @staticmethod + @deprecated(since='2.77.0', current='generate_secret_bytes') + def generate_secret_bytes() -> bytes: + """Generates a new secret key. + + Deprecated: Use global :func:`generate_secret_bytes` instead. + """ + return generate_secret_bytes() + + @classmethod + @deprecated(since='2.77.0', current='from_option_string') + def parse_secret_option(cls, secret: str) -> 'Secret': + """Parses a secret string and returns the appropriate secret type. + + Deprecated: Use :meth:`from_option_string` instead. + """ + return cls.from_option_string(secret) + def __getstate__(self): """Strip cached secrets before pickling for pipeline submission/transmission.""" state = self.__dict__.copy() @@ -151,10 +176,13 @@ def from_option_string(cls, option: str) -> 'Secret': param_map = {} for param in option.split(';'): parts = param.split(':') - param_map[parts[0]] = parts[1] + if len(parts) == 2: + param_map[parts[0]] = parts[1] + + if 'type' not in param_map: + raise ValueError('Secret string must contain a valid type parameter') - secret_manager = param_map.get('type', None) - del param_map['type'] + secret_manager = param_map.pop('type') return cls.from_spec(json.dumps(param_map), secret_manager) @@ -185,6 +213,11 @@ def __init__(self, version_name: str): @classmethod def from_dict(cls, spec_dict: Dict[str, str]) -> 'GcpSecret': """Initialize GcpSecret from a dictionary specification.""" + allowed_keys = {'version_name', 'name', 'project', 'version'} + invalid_keys = set(spec_dict.keys()) - allowed_keys + if invalid_keys: + raise ValueError( + f"Invalid secret parameter {', '.join(sorted(invalid_keys))}") version_name = cls._parse_version_name(spec_dict) return cls(version_name) @@ -277,13 +310,17 @@ def __init__( @classmethod def from_dict(cls, spec_dict: Dict[str, str]) -> 'GcpHsmGeneratedSecret': """Initialize GcpHsmGeneratedSecret from a dictionary specification.""" - required_keys = [ + allowed_keys = { 'project_id', 'location_id', 'key_ring_id', 'key_id', 'job_name' - ] - missing = [k for k in required_keys if k not in spec_dict] + } + missing = allowed_keys - set(spec_dict.keys()) if missing: raise ValueError( - f"Missing required parameter(s) for GcpHsmGeneratedSecret: {missing}") + f"Missing required parameter(s) for GcpHsmGeneratedSecret: {sorted(list(missing))}") + invalid_keys = set(spec_dict.keys()) - allowed_keys + if invalid_keys: + raise ValueError( + f"Invalid secret parameter {', '.join(sorted(invalid_keys))}") return cls( project_id=spec_dict['project_id'], location_id=spec_dict['location_id'], @@ -417,9 +454,3 @@ def generate_dek(self, dek_size: int = 32) -> bytes: "gcpsecret": GcpSecret, "gcphsmgeneratedsecret": GcpHsmGeneratedSecret, } - - -def generate_secret_bytes() -> bytes: - """Generates a new secret key using Fernet.""" - from cryptography.fernet import Fernet - return Fernet.generate_key() diff --git a/sdks/python/apache_beam/utils/secret_test.py b/sdks/python/apache_beam/utils/secret_test.py index 1cc2bda412fd..2bd913d33a19 100644 --- a/sdks/python/apache_beam/utils/secret_test.py +++ b/sdks/python/apache_beam/utils/secret_test.py @@ -268,6 +268,19 @@ def test_generate_secret_bytes(self): self.assertIsInstance(key, bytes) self.assertTrue(len(key) > 0) + def test_deprecated_parse_secret_option(self): + option_str = "type:gcpsecreT;version_name:my_secret/versions/latest" + with self.assertWarns(BeamDeprecationWarning): + secret = Secret.parse_secret_option(option_str) + self.assertIsInstance(secret, GcpSecret) + self.assertEqual(secret.version_name, "my_secret/versions/latest") + + def test_deprecated_generate_secret_bytes(self): + with self.assertWarns(BeamDeprecationWarning): + key = Secret.generate_secret_bytes() + self.assertIsInstance(key, bytes) + self.assertTrue(len(key) > 0) + def test_equality(self): raw1 = RawSecret("secret_value") raw2 = RawSecret("secret_value") From fdb7ec3237fc2244fbeb0807c11ba6383a41ac7e Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Wed, 5 Aug 2026 11:52:27 -0400 Subject: [PATCH 03/11] Make some attributes private --- sdks/python/apache_beam/utils/secret.py | 40 ++++++++++---------- sdks/python/apache_beam/utils/secret_test.py | 18 ++++----- 2 files changed, 29 insertions(+), 29 deletions(-) diff --git a/sdks/python/apache_beam/utils/secret.py b/sdks/python/apache_beam/utils/secret.py index dfbc3a52daa9..9db6faca045d 100644 --- a/sdks/python/apache_beam/utils/secret.py +++ b/sdks/python/apache_beam/utils/secret.py @@ -191,24 +191,24 @@ class RawSecret(Secret): def __init__(self, secret: Union[str, bytes]): super().__init__() if isinstance(secret, str): - self.secret = secret.encode("utf-8") + self._secret = secret.encode("utf-8") else: - self.secret = secret + self._secret = secret def get_secret_bytes(self) -> bytes: - return self.secret + return self._secret def __eq__(self, other: Any) -> bool: if not isinstance(other, RawSecret): return False - return self.secret == other.secret + return self._secret == other._secret class GcpSecret(Secret): """Secret implementation using Google Cloud Secret Manager.""" def __init__(self, version_name: str): super().__init__() - self.version_name = version_name + self._version_name = version_name @classmethod def from_dict(cls, spec_dict: Dict[str, str]) -> 'GcpSecret': @@ -255,7 +255,7 @@ def _parse_version_name(cls, spec_dict: Dict[str, str]) -> str: def __eq__(self, other: Any) -> bool: if not isinstance(other, GcpSecret): return False - return self.version_name == other.version_name + return self._version_name == other._version_name def get_secret_bytes(self) -> bytes: """Get the secret value as bytes from GCP Secret Manager. @@ -266,16 +266,16 @@ def get_secret_bytes(self) -> bytes: try: client = secretmanager.SecretManagerServiceClient() response = client.access_secret_version( - request={"name": self.version_name}) + request={"name": self._version_name}) secret_val_bytes = response.payload.data logging.info( "Successfully fetched secret from GCP Secret Manager (version_name '%s')", - self.version_name) + self._version_name) return secret_val_bytes except Exception as e: raise RuntimeError( f'Failed to retrieve secret bytes for secret ' - f'{self.version_name} with exception {e}') + f'{self._version_name} with exception {e}') class GcpHsmGeneratedSecret(Secret): @@ -300,11 +300,11 @@ def __init__( job_name: The name of the job, used to generate a unique secret name. """ super().__init__() - self.project_id = project_id - self.location_id = location_id - self.key_ring_id = key_ring_id - self.key_id = key_id - self.job_name = job_name + self._project_id = project_id + self._location_id = location_id + self._key_ring_id = key_ring_id + self._key_id = key_id + self._job_name = job_name self._secret_version_name = f'HsmGeneratedSecret_{job_name}' @classmethod @@ -333,10 +333,10 @@ def __eq__(self, other: Any) -> bool: if not isinstance(other, GcpHsmGeneratedSecret): return False return ( - self.project_id == other.project_id and - self.location_id == other.location_id and - self.key_ring_id == other.key_ring_id and - self.key_id == other.key_id and self.job_name == other.job_name) + self._project_id == other._project_id and + self._location_id == other._location_id and + self._key_ring_id == other._key_ring_id and + self._key_id == other._key_id and self._job_name == other._job_name) def get_secret_bytes(self) -> bytes: """Retrieves the secret bytes from GCP Secret Manager, creating it if needed. @@ -349,7 +349,7 @@ def get_secret_bytes(self) -> bytes: client = secretmanager.SecretManagerServiceClient() - project_path = f"projects/{self.project_id}" + project_path = f"projects/{self._project_id}" secret_path = f"{project_path}/secrets/{self._secret_version_name}" secret_version_path = f"{secret_path}/versions/1" @@ -426,7 +426,7 @@ def generate_dek(self, dek_size: int = 32) -> bytes: # 2. Use the HSM-backed key to encrypt nonce_one to create nonce_two kms_client = kms.KeyManagementServiceClient() key_path = kms_client.crypto_key_path( - self.project_id, self.location_id, self.key_ring_id, self.key_id) + self._project_id, self._location_id, self._key_ring_id, self._key_id) response = kms_client.encrypt( request={ 'name': key_path, 'plaintext': nonce_one diff --git a/sdks/python/apache_beam/utils/secret_test.py b/sdks/python/apache_beam/utils/secret_test.py index 2bd913d33a19..c54d3c94adb1 100644 --- a/sdks/python/apache_beam/utils/secret_test.py +++ b/sdks/python/apache_beam/utils/secret_test.py @@ -147,11 +147,11 @@ def test_from_dict_success(self): "job_name": "my-job" } secret = GcpHsmGeneratedSecret.from_dict(spec) - self.assertEqual(secret.project_id, "test-proj") - self.assertEqual(secret.location_id, "global") - self.assertEqual(secret.key_ring_id, "ring") - self.assertEqual(secret.key_id, "key") - self.assertEqual(secret.job_name, "my-job") + self.assertEqual(secret._project_id, "test-proj") + self.assertEqual(secret._location_id, "global") + self.assertEqual(secret._key_ring_id, "ring") + self.assertEqual(secret._key_id, "key") + self.assertEqual(secret._job_name, "my-job") self.assertEqual(secret._secret_version_name, "HsmGeneratedSecret_my-job") def test_from_dict_missing_params_raises_value_error(self): @@ -252,7 +252,7 @@ def test_from_option_string(self): option_str = "type:gcpsecreT;version_name:my_secret/versions/latest" secret = Secret.from_option_string(option_str) self.assertIsInstance(secret, GcpSecret) - self.assertEqual(secret.version_name, "my_secret/versions/latest") + self.assertEqual(secret._version_name, "my_secret/versions/latest") def test_from_option_string_hsm(self): option_str = ( @@ -260,8 +260,8 @@ def test_from_option_string_hsm(self): "key_ring_id:r;key_id:k;job_name:j") secret = Secret.from_option_string(option_str) self.assertIsInstance(secret, GcpHsmGeneratedSecret) - self.assertEqual(secret.project_id, "p") - self.assertEqual(secret.job_name, "j") + self.assertEqual(secret._project_id, "p") + self.assertEqual(secret._job_name, "j") def test_generate_secret_bytes(self): key = generate_secret_bytes() @@ -273,7 +273,7 @@ def test_deprecated_parse_secret_option(self): with self.assertWarns(BeamDeprecationWarning): secret = Secret.parse_secret_option(option_str) self.assertIsInstance(secret, GcpSecret) - self.assertEqual(secret.version_name, "my_secret/versions/latest") + self.assertEqual(secret._version_name, "my_secret/versions/latest") def test_deprecated_generate_secret_bytes(self): with self.assertWarns(BeamDeprecationWarning): From 5f2f84f2ba8b939a6854c400e61f7ebe903e13e1 Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Wed, 5 Aug 2026 12:30:52 -0400 Subject: [PATCH 04/11] Make the factory method more tolerant to different constructors. --- sdks/python/apache_beam/utils/secret.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/sdks/python/apache_beam/utils/secret.py b/sdks/python/apache_beam/utils/secret.py index 9db6faca045d..13224e445717 100644 --- a/sdks/python/apache_beam/utils/secret.py +++ b/sdks/python/apache_beam/utils/secret.py @@ -156,6 +156,10 @@ def from_spec( if secret_cls: if isinstance(spec_dict, dict) and hasattr(secret_cls, 'from_dict'): return secret_cls.from_dict(spec_dict) + elif isinstance(spec_dict, dict): + return secret_cls(**spec_dict) + else: + return secret_cls(spec) else: raise ValueError( f"Unsupported secret provider: '{provider_str}'. Currently supported options: 'GoogleCloudSecretManager', 'gcpsecret'." From 3534d367a449ee20f3a09683b4d1d5be057185e0 Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Wed, 5 Aug 2026 12:32:20 -0400 Subject: [PATCH 05/11] Remove duplicate Secret classes from util.py and update util_test mocks --- sdks/python/apache_beam/transforms/util.py | 245 +----------------- .../apache_beam/transforms/util_test.py | 7 +- 2 files changed, 8 insertions(+), 244 deletions(-) diff --git a/sdks/python/apache_beam/transforms/util.py b/sdks/python/apache_beam/transforms/util.py index 2ea9df9399cb..30f8dc01c458 100644 --- a/sdks/python/apache_beam/transforms/util.py +++ b/sdks/python/apache_beam/transforms/util.py @@ -82,6 +82,9 @@ 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 GcpHsmGeneratedSecret +from apache_beam.utils.secret import GcpSecret +from apache_beam.utils.secret import Secret from apache_beam.utils.sharded_key import ShardedKey from apache_beam.utils.timestamp import Timestamp @@ -94,6 +97,7 @@ 'BatchElements', 'CoGroupByKey', 'Distinct', + 'GcpHsmGeneratedSecret', 'GcpSecret', 'GroupByEncryptedKey', 'Keys', @@ -327,247 +331,6 @@ def RemoveDuplicates(pcoll): return pcoll | 'RemoveDuplicates' >> Distinct() -class Secret(): - """A secret management class used for handling sensitive data. - - This class provides a generic interface for secret management. Implementations - of this class should handle fetching secrets from a secret management system. - """ - def get_secret_bytes(self) -> bytes: - """Returns the secret as a byte string.""" - raise NotImplementedError() - - @staticmethod - def generate_secret_bytes() -> bytes: - """Generates a new secret key.""" - return Fernet.generate_key() - - @staticmethod - def parse_secret_option(secret) -> 'Secret': - """Parses a secret string and returns the appropriate secret type. - - The secret string should be formatted like: - 'type:;:' - - For example, 'type:GcpSecret;version_name:my_secret/versions/latest' - would return a GcpSecret initialized with 'my_secret/versions/latest'. - """ - param_map = {} - for param in secret.split(';'): - parts = param.split(':') - param_map[parts[0]] = parts[1] - - if 'type' not in param_map: - raise ValueError('Secret string must contain a valid type parameter') - - secret_type = param_map['type'].lower() - del param_map['type'] - secret_class = Secret - secret_params = None - if secret_type == 'gcpsecret': - secret_class = GcpSecret # type: ignore[assignment] - secret_params = ['version_name'] - elif secret_type == 'gcphsmgeneratedsecret': - secret_class = GcpHsmGeneratedSecret # type: ignore[assignment] - secret_params = [ - 'project_id', 'location_id', 'key_ring_id', 'key_id', 'job_name' - ] - else: - raise ValueError( - f'Invalid secret type {secret_type}, currently only ' - 'GcpSecret and GcpHsmGeneratedSecret are supported') - - for param_name in param_map.keys(): - if param_name not in secret_params: - raise ValueError( - f'Invalid secret parameter {param_name}, ' - f'{secret_type} only supports the following ' - f'parameters: {secret_params}') - return secret_class(**param_map) - - -class GcpSecret(Secret): - """A secret manager implementation that retrieves secrets from Google Cloud - Secret Manager. - """ - def __init__(self, version_name: str): - """Initializes a GcpSecret object. - - Args: - version_name: The full version name of the secret in Google Cloud Secret - Manager. For example: - projects//secrets//versions/1. - For more info, see - https://cloud.google.com/python/docs/reference/secretmanager/latest/google.cloud.secretmanager_v1beta1.services.secret_manager_service.SecretManagerServiceClient#google_cloud_secretmanager_v1beta1_services_secret_manager_service_SecretManagerServiceClient_access_secret_version - """ - self._version_name = version_name - - def get_secret_bytes(self) -> bytes: - try: - from google.cloud import secretmanager - client = secretmanager.SecretManagerServiceClient() - response = client.access_secret_version( - request={"name": self._version_name}) - secret = response.payload.data - return secret - except Exception as e: - raise RuntimeError( - 'Failed to retrieve secret bytes for secret ' - f'{self._version_name} with exception {e}') - - def __eq__(self, secret): - return self._version_name == getattr(secret, '_version_name', None) - - -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, - it will be retrieved. - """ - def __init__( - self, - project_id: str, - location_id: str, - key_ring_id: str, - key_id: str, - job_name: str): - """Initializes a GcpHsmGeneratedSecret object. - - Args: - project_id: The GCP project ID. - location_id: The GCP location ID for the HSM key. - key_ring_id: The ID of the KMS key ring. - key_id: The ID of the KMS key. - job_name: The name of the job, used to generate a unique secret name. - """ - self._project_id = project_id - self._location_id = location_id - self._key_ring_id = key_ring_id - self._key_id = key_id - self._secret_version_name = f'HsmGeneratedSecret_{job_name}' - - def get_secret_bytes(self) -> bytes: - """Retrieves the secret bytes. - - If the secret version already exists in Secret Manager, it is retrieved. - Otherwise, a new secret and version are created. The new secret is - generated using the HSM key. - - Returns: - The secret as a byte string. - """ - try: - from google.api_core import exceptions as api_exceptions - from google.cloud import secretmanager - client = secretmanager.SecretManagerServiceClient() - - project_path = f"projects/{self._project_id}" - secret_path = f"{project_path}/secrets/{self._secret_version_name}" - # Since we may generate multiple versions when doing this on workers, - # just always take the first version added to maintain consistency. - secret_version_path = f"{secret_path}/versions/1" - - try: - response = client.access_secret_version( - request={"name": secret_version_path}) - return response.payload.data - except api_exceptions.NotFound: - # Don't bother logging yet, we'll only log if we actually add the - # secret version below - pass - - try: - client.create_secret( - request={ - "parent": project_path, - "secret_id": self._secret_version_name, - "secret": { - "replication": { - "automatic": {} - } - }, - }) - except api_exceptions.AlreadyExists: - # Don't bother logging yet, we'll only log if we actually add the - # secret version below - pass - - new_key = self.generate_dek() - try: - # Try one more time in case it was created while we were generating the - # DEK. - response = client.access_secret_version( - request={"name": secret_version_path}) - return response.payload.data - except api_exceptions.NotFound: - _LOGGER.info( - "Secret version %s not found. " - "Creating new secret and version.", - secret_version_path) - client.add_secret_version( - request={ - "parent": secret_path, "payload": { - "data": new_key - } - }) - response = client.access_secret_version( - request={"name": secret_version_path}) - return response.payload.data - - except Exception as e: - raise RuntimeError( - f'Failed to retrieve or create secret bytes for secret ' - f'{self._secret_version_name} with exception {e}') - - def generate_dek(self, dek_size: int = 32) -> bytes: - """Generates a new Data Encryption Key (DEK) using an HSM-backed key. - - This function follows a key derivation process that incorporates entropy - from the HSM-backed key into the nonce used for key derivation. - - Args: - dek_size: The size of the DEK to generate. - - Returns: - A new DEK of the specified size, url-safe base64-encoded. - """ - try: - import base64 - import os - - from cryptography.hazmat.primitives import hashes - from cryptography.hazmat.primitives.kdf.hkdf import HKDF - from google.cloud import kms - - # 1. Generate a random nonce (nonce_one) - nonce_one = os.urandom(dek_size) - - # 2. Use the HSM-backed key to encrypt nonce_one to create nonce_two - kms_client = kms.KeyManagementServiceClient() - key_path = kms_client.crypto_key_path( - self._project_id, self._location_id, self._key_ring_id, self._key_id) - response = kms_client.encrypt( - request={ - 'name': key_path, 'plaintext': nonce_one - }) - nonce_two = response.ciphertext - - # 3. Generate a Derivation Key (DK) - dk = os.urandom(dek_size) - - # 4. Use a KDF to derive the DEK using DK and nonce_two - hkdf = HKDF( - algorithm=hashes.SHA256(), - length=dek_size, - salt=nonce_two, - info=None, - ) - dek = hkdf.derive(dk) - return base64.urlsafe_b64encode(dek) - except Exception as e: - raise RuntimeError(f'Failed to generate DEK with exception {e}') - - class _EncryptMessage(DoFn): """A DoFn that encrypts the key and value of each element.""" def __init__( diff --git a/sdks/python/apache_beam/transforms/util_test.py b/sdks/python/apache_beam/transforms/util_test.py index 63ce42726c1f..38086074c100 100644 --- a/sdks/python/apache_beam/transforms/util_test.py +++ b/sdks/python/apache_beam/transforms/util_test.py @@ -243,7 +243,7 @@ def test_co_group_by_key_on_unpickled(self): assert_that(pcoll, equal_to(expected)) -class FakeSecret(beam.Secret): +class FakeSecret(beam.utils.secret.Secret): def __init__(self, version_name=None, should_throw=False): self._secret = b'aKwI2PmqYFt2p5tNKCyBS5qYmHhHsGZcyZrnZQiQ-uE=' self._should_throw = should_throw @@ -308,7 +308,7 @@ def test_secret_manager_parses_correctly(self, secret_string, secret): exception_str='must contain a valid type parameter'), param( secret_string='type:gcpsecreT', - exception_str='missing 1 required positional argument'), + exception_str=r"Secret name \('name'\) must be specified in secret spec."), param( secret_string='type:gcpsecreT;version_name:foo;extra:val', exception_str='Invalid secret parameter extra'), @@ -387,7 +387,8 @@ def test_gbek_fake_secret_manager_actually_does_encryption(self): result, equal_to([('a', ([1, 2])), ('b', ([3])), ('c', ([4]))])) @mock.patch('apache_beam.transforms.util._DecryptMessage', MockNoOpDecrypt) - @mock.patch('apache_beam.transforms.util.GcpSecret', FakeSecret) + @mock.patch.dict( + 'apache_beam.utils.secret._SECRET_CLASSES', {'gcpsecret': FakeSecret}) def test_gbk_actually_does_encryption(self): options = PipelineOptions() # Version of GcpSecret doesn't matter since it is replaced by FakeSecret From ff7426aeb11586af64146a1062bebcb027d09729 Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Wed, 5 Aug 2026 13:19:43 -0400 Subject: [PATCH 06/11] Move the import in the function so it won't fail on module import --- sdks/python/apache_beam/utils/secret.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/sdks/python/apache_beam/utils/secret.py b/sdks/python/apache_beam/utils/secret.py index 13224e445717..a3498b0ae3aa 100644 --- a/sdks/python/apache_beam/utils/secret.py +++ b/sdks/python/apache_beam/utils/secret.py @@ -24,8 +24,6 @@ import warnings from typing import Any, Dict, Optional, Union -from google.cloud import secretmanager - from apache_beam.utils.annotations import deprecated __all__ = [ @@ -267,6 +265,8 @@ def get_secret_bytes(self) -> bytes: Returns: The secret bytes. """ + from google.cloud import secretmanager + try: client = secretmanager.SecretManagerServiceClient() response = client.access_secret_version( From be79feffd04ac4a55ee0312e58aafb25c7674a8c Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Wed, 5 Aug 2026 13:20:00 -0400 Subject: [PATCH 07/11] Reformat --- sdks/python/apache_beam/transforms/util_test.py | 3 ++- sdks/python/apache_beam/utils/secret.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/sdks/python/apache_beam/transforms/util_test.py b/sdks/python/apache_beam/transforms/util_test.py index 38086074c100..c5c4e63e4d04 100644 --- a/sdks/python/apache_beam/transforms/util_test.py +++ b/sdks/python/apache_beam/transforms/util_test.py @@ -308,7 +308,8 @@ def test_secret_manager_parses_correctly(self, secret_string, secret): exception_str='must contain a valid type parameter'), param( secret_string='type:gcpsecreT', - exception_str=r"Secret name \('name'\) must be specified in secret spec."), + exception_str= + r"Secret name \('name'\) must be specified in secret spec."), param( secret_string='type:gcpsecreT;version_name:foo;extra:val', exception_str='Invalid secret parameter extra'), diff --git a/sdks/python/apache_beam/utils/secret.py b/sdks/python/apache_beam/utils/secret.py index a3498b0ae3aa..50dce10176d8 100644 --- a/sdks/python/apache_beam/utils/secret.py +++ b/sdks/python/apache_beam/utils/secret.py @@ -320,7 +320,8 @@ def from_dict(cls, spec_dict: Dict[str, str]) -> 'GcpHsmGeneratedSecret': missing = allowed_keys - set(spec_dict.keys()) if missing: raise ValueError( - f"Missing required parameter(s) for GcpHsmGeneratedSecret: {sorted(list(missing))}") + f"Missing required parameter(s) for GcpHsmGeneratedSecret: {sorted(list(missing))}" + ) invalid_keys = set(spec_dict.keys()) - allowed_keys if invalid_keys: raise ValueError( From 401ce4bba751caf5a87f3fd77598ee168f98f67d Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Wed, 5 Aug 2026 14:13:15 -0400 Subject: [PATCH 08/11] Add back comments. --- sdks/python/apache_beam/utils/secret.py | 39 ++++++++++++++++++++++--- 1 file changed, 35 insertions(+), 4 deletions(-) diff --git a/sdks/python/apache_beam/utils/secret.py b/sdks/python/apache_beam/utils/secret.py index 50dce10176d8..b4201b265934 100644 --- a/sdks/python/apache_beam/utils/secret.py +++ b/sdks/python/apache_beam/utils/secret.py @@ -42,7 +42,11 @@ def generate_secret_bytes() -> bytes: class Secret(abc.ABC): - """Abstract base interface for Secrets in Apache Beam.""" + """A secret management base interface used for handling sensitive data. + + This class provides a generic interface for secret management. Implementations + of this class should handle fetching secrets from a secret management system. + """ def __init__(self): self._cached_secret_bytes: Optional[bytes] = None @@ -99,6 +103,12 @@ def generate_secret_bytes() -> bytes: def parse_secret_option(cls, secret: str) -> 'Secret': """Parses a secret string and returns the appropriate secret type. + The secret string should be formatted like: + 'type:;:' + + For example, 'type:GcpSecret;version_name:my_secret/versions/latest' + would return a GcpSecret initialized with 'my_secret/versions/latest'. + Deprecated: Use :meth:`from_option_string` instead. """ return cls.from_option_string(secret) @@ -207,8 +217,19 @@ def __eq__(self, other: Any) -> bool: class GcpSecret(Secret): - """Secret implementation using Google Cloud Secret Manager.""" + """A secret manager implementation that retrieves secrets from Google Cloud + Secret Manager. + """ def __init__(self, version_name: str): + """Initializes a GcpSecret object. + + Args: + version_name: The full version name of the secret in Google Cloud Secret + Manager. For example: + projects//secrets//versions/1. + For more info, see + https://cloud.google.com/python/docs/reference/secretmanager/latest/google.cloud.secretmanager_v1beta1.services.secret_manager_service.SecretManagerServiceClient#google_cloud_secretmanager_v1beta1_services_secret_manager_service_SecretManagerServiceClient_access_secret_version + """ super().__init__() self._version_name = version_name @@ -344,10 +365,14 @@ def __eq__(self, other: Any) -> bool: self._key_id == other._key_id and self._job_name == other._job_name) def get_secret_bytes(self) -> bytes: - """Retrieves the secret bytes from GCP Secret Manager, creating it if needed. + """Retrieves the secret bytes. + + If the secret version already exists in Secret Manager, it is retrieved. + Otherwise, a new secret and version are created. The new secret is + generated using the HSM key. Returns: - The secret bytes. + The secret as a byte string. """ from google.api_core import exceptions as api_exceptions from google.cloud import secretmanager @@ -356,6 +381,8 @@ def get_secret_bytes(self) -> bytes: project_path = f"projects/{self._project_id}" secret_path = f"{project_path}/secrets/{self._secret_version_name}" + # Since we may generate multiple versions when doing this on workers, + # just always take the first version added to maintain consistency. secret_version_path = f"{secret_path}/versions/1" try: @@ -364,6 +391,8 @@ def get_secret_bytes(self) -> bytes: request={"name": secret_version_path}) return response.payload.data except api_exceptions.NotFound: + # Don't bother logging yet, we'll only log if we actually add the + # secret version below pass try: @@ -378,6 +407,8 @@ def get_secret_bytes(self) -> bytes: }, }) except api_exceptions.AlreadyExists: + # Don't bother logging yet, we'll only log if we actually add the + # secret version below pass new_key = self.generate_dek() From 091766674379e27ae860bf1d981deba4d972ce17 Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Wed, 5 Aug 2026 14:26:12 -0400 Subject: [PATCH 09/11] Skip unit tests if secretmanager module not installed --- sdks/python/apache_beam/utils/secret_test.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/sdks/python/apache_beam/utils/secret_test.py b/sdks/python/apache_beam/utils/secret_test.py index c54d3c94adb1..e43d51f41be2 100644 --- a/sdks/python/apache_beam/utils/secret_test.py +++ b/sdks/python/apache_beam/utils/secret_test.py @@ -26,7 +26,13 @@ from apache_beam.utils.secret import Secret from apache_beam.utils.secret import generate_secret_bytes +try: + from google.cloud import secretmanager +except ImportError: + secretmanager = None # type: ignore[assignment] + +@unittest.skipIf(secretmanager is None, 'GCP dependencies are not installed') class GcpSecretTest(unittest.TestCase): @patch("google.cloud.secretmanager.SecretManagerServiceClient") def test_gcp_secret_success(self, mock_client_cls): @@ -137,6 +143,7 @@ def test_ill_formed_missing_secret_name_raises_value_error(self): self.assertIn("Secret name ('name') must be specified", str(ctx.exception)) +@unittest.skipIf(secretmanager is None, 'GCP dependencies are not installed') class GcpHsmGeneratedSecretTest(unittest.TestCase): def test_from_dict_success(self): spec = { From 6b1fa899c99f0517400ce60b4e06713f79daa52c Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Wed, 5 Aug 2026 14:32:12 -0400 Subject: [PATCH 10/11] Minor edits on docstring --- sdks/python/apache_beam/utils/secret.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/sdks/python/apache_beam/utils/secret.py b/sdks/python/apache_beam/utils/secret.py index b4201b265934..f80aaaee15b9 100644 --- a/sdks/python/apache_beam/utils/secret.py +++ b/sdks/python/apache_beam/utils/secret.py @@ -42,7 +42,7 @@ def generate_secret_bytes() -> bytes: class Secret(abc.ABC): - """A secret management base interface used for handling sensitive data. + """A secret management class used for handling sensitive data. This class provides a generic interface for secret management. Implementations of this class should handle fetching secrets from a secret management system. @@ -372,7 +372,7 @@ def get_secret_bytes(self) -> bytes: generated using the HSM key. Returns: - The secret as a byte string. + The secret bytes. """ from google.api_core import exceptions as api_exceptions from google.cloud import secretmanager From 32d4cb276672f343160950e421e7a7f63092b732 Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Wed, 5 Aug 2026 14:43:27 -0400 Subject: [PATCH 11/11] Minor refactoring. --- sdks/python/apache_beam/utils/secret.py | 29 +++++++++++++------------ 1 file changed, 15 insertions(+), 14 deletions(-) diff --git a/sdks/python/apache_beam/utils/secret.py b/sdks/python/apache_beam/utils/secret.py index f80aaaee15b9..608aba77c400 100644 --- a/sdks/python/apache_beam/utils/secret.py +++ b/sdks/python/apache_beam/utils/secret.py @@ -34,6 +34,8 @@ 'generate_secret_bytes', ] +_LOGGER = logging.getLogger(__name__) + def generate_secret_bytes() -> bytes: """Generates a new secret key using Fernet.""" @@ -178,7 +180,7 @@ def from_spec( msg = ( "The 'spec' parameter appears to be a JSON specification, but " "'secret_manager' is not set. Defaulting to Raw.") - logging.warning(msg) + _LOGGER.warning(msg) warnings.warn(msg, UserWarning) return RawSecret(spec) @@ -286,14 +288,13 @@ def get_secret_bytes(self) -> bytes: Returns: The secret bytes. """ - from google.cloud import secretmanager - try: + from google.cloud import secretmanager client = secretmanager.SecretManagerServiceClient() response = client.access_secret_version( request={"name": self._version_name}) secret_val_bytes = response.payload.data - logging.info( + _LOGGER.info( "Successfully fetched secret from GCP Secret Manager (version_name '%s')", self._version_name) return secret_val_bytes @@ -374,18 +375,18 @@ def get_secret_bytes(self) -> bytes: Returns: The secret bytes. """ - from google.api_core import exceptions as api_exceptions - from google.cloud import secretmanager + try: + from google.api_core import exceptions as api_exceptions + from google.cloud import secretmanager - client = secretmanager.SecretManagerServiceClient() + client = secretmanager.SecretManagerServiceClient() - project_path = f"projects/{self._project_id}" - secret_path = f"{project_path}/secrets/{self._secret_version_name}" - # Since we may generate multiple versions when doing this on workers, - # just always take the first version added to maintain consistency. - secret_version_path = f"{secret_path}/versions/1" + project_path = f"projects/{self._project_id}" + secret_path = f"{project_path}/secrets/{self._secret_version_name}" + # Since we may generate multiple versions when doing this on workers, + # just always take the first version added to maintain consistency. + secret_version_path = f"{secret_path}/versions/1" - try: try: response = client.access_secret_version( request={"name": secret_version_path}) @@ -418,7 +419,7 @@ def get_secret_bytes(self) -> bytes: request={"name": secret_version_path}) return response.payload.data except api_exceptions.NotFound: - logging.info( + _LOGGER.info( "Secret version %s not found. Creating new secret and version.", secret_version_path) client.add_secret_version(