This is an automated email from the ASF dual-hosted git repository. jedcunningham pushed a commit to branch v2-2-test in repository https://gitbox.apache.org/repos/asf/airflow.git
commit bef01d94d46c8beb2cc0f535158c9bb24bf92a3d Author: raphaelauv <raphael...@users.noreply.github.com> AuthorDate: Sat Nov 6 00:27:36 2021 +0100 Fix bug when checking for existence of a Variable (#19395) `check_for_write_conflict` was a `staticmethod` but for some reason it was ignored (cherry picked from commit 93d2a1626d4da4ae372f5c5edb47a12afc388d33) --- airflow/models/variable.py | 3 ++- tests/secrets/test_local_filesystem.py | 24 +++++++++++++++++-- tests/secrets/test_secrets.py | 44 ++++++++++++++++++++++++++++++++-- 3 files changed, 66 insertions(+), 5 deletions(-) diff --git a/airflow/models/variable.py b/airflow/models/variable.py index b724686..6da65c5 100644 --- a/airflow/models/variable.py +++ b/airflow/models/variable.py @@ -197,7 +197,7 @@ class Variable(Base, LoggingMixin): """ cls.check_for_write_conflict(key) - if cls.get_variable_from_secrets(key) is None: + if cls.get_variable_from_secrets(key=key) is None: raise KeyError(f'Variable {key} does not exist') obj = session.query(cls).filter(cls.key == key).first() @@ -223,6 +223,7 @@ class Variable(Base, LoggingMixin): if self._val and self.is_encrypted: self._val = fernet.rotate(self._val.encode('utf-8')).decode() + @staticmethod def check_for_write_conflict(key: str) -> None: """ Logs a warning if a variable exists outside of the metastore. diff --git a/tests/secrets/test_local_filesystem.py b/tests/secrets/test_local_filesystem.py index 1c6c5ae..85f0aaa 100644 --- a/tests/secrets/test_local_filesystem.py +++ b/tests/secrets/test_local_filesystem.py @@ -25,9 +25,12 @@ from unittest import mock import pytest from parameterized import parameterized +from airflow.configuration import ensure_secrets_loaded from airflow.exceptions import AirflowException, AirflowFileParseException, ConnectionNotUnique +from airflow.models import Variable from airflow.secrets import local_filesystem from airflow.secrets.local_filesystem import LocalFilesystemBackend +from tests.test_utils.config import conf_vars @contextmanager @@ -380,15 +383,32 @@ class TestLocalFileBackend(unittest.TestCase): assert "VAL_A" == backend.get_variable("KEY_A") assert backend.get_variable("KEY_B") is None + @conf_vars( + { + ( + "secrets", + "backend", + ): "airflow.secrets.local_filesystem.LocalFilesystemBackend", + ("secrets", "backend_kwargs"): '{"variables_file_path": "var.env"}', + } + ) + def test_load_secret_backend_LocalFilesystemBackend(self): + with mock_local_file("KEY_A=VAL_A"): + backends = ensure_secrets_loaded() + + backend_classes = [backend.__class__.__name__ for backend in backends] + assert 'LocalFilesystemBackend' in backend_classes + assert Variable.get("KEY_A") == "VAL_A" + def test_should_read_connection(self): with NamedTemporaryFile(suffix=".env") as tmp_file: tmp_file.write(b"CONN_A=mysql://host_a") tmp_file.flush() backend = LocalFilesystemBackend(connections_file_path=tmp_file.name) - assert ["mysql://host_a"] == [conn.get_uri() for conn in backend.get_connections("CONN_A")] + assert "mysql://host_a" == backend.get_connection("CONN_A").get_uri() assert backend.get_variable("CONN_B") is None def test_files_are_optional(self): backend = LocalFilesystemBackend() - assert [] == backend.get_connections("CONN_A") + assert None is backend.get_connection("CONN_A") assert backend.get_variable("VAR_A") is None diff --git a/tests/secrets/test_secrets.py b/tests/secrets/test_secrets.py index 31baf38..83a7fe3 100644 --- a/tests/secrets/test_secrets.py +++ b/tests/secrets/test_secrets.py @@ -137,10 +137,10 @@ class TestVariableFromSecrets(unittest.TestCase): Test if Variable is present in Environment Variable, it does not look for it in Metastore DB """ - mock_env_get.return_value = [["something"]] # returns nonempty list + mock_env_get.return_value = "something" Variable.get_variable_from_secrets("fake_var_key") mock_env_get.assert_called_once_with(key="fake_var_key") - mock_meta_get.not_called() + mock_meta_get.assert_not_called() def test_backend_fallback_to_default_var(self): """ @@ -149,3 +149,43 @@ class TestVariableFromSecrets(unittest.TestCase): """ variable_value = Variable.get(key="test_var", default_var="new") assert "new" == variable_value + + @conf_vars( + { + ( + "secrets", + "backend", + ): "airflow.providers.amazon.aws.secrets.systems_manager.SystemsManagerParameterStoreBackend", + ("secrets", "backend_kwargs"): '{"variables_prefix": "/airflow", "profile_name": null}', + } + ) + @mock.patch.dict( + 'os.environ', + { + 'AIRFLOW_VAR_MYVAR': 'a_venv_value', + }, + ) + @mock.patch("airflow.secrets.metastore.MetastoreBackend.get_variable") + @mock.patch( + "airflow.providers.amazon.aws.secrets.systems_manager." + "SystemsManagerParameterStoreBackend.get_variable" + ) + def test_backend_variable_order(self, mock_secret_get, mock_meta_get): + backends = ensure_secrets_loaded() + backend_classes = [backend.__class__.__name__ for backend in backends] + assert 'SystemsManagerParameterStoreBackend' in backend_classes + + mock_secret_get.return_value = None + mock_meta_get.return_value = None + + assert "a_venv_value" == Variable.get(key="MYVAR") + mock_secret_get.assert_called_with(key="MYVAR") + mock_meta_get.assert_not_called() + + mock_secret_get.return_value = None + mock_meta_get.return_value = "a_metastore_value" + assert "a_metastore_value" == Variable.get(key="not_myvar") + mock_meta_get.assert_called_once_with(key="not_myvar") + + mock_secret_get.return_value = "a_secret_value" + assert "a_secret_value" == Variable.get(key="not_myvar")