Add AWSBaseSettings

pydantic_settings_aws/__init__.py

@@ -1,6 +1,14 @@
-from .settings import ParameterStoreBaseSettings, SecretsManagerBaseSettings
+from .settings import (
+    AWSBaseSettings,
+    ParameterStoreBaseSettings,
+    SecretsManagerBaseSettings,
+)
 from .version import VERSION
 
-__all__ = ["ParameterStoreBaseSettings", "SecretsManagerBaseSettings"]
+__all__ = [
+    "AWSBaseSettings",
+    "ParameterStoreBaseSettings",
+    "SecretsManagerBaseSettings",
+]
 
 __version__ = VERSION

pydantic_settings_aws/aws.py

@@ -130,6 +130,7 @@ def _create_client_from_settings(  # type: ignore[no-untyped-def]
     client = settings.model_config.get(client_param)
 
     if client:
+        logger.debug(f"Will use client from model config {client_param}")
         return client
 
     logger.debug("Extracting settings prefixed with aws_")

pydantic_settings_aws/settings.py

@@ -5,7 +5,30 @@
     PydanticBaseSettingsSource,
 )
 
-from .sources import ParameterStoreSettingsSource, SecretsManagerSettingsSource
+from .sources import (
+    AWSSettingsSource,
+    ParameterStoreSettingsSource,
+    SecretsManagerSettingsSource,
+)
+
+
+class AWSBaseSettings(BaseSettings):
+    @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 (
+            init_settings,
+            AWSSettingsSource(settings_cls),
+            env_settings,
+            dotenv_settings,
+            file_secret_settings,
+        )
 
 
 class ParameterStoreBaseSettings(BaseSettings):

pydantic_settings_aws/sources.py

@@ -7,19 +7,82 @@
 )
 
 from pydantic_settings_aws import aws, utils
+from pydantic_settings_aws.logger import logger
+
+
+class AWSSettingsSource(PydanticBaseSettingsSource):
+    def __init__(self, settings_cls: Type[BaseSettings]):
+        super().__init__(settings_cls)
+
+    def get_field_value(
+        self, field: FieldInfo, field_name: str
+    ) -> Tuple[Any, str, bool]:
+        logger.debug(f"Getting {field_name} value from AWS service")
+        field_value = None
+
+        service_metadata = utils.get_annotated_service_metadata(field.metadata)
+
+        if not service_metadata:
+            logger.info("No information about AWS service was found")
+            return None, field_name, False
+
+        service = service_metadata.get("service")
+
+        if service == "ssm":
+            ssm_info = utils.get_ssm_name_from_annotated_field(field.metadata)
+            field_value = aws.get_ssm_content(
+                self.settings_cls, field_name, ssm_info
+            )
+
+        elif service == "secrets":
+            logger.debug(
+                f"Getting value from secrets manager for filed {field_name}"
+            )
+            json_content = aws.get_secrets_content(self.settings_cls)
+            field_value = json_content.get(field_name)
+
+        logger.info(f"field value={field_value}")
+
+        return field_value, field_name, False
+
+    def prepare_field_value(
+        self,
+        field_name: str,
+        field: FieldInfo,
+        value: Any,
+        value_is_complex: bool,
+    ) -> Any:
+        return value
+
+    def __call__(self) -> Dict[str, Any]:
+        d: Dict[str, Any] = {}
+
+        for field_name, field in self.settings_cls.model_fields.items():
+            field_value, field_key, value_is_complex = self.get_field_value(
+                field, field_name
+            )
+            field_value = self.prepare_field_value(
+                field_name, field, field_value, value_is_complex
+            )
+            if field_value is not None:
+                d[field_key] = field_value
+
+        return d
 
 
 class ParameterStoreSettingsSource(PydanticBaseSettingsSource):
-    """Source class for loading settings from AWS Parameter Store.
-    """
+    """Source class for loading settings from AWS Parameter Store."""
+
     def __init__(self, settings_cls: Type[BaseSettings]):
         super().__init__(settings_cls)
 
     def get_field_value(
         self, field: FieldInfo, field_name: str
     ) -> Tuple[Any, str, bool]:
         ssm_info = utils.get_ssm_name_from_annotated_field(field.metadata)
