diff --git a/.github/workflows/test_and_lint.yml b/.github/workflows/test_and_lint.yml index e3f6a55..a10ce3c 100644 --- a/.github/workflows/test_and_lint.yml +++ b/.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 }} diff --git a/pyproject.toml b/pyproject.toml index 5c32ead..d60bf33 100644 --- a/pyproject.toml +++ b/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', ] diff --git a/src/aind_settings_utils/aws.py b/src/aind_settings_utils/aws.py index cbdcf27..9ce9267 100644 --- a/src/aind_settings_utils/aws.py +++ b/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, + ) diff --git a/tests/test_aws.py b/tests/test_aws.py index 2519623..dc9490f 100644 --- a/tests/test_aws.py +++ b/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()