diff --git a/airflow-core/newsfragments/71968.bugfix.rst b/airflow-core/newsfragments/71968.bugfix.rst new file mode 100644 index 0000000000000..fab9a98070e2d --- /dev/null +++ b/airflow-core/newsfragments/71968.bugfix.rst @@ -0,0 +1,7 @@ +``Variable.get``, ``Variable.get_variable_from_secrets`` and ``Variable.setdefault`` now accept an +optional keyword-only ``session``, which is forwarded to the metastore secrets backend so the lookup +reuses the caller's transaction instead of committing it. Callers that read a Variable while holding +an open session (most notably code running inside the scheduler's ``prohibit_commit`` guard) should +pass it; previously the backend opened the same scoped session and committed it, detaching the caller's +objects or failing the lookup outright. ``Variable.update`` now forwards the session it was given to +its own existence check, which was affected by the same problem. diff --git a/airflow-core/src/airflow/models/variable.py b/airflow-core/src/airflow/models/variable.py index 9493cab60aa38..0db7c5102f382 100644 --- a/airflow-core/src/airflow/models/variable.py +++ b/airflow-core/src/airflow/models/variable.py @@ -126,7 +126,7 @@ def val(cls): return synonym("_val", descriptor=property(cls.get_val, cls.set_val)) @classmethod - def setdefault(cls, key, default, description=None, deserialize_json=False): + def setdefault(cls, key, default, description=None, deserialize_json=False, *, session=None): """ Return the current value for a key or store the default value and return it. @@ -138,13 +138,20 @@ def setdefault(cls, key, default, description=None, deserialize_json=False): :param description: Default value to set Description of the Variable :param deserialize_json: Store this as a JSON encoded value in the DB and un-encode it when retrieving a value - :param session: Session + :param session: Existing session to reuse for the metadata database read and write. + Callers holding an open transaction must pass it. :return: Mixed """ - obj = Variable.get(key, default_var=None, deserialize_json=deserialize_json) + obj = Variable.get(key, default_var=None, deserialize_json=deserialize_json, session=session) if obj is None: if default is not None: - Variable.set(key=key, value=default, description=description, serialize_json=deserialize_json) + Variable.set( + key=key, + value=default, + description=description, + serialize_json=deserialize_json, + session=session, + ) return default raise ValueError("Default Value must be set") return obj @@ -156,6 +163,8 @@ def get( default_var: Any = __NO_DEFAULT_SENTINEL, deserialize_json: bool = False, team_name: str | None = None, + *, + session: Session | None = None, ) -> Any: """ Get a value for an Airflow Variable Key. @@ -164,6 +173,8 @@ def get( :param default_var: Default value of the Variable if the Variable doesn't exist :param deserialize_json: Deserialize the value to a Python dict :param team_name: Team name associated to the task trying to access the variable (if any) + :param session: Existing session to reuse for the metadata database lookup. Callers holding an + open transaction (the scheduler under ``prohibit_commit``, for example) must pass it. """ # TODO: This is not the best way of having compat, but it's "better than erroring" for now. This still # means SQLA etc is loaded, but we can't avoid that unless/until we add import shims as a big @@ -172,6 +183,11 @@ def get( # If this is set it means we are in some kind of execution context (Task, Dag Parse or Triggerer perhaps) # and should use the Task SDK API server path if hasattr(sys.modules.get("airflow.sdk.execution_time.task_runner"), "SUPERVISOR_COMMS"): + if session is not None: + raise ValueError( + "Variable.get() cannot use a metadata database session from an execution context; " + "reads there go through the Execution API. Use airflow.sdk.Variable.get() instead." + ) warnings.warn( "Using Variable.get from `airflow.models` is deprecated." "Please use `get` on Variable from sdk(`airflow.sdk.Variable`) instead", @@ -192,7 +208,7 @@ def get( "Multi-team mode is not configured in the Airflow environment but the task trying to access the variable belongs to a team" ) - var_val = Variable.get_variable_from_secrets(key=key, team_name=team_name) + var_val = Variable.get_variable_from_secrets(key=key, team_name=team_name, session=session) if var_val is None: if default_var is not cls.__NO_DEFAULT_SENTINEL: return default_var @@ -344,7 +360,7 @@ def update( Variable.check_for_write_conflict(key=key) - if Variable.get_variable_from_secrets(key=key, team_name=team_name) is None: + if Variable.get_variable_from_secrets(key=key, team_name=team_name, session=session) is None: raise KeyError(f"Variable {key} does not exist") ctx: contextlib.AbstractContextManager @@ -465,12 +481,18 @@ def check_for_write_conflict(key: str) -> None: return None @staticmethod - def get_variable_from_secrets(key: str, team_name: str | None = None) -> str | None: + def get_variable_from_secrets( + key: str, team_name: str | None = None, *, session: Session | None = None + ) -> str | None: """ Get Airflow Variable by iterating over all Secret Backends. :param key: Variable Key :param team_name: Team name associated to the task trying to access the variable (if any) + :param session: Existing session to reuse for the metadata database lookup. Callers that + already hold a transaction must pass it, otherwise ``MetastoreBackend`` opens the same + scoped session and commits it, which detaches the caller's objects and is rejected + outright under ``prohibit_commit``. :return: Variable Value """ from airflow.sdk import SecretCache @@ -487,7 +509,16 @@ def get_variable_from_secrets(key: str, team_name: str | None = None) -> str | N for secrets_backend in ensure_secrets_loaded(): try: var_val = call_secrets_backend_method( - secrets_backend.get_variable, team_name=team_name, key=key + secrets_backend.get_variable, + team_name=team_name, + key=key, + # Only the metastore backend touches the metadata database, and it is the only + # one whose signature accepts a session. + **( + {"session": session} + if session is not None and isinstance(secrets_backend, MetastoreBackend) + else {} + ), ) if var_val is not None: break diff --git a/airflow-core/tests/unit/models/test_variable.py b/airflow-core/tests/unit/models/test_variable.py index 0bc32373a305d..c1db11df9f239 100644 --- a/airflow-core/tests/unit/models/test_variable.py +++ b/airflow-core/tests/unit/models/test_variable.py @@ -29,6 +29,7 @@ from airflow.models import Variable, crypto, variable from airflow.sdk import SecretCache from airflow.secrets.metastore import MetastoreBackend +from airflow.utils.sqlalchemy import prohibit_commit from tests_common.test_utils import db from tests_common.test_utils.config import conf_vars @@ -192,6 +193,84 @@ def test_variable_set_with_extra_secret_backend(self, mock_ensure_secrets, caplo ) Variable.delete(key="key", session=session) + @mock.patch.object(MetastoreBackend, "get_variable", autospec=True) + @mock.patch("airflow.models.variable.ensure_secrets_loaded") + def test_get_forwards_session_to_metastore_backend(self, mock_ensure_secrets, mock_get_variable, session): + mock_get_variable.return_value = "from_db" + mock_ensure_secrets.return_value = [MetastoreBackend()] + + assert Variable.get("some_key", session=session) == "from_db" + assert mock_get_variable.call_args.kwargs["session"] is session + + @mock.patch.object(MetastoreBackend, "get_variable", autospec=True) + @mock.patch("airflow.models.variable.ensure_secrets_loaded") + def test_get_without_session_omits_session_kwarg(self, mock_ensure_secrets, mock_get_variable): + mock_get_variable.return_value = "from_db" + mock_ensure_secrets.return_value = [MetastoreBackend()] + + assert Variable.get("some_key") == "from_db" + assert "session" not in mock_get_variable.call_args.kwargs + + @mock.patch("airflow.models.variable.ensure_secrets_loaded") + def test_get_does_not_forward_session_to_other_backends(self, mock_ensure_secrets, session): + """Only the metastore backend reads the metadata database, so only it accepts a session.""" + mock_backend = mock.Mock() + mock_backend.get_variable.return_value = "from_backend" + mock_backend.__class__.__name__ = "MockSecretsBackend" + mock_ensure_secrets.return_value = [mock_backend] + + assert Variable.get("some_key", session=session) == "from_backend" + assert "session" not in mock_backend.get_variable.call_args.kwargs + + def test_get_with_session_does_not_commit_under_prohibit_commit(self, session): + """ + A caller holding an open transaction can read a Variable without its session being committed. + + Without the session being forwarded, ``MetastoreBackend.get_variable``'s ``provide_session`` + takes the same scoped session and commits it, which the guard rejects. + """ + Variable.set(key="interval_key", value="60", session=session) + session.commit() + SecretCache.invalidate_variable("interval_key") + + with prohibit_commit(session): + assert Variable.get("interval_key", session=session) == "60" + + def test_update_with_session_does_not_commit_under_prohibit_commit(self, session): + """``update`` verifies existence through the secrets chain, which must reuse the session too.""" + Variable.set(key="interval_key", value="60", session=session) + session.commit() + SecretCache.invalidate_variable("interval_key") + + with prohibit_commit(session): + Variable.update(key="interval_key", value="120", session=session) + + def test_setdefault_with_session_does_not_commit_under_prohibit_commit(self, session): + """``setdefault`` reads through the secrets chain before deciding whether to write.""" + Variable.set(key="interval_key", value="60", session=session) + session.commit() + SecretCache.invalidate_variable("interval_key") + + with prohibit_commit(session): + assert Variable.setdefault("interval_key", "120", session=session) == "60" + + def test_setdefault_writes_default_with_session_under_prohibit_commit(self, session): + """The write half must reuse the session too, so the miss path stays inside the transaction.""" + with prohibit_commit(session): + assert Variable.setdefault("absent_key", "30", session=session) == "30" + session.commit() + + assert Variable.get("absent_key", session=session) == "30" + + def test_get_rejects_session_in_execution_context(self): + """Reads from an execution context go via the Execution API, where a session is meaningless.""" + task_runner = mock.Mock(SUPERVISOR_COMMS=mock.Mock()) + with ( + mock.patch.dict("sys.modules", {"airflow.sdk.execution_time.task_runner": task_runner}), + pytest.raises(ValueError, match="cannot use a metadata database session"), + ): + Variable.get("some_key", session=mock.Mock()) + def test_variable_set_get_round_trip_json(self): value = {"a": 17, "b": 47} Variable.set(key="tested_var_set_id", value=value, serialize_json=True)