-        field_value = aws.get_ssm_content(self.settings_cls, field_name, ssm_info)
+        field_value = aws.get_ssm_content(
+            self.settings_cls, field_name, ssm_info
+        )
 
         return field_value, field_name, False
 

pydantic_settings_aws/utils.py

@@ -1,10 +1,16 @@
-from typing import Any, List, Optional
+from typing import Any, Dict, List, Optional
+
+
+def get_annotated_service_metadata(
+    metadata: List[Any],
+) -> Optional[Dict[str, Any]]:
+    service_metadata = list(filter(_get_service_metadata, metadata))
+
+    return service_metadata[0] if service_metadata else None
 
 
 def get_ssm_name_from_annotated_field(metadata: List[Any]) -> Optional[str]:
-    ssm_metadata = list(
-        filter(_get_ssm_info_from_metadata, metadata)
-    )
+    ssm_metadata = list(filter(_get_ssm_info_from_metadata, metadata))
 
     if ssm_metadata:
         return ssm_metadata[0]
@@ -16,7 +22,14 @@ def _get_ssm_info_from_metadata(metadata: Any) -> Optional[Any]:
     if isinstance(metadata, str):
         return metadata
 
-    if isinstance(metadata, dict) and "ssm" in metadata.keys():
+    if isinstance(metadata, dict) and "ssm" in metadata:
+        return metadata
+
+    return None
+
+
+def _get_service_metadata(metadata: Any) -> Optional[Dict[str, Any]]:
+    if isinstance(metadata, dict) and "service" in metadata:
         return metadata
 
     return None

pyproject.toml

@@ -51,11 +51,11 @@ Repository = "https://github.com/ceb10n/pydantic-settings-aws"
 [tool.pytest.ini_options]
 testpaths = 'tests'
 log_cli = true
-log_cli_level = "CRITICAL"
+log_cli_level = "ERROR"
 log_cli_format = "%(message)s"
 
 log_file = "pytest.log"
-log_file_level = "DEBUG"
+log_file_level = "ERROR"
 log_file_format = "%(asctime)s [%(levelname)8s] %(message)s (%(filename)s:%(lineno)s)"
 log_file_date_format = "%Y-%m-%d %H:%M:%S"
 

tests/settings_mocks.py

@@ -6,14 +6,21 @@
 from typing_extensions import Annotated
 
 from pydantic_settings_aws import (
+    AWSBaseSettings,
     ParameterStoreBaseSettings,
     SecretsManagerBaseSettings,
 )
 
 from .boto3_mocks import ClientMock
 
