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

airflow-core/newsfragments/71968.bugfix.rst

@@ -0,0 +1 @@
+``Variable.get``, ``Variable.get_variable_from_secrets`` and ``Variable.setdefault`` now accept an optional keyword-only ``session`` that is forwarded to the metastore secrets backend.  A lookup made while holding an open session (most notably inside the scheduler's ``prohibit_commit`` guard) now reuses the caller's transaction instead of opening a scoped session and committing it.  ``Variable.update`` now forwards its session to its own existence check, which had the same problem.

airflow-core/src/airflow/models/variable.py

@@ -28,7 +28,7 @@
 from sqlalchemy.dialects.mysql import MEDIUMTEXT
 from sqlalchemy.orm import Mapped, declared_attr, mapped_column, reconstructor, synonym
 
-from airflow._shared.secrets_backend.base import call_secrets_backend_method
+from airflow._shared.secrets_backend.base import accepts_kwarg, call_secrets_backend_method
 from airflow._shared.secrets_masker import mask_secret
 from airflow.configuration import conf, ensure_secrets_loaded
 from airflow.models.base import ID_LEN, Base
@@ -126,7 +126,15 @@ 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: str,
+        default: Any,
+        description: str | None = None,
+        deserialize_json: bool = False,
+        *,
+        session: Session | None = None,
+    ) -> Any:
         """
         Return the current value for a key or store the default value and return it.
 
@@ -138,13 +146,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 +171,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 +181,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 +191,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 +216,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 +368,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 +489,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
@@ -485,9 +515,23 @@ def get_variable_from_secrets(key: str, team_name: str | None = None) -> str | N
         var_val = None
         # iterate over backends if not in cache (or expired)
         for secrets_backend in ensure_secrets_loaded():
+            # Only the metastore Variable backend touches the metadata database, so it is the only
+            # one offered a session, and only when its own override accepts one.  A subclass that
+            # overrides get_variable without the parameter raises TypeError, which the handler
+            # below swallows into a false "not found" that then gets cached.
+            session_kwargs: dict[str, Session] = {}
+            if (
+                session is not None
+                and isinstance(secrets_backend, MetastoreBackend)
+                and accepts_kwarg(secrets_backend.get_variable, "session")
+            ):
+                session_kwargs["session"] = session
             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,
+                    **session_kwargs,
                 )
                 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
@@ -41,6 +42,29 @@
 pytestmark = pytest.mark.db_test
 
 
+class _SessionUnawareMetastoreBackend(MetastoreBackend):
+    """A custom backend whose ``get_variable`` override predates the ``session`` keyword."""
+
+    # The signature mismatch with the base class is the point of this fixture, so mypy's
+    # override check has to be waived here rather than fixed.
+    def get_variable(self, key: str, team_name: str | None = None) -> str | None:  # type: ignore[override]
+        return "from_subclass"
+
+
+class _SessionAwareMetastoreBackend(MetastoreBackend):
+    """A custom backend whose ``get_variable`` override does accept ``session``."""
+
+    def __init__(self):
+        super().__init__()
+        self.received_session: Session | None = None
+
+    def get_variable(
+        self, key: str, team_name: str | None = None, *, session: Session | None = None
+    ) -> str | None:
+        self.received_session = session
+        return "from_subclass"
+
+
 class TestVariable:
     @pytest.fixture(autouse=True)
     def setup_test_cases(self):
@@ -192,6 +216,100 @@ 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
+
+    @mock.patch("airflow.models.variable.ensure_secrets_loaded")
+    def test_get_omits_session_for_session_unaware_metastore_subclass(self, mock_ensure_secrets, session):
+        """Forwarding a session to an override that predates it would read as a missing Variable."""
+        mock_ensure_secrets.return_value = [_SessionUnawareMetastoreBackend()]
+
+        assert Variable.get("some_key", session=session) == "from_subclass"
+
+    @mock.patch("airflow.models.variable.ensure_secrets_loaded")
+    def test_get_forwards_session_to_session_aware_metastore_subclass(self, mock_ensure_secrets, session):
+        """A subclass that does accept a session still receives it, so it reuses the transaction."""
+        backend = _SessionAwareMetastoreBackend()
+        mock_ensure_secrets.return_value = [backend]
+
+        assert Variable.get("some_key", session=session) == "from_subclass"
+        assert backend.received_session is session
+
+    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)

shared/secrets_backend/src/airflow_shared/secrets_backend/base.py

@@ -21,24 +21,23 @@
 from collections.abc import Callable
 
 
-def _accepts_team_name(method: Callable) -> bool:
+def accepts_kwarg(method: Callable, name: str) -> bool:
     """
-    Return whether a secrets-backend method accepts the ``team_name`` keyword.
-
-    Backends written before Airflow 3.2 override ``get_conn_value`` / ``get_variable`` /
-    ``get_connection`` with the legacy ``(self, conn_id)`` / ``(self, key)`` signature.
-    AIP-67 (multi-team) added a ``team_name`` keyword; forwarding it to those raises
-    ``TypeError``. A method accepts it if it declares a ``team_name`` parameter or a
-    ``**kwargs`` catch-all.
+    Return whether a secrets-backend method accepts the keyword *name*.
+
+    Backends override ``get_conn_value`` / ``get_variable`` / ``get_connection`` with
+    whatever signature was current when they were written, so a keyword added later cannot
+    be forwarded blindly: it raises ``TypeError`` inside the backend, which callers swallow
+    per-backend and report as a missing secret.  ``team_name`` (added by AIP-67 in 3.2) is
+    one example.  A method is considered to accept *name* if it explicitly declares the
+    parameter or has a ``**kwargs`` catch-all.
     """
     try:
         parameters = inspect.signature(method).parameters
     except (TypeError, ValueError):
-        # Un-introspectable callable (e.g. C-implemented): assume the 3.2+ signature.
+        # Un-introspectable callable (e.g. C-implemented): assume the current signature.
         return True
-    return "team_name" in parameters or any(
-        p.kind is inspect.Parameter.VAR_KEYWORD for p in parameters.values()
-    )
+    return name in parameters or any(p.kind is inspect.Parameter.VAR_KEYWORD for p in parameters.values())
 
 
 def call_secrets_backend_method(method: Callable, *, team_name: str | None, **kwargs):
@@ -52,7 +51,7 @@ def call_secrets_backend_method(method: Callable, *, team_name: str | None, **kw
     rather than retried without ``team_name``, which could mask the error and resolve a
     team-scoped lookup against the global scope.
     """
-    if _accepts_team_name(method):
+    if accepts_kwarg(method, "team_name"):
         return method(team_name=team_name, **kwargs)
     return method(**kwargs)
 

shared/secrets_backend/tests/secrets_backend/test_base.py

@@ -19,7 +19,7 @@
 
 import pytest
 
-from airflow_shared.secrets_backend.base import BaseSecretsBackend
+from airflow_shared.secrets_backend.base import BaseSecretsBackend, accepts_kwarg
 
 
 class MockConnection:
@@ -243,3 +243,26 @@ def test_legacy_backend_missing_conn_returns_none(self, team_name):
         backend = _LegacyConnValueBackend(conn_values={})
 
         assert backend.get_connection(conn_id="missing", team_name=team_name) is None
+
+
+class TestAcceptsKwarg:
+    @pytest.mark.parametrize(
+        ("method", "name", "expected"),
+        [
+            (lambda key: None, "session", False),
+            (lambda key, session=None: None, "session", True),
+            (lambda key, **kwargs: None, "session", True),
+            (lambda conn_id: None, "team_name", False),
+            (lambda conn_id, team_name=None: None, "team_name", True),
+            (lambda conn_id, **kwargs: None, "team_name", True),
+        ],
+    )
+    def test_declared_parameter_or_kwargs_catch_all(self, method, name, expected):
+        assert accepts_kwarg(method, name) is expected
+
+    def test_each_keyword_judged_independently(self):
+        def get_variable(key, team_name=None):
+            return None
+
+        assert accepts_kwarg(get_variable, "team_name") is True
+        assert accepts_kwarg(get_variable, "session") is False