feat: support custom boto3 client

Author: alukach

pydantic_ssm_settings/settings.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Any, Optional, Tuple, Type
+from typing import TYPE_CHECKING, Any, Optional, Tuple, Type
 
 from pydantic_settings import (
     BaseSettings,
@@ -12,11 +12,15 @@
 
 from .source import AwsSsmSettingsSource
 
+if TYPE_CHECKING:
+    from mypy_boto3_ssm.client import SSMClient
+
 logger = logging.getLogger(__name__)
 
 
 class SsmSettingsConfigDict(SettingsConfigDict):
     ssm_prefix: str
+    ssm_client: Optional["SSMClient"]
 
 
 class AwsSsmBaseSettings(BaseSettings):
@@ -29,6 +33,7 @@ def __init__(
         self,
         *args,
         _ssm_prefix: Optional[str] = None,
+        _ssm_client: Optional["SSMClient"] = None,
         **kwargs: Any,
     ) -> None:
         """
@@ -37,9 +42,12 @@ def __init__(
             separated by "/". NB:unlike its _env_prefix counterpart, _ssm_prefix
             is treated case sensitively regardless of the _case_sensitive
             parameter value.
+            _ssm_client: Optional boto3 SSM client. If not provided, a new client
+            will be created.
         """
         # NOTE: Need a direct access to the attributes dictionary to avoid raising an AttributeError: __pydantic_private__ exception
         self.__dict__["__ssm_prefix"] = _ssm_prefix
+        self.__dict__["__ssm_client"] = _ssm_client
         super().__init__(self, *args, **kwargs)
 
     def settings_customise_sources(
@@ -53,6 +61,7 @@ def settings_customise_sources(
         ssm_settings = AwsSsmSettingsSource(
             settings_cls=settings_cls,
             ssm_prefix=self.__dict__["__ssm_prefix"],
+            ssm_client=self.__dict__["__ssm_client"],
         )
 
         return (

pydantic_ssm_settings/source.py

@@ -31,31 +31,31 @@ def __init__(
         settings_cls: type[BaseSettings],
         case_sensitive: Optional[bool] = None,
         ssm_prefix: Optional[str] = None,
+        ssm_client: Optional["SSMClient"] = None,
     ):
-        # Ideally would retrieve ssm_prefix from self.config
-        # but need the superclass to be initialized for that
-        ssm_prefix_ = (
+        ssm_prefix = (
             ssm_prefix
             if ssm_prefix is not None
             else settings_cls.model_config.get("ssm_prefix", "/")
         )
+        self.ssm_client = (
+            ssm_client
+            if ssm_client
+            else settings_cls.model_config.get("ssm_client", self._build_client())
+        )
         super().__init__(
             settings_cls,
             case_sensitive=case_sensitive,
-            env_prefix=ssm_prefix_,
+            env_prefix=ssm_prefix,
             env_nested_delimiter="/",  # SSM only accepts / as a delimiter
         )
-        self.ssm_prefix = ssm_prefix_
-        assert self.ssm_prefix == self.env_prefix
-
-    @property
-    def client(self) -> "SSMClient":
-        return boto3.client("ssm", config=self.client_config)
+        assert ssm_prefix == self.env_prefix
 
-    @property
-    def client_config(self) -> Config:
+    def _build_client(self) -> "SSMClient":
         timeout = float(os.environ.get("SSM_TIMEOUT", 0.5))
-        return Config(connect_timeout=timeout, read_timeout=timeout)
+        return boto3.client(
+            "ssm", config=Config(connect_timeout=timeout, read_timeout=timeout)
+        )
 
     def _load_env_vars(
         self,
@@ -68,7 +68,7 @@ def _load_env_vars(
 
         output = {}
         try:
-            paginator = self.client.get_paginator("get_parameters_by_path")
+            paginator = self.ssm_client.get_paginator("get_parameters_by_path")
             response_iterator = paginator.paginate(
                 Path=self.env_prefix, WithDecryption=True, Recursive=True
             )

tests/test_main.py

@@ -1,3 +1,5 @@
+import boto3
+import moto
 import pytest
 from pydantic import BaseModel
 from pydantic_settings import SettingsConfigDict
@@ -116,3 +118,13 @@ def test_parameters_from_model_config(ssm):
     ssm.put_parameter(Name="/asdf/foo", Value="bar", Type="String")
     settings = CustomConfigDict()
     assert settings.foo == "bar"
+
+
+@pytest.mark.parametrize("region", ["us-east-1", "us-west-2"])
+def test_custom_client(region: str):
+    with moto.mock_aws():
+        client = boto3.client("ssm", region_name=region)
+        client.put_parameter(Name="/foo", Value=region, Type="String")
+
+        settings = SimpleSettings(_ssm_client=client)
+        assert settings.foo == region