diff --git a/airflow/config_templates/config.yml b/airflow/config_templates/config.yml index b7eb88d62199a..6d6bba2493166 100644 --- a/airflow/config_templates/config.yml +++ b/airflow/config_templates/config.yml @@ -782,6 +782,17 @@ type: string example: ~ default: "" + - name: backends_config + description: | + Advanced secrets backend configuration, allow to user configure more than one secret backend, + order of secrets backend and turn off built-in backends. + Expected JSON list of objects. See + https://airflow.apache.org/docs/apache-airflow/stable/security/secrets/secrets-backend/index.html#advanced-configuration + for more details. + version_added: 2.4.0 + example: ~ + type: string + default: "" - name: cli description: ~ options: diff --git a/airflow/config_templates/default_airflow.cfg b/airflow/config_templates/default_airflow.cfg index f47d0a82a2e6b..44f1c606a9569 100644 --- a/airflow/config_templates/default_airflow.cfg +++ b/airflow/config_templates/default_airflow.cfg @@ -422,6 +422,13 @@ backend = # ``{{"connections_prefix": "/airflow/connections", "profile_name": "default"}}`` backend_kwargs = +# Advanced secrets backend configuration, allow to user configure more than one secret backend, +# order of secrets backend and turn off built-in backends. +# Expected JSON list of objects. See +# https://airflow.apache.org/docs/apache-airflow/stable/security/secrets/secrets-backend/index.html#advanced-configuration +# for more details. +backends_config = + [cli] # In what way should the cli access the API. The LocalClient will use the # database directly, while the json_client will use the api running on the diff --git a/airflow/configuration.py b/airflow/configuration.py index 93612160c8290..d73dd123e3c2c 100644 --- a/airflow/configuration.py +++ b/airflow/configuration.py @@ -27,20 +27,21 @@ import sys import warnings from base64 import b64encode -from collections import OrderedDict +from collections import OrderedDict, UserList # Ignored Mypy on configparser because it thinks the configparser module has no _UNSET attribute from configparser import _UNSET, ConfigParser, NoOptionError, NoSectionError # type: ignore from contextlib import suppress +from dataclasses import dataclass, field from json.decoder import JSONDecodeError from re import Pattern -from typing import IO, Any, Dict, Iterable, List, Optional, Set, Tuple, Union +from typing import IO, Any, Dict, Iterable, List, Optional, Sequence, Set, Tuple, Union from urllib.parse import urlparse from typing_extensions import overload from airflow.exceptions import AirflowConfigException -from airflow.secrets import DEFAULT_SECRETS_SEARCH_PATH, BaseSecretsBackend +from airflow.secrets import BaseSecretsBackend from airflow.utils import yaml from airflow.utils.module_loading import import_string from airflow.utils.weight_rule import WeightRule @@ -62,6 +63,16 @@ ENV_VAR_PREFIX = 'AIRFLOW__' +DEFAULT_SECRETS_SEARCH_PATH: List[Dict[str, Any]] = [ + { + "backend": "airflow.secrets.environment_variables.EnvironmentVariablesBackend", + }, + { + "backend": "airflow.secrets.metastore.MetastoreBackend", + }, +] + + def _parse_sqlite_version(s: str) -> Tuple[int, ...]: match = _SQLITE3_VERSION_PATTERN.match(s) if match is None: @@ -113,18 +124,20 @@ def run_command(command: str) -> str: def _get_config_value_from_secret_backend(config_key: str) -> Optional[str]: """Get Config option values from Secret Backend""" - try: - secrets_client = get_custom_secret_backend() - if not secrets_client: - return None - return secrets_client.get_config(config_key) - except Exception as e: - raise AirflowConfigException( - 'Cannot retrieve config from alternative secrets backend. ' - 'Make sure it is configured properly and that the Backend ' - 'is accessible.\n' - f'{e}' - ) + config_value = None + for secrets_client in ensure_secrets_loaded(): + try: + config_value = secrets_client.get_config(config_key) + if config_value: + break + except Exception as e: + raise AirflowConfigException( + f'Cannot retrieve config from secrets backend `{secrets_client.__class__.__name__}`. ' + f'Make sure it is configured properly and that the Backend is accessible.\n' + f'{e}' + ) + + return config_value def _default_config_file_path(file_name: str) -> str: @@ -1487,47 +1500,118 @@ def set(*args, **kwargs) -> None: conf.set(*args, **kwargs) -def ensure_secrets_loaded() -> List[BaseSecretsBackend]: - """ - Ensure that all secrets backends are loaded. - If the secrets_backend_list contains only 2 default backends, reload it. - """ - # Check if the secrets_backend_list contains only 2 default backends - if len(secrets_backend_list) == 2: - return initialize_secrets_backends() - return secrets_backend_list +class DefaultSecretsBackend(UserList): + """List Container which use for store default secrets backends.""" + + +@dataclass(frozen=True) +class SecretsBackendConfig: + """Secrets Backend Config dataclass helper.""" + + backend: str + backend_kwargs: Dict[str, Any] = field(default_factory=dict) + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> 'SecretsBackendConfig': + """ + Read Secret Backend Config from dictionary + + Ignores all unexpected keywords + """ + return cls(d['backend'], d.get('backend_kwargs', {})) + @classmethod + def from_config(cls) -> Optional['SecretsBackendConfig']: + """Try to get ``SecretsBackendConfig`` from airflow config [secrets] section""" + secrets_backend = conf.get(section='secrets', key='backend', fallback=None) -def get_custom_secret_backend() -> Optional[BaseSecretsBackend]: - """Get Secret Backend if defined in airflow.cfg""" - secrets_backend_cls = conf.getimport(section='secrets', key='backend') + if not secrets_backend: + return None - if secrets_backend_cls: try: - backends: Any = conf.get(section='secrets', key='backend_kwargs', fallback='{}') - alternative_secrets_config_dict = json.loads(backends) + secrets_backend_kwargs: Any = conf.get(section='secrets', key='backend_kwargs', fallback='{}') + secrets_backend_kwargs = json.loads(secrets_backend_kwargs) except JSONDecodeError: - alternative_secrets_config_dict = {} + secrets_backend_kwargs = {} + + return cls(secrets_backend, secrets_backend_kwargs) + + def initialize(self) -> BaseSecretsBackend: + """Initialize Secrets Backend.""" + return import_string(self.backend)(**self.backend_kwargs) + + +class UniqueSecretsBackendsConfigs(UserList): + """List Container which use for store unique secrets backends configs.""" - return secrets_backend_cls(**alternative_secrets_config_dict) - return None + def append(self, config) -> None: + """Append item to a list if it not exists yet""" + if config in self.data: + log.warning("%r already exists.", config) + return + self.data.append(config) + + def extend(self, configs) -> None: + """Extends item if it not exists yet""" + for config in configs: + self.append(config) + + +def ensure_secrets_loaded() -> Sequence[BaseSecretsBackend]: + """Ensure that all secrets backends are loaded.""" + if isinstance(secrets_backend_list, DefaultSecretsBackend): + return initialize_secrets_backends() + return secrets_backend_list -def initialize_secrets_backends() -> List[BaseSecretsBackend]: +def initialize_secrets_backends() -> Sequence[BaseSecretsBackend]: """ + * read secrets backend configurations + * keep only unique secret backend configurations * import secrets backend classes * instantiate them and return them in a list """ - backend_list = [] + backend_configs = UniqueSecretsBackendsConfigs() + default_configs = [SecretsBackendConfig.from_dict(config) for config in DEFAULT_SECRETS_SEARCH_PATH] + + secrets_backend_config = conf.getjson(section='secrets', key='backends_config', fallback=None) + if secrets_backend_config: + # If [secrets] 'backend_config' exists and defined than use this configuration + # instead of [secrets] 'backend' and [secrets] 'backend_kwargs'. + if not isinstance(secrets_backend_config, list): + raise AirflowConfigException( + f"[secrets] 'backends_config' expected list of backends configurations but got" + f" {secrets_backend_config!r}." + ) + + for config in secrets_backend_config: + try: + backend_configs.append(SecretsBackendConfig.from_dict(config)) + except Exception as e: + raise AirflowConfigException( + f"Cannot read config: {config!r} from [secrets] 'backends_config'.\n{e}" + ) from e - custom_secret_backend = get_custom_secret_backend() + else: + custom_secrets_backend_config = SecretsBackendConfig.from_config() + if not custom_secrets_backend_config: + # Returns default secrets backend list for further checks in `ensure_secrets_loaded()`. + return DefaultSecretsBackend(config.initialize() for config in default_configs) - if custom_secret_backend is not None: - backend_list.append(custom_secret_backend) + backend_configs.append(custom_secrets_backend_config) + backend_configs.extend(default_configs) - for class_name in DEFAULT_SECRETS_SEARCH_PATH: - secrets_backend_cls = import_string(class_name) - backend_list.append(secrets_backend_cls()) + # Initialize secrets backends + backend_list = [] + for config in backend_configs: + try: + backend_list.append(config.initialize()) + except Exception as e: + raise AirflowConfigException( + f"Cannot initialize secrets backend {config.backend!r} " + f"with keyword arguments {config.backend_kwargs}.\n" + f"{e}" + ) from e return backend_list diff --git a/airflow/secrets/__init__.py b/airflow/secrets/__init__.py index 5b12c8a300137..ab6242c59b558 100644 --- a/airflow/secrets/__init__.py +++ b/airflow/secrets/__init__.py @@ -22,11 +22,6 @@ * Metastore database * AWS SSM Parameter store """ -__all__ = ['BaseSecretsBackend', 'DEFAULT_SECRETS_SEARCH_PATH'] +__all__ = ['BaseSecretsBackend'] from airflow.secrets.base_secrets import BaseSecretsBackend - -DEFAULT_SECRETS_SEARCH_PATH = [ - "airflow.secrets.environment_variables.EnvironmentVariablesBackend", - "airflow.secrets.metastore.MetastoreBackend", -] diff --git a/docs/apache-airflow/security/secrets/secrets-backend/index.rst b/docs/apache-airflow/security/secrets/secrets-backend/index.rst index fa92e1d61e4b7..0aec99c364592 100644 --- a/docs/apache-airflow/security/secrets/secrets-backend/index.rst +++ b/docs/apache-airflow/security/secrets/secrets-backend/index.rst @@ -38,8 +38,8 @@ Search path When looking up a connection/variable, by default Airflow will search environment variables first and metastore database second. -If you enable an alternative secrets backend, it will be searched first, followed by environment variables, -then metastore. This search ordering is not configurable. +If you enable an alternative secrets backend, it will be searched first, followed by environment +variables, then metastore. This search ordering is not configurable in Airflow <2.4.0. .. warning:: @@ -59,6 +59,7 @@ The ``[secrets]`` section has the following options: [secrets] backend = backend_kwargs = + backends_config = Set ``backend`` to the fully qualified class name of the backend you want to enable. @@ -73,6 +74,47 @@ the example below. $ airflow config get-value secrets backend airflow.providers.google.cloud.secrets.secret_manager.CloudSecretManagerBackend +Advanced Configuration +"""""""""""""""""""""" + +.. versionadded:: 2.4.0 + +You can provide ``backends_config`` with json-list which contains json-objects of required backends for: + +* Configure more than one alternative secrets backend. +* Change search ordering of secrets backends. +* Turn off built-in backends (environment variables or the metastore database). + +Object should contains mandatory **backend** attribute with fully qualified class name of the backend you want to enable +and optionally contains **backend_kwargs** with json-object which will be passed as kwargs to the ``__init__`` +method of your secrets backend. + +Here's a basic example of how to configure to use Environment Variables first, +next :ref:`SSM Parameter Store `, and then metastore.: + +.. code-block:: json + + [ + { + "backend": "airflow.secrets.environment_variables.EnvironmentVariablesBackend" + }, + { + "backend": "airflow.providers.amazon.aws.secrets.systems_manager.SystemsManagerParameterStoreBackend", + "backend_kwargs": { + "connections_prefix": "/airflow/connections", + "profile_name": "default" + } + }, + { + "backend": "airflow.secrets.metastore.MetastoreBackend" + } + ] + +.. warning:: + + If you specify ``backends_config`` options than Airflow will ignore values + in ``backend`` and ``backend_kwargs`` options. + Supported core backends ^^^^^^^^^^^^^^^^^^^^^^^ @@ -101,8 +143,8 @@ Roll your own secrets backend A secrets backend is a subclass of :py:class:`airflow.secrets.BaseSecretsBackend` and must implement either :py:meth:`~airflow.secrets.BaseSecretsBackend.get_connection` or :py:meth:`~airflow.secrets.BaseSecretsBackend.get_conn_value`. -After writing your backend class, provide the fully qualified class name in the ``backend`` key in the ``[secrets]`` -section of ``airflow.cfg``. +After writing your backend class, provide the fully qualified class name in the ``backend`` key or provide +configurations in the ``backends_config`` in the ``[secrets]`` section of ``airflow.cfg``. Additional arguments to your SecretsBackend can be configured in ``airflow.cfg`` by supplying a JSON string to ``backend_kwargs``, which will be passed to the ``__init__`` of your SecretsBackend. See :ref:`Configuration ` for more details, and :ref:`SSM Parameter Store ` for an example. diff --git a/tests/core/test_configuration.py b/tests/core/test_configuration.py index 991695a969d08..90cf93a70dc05 100644 --- a/tests/core/test_configuration.py +++ b/tests/core/test_configuration.py @@ -290,10 +290,10 @@ def test_config_raise_exception_from_secret_backend_connection_error(self, mock_ with pytest.raises( AirflowConfigException, - match=re.escape( - 'Cannot retrieve config from alternative secrets backend. ' - 'Make sure it is configured properly and that the Backend ' - 'is accessible.' + match=( + r'Cannot retrieve config from secrets backend `.*`. ' + r'Make sure it is configured properly and that the Backend ' + r'is accessible.' ), ): test_conf.get('test', 'sql_alchemy_conn') @@ -1184,11 +1184,9 @@ def test_conf_as_dict_when_deprecated_value_in_cmd_disabled_config(self, display @pytest.mark.parametrize("display_source", [True, False]) @mock.patch.dict('os.environ', {"AIRFLOW__CORE__SQL_ALCHEMY_CONN_SECRET": "secret_path'"}, clear=True) - @mock.patch("airflow.configuration.get_custom_secret_backend") - def test_conf_as_dict_when_deprecated_value_in_secrets( - self, get_custom_secret_backend, display_source: bool - ): - get_custom_secret_backend.return_value.get_config.return_value = "postgresql://" + @mock.patch("airflow.secrets.BaseSecretsBackend.get_config") + def test_conf_as_dict_when_deprecated_value_in_secrets(self, mock_get_config, display_source: bool): + mock_get_config.return_value = "postgresql://" with use_config(config="empty.cfg"): cfg_dict = conf.as_dict( display_source=display_source, @@ -1210,11 +1208,11 @@ def test_conf_as_dict_when_deprecated_value_in_secrets( @pytest.mark.parametrize("display_source", [True, False]) @mock.patch.dict('os.environ', {"AIRFLOW__CORE__SQL_ALCHEMY_CONN_SECRET": "secret_path'"}, clear=True) - @mock.patch("airflow.configuration.get_custom_secret_backend") + @mock.patch("airflow.secrets.BaseSecretsBackend.get_config") def test_conf_as_dict_when_deprecated_value_in_secrets_disabled_env( - self, get_custom_secret_backend, display_source: bool + self, mock_get_config, display_source: bool ): - get_custom_secret_backend.return_value.get_config.return_value = "postgresql://" + mock_get_config.return_value = "postgresql://" with use_config(config="empty.cfg"): cfg_dict = conf.as_dict( display_source=display_source, @@ -1236,12 +1234,12 @@ def test_conf_as_dict_when_deprecated_value_in_secrets_disabled_env( assert conf.get('database', 'sql_alchemy_conn') == f'sqlite:///{HOME_DIR}/airflow/airflow.db' @pytest.mark.parametrize("display_source", [True, False]) - @mock.patch("airflow.configuration.get_custom_secret_backend") + @mock.patch("airflow.secrets.BaseSecretsBackend.get_config") @mock.patch.dict('os.environ', {}, clear=True) def test_conf_as_dict_when_deprecated_value_in_secrets_disabled_config( - self, get_custom_secret_backend, display_source: bool + self, mock_get_config, display_source: bool ): - get_custom_secret_backend.return_value.get_config.return_value = "postgresql://" + mock_get_config.return_value = "postgresql://" with use_config(config="deprecated_secret.cfg"): cfg_dict = conf.as_dict( display_source=display_source, diff --git a/tests/secrets/test_secrets.py b/tests/secrets/test_secrets.py index f8d8a6717cd63..8276dea8de32e 100644 --- a/tests/secrets/test_secrets.py +++ b/tests/secrets/test_secrets.py @@ -16,14 +16,23 @@ # specific language governing permissions and limitations # under the License. # +import json import unittest from unittest import mock +import pytest + from airflow.configuration import ensure_secrets_loaded, initialize_secrets_backends +from airflow.exceptions import AirflowConfigException, AirflowNotFoundException from airflow.models import Connection, Variable from tests.test_utils.config import conf_vars from tests.test_utils.db import clear_db_variables +SSM_SB = "airflow.providers.amazon.aws.secrets.systems_manager.SystemsManagerParameterStoreBackend" +LOCAL_SB = "airflow.secrets.local_filesystem.LocalFilesystemBackend" +ENV_SB = "airflow.secrets.environment_variables.EnvironmentVariablesBackend" +METASTORE_DB_SB = "airflow.secrets.metastore.MetastoreBackend" + class TestConnectionsFromSecrets(unittest.TestCase): @mock.patch("airflow.secrets.metastore.MetastoreBackend.get_connection") @@ -44,10 +53,7 @@ def test_get_connection_first_try(self, mock_env_get, mock_meta_get): @conf_vars( { - ( - "secrets", - "backend", - ): "airflow.providers.amazon.aws.secrets.systems_manager.SystemsManagerParameterStoreBackend", + ("secrets", "backend"): SSM_SB, ("secrets", "backend_kwargs"): '{"connections_prefix": "/airflow", "profile_name": null}', } ) @@ -60,10 +66,7 @@ def test_initialize_secrets_backends(self): @conf_vars( { - ( - "secrets", - "backend", - ): "airflow.providers.amazon.aws.secrets.systems_manager.SystemsManagerParameterStoreBackend", + ("secrets", "backend"): SSM_SB, ("secrets", "backend_kwargs"): '{"use_ssl": false}', } ) @@ -79,10 +82,7 @@ def test_backends_kwargs(self): @conf_vars( { - ( - "secrets", - "backend", - ): "airflow.providers.amazon.aws.secrets.systems_manager.SystemsManagerParameterStoreBackend", + ("secrets", "backend"): SSM_SB, ("secrets", "backend_kwargs"): '{"connections_prefix": "/airflow", "profile_name": null}', } ) @@ -152,10 +152,7 @@ def test_backend_fallback_to_default_var(self): @conf_vars( { - ( - "secrets", - "backend", - ): "airflow.providers.amazon.aws.secrets.systems_manager.SystemsManagerParameterStoreBackend", + ("secrets", "backend"): SSM_SB, ("secrets", "backend_kwargs"): '{"variables_prefix": "/airflow", "profile_name": null}', } ) @@ -189,3 +186,217 @@ def test_backend_variable_order(self, mock_secret_get, mock_meta_get): mock_secret_get.return_value = "a_secret_value" assert "a_secret_value" == Variable.get(key="not_myvar") + + +class TestSecretsBackendsConfig(unittest.TestCase): + def setUp(self) -> None: + clear_db_variables() + + def tearDown(self) -> None: + clear_db_variables() + + def test_default_secrets_backend(self): + """ + Test expected default order of backends which compatible with previous versions of Apache Airflow + """ + backends = ensure_secrets_loaded() + backend_classes = [backend.__class__.__name__ for backend in backends] + + assert backend_classes == ["EnvironmentVariablesBackend", "MetastoreBackend"] + + @conf_vars({("secrets", "backends_config"): "[]"}) + def test_empty_secrets_backend_config(self): + """Test fallback if [secrets] 'backend_config' is empty list""" + backends = ensure_secrets_loaded() + backend_classes = [backend.__class__.__name__ for backend in backends] + + assert backend_classes == ["EnvironmentVariablesBackend", "MetastoreBackend"] + + def test_custom_backend_order(self): + secrets_backend_config = [ + {"backend": LOCAL_SB}, + {"backend": SSM_SB}, + {"backend": ENV_SB}, + {"backend": METASTORE_DB_SB}, + ] + + with conf_vars({("secrets", "backends_config"): json.dumps(secrets_backend_config)}): + backends = ensure_secrets_loaded() + backend_classes = [backend.__class__.__name__ for backend in backends] + + assert backend_classes == [ + "LocalFilesystemBackend", + "SystemsManagerParameterStoreBackend", + "EnvironmentVariablesBackend", + "MetastoreBackend", + ] + + def test_multiple_same_backends_classes(self): + secrets_backend_config = [ + {"backend": SSM_SB, "backend_kwargs": {"use_ssl": False}}, + {"backend": ENV_SB}, + {"backend": METASTORE_DB_SB}, + {"backend": SSM_SB, "backend_kwargs": {"use_ssl": True}}, + ] + + with conf_vars({("secrets", "backends_config"): json.dumps(secrets_backend_config)}): + backends = ensure_secrets_loaded() + backend_classes = [backend.__class__.__name__ for backend in backends] + + assert backend_classes == [ + "SystemsManagerParameterStoreBackend", + "EnvironmentVariablesBackend", + "MetastoreBackend", + "SystemsManagerParameterStoreBackend", + ] + + assert backends[0] is not backends[-1] + assert backends[0].kwargs == {'use_ssl': False} + assert backends[-1].kwargs == {'use_ssl': True} + + def test_multiple_duplicates_backends_classes(self): + secrets_backend_config = [ + {"backend": ENV_SB}, + {"backend": ENV_SB}, + {"backend": METASTORE_DB_SB}, + {"backend": METASTORE_DB_SB}, + {"backend": ENV_SB}, + {"backend": METASTORE_DB_SB}, + ] + + with conf_vars({("secrets", "backends_config"): json.dumps(secrets_backend_config)}): + with self.assertLogs(level='WARNING') as capture_logs: + backends = ensure_secrets_loaded() + assert len(capture_logs.records) == 4 + + backend_classes = [backend.__class__.__name__ for backend in backends] + assert backend_classes == [ + "EnvironmentVariablesBackend", + "MetastoreBackend", + ] + + @conf_vars({("secrets", "backends_config"): '"INVALID"'}) + def test_invalid_backend_config_type(self): + error_match = r"\[secrets\] 'backends_config' expected list of backends.*" + with pytest.raises(AirflowConfigException, match=error_match): + ensure_secrets_loaded() + + def test_invalid_backend_config(self): + secrets_backend_config = [ + {"backend": LOCAL_SB, "backend_kwargs": {"use_ssl": False}}, + {"backend": ENV_SB}, + {"backend": METASTORE_DB_SB}, + ] + + with conf_vars({("secrets", "backends_config"): json.dumps(secrets_backend_config)}): + error_match = r"Cannot initialize secrets backend .* with keyword arguments .*" + with pytest.raises(AirflowConfigException, match=error_match): + ensure_secrets_loaded() + + @mock.patch("airflow.secrets.local_filesystem.LocalFilesystemBackend.get_variable") + @mock.patch("airflow.secrets.metastore.MetastoreBackend.get_variable") + @mock.patch("airflow.secrets.environment_variables.EnvironmentVariablesBackend.get_variable") + @mock.patch( + "airflow.providers.amazon.aws.secrets.systems_manager." + "SystemsManagerParameterStoreBackend.get_variable" + ) + def test_backend_config_variable_order(self, mock_ssm_get, mock_env_get, mock_meta_get, mock_local_get): + secrets_backend_config = [ + {"backend": LOCAL_SB}, + {"backend": METASTORE_DB_SB}, + {"backend": SSM_SB}, + {"backend": ENV_SB}, + ] + + with conf_vars({("secrets", "backends_config"): json.dumps(secrets_backend_config)}): + ensure_secrets_loaded() + + # Test Variable exists in first backend. + mock_local_get.return_value = "local" + assert "local" == Variable.get(key="test_var_1") + mock_local_get.assert_called_once_with(key="test_var_1") + mock_meta_get.assert_not_called() + mock_ssm_get.assert_not_called() + mock_env_get.assert_not_called() + + # Test Variable exists in second backend. + mock_local_get.return_value = None + mock_meta_get.return_value = "meta" + assert "meta" == Variable.get(key="test_var_2") + mock_meta_get.assert_called_once_with(key="test_var_2") + mock_ssm_get.assert_not_called() + mock_env_get.assert_not_called() + + # Test Variable exists in third backend. + mock_meta_get.return_value = None + mock_ssm_get.return_value = "ssm" + assert "ssm" == Variable.get(key="test_var_3") + mock_ssm_get.assert_called_once_with(key="test_var_3") + mock_env_get.assert_not_called() + + # Test Variable exists in forth backend. + mock_ssm_get.return_value = None + mock_env_get.return_value = "env" + assert "env" == Variable.get(key="test_var_4") + mock_env_get.assert_called_once_with(key="test_var_4") + + # Test Variable not exists in any backends. + mock_env_get.return_value = None + with pytest.raises(KeyError, match="Variable test_var_5 does not exist"): + Variable.get(key="test_var_5") + + @mock.patch("airflow.secrets.local_filesystem.LocalFilesystemBackend.get_connection") + @mock.patch("airflow.secrets.metastore.MetastoreBackend.get_connection") + @mock.patch("airflow.secrets.environment_variables.EnvironmentVariablesBackend.get_connection") + @mock.patch( + "airflow.providers.amazon.aws.secrets.systems_manager." + "SystemsManagerParameterStoreBackend.get_connection" + ) + def test_backend_config_conn_order(self, mock_ssm_get, mock_env_get, mock_meta_get, mock_local_get): + secrets_backend_config = [ + {"backend": LOCAL_SB}, + {"backend": METASTORE_DB_SB}, + {"backend": SSM_SB}, + {"backend": ENV_SB}, + ] + + with conf_vars({("secrets", "backends_config"): json.dumps(secrets_backend_config)}): + ensure_secrets_loaded() + + # Test Connection exists in first backend. + mock_local_conn = mock.MagicMock() + mock_local_get.return_value = mock_local_conn + assert mock_local_conn == Connection.get_connection_from_secrets("test_conn_1") + mock_local_get.assert_called_once_with(conn_id="test_conn_1") + mock_meta_get.assert_not_called() + mock_ssm_get.assert_not_called() + mock_env_get.assert_not_called() + + # Test Connection exists in second backend. + mock_meta_conn = mock.MagicMock() + mock_local_get.return_value = None + mock_meta_get.return_value = mock_meta_conn + assert mock_meta_conn == Connection.get_connection_from_secrets("test_conn_2") + mock_meta_get.assert_called_once_with(conn_id="test_conn_2") + mock_ssm_get.assert_not_called() + mock_env_get.assert_not_called() + + # Test Connection exists in third backend. + mock_ssm_conn = mock.MagicMock() + mock_meta_get.return_value = None + mock_ssm_get.return_value = mock_ssm_conn + assert mock_ssm_conn == Connection.get_connection_from_secrets("test_conn_3") + mock_ssm_get.assert_called_once_with(conn_id="test_conn_3") + mock_env_get.assert_not_called() + + # Test Connection exists in fourth backend. + mock_env_conn = mock.MagicMock() + mock_ssm_get.return_value = None + mock_env_get.return_value = mock_env_conn + assert mock_env_conn == Connection.get_connection_from_secrets("test_conn_4") + mock_env_get.assert_called_once_with(conn_id="test_conn_4") + + # Test Connection not exists in any backends. + mock_env_get.return_value = None + with pytest.raises(AirflowNotFoundException, match="The conn_id `test_conn_5` isn't defined"): + Connection.get_connection_from_secrets("test_conn_5")