Add Azure Key Vault settings source (#272)

Co-authored-by: Hasan Ramezani <[email protected]>

Makefile

@@ -13,7 +13,7 @@ refresh-lockfiles:
 	find requirements/ -name '*.txt' ! -name 'all.txt' -type f -delete
 	pip-compile -q --no-emit-index-url --resolver backtracking -o requirements/linting.txt requirements/linting.in
 	pip-compile -q --no-emit-index-url --resolver backtracking -o requirements/testing.txt requirements/testing.in
-	pip-compile -q --no-emit-index-url --resolver backtracking --extra toml --extra yaml -o requirements/pyproject.txt pyproject.toml
+	pip-compile -q --no-emit-index-url --resolver backtracking --extra toml --extra yaml --extra azure-key-vault -o requirements/pyproject.txt pyproject.toml
 	pip install --dry-run -r requirements/all.txt
 
 .PHONY: format

docs/index.md

@@ -1173,6 +1173,65 @@ Last, run your application inside a Docker container and supply your newly creat
 docker service create --name pydantic-with-secrets --secret my_secret_data pydantic-app:latest
 ```
 
+## Azure Key Vault
+
+You must set two parameters:
+
+- `url`: For example, `https://my-resource.vault.azure.net/`.
+- `credential`: If you use `DefaultAzureCredential`, in local you can execute `az login` to get your identity credentials. The identity must have a role assignment (the recommended one is `Key Vault Secrets User`), so you can access the secrets.
+
+You must have the same naming convention in the field name as in the Key Vault secret name. For example, if the secret is named `SqlServerPassword`, the field name must be the same. You can use an alias too.
+
+In Key Vault, nested models are supported with the `--` separator. For example, `SqlServer--Password`.
+
+Key Vault arrays (e.g. `MySecret--0`, `MySecret--1`) are not supported.
+
+```py
+import os
+from typing import Tuple, Type
+
+from azure.identity import DefaultAzureCredential
+from pydantic import BaseModel
+
+from pydantic_settings import (
+    AzureKeyVaultSettingsSource,
+    BaseSettings,
+    PydanticBaseSettingsSource,
+)
+
+
+class SubModel(BaseModel):
+    a: str
+
+
+class AzureKeyVaultSettings(BaseSettings):
+    foo: str
+    bar: int
+    sub: SubModel
+
+    @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, ...]:
+        az_key_vault_settings = AzureKeyVaultSettingsSource(
+            settings_cls,
+            os.environ['AZURE_KEY_VAULT_URL'],
+            DefaultAzureCredential(),
+        )
+        return (
+            init_settings,
+            env_settings,
+            dotenv_settings,
+            file_secret_settings,
+            az_key_vault_settings,
+        )
+```
+
 ## Other settings source
 
 Other settings sources are available for common configuration files:

pydantic_settings/__init__.py

@@ -1,5 +1,6 @@
 from .main import BaseSettings, SettingsConfigDict
 from .sources import (
+    AzureKeyVaultSettingsSource,
     CliPositionalArg,
     CliSettingsSource,
     CliSubCommand,
@@ -30,6 +31,7 @@
     'SettingsConfigDict',
     'TomlConfigSettingsSource',
     'YamlConfigSettingsSource',
+    'AzureKeyVaultSettingsSource',
     '__version__',
 )
 

pydantic_settings/sources.py

@@ -19,8 +19,10 @@
     Any,
     Callable,
     Generic,
+    Iterator,
     List,
     Mapping,
+    Optional,
     Sequence,
     Tuple,
     TypeVar,
@@ -83,6 +85,21 @@ def import_toml() -> None:
         import tomllib
 
 
+def import_azure_key_vault() -> None:
+    global TokenCredential
+    global SecretClient
+    global ResourceNotFoundError
+
+    try:
+        from azure.core.credentials import TokenCredential
+        from azure.core.exceptions import ResourceNotFoundError
+        from azure.keyvault.secrets import SecretClient
+    except ImportError as e:
+        raise ImportError(
+            'Azure Key Vault dependencies are not installed, run `pip install pydantic-settings[azure-key-vault]`'
+        ) from e
+
+
 DotenvType = Union[Path, str, List[Union[Path, str]], Tuple[Union[Path, str], ...]]
 PathType = Union[Path, str, List[Union[Path, str]], Tuple[Union[Path, str], ...]]
 DEFAULT_PATH: PathType = Path('')
@@ -1725,6 +1742,70 @@ def _read_file(self, file_path: Path) -> dict[str, Any]:
             return yaml.safe_load(yaml_file) or {}
 
 
+class AzureKeyVaultMapping(Mapping[str, Optional[str]]):
+    _loaded_secrets: dict[str, str | None]
+    _secret_client: SecretClient  # type: ignore
+    _secret_names: list[str]
+
+    def __init__(
+        self,
+        secret_client: SecretClient,  # type: ignore
+    ) -> None:
+        self._loaded_secrets = {}
+        self._secret_client = secret_client
+        self._secret_names: list[str] = [secret.name for secret in self._secret_client.list_properties_of_secrets()]
+
+    def __getitem__(self, key: str) -> str | None:
+        if key not in self._loaded_secrets:
+            try:
+                self._loaded_secrets[key] = self._secret_client.get_secret(key).value
+            except ResourceNotFoundError:  # type: ignore
+                raise KeyError(key)
+
+        return self._loaded_secrets[key]
+
+    def __len__(self) -> int:
+        return len(self._secret_names)
+
+    def __iter__(self) -> Iterator[str]:
+        return iter(self._secret_names)
+
+
+class AzureKeyVaultSettingsSource(EnvSettingsSource):
+    _url: str
+    _credential: TokenCredential  # type: ignore
+    _secret_client: SecretClient  # type: ignore
+
+    def __init__(
+        self,
+        settings_cls: type[BaseSettings],
+        url: str,
+        credential: TokenCredential,  # type: ignore
+        env_prefix: str | None = None,
+        env_parse_none_str: str | None = None,
+        env_parse_enums: bool | None = None,
+    ) -> None:
+        import_azure_key_vault()
+        self._url = url
+        self._credential = credential
+        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]]:
+        secret_client = SecretClient(vault_url=self._url, credential=self._credential)  # type: ignore
+        return AzureKeyVaultMapping(secret_client)
+
+    def __repr__(self) -> str:
+        return f'AzureKeyVaultSettingsSource(url={self._url!r}, ' f'env_nested_delimiter={self.env_nested_delimiter!r})'
+
+
 def _get_env_var_key(key: str, case_sensitive: bool = False) -> str:
     return key if case_sensitive else key.lower()
 

