-
-
Notifications
You must be signed in to change notification settings - Fork 108
Add support for AWS Secrets Manager #532
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from 2 commits
fa6f3a0
5a76f48
a2f7478
6134b0a
e00cd49
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,69 @@ | ||
| from __future__ import annotations as _annotations # important for BaseSettings import to work | ||
|
|
||
| import json | ||
| from collections.abc import Mapping | ||
| from typing import TYPE_CHECKING, Optional | ||
|
|
||
| from .env import EnvSettingsSource | ||
|
|
||
| if TYPE_CHECKING: | ||
| from pydantic_settings.main import BaseSettings | ||
|
|
||
|
|
||
| boto3_client = None | ||
| SecretsManagerClient = None | ||
|
|
||
|
|
||
| def import_aws_secrets_manager() -> None: | ||
| global boto3_client | ||
| global SecretsManagerClient | ||
|
|
||
| try: | ||
| from boto3 import client as boto3_client | ||
| from mypy_boto3_secretsmanager.client import SecretsManagerClient | ||
| except ImportError as e: | ||
| raise ImportError( | ||
| 'AWS Secrets Manager dependencies are not installed, run `pip install pydantic-settings[aws-secrets-manager]`' | ||
| ) from e | ||
|
|
||
|
|
||
| class AWSSecretsManagerSettingsSource(EnvSettingsSource): | ||
| _secret_id: str | ||
| _secretsmanager_client: SecretsManagerClient # type: ignore | ||
|
|
||
| def __init__( | ||
| self, | ||
| settings_cls: type[BaseSettings], | ||
| secret_id: str, | ||
| env_prefix: str | None = None, | ||
| env_parse_none_str: str | None = None, | ||
| env_parse_enums: bool | None = None, | ||
| ) -> None: | ||
| import_aws_secrets_manager() | ||
| self._secretsmanager_client = boto3_client('secretsmanager') # type: ignore | ||
| self._secret_id = secret_id | ||
| super().__init__( | ||
| settings_cls, | ||
| case_sensitive=True, | ||
| env_prefix=env_prefix, | ||
| env_nested_delimiter='--', | ||
| env_ignore_empty=False, | ||
| env_parse_none_str=env_parse_none_str, | ||
| env_parse_enums=env_parse_enums, | ||
| ) | ||
|
|
||
| def _load_env_vars(self) -> Mapping[str, Optional[str]]: | ||
| response = self._secretsmanager_client.get_secret_value(SecretId=self._secret_id) # type: ignore | ||
|
|
||
| return json.loads(response['SecretString']) | ||
|
|
||
| def __repr__(self) -> str: | ||
| return ( | ||
| f'{self.__class__.__name__}(secret_id={self._secret_id!r}, ' | ||
| f'env_nested_delimiter={self.env_nested_delimiter!r})' | ||
| ) | ||
|
|
||
|
|
||
| __all__ = [ | ||
| 'AWSSecretsManagerSettingsSource', | ||
| ] |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,120 @@ | ||
| """ | ||
| Test pydantic_settings.AWSSecretsManagerSettingsSource. | ||
| """ | ||
|
|
||
| import json | ||
| import os | ||
|
|
||
| import pytest | ||
|
|
||
| try: | ||
| import yaml | ||
| from moto import mock_aws | ||
| except ImportError: | ||
| yaml = None | ||
| mock_aws = None | ||
|
|
||
| from pydantic import BaseModel, Field | ||
|
|
||
| from pydantic_settings import ( | ||
| AWSSecretsManagerSettingsSource, | ||
| BaseSettings, | ||
| PydanticBaseSettingsSource, | ||
| ) | ||
| from pydantic_settings.sources.providers.aws import import_aws_secrets_manager | ||
|
|
||
| try: | ||
| aws_secrets_manager = True | ||
| import_aws_secrets_manager() | ||
| import boto3 | ||
|
|
||
| os.environ['AWS_DEFAULT_REGION'] = os.environ.get('AWS_DEFAULT_REGION', 'us-east-1') | ||
| except ImportError: | ||
| aws_secrets_manager = False | ||
|
|
||
|
|
||
| MODULE = 'pydantic_settings.sources' | ||
|
|
||
| if not yaml: | ||
| pytest.skip('PyYAML is not installed', allow_module_level=True) | ||
|
|
||
|
|
||
| @pytest.mark.skipif(not aws_secrets_manager, reason='pydantic-settings[aws-secrets-manager] is not installed') | ||
| class TestAWSSecretsManagerSettingsSource: | ||
| """Test AWSSecretsManagerSettingsSource.""" | ||
|
|
||
| @mock_aws | ||
| def test___init__(self) -> None: | ||
| """Test __init__.""" | ||
|
|
||
| class AWSSecretsManagerSettings(BaseSettings): | ||
| """AWSSecretsManager settings.""" | ||
|
|
||
| client = boto3.client('secretsmanager') | ||
| client.create_secret(Name='test-secret', SecretString='{}') | ||
|
|
||
| AWSSecretsManagerSettingsSource(AWSSecretsManagerSettings, 'test-secret') | ||
|
|
||
| @mock_aws | ||
| def test___call__(self) -> None: | ||
| """Test __call__.""" | ||
|
|
||
| class SqlServer(BaseModel): | ||
| password: str = Field(..., alias='Password') | ||
|
|
||
| class AWSSecretsManagerSettings(BaseSettings): | ||
| """AWSSecretsManager settings.""" | ||
|
|
||
| SqlServerUser: str | ||
|
||
| sql_server_user: str = Field(..., alias='SqlServerUser') | ||
| sql_server: SqlServer = Field(..., alias='SqlServer') | ||
|
|
||
| expected_secret_value = 'SecretValue' | ||
| secret_data = {'SqlServerUser': expected_secret_value, 'SqlServer--Password': expected_secret_value} | ||
hramezani marked this conversation as resolved.
Outdated
Show resolved
Hide resolved
|
||
|
|
||
| client = boto3.client('secretsmanager') | ||
| client.create_secret(Name='test-secret', SecretString=json.dumps(secret_data)) | ||
|
|
||
| obj = AWSSecretsManagerSettingsSource(AWSSecretsManagerSettings, 'test-secret') | ||
|
|
||
| settings = obj() | ||
|
|
||
| assert settings['SqlServerUser'] == expected_secret_value | ||
| assert settings['SqlServer']['Password'] == expected_secret_value | ||
hramezani marked this conversation as resolved.
Outdated
Show resolved
Hide resolved
|
||
|
|
||
| @mock_aws | ||
| def test_aws_secrets_manager_settings_source(self) -> None: | ||
| """Test AWSSecretsManagerSettingsSource.""" | ||
|
|
||
| class SqlServer(BaseModel): | ||
| password: str = Field(..., alias='Password') | ||
|
|
||
| class AWSSecretsManagerSettings(BaseSettings): | ||
| """AWSSecretsManager settings.""" | ||
|
|
||
| SqlServerUser: str | ||
| sql_server_user: str = Field(..., alias='SqlServerUser') | ||
| sql_server: SqlServer = Field(..., alias='SqlServer') | ||
|
|
||
| @classmethod | ||
| def settings_customise_sources( | ||
| cls, | ||
| settings_cls: type[BaseSettings], | ||
| init_settings: PydanticBaseSettingsSource, | ||
| env_settings: PydanticBaseSettingsSource, | ||
| dotenv_settings: PydanticBaseSettingsSource, | ||
| file_secret_settings: PydanticBaseSettingsSource, | ||
| ) -> tuple[PydanticBaseSettingsSource, ...]: | ||
| return (AWSSecretsManagerSettingsSource(settings_cls, 'test-secret'),) | ||
|
|
||
| expected_secret_value = 'SecretValue' | ||
| secret_data = {'SqlServerUser': expected_secret_value, 'SqlServer--Password': expected_secret_value} | ||
|
|
||
| client = boto3.client('secretsmanager') | ||
| client.create_secret(Name='test-secret', SecretString=json.dumps(secret_data)) | ||
|
|
||
| settings = AWSSecretsManagerSettings() # type: ignore | ||
|
|
||
| assert settings.SqlServerUser == expected_secret_value | ||
| assert settings.sql_server_user == expected_secret_value | ||
| assert settings.sql_server.password == expected_secret_value | ||
hramezani marked this conversation as resolved.
Outdated
Show resolved
Hide resolved
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Do we need
pyyamlto be installed for this test?There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yes. Moto the mocking library for boto3 has a dependency on responses which imports yaml. The second tests run in ci was failing and tests won't even start.
Hence the import moto is inside try catch after trying to import yaml.
And the skip module is because it fails for the @mock_aws decorator.