feat: uses latest pydantic settings with aws

.github/workflows/test_and_lint.yml

@@ -10,7 +10,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [ '3.8', '3.9', '3.10', '3.11', '3.12' ]
+        python-version: [ '3.9', '3.10', '3.11', '3.12' ]
     steps:
     - uses: actions/checkout@v4
     - name: Set up Python ${{ matrix.python-version }}

pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "setuptools.build_meta"
 name = "aind-settings-utils"
 description = "Utility package to add custom pydantic settings sources"
 license = {text = "MIT"}
-requires-python = ">=3.8"
+requires-python = ">=3.9"
 authors = [
     {name = "Allen Institute for Neural Dynamics"}
 ]
@@ -17,8 +17,7 @@ readme = "README.md"
 dynamic = ["version"]
 
 dependencies = [
-    'boto3',
-    'pydantic-settings>=2.0',
+    'pydantic-settings[aws-secrets-manager]>=2.9.0',
     'pydantic>=2.0',
 ]
 

src/aind_settings_utils/aws.py

@@ -3,11 +3,16 @@
 import functools
 import json
 import logging
+import os
 from typing import Any, Dict, Optional, Tuple, Type
 
 import boto3
 from pydantic.fields import FieldInfo
-from pydantic_settings import BaseSettings, PydanticBaseSettingsSource
+from pydantic_settings import (
+    AWSSecretsManagerSettingsSource,
+    BaseSettings,
+    PydanticBaseSettingsSource,
+)
 from pydantic_settings.sources import PydanticBaseEnvSettingsSource
 
 
@@ -176,3 +181,41 @@ def settings_customise_sources(
                 dotenv_settings,
                 file_secret_settings,
             )
+
+
+class SecretsManagerBaseSettings(BaseSettings):
+    """Base Settings that will fall back to AWS Secrets Manager."""
+
+    @classmethod
+    def settings_customise_sources(
+        cls,
+        settings_cls: type[BaseSettings],  # noqa
+        init_settings: PydanticBaseSettingsSource,
+        env_settings: PydanticBaseSettingsSource,
+        dotenv_settings: PydanticBaseSettingsSource,
+        file_secret_settings: PydanticBaseSettingsSource,
+    ) -> tuple[PydanticBaseSettingsSource, ...]:
+        """If env var is not set, then use standard fall backs."""
+
+        secret_manager_id = os.getenv("AWS_SECRETS_MANAGER_SECRET_ID")
+        if secret_manager_id:
+            aws_secrets_manager_settings = AWSSecretsManagerSettingsSource(
+                settings_cls,
+                secret_manager_id,
+                case_sensitive=False,
+                region_name=os.getenv("AWS_REGION"),
+            )
+            return (
+                init_settings,
+                env_settings,
+                dotenv_settings,
+                file_secret_settings,
+                aws_secrets_manager_settings,
+            )
+        else:
+            return (
+                init_settings,
+                env_settings,
+                dotenv_settings,
+                file_secret_settings,
+            )

tests/test_aws.py

@@ -1,11 +1,13 @@
 """Module to test custom settings"""
 
+import os
 import unittest
 from unittest.mock import MagicMock, PropertyMock, patch
 
 from aind_settings_utils.aws import (
     AWSParamStoreAppSource,
     ParameterStoreAppBaseSettings,
+    SecretsManagerBaseSettings,
 )
 
 
@@ -20,6 +22,13 @@ class ExampleSettings(ParameterStoreAppBaseSettings):
     }
 
 
+class ExampleSettingsForSecretsManager(SecretsManagerBaseSettings):
+    """Example settings class for testing"""
+
+    my_param_1: str
+    my_param_2: int
+
+
 class TestAWSParamStoreSource(unittest.TestCase):
     """Test AWSParamStoreAppSource class"""
 
@@ -133,5 +142,47 @@ class NoParamSettings(ParameterStoreAppBaseSettings):
         self.assertEqual(len(sources), 4)
 
 
+class TestSecretsManagerBaseSettings(unittest.TestCase):
+    """Test SecretsManagerBaseSettings class"""
+
+    @patch.dict(
+        os.environ,
+        dict(),
+        clear=True,
+    )
+    @patch("boto3.client")
+    def test_settings_no_secrets_manager(self, mock_boto: MagicMock):
+        """Tests settings with no secret manager id"""
+        settings = ExampleSettingsForSecretsManager(
+            my_param_1="a", my_param_2=1
+        )
+        mock_boto.assert_not_called()
+        self.assertIsNotNone(settings)
+
+    @patch.dict(
+        os.environ,
+        {
+            "AWS_SECRETS_MANAGER_SECRET_ID": "abc/def",
+            "AWS_REGION": "us-west-2",
+        },
+        clear=True,
+    )
+    @patch("boto3.client")
+    def test_settings_with_secrets_manager(self, mock_boto: MagicMock):
+        """Tests settings with no secret manager id"""
+
+        mock_sm = MagicMock()
+        mock_sm.get_secret_value.return_value = {
+            "SecretString": '{"MY_PARAM_1":"a","MY_PARAM_2":"1"}'
+        }
+        mock_boto.return_value = mock_sm
+
+        settings = ExampleSettingsForSecretsManager()
+        expected_settings = ExampleSettingsForSecretsManager(
+            my_param_1="a", my_param_2=1
+        )
+        self.assertEqual(settings, expected_settings)
+
+
 if __name__ == "__main__":
     unittest.main()