pyproject.toml

@@ -48,6 +48,7 @@ dynamic = ['version']
 [project.optional-dependencies]
 yaml = ["pyyaml>=6.0.1"]
 toml = ["tomli>=2.0.1"]
+azure-key-vault = ["azure-keyvault-secrets>=4.8.0", "azure-identity>=1.16.0"]
 
 [project.urls]
 Homepage = 'https://github.com/pydantic/pydantic-settings'

requirements/linting.txt

@@ -4,43 +4,43 @@
 #
 #    pip-compile --no-emit-index-url --output-file=requirements/linting.txt requirements/linting.in
 #
-black==24.4.0
+black==24.4.2
     # via -r requirements/linting.in
 cfgv==3.4.0
     # via pre-commit
 click==8.1.7
     # via black
 distlib==0.3.8
     # via virtualenv
-filelock==3.13.4
+filelock==3.15.3
     # via virtualenv
 identify==2.5.36
     # via pre-commit
-mypy==1.9.0
+mypy==1.10.0
     # via -r requirements/linting.in
 mypy-extensions==1.0.0
     # via
     #   black
     #   mypy
-nodeenv==1.8.0
+nodeenv==1.9.1
     # via pre-commit
-packaging==24.0
+packaging==24.1
     # via black
 pathspec==0.12.1
     # via black
-platformdirs==4.2.0
+platformdirs==4.2.2
     # via
     #   black
     #   virtualenv
 pre-commit==3.5.0
     # via -r requirements/linting.in
-pyupgrade==3.15.2
+pyupgrade==3.16.0
     # via -r requirements/linting.in
 pyyaml==6.0.1
     # via
     #   -r requirements/linting.in
     #   pre-commit
-ruff==0.4.1
+ruff==0.4.10
     # via -r requirements/linting.in
 tokenize-rt==5.2.0
     # via pyupgrade
@@ -50,12 +50,9 @@ tomli==2.0.1
     #   mypy
 types-pyyaml==6.0.12.20240311
     # via -r requirements/linting.in
-typing-extensions==4.11.0
+typing-extensions==4.12.2
     # via
     #   black
     #   mypy
-virtualenv==20.25.3
+virtualenv==20.26.2
     # via pre-commit
-
-# The following packages are considered to be unsafe in a requirements file:
-# setuptools

requirements/pyproject.txt

@@ -2,22 +2,74 @@
 # This file is autogenerated by pip-compile with Python 3.8
 # by the following command:
 #
-#    pip-compile --extra=toml --extra=yaml --no-emit-index-url --output-file=requirements/pyproject.txt pyproject.toml
+#    pip-compile --extra=azure-key-vault --extra=toml --extra=yaml --no-emit-index-url --output-file=requirements/pyproject.txt pyproject.toml
 #
-annotated-types==0.6.0
+annotated-types==0.7.0
     # via pydantic
-pydantic==2.7.0
+azure-core==1.30.2
+    # via
+    #   azure-identity
+    #   azure-keyvault-secrets
+azure-identity==1.17.0
+    # via pydantic-settings (pyproject.toml)
+azure-keyvault-secrets==4.8.0
+    # via pydantic-settings (pyproject.toml)
+certifi==2024.6.2
+    # via requests
+cffi==1.16.0
+    # via cryptography
+charset-normalizer==3.3.2
+    # via requests
+cryptography==42.0.8
+    # via
+    #   azure-identity
+    #   msal
+    #   pyjwt
+idna==3.7
+    # via requests
+isodate==0.6.1
+    # via azure-keyvault-secrets
+msal==1.28.1
+    # via
+    #   azure-identity
+    #   msal-extensions
+msal-extensions==1.1.0
+    # via azure-identity
+packaging==24.1
+    # via msal-extensions
+portalocker==2.8.2
+    # via msal-extensions
+pycparser==2.22
+    # via cffi
+pydantic==2.7.4
     # via pydantic-settings (pyproject.toml)
-pydantic-core==2.18.1
+pydantic-core==2.18.4
     # via pydantic
+pyjwt[crypto]==2.8.0
+    # via
+    #   msal
+    #   pyjwt
 python-dotenv==1.0.1
     # via pydantic-settings (pyproject.toml)
 pyyaml==6.0.1
     # via pydantic-settings (pyproject.toml)
+requests==2.32.3
+    # via
+    #   azure-core
+    #   msal
+six==1.16.0
+    # via
+    #   azure-core
+    #   isodate
 tomli==2.0.1
     # via pydantic-settings (pyproject.toml)
-typing-extensions==4.11.0
+typing-extensions==4.12.2
     # via
     #   annotated-types
+    #   azure-core
+    #   azure-identity
+    #   azure-keyvault-secrets
     #   pydantic
     #   pydantic-core
+urllib3==2.2.2
+    # via requests