+dict_secrets_with_username_and_password = {
+    "username": "myusername",
+    "password": "password1234",
+    "name": None,
+}
+
 secrets_with_username_and_password = json.dumps(
-    {"username": "myusername", "password": "password1234", "name": None}
+    dict_secrets_with_username_and_password
 )
 
 mock_secrets_with_username_and_pwd = ClientMock(
@@ -61,22 +68,60 @@ class SecretsWithNestedContent(SecretsManagerBaseSettings):
 
 
 class ParameterSettings(ParameterStoreBaseSettings):
-    my_ssm: Annotated[str, {"ssm": "my/parameter", "ssm_client": ClientMock(ssm_value="value")}]
+    my_ssm: Annotated[
+        str,
+        {"ssm": "my/parameter", "ssm_client": ClientMock(ssm_value="value")},
+    ]
 
 
 class ParameterWithTwoSSMClientSettings(ParameterStoreBaseSettings):
+    model_config = SettingsConfigDict(ssm_client=ClientMock(ssm_value="value"))
+
+    my_ssm: Annotated[
+        str,
+        {"ssm": "my/parameter", "ssm_client": ClientMock(ssm_value="value")},
+    ]
+    my_ssm_1: Annotated[
+        str,
+        {"ssm": "my/parameter", "ssm_client": ClientMock(ssm_value="value1")},
+    ]
+    my_ssm_2: Annotated[str, "my/ssm/2/parameter"]
+
+
+class ParameterWithOptionalValueSettings(ParameterStoreBaseSettings):
+    model_config = SettingsConfigDict(ssm_client=ClientMock())
+
+    my_ssm: Annotated[Optional[str], "my/ssm/2/parameter"] = None
+
+
+class AWSWithParameterAndSecretsWithDefaultBoto3Client(AWSBaseSettings):
     model_config = SettingsConfigDict(
-        ssm_client=ClientMock(ssm_value="value")
+        ssm_client=ClientMock(ssm_value="value"),
+        secrets_client=mock_secrets_with_username_and_pwd,
+        secrets_name="my/secret",
     )
 
-    my_ssm: Annotated[str, {"ssm": "my/parameter", "ssm_client": ClientMock(ssm_value="value")}]
-    my_ssm_1: Annotated[str, {"ssm": "my/parameter", "ssm_client": ClientMock(ssm_value="value1")}]
-    my_ssm_2: Annotated[str, "my/ssm/2/parameter"]
+    username: Annotated[str, {"service": "secrets"}]
+    password: Annotated[str, {"service": "secrets"}]
+    host: Annotated[str, {"service": "ssm"}]
 
 
-class ParameterWithOptionalValueSettings(ParameterStoreBaseSettings):
+class AWSWithParameterSecretsAndEnvironmentWithDefaultBoto3Client(AWSBaseSettings):
     model_config = SettingsConfigDict(
-        ssm_client=ClientMock()
+        ssm_client=ClientMock(ssm_value="value"),
+        secrets_client=mock_secrets_with_username_and_pwd,
+        secrets_name="my/secret",
     )
 
-    my_ssm: Annotated[Optional[str], "my/ssm/2/parameter"] = None
+    username: Annotated[str, {"service": "secrets"}]
+    password: Annotated[str, {"service": "secrets"}]
+    host: Annotated[str, {"service": "ssm"}]
+    server_name: str
+
+
+class AWSWithUnknownService(AWSBaseSettings):
+    my_name: Annotated[Optional[str], {"service": "s3"}] = None
+
+
+class AWSWithNonDictMetadata(AWSBaseSettings):
+    my_name: Annotated[Optional[str], "my irrelevant metadata"] = None

tests/settings_test.py

@@ -1,9 +1,16 @@
+import os
+
 from .settings_mocks import (
+    AWSWithNonDictMetadata,
+    AWSWithParameterAndSecretsWithDefaultBoto3Client,
+    AWSWithParameterSecretsAndEnvironmentWithDefaultBoto3Client,
+    AWSWithUnknownService,
     MySecretsWithClientConfig,
     ParameterSettings,
     ParameterWithOptionalValueSettings,
     ParameterWithTwoSSMClientSettings,
     SecretsWithNestedContent,
+    dict_secrets_with_username_and_password,
 )
 
 
@@ -47,3 +54,41 @@ def test_ssm_with_none_in_optional_values():
 
     assert my_config is not None
     assert my_config.my_ssm is None
+
+
+def test_aws_with_secrets_and_parameters():
+    my_config = AWSWithParameterAndSecretsWithDefaultBoto3Client()
+
+    assert my_config is not None
+    assert (
+        my_config.username
+        == dict_secrets_with_username_and_password["username"]
+    )
+    assert (
+        my_config.password
+        == dict_secrets_with_username_and_password["password"]
+    )
+    assert my_config.host is not None
+
+
+def test_aws_settings_should_get_value_from_environment_if_not_found_in_ssm_or_secrets():
+    os.environ["server_name"] = "test-server"
+
+    my_config = AWSWithParameterSecretsAndEnvironmentWithDefaultBoto3Client()
+    assert my_config is not None
+    assert my_config.username is not None
+    assert my_config.password is not None
+    assert my_config.host is not None
+    assert my_config.server_name == "test-server"
+
+
+def test_aws_settings_should_ignore_value_if_service_is_unknown():
+    my_config = AWSWithUnknownService()
+    assert my_config is not None
+    assert my_config.my_name is None
+
+
+def test_aws_settings_should_ignore_value_if_metadata_is_not_a_dict():
+    my_config = AWSWithNonDictMetadata()
+    assert my_config is not None
+    assert my_config.my_name is None