codekingpro/portable-devtools

1from __future__ import annotations as _annotations2 3import warnings4from collections.abc import Iterator, Mapping5from functools import cached_property6from typing import TYPE_CHECKING, Any7 8from pydantic.fields import FieldInfo9 10from ..types import SecretVersion11from .env import EnvSettingsSource12 13if TYPE_CHECKING:14    from google.auth import default as google_auth_default15    from google.auth.credentials import Credentials16    from google.cloud.secretmanager import SecretManagerServiceClient17 18    from pydantic_settings.main import BaseSettings19else:20    Credentials = None21    SecretManagerServiceClient = None22    google_auth_default = None23 24 25def import_gcp_secret_manager() -> None:26    global Credentials27    global SecretManagerServiceClient28    global google_auth_default29 30    try:31        from google.auth import default as google_auth_default32        from google.auth.credentials import Credentials33 34        with warnings.catch_warnings():35            warnings.filterwarnings('ignore', category=FutureWarning)36            from google.cloud.secretmanager import SecretManagerServiceClient37    except ImportError as e:  # pragma: no cover38        raise ImportError(39            'GCP Secret Manager dependencies are not installed, run `pip install pydantic-settings[gcp-secret-manager]`'40        ) from e41 42 43class GoogleSecretManagerMapping(Mapping[str, str | None]):44    _loaded_secrets: dict[str, str | None]45    _secret_client: SecretManagerServiceClient46 47    def __init__(self, secret_client: SecretManagerServiceClient, project_id: str, case_sensitive: bool) -> None:48        self._loaded_secrets = {}49        self._secret_client = secret_client50        self._project_id = project_id51        self._case_sensitive = case_sensitive52 53    @property54    def _gcp_project_path(self) -> str:55        return self._secret_client.common_project_path(self._project_id)56 57    def _select_case_insensitive_secret(self, lower_name: str, candidates: list[str]) -> str:58        if len(candidates) == 1:59            return candidates[0]60 61        # Sort to ensure deterministic selection (prefer lowercase / ASCII last)62        candidates.sort()63        winner = candidates[-1]64        warnings.warn(65            f"Secret collision: Found multiple secrets {candidates} normalizing to '{lower_name}'. "66            f"Using '{winner}' for case-insensitive lookup.",67            UserWarning,68            stacklevel=2,69        )70        return winner71 72    @cached_property73    def _secret_name_map(self) -> dict[str, str]:74        mapping: dict[str, str] = {}75        # Group secrets by normalized name to detect collisions76        normalized_groups: dict[str, list[str]] = {}77 78        secrets = self._secret_client.list_secrets(parent=self._gcp_project_path)79        for secret in secrets:80            name = self._secret_client.parse_secret_path(secret.name).get('secret', '')81            mapping[name] = name82 83            if not self._case_sensitive:84                lower_name = name.lower()85                if lower_name not in normalized_groups:86                    normalized_groups[lower_name] = []87                normalized_groups[lower_name].append(name)88 89        if not self._case_sensitive:90            for lower_name, candidates in normalized_groups.items():91                mapping[lower_name] = self._select_case_insensitive_secret(lower_name, candidates)92 93        return mapping94 95    @property96    def _secret_names(self) -> list[str]:97        return list(self._secret_name_map.keys())98 99    def _secret_version_path(self, key: str, version: str = 'latest') -> str:100        return self._secret_client.secret_version_path(self._project_id, key, version)101 102    def _get_secret_value(self, gcp_secret_name: str, version: str = 'latest') -> str | None:103        try:104            return self._secret_client.access_secret_version(105                name=self._secret_version_path(gcp_secret_name, version)106            ).payload.data.decode('UTF-8')107        except Exception:108            return None109 110    def __getitem__(self, key: str) -> str | None:111        if key in self._loaded_secrets:112            return self._loaded_secrets[key]113 114        gcp_secret_name = self._secret_name_map.get(key)115        if gcp_secret_name is None and not self._case_sensitive:116            gcp_secret_name = self._secret_name_map.get(key.lower())117 118        if gcp_secret_name:119            self._loaded_secrets[key] = self._get_secret_value(gcp_secret_name)120        else:121            raise KeyError(key)122 123        return self._loaded_secrets[key]124 125    def __len__(self) -> int:126        return len(self._secret_names)127 128    def __iter__(self) -> Iterator[str]:129        return iter(self._secret_names)130 131 132class GoogleSecretManagerSettingsSource(EnvSettingsSource):133    _credentials: Credentials134    _secret_client: SecretManagerServiceClient135    _project_id: str136 137    def __init__(138        self,139        settings_cls: type[BaseSettings],140        credentials: Credentials | None = None,141        project_id: str | None = None,142        env_prefix: str | None = None,143        env_parse_none_str: str | None = None,144        env_parse_enums: bool | None = None,145        secret_client: SecretManagerServiceClient | None = None,146        case_sensitive: bool | None = True,147    ) -> None:148        # Import Google Packages if they haven't already been imported149        if SecretManagerServiceClient is None or Credentials is None or google_auth_default is None:150            import_gcp_secret_manager()151 152        # If credentials or project_id are not passed, then153        # try to get them from the default function154        if not credentials or not project_id:155            _creds, _project_id = google_auth_default()156 157        # Set the credentials and/or project id if they weren't specified158        if credentials is None:159            credentials = _creds160 161        if project_id is None:162            if isinstance(_project_id, str):163                project_id = _project_id164            else:165                raise AttributeError(166                    'project_id is required to be specified either as an argument or from the google.auth.default. See https://google-auth.readthedocs.io/en/master/reference/google.auth.html#google.auth.default'167                )168 169        self._credentials: Credentials = credentials170        self._project_id: str = project_id171 172        if secret_client:173            self._secret_client = secret_client174        else:175            self._secret_client = SecretManagerServiceClient(credentials=self._credentials)176 177        super().__init__(178            settings_cls,179            case_sensitive=case_sensitive,180            env_prefix=env_prefix,181            env_ignore_empty=False,182            env_parse_none_str=env_parse_none_str,183            env_parse_enums=env_parse_enums,184        )185 186    def get_field_value(self, field: FieldInfo, field_name: str) -> tuple[Any, str, bool]:187        """Override get_field_value to get the secret value from GCP Secret Manager.188        Look for a SecretVersion metadata field to specify a particular SecretVersion.189 190        Args:191            field: The field to get the value for192            field_name: The declared name of the field193 194        Returns:195            A tuple of (value, key, value_is_complex), where `key` is the identifier used196            to populate the model (either the field name or an alias, depending on197            configuration).198        """199 200        secret_version = next((m.version for m in field.metadata if isinstance(m, SecretVersion)), None)201 202        # If a secret version is specified, try to get that specific version of the secret from203        # GCP Secret Manager via the GoogleSecretManagerMapping. This allows different versions204        # of the same secret name to be retrieved independently and cached in the GoogleSecretManagerMapping205        if secret_version and isinstance(self.env_vars, GoogleSecretManagerMapping):206            for field_key, env_name, value_is_complex in self._extract_field_info(field, field_name):207                gcp_secret_name = self.env_vars._secret_name_map.get(env_name)208                if gcp_secret_name is None and not self.case_sensitive:209                    gcp_secret_name = self.env_vars._secret_name_map.get(env_name.lower())210 211                if gcp_secret_name:212                    env_val = self.env_vars._get_secret_value(gcp_secret_name, secret_version)213                    if env_val is not None:214                        # If populate_by_name is enabled, return field_name to allow multiple fields215                        # with the same alias but different versions to be distinguished216                        if self.settings_cls.model_config.get('populate_by_name'):217                            return env_val, field_name, value_is_complex218                        return env_val, field_key, value_is_complex219 220            # If a secret version is specified but not found, we should not fall back to "latest" (default behavior)221            # as that would be incorrect. We return None to indicate the value was not found.222            return None, field_name, False223 224        val, key, is_complex = super().get_field_value(field, field_name)225 226        # If populate_by_name is enabled, we need to return the field_name as the key227        # without this being enabled, you cannot load two secrets with the same name but different versions228        if self.settings_cls.model_config.get('populate_by_name') and val is not None:229            return val, field_name, is_complex230        return val, key, is_complex231 232    def _load_env_vars(self) -> Mapping[str, str | None]:233        return GoogleSecretManagerMapping(234            self._secret_client, project_id=self._project_id, case_sensitive=self.case_sensitive235        )236 237    def __repr__(self) -> str:238        return f'{self.__class__.__name__}(project_id={self._project_id!r}, env_nested_delimiter={self.env_nested_delimiter!r})'239 240 241__all__ = ['GoogleSecretManagerSettingsSource', 'GoogleSecretManagerMapping']242