Allow Variable.get to reuse the caller's database session

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.

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

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)