From cade49845459af57f2ba32dfe6a6ab144bbb8f98 Mon Sep 17 00:00:00 2001 From: pdebelak Date: Fri, 5 Aug 2022 11:14:30 -0500 Subject: [PATCH 1/3] Cache the custom secrets backend so the same instance gets re-used Fixes #25555 This uses `functools.lru_cache` to re-use the same secrets backend instance between the `conf` global when it loads configuration from secrets and uses outside the `configuration` module like variables and connections. Previously, each fetch of a configuration value from secrets would use its own secrets backend instance. --- airflow/configuration.py | 5 +++++ tests/core/test_configuration.py | 2 ++ 2 files changed, 7 insertions(+) diff --git a/airflow/configuration.py b/airflow/configuration.py index 139172d0f3c7..978aeec38d38 100644 --- a/airflow/configuration.py +++ b/airflow/configuration.py @@ -1535,6 +1535,11 @@ def get_custom_secret_backend() -> Optional[BaseSecretsBackend]: """Get Secret Backend if defined in airflow.cfg""" secrets_backend_cls = conf.getimport(section='secrets', key='backend') + return _custom_secrets_backend(secrets_backend_cls) + + +@functools.lru_cache(maxsize=2) +def _custom_secrets_backend(secrets_backend_cls) -> Optional[BaseSecretsBackend]: if secrets_backend_cls: try: backends: Any = conf.get(section='secrets', key='backend_kwargs', fallback='{}') diff --git a/tests/core/test_configuration.py b/tests/core/test_configuration.py index 15c8bbffc149..47914eb6e05d 100644 --- a/tests/core/test_configuration.py +++ b/tests/core/test_configuration.py @@ -33,6 +33,7 @@ from airflow.configuration import ( AirflowConfigException, AirflowConfigParser, + _custom_secrets_backend, conf, expand_env_var, get_airflow_config, @@ -296,6 +297,7 @@ def test_hidding_of_sensitive_config_values(self): def test_config_raise_exception_from_secret_backend_connection_error(self, mock_hvac): """Get Config Value from a Secret Backend""" + _custom_secrets_backend.cache_clear() # so secrets backend gets reloaded mock_client = mock.MagicMock() # mock_client.side_effect = AirflowConfigException mock_hvac.Client.return_value = mock_client From 9097b53efd4dbf968ba568d14f7b38258cf9e2fb Mon Sep 17 00:00:00 2001 From: pdebelak Date: Sat, 6 Aug 2022 08:05:51 -0500 Subject: [PATCH 2/3] Add secrets backend kwargs to cache key Also add unit test to confirm that only one secrets backend instance gets created. --- airflow/configuration.py | 14 +++--- tests/core/test_configuration.py | 81 +++++++++++++++++++++++++++++++- 2 files changed, 87 insertions(+), 8 deletions(-) diff --git a/airflow/configuration.py b/airflow/configuration.py index 978aeec38d38..7e1974a7b817 100644 --- a/airflow/configuration.py +++ b/airflow/configuration.py @@ -1534,12 +1534,6 @@ def ensure_secrets_loaded() -> List[BaseSecretsBackend]: def get_custom_secret_backend() -> Optional[BaseSecretsBackend]: """Get Secret Backend if defined in airflow.cfg""" secrets_backend_cls = conf.getimport(section='secrets', key='backend') - - return _custom_secrets_backend(secrets_backend_cls) - - -@functools.lru_cache(maxsize=2) -def _custom_secrets_backend(secrets_backend_cls) -> Optional[BaseSecretsBackend]: if secrets_backend_cls: try: backends: Any = conf.get(section='secrets', key='backend_kwargs', fallback='{}') @@ -1547,10 +1541,16 @@ def _custom_secrets_backend(secrets_backend_cls) -> Optional[BaseSecretsBackend] except JSONDecodeError: alternative_secrets_config_dict = {} - return secrets_backend_cls(**alternative_secrets_config_dict) + return _custom_secrets_backend(secrets_backend_cls, **alternative_secrets_config_dict) return None +@functools.lru_cache +def _custom_secrets_backend(secrets_backend_cls, **alternative_secrets_config_dict): + """Separate function to create secrets backend instance to allow caching""" + return secrets_backend_cls(**alternative_secrets_config_dict) + + def initialize_secrets_backends() -> List[BaseSecretsBackend]: """ * import secrets backend classes diff --git a/tests/core/test_configuration.py b/tests/core/test_configuration.py index 47914eb6e05d..175978865496 100644 --- a/tests/core/test_configuration.py +++ b/tests/core/test_configuration.py @@ -297,7 +297,7 @@ def test_hidding_of_sensitive_config_values(self): def test_config_raise_exception_from_secret_backend_connection_error(self, mock_hvac): """Get Config Value from a Secret Backend""" - _custom_secrets_backend.cache_clear() # so secrets backend gets reloaded + _custom_secrets_backend.cache_clear() mock_client = mock.MagicMock() # mock_client.side_effect = AirflowConfigException mock_hvac.Client.return_value = mock_client @@ -324,6 +324,7 @@ def test_config_raise_exception_from_secret_backend_connection_error(self, mock_ ), ): test_conf.get('test', 'sql_alchemy_conn') + _custom_secrets_backend.cache_clear() def test_getboolean(self): """Test AirflowConfigParser.getboolean""" @@ -1299,3 +1300,81 @@ def test_conf_as_dict_when_deprecated_value_in_secrets_disabled_config( conf.read_dict(dictionary=cfg_dict) os.environ.clear() assert conf.get('database', 'sql_alchemy_conn') == f'sqlite:///{HOME_DIR}/airflow/airflow.db' + + @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac") + @conf_vars( + { + ("secrets", "backend"): "airflow.providers.hashicorp.secrets.vault.VaultBackend", + ("secrets", "backend_kwargs"): '{"url": "http://127.0.0.1:8200", "token": "token"}', + } + ) + def test_config_from_secret_backend_caches_instance(self, mock_hvac): + """Get Config Value from a Secret Backend""" + _custom_secrets_backend.cache_clear() + + test_config = '''[test] +sql_alchemy_conn_secret = sql_alchemy_conn +secret_key_secret = secret_key +''' + test_config_default = '''[test] +sql_alchemy_conn = airflow +secret_key = airflow +''' + + mock_client = mock.MagicMock() + mock_hvac.Client.return_value = mock_client + + def fake_read_secret(path, mount_point, version): + if path.endswith('sql_alchemy_conn'): + return { + 'request_id': '2d48a2ad-6bcb-e5b6-429d-da35fdf31f56', + 'lease_id': '', + 'renewable': False, + 'lease_duration': 0, + 'data': { + 'data': {'value': 'fake_conn'}, + 'metadata': { + 'created_time': '2020-03-28T02:10:54.301784Z', + 'deletion_time': '', + 'destroyed': False, + 'version': 1, + }, + }, + 'wrap_info': None, + 'warnings': None, + 'auth': None, + } + if path.endswith('secret_key'): + return { + 'request_id': '2d48a2ad-6bcb-e5b6-429d-da35fdf31f56', + 'lease_id': '', + 'renewable': False, + 'lease_duration': 0, + 'data': { + 'data': {'value': 'fake_key'}, + 'metadata': { + 'created_time': '2020-03-28T02:10:54.301784Z', + 'deletion_time': '', + 'destroyed': False, + 'version': 1, + }, + }, + 'wrap_info': None, + 'warnings': None, + 'auth': None, + } + + mock_client.secrets.kv.v2.read_secret_version.side_effect = fake_read_secret + + test_conf = AirflowConfigParser(default_config=parameterized_config(test_config_default)) + test_conf.read_string(test_config) + test_conf.sensitive_config_values = test_conf.sensitive_config_values | { + ('test', 'sql_alchemy_conn'), + ('test', 'secret_key'), + } + + assert 'fake_conn' == test_conf.get('test', 'sql_alchemy_conn') + mock_hvac.Client.assert_called_once() + assert 'fake_key' == test_conf.get('test', 'secret_key') + mock_hvac.Client.assert_called_once() + _custom_secrets_backend.cache_clear() From db6b8ab44ebe0ffee216ab55d119979c092960bc Mon Sep 17 00:00:00 2001 From: pdebelak Date: Sat, 6 Aug 2022 08:27:13 -0500 Subject: [PATCH 3/3] Set maxsize for cache on _custom_secrets_backend --- airflow/configuration.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/configuration.py b/airflow/configuration.py index 7e1974a7b817..5e84de6ea4b4 100644 --- a/airflow/configuration.py +++ b/airflow/configuration.py @@ -1545,7 +1545,7 @@ def get_custom_secret_backend() -> Optional[BaseSecretsBackend]: return None -@functools.lru_cache +@functools.lru_cache(maxsize=2) def _custom_secrets_backend(secrets_backend_cls, **alternative_secrets_config_dict): """Separate function to create secrets backend instance to allow caching""" return secrets_backend_cls(**alternative_secrets_config_dict)