This is an automated email from the ASF dual-hosted git repository.
vincbeck pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 2c9088e078d Allow Variable.get to reuse the caller's database session
(#71968)
2c9088e078d is described below
commit 2c9088e078de7acc297a998a611490632cd1fa29
Author: D. Ferruzzi <[email protected]>
AuthorDate: Mon Sep 14 09:54:11 2026 -0700
Allow Variable.get to reuse the caller's database session (#71968)
`MetastoreBackend.get_variable` is decorated with `@provide_session`. Called
without a session it goes through `create_session()`, which for a scoped
session
returns *the caller's own session* and commits it on exit. So any code that
reads
a Variable while holding a transaction gets that transaction committed
underneath
it — detaching its objects, or raising `UNEXPECTED COMMIT` under the
scheduler's
`prohibit_commit` guard, where the error is then swallowed per-backend and
surfaces as a missing Variable.
`Variable.get` and `Variable.get_variable_from_secrets` now take an optional
keyword-only `session`, forwarded only to `MetastoreBackend`.
`Variable.update`
forwards its own.
---
airflow-core/newsfragments/71968.bugfix.rst | 1 +
airflow-core/src/airflow/models/variable.py | 60 +++++++++--
.../src/airflow/serialization/definitions/dag.py | 2 +-
.../airflow/serialization/definitions/deadline.py | 8 +-
airflow-core/tests/unit/models/test_variable.py | 118 +++++++++++++++++++++
.../serialization/definitions/test_deadline.py | 12 +++
.../src/airflow_shared/secrets_backend/base.py | 25 +++--
.../tests/secrets_backend/test_base.py | 25 ++++-
8 files changed, 225 insertions(+), 26 deletions(-)
diff --git a/airflow-core/newsfragments/71968.bugfix.rst
b/airflow-core/newsfragments/71968.bugfix.rst
new file mode 100644
index 00000000000..a312ab29091
--- /dev/null
+++ b/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.
diff --git a/airflow-core/src/airflow/models/variable.py
b/airflow-core/src/airflow/models/variable.py
index bcbc07a9703..1c53237aa16 100644
--- a/airflow-core/src/airflow/models/variable.py
+++ b/airflow-core/src/airflow/models/variable.py
@@ -28,7 +28,7 @@ from sqlalchemy import Boolean, ForeignKey, Integer, String,
Text, delete, or_,
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 @@ class Variable(Base, LoggingMixin):
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,19 @@ class Variable(Base, LoggingMixin):
: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 SQLAlchemy Session. 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 +170,8 @@ class Variable(Base, LoggingMixin):
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 +180,7 @@ class Variable(Base, LoggingMixin):
: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 SQLAlchemy Session. Callers holding an open
transaction 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 +189,11 @@ class Variable(Base, LoggingMixin):
# 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 +214,7 @@ class Variable(Base, LoggingMixin):
"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 +366,7 @@ class Variable(Base, LoggingMixin):
Variable.check_for_write_conflict(key=key, team_name=team_name)
- 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
@@ -469,12 +491,18 @@ class Variable(Base, LoggingMixin):
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
@@ -489,9 +517,23 @@ class Variable(Base, LoggingMixin):
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
diff --git a/airflow-core/src/airflow/serialization/definitions/dag.py
b/airflow-core/src/airflow/serialization/definitions/dag.py
index 5d5c4763622..e6979fc44f3 100644
--- a/airflow-core/src/airflow/serialization/definitions/dag.py
+++ b/airflow-core/src/airflow/serialization/definitions/dag.py
@@ -763,7 +763,7 @@ class SerializedDAG:
interval = deserialized_deadline_alert.interval
if isinstance(interval, SerializedVariableInterval):
- interval = interval.resolve()
+ interval = interval.resolve(session=session)
if isinstance(deserialized_deadline_alert.reference,
SerializedReferenceModels.TYPES.DAGRUN):
deadline_time =
deserialized_deadline_alert.reference.evaluate_with(
diff --git a/airflow-core/src/airflow/serialization/definitions/deadline.py
b/airflow-core/src/airflow/serialization/definitions/deadline.py
index 5f462c63d1a..e4362338cc0 100644
--- a/airflow-core/src/airflow/serialization/definitions/deadline.py
+++ b/airflow-core/src/airflow/serialization/definitions/deadline.py
@@ -388,10 +388,14 @@ class SerializedVariableInterval:
key: str
- def resolve(self) -> timedelta:
+ def resolve(self, *, session: Session | None = None) -> timedelta:
+ """
+ Get the Airflow Variable and return it as a ``timedelta``.
+ :param session: Existing SQLAlchemy Session. Callers holding an open
transaction must pass it.
+ """
try:
- value = Variable.get(self.key)
+ value = Variable.get(self.key, session=session)
except KeyError as e:
raise ValueError(f"VariableInterval '{self.key}' not found") from e
diff --git a/airflow-core/tests/unit/models/test_variable.py
b/airflow-core/tests/unit/models/test_variable.py
index cc2f9f94b93..57742aff4fa 100644
--- a/airflow-core/tests/unit/models/test_variable.py
+++ b/airflow-core/tests/unit/models/test_variable.py
@@ -30,6 +30,7 @@ from airflow.models import Variable, crypto, variable
from airflow.sdk import SecretCache
from airflow.secrets import BaseSecretsBackend
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
@@ -66,6 +67,29 @@ class _TeamAwareVariableBackend(BaseSecretsBackend):
return "secret_val"
+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):
@@ -254,6 +278,100 @@ class TestVariable:
assert mock_check.call_args.kwargs["team_name"] == testing_team.name
+ @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)
diff --git a/airflow-core/tests/unit/serialization/definitions/test_deadline.py
b/airflow-core/tests/unit/serialization/definitions/test_deadline.py
index 3e12fe0830c..f0fad39a37e 100644
--- a/airflow-core/tests/unit/serialization/definitions/test_deadline.py
+++ b/airflow-core/tests/unit/serialization/definitions/test_deadline.py
@@ -43,6 +43,18 @@ class TestVariableInterval:
assert interval.resolve() == expected
+ def test_resolve_forwards_session(self, mocker):
+ """The scheduler resolves intervals while holding an open transaction,
so the caller's
+ session has to reach ``Variable.get`` rather than open a new
session."""
+ expected_seconds = 42
+ mock_get = mocker.patch.object(Variable, "get",
return_value=str(expected_seconds))
+ session = mocker.MagicMock()
+
+ interval = SerializedVariableInterval(key="test_interval")
+
+ assert interval.resolve(session=session) ==
timedelta(seconds=expected_seconds)
+ mock_get.assert_called_once_with("test_interval", session=session)
+
@pytest.mark.parametrize(
("value", "raise_missing", "match"),
[
diff --git a/shared/secrets_backend/src/airflow_shared/secrets_backend/base.py
b/shared/secrets_backend/src/airflow_shared/secrets_backend/base.py
index 77c08be45f1..d98829c118e 100644
--- a/shared/secrets_backend/src/airflow_shared/secrets_backend/base.py
+++ b/shared/secrets_backend/src/airflow_shared/secrets_backend/base.py
@@ -21,24 +21,23 @@ from abc import ABC
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)
diff --git a/shared/secrets_backend/tests/secrets_backend/test_base.py
b/shared/secrets_backend/tests/secrets_backend/test_base.py
index e374a883cc4..60587892d8e 100644
--- a/shared/secrets_backend/tests/secrets_backend/test_base.py
+++ b/shared/secrets_backend/tests/secrets_backend/test_base.py
@@ -19,7 +19,7 @@ from __future__ import annotations
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 @@ class TestTeamNameBackwardCompat:
backend = _TeamUnawareConnValueBackend(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
