ferruzzi commented on code in PR #68917:
URL: https://github.com/apache/airflow/pull/68917#discussion_r3777965163
##########
task-sdk/src/airflow/sdk/definitions/deadline.py:
##########
@@ -423,9 +423,19 @@ def resolve(self) -> timedelta:
value = Variable.get(self.key)
except AirflowRuntimeError as e:
raise ValueError(f"VariableInterval '{self.key}' not found") from e
+ return self.coerce_to_timedelta(value)
+ def coerce_to_timedelta(self, value: str | int | float | None) ->
timedelta:
+ """
Review Comment:
Nit: Your docstring in SDK is explicitly referring to core internals, I
don't think we're meant to be doing that. I'd just leave it as something like
"Split out so a caller that already holds the value can reuse the validation"
since that's all the SDK side knows.
##########
airflow-core/tests/unit/models/test_dagrun.py:
##########
@@ -1509,71 +1511,121 @@ def
test_dagrun_success_handles_empty_deadline_list(self, mock_prune, dag_maker,
mock_prune.assert_not_called()
assert dag_run.state == DagRunState.SUCCESS
- @mock.patch.object(Variable, "get")
+ @pytest.mark.parametrize(
+ ("interval", "failure"),
+ [
+ pytest.param(VariableInterval("missing_key"), nullcontext(),
id="missing_variable"),
+ pytest.param(
+ datetime.timedelta(hours=1),
+ mock.patch(
+
"airflow.serialization.definitions.dag.decode_deadline_alert",
+ autospec=True,
+ side_effect=ValueError("corrupt deadline alert blob"),
+ ),
+ id="decode_failure",
+ ),
+ pytest.param(
+ datetime.timedelta(hours=1),
+ mock.patch.object(
+ SerializedReferenceModels.FixedDatetimeDeadline,
+ "evaluate_with",
+ autospec=True,
+ side_effect=RuntimeError("evaluate_with failed"),
+ ),
+ id="evaluate_with_failure",
+ ),
+ ],
+ )
@mock.patch.object(Deadline, "prune_deadlines")
- def test_dagrun_deadline_variable_interval_stable(self, _, mock_get,
session, deadline_test_dag):
Review Comment:
Looks like a dropped test, is this covered by another or am I just
misreading the diff and it's just a rename?
##########
airflow-core/src/airflow/serialization/definitions/dag.py:
##########
@@ -741,52 +743,100 @@ def _process_dagrun_deadline_alerts(
if not deadline_alert:
continue
- deserialized_deadline_alert = decode_deadline_alert(
- {
- Encoding.TYPE: DAT.DEADLINE_ALERT,
- Encoding.VAR: {
- DeadlineAlertFields.REFERENCE:
deadline_alert.reference,
- DeadlineAlertFields.INTERVAL: deadline_alert.interval,
- DeadlineAlertFields.CALLBACK:
deadline_alert.callback_def,
- },
- }
- )
-
- interval = deserialized_deadline_alert.interval
+ # Deadline creation is best-effort. A failure here must not
prevent the DagRun
+ # itself from being created. Use a plain try/except rather than
+ # ``session.begin_nested()`` since ``create_dagrun`` runs under
+ # ``prohibit_commit`` and releasing a SAVEPOINT would trip that
guard.
+ try:
+ deserialized_deadline_alert = decode_deadline_alert(
+ {
+ Encoding.TYPE: DAT.DEADLINE_ALERT,
+ Encoding.VAR: {
+ DeadlineAlertFields.REFERENCE:
deadline_alert.reference,
+ DeadlineAlertFields.INTERVAL:
deadline_alert.interval,
+ DeadlineAlertFields.CALLBACK:
deadline_alert.callback_def,
+ },
+ }
+ )
- if isinstance(interval, VariableInterval):
- interval = interval.resolve()
+ interval = deserialized_deadline_alert.interval
- if isinstance(deserialized_deadline_alert.reference,
SerializedReferenceModels.TYPES.DAGRUN):
- deadline_time =
deserialized_deadline_alert.reference.evaluate_with(
- session=session,
- interval=interval,
- # TODO : Pretty sure we can drop these last two; verify
after testing is complete
- dag_id=self.dag_id,
- run_id=orm_dagrun.run_id,
+ # Resolve the DagRun's team once, so a team-scoped
VariableInterval is looked up
+ # against the right team (not the global scope) and the stats
tag is consistent.
+ team_name = (
+ DagModel.get_team_name(self.dag_id, session=session)
+ if airflow_conf.getboolean("core", "multi_team")
+ else None
)
- if deadline_time is not None:
- session.add(
- Deadline(
- deadline_time=deadline_time,
- callback=deserialized_deadline_alert.callback,
- dagrun_id=orm_dagrun.id,
- deadline_alert_id=deadline_alert.id,
- dag_id=orm_dagrun.dag_id,
- bundle_name=orm_dagrun.dag_model.bundle_name,
- )
- )
- team_name = (
- DagModel.get_team_name(self.dag_id, session=session)
- if airflow_conf.getboolean("core", "multi_team")
- else None
- )
- stats.incr(
- "deadline_alerts.deadline_created",
- tags=prune_dict({"dag_id": self.dag_id, "team_name":
team_name}),
+ if isinstance(interval, VariableInterval):
+ interval = self._resolve_variable_interval(interval,
team_name=team_name, session=session)
+
+ if isinstance(deserialized_deadline_alert.reference,
SerializedReferenceModels.TYPES.DAGRUN):
+ deadline_time =
deserialized_deadline_alert.reference.evaluate_with(
+ session=session,
+ interval=interval,
+ # TODO : Pretty sure we can drop these last two;
verify after testing is complete
+ dag_id=self.dag_id,
+ run_id=orm_dagrun.run_id,
)
+ if deadline_time is not None:
+ session.add(
+ Deadline(
+ deadline_time=deadline_time,
+ callback=deserialized_deadline_alert.callback,
+ dagrun_id=orm_dagrun.id,
+ deadline_alert_id=deadline_alert.id,
+ dag_id=orm_dagrun.dag_id,
+ bundle_name=orm_dagrun.dag_model.bundle_name,
+ )
+ )
+ stats.incr(
+ "deadline_alerts.deadline_created",
+ tags=prune_dict({"dag_id": self.dag_id,
"team_name": team_name}),
+ )
+ except Exception:
+ log.exception(
+ "Failed to create deadline for alert %s on DagRun %s
(dag_id=%s); "
+ "skipping this deadline, the DagRun is unaffected",
+ getattr(deadline_alert, "id", "<unknown>"),
+ orm_dagrun.run_id,
+ self.dag_id,
+ )
+ stats.incr("deadline_alerts.deadline_creation_failed",
tags={"dag_id": self.dag_id})
+
+ @staticmethod
+ def _resolve_variable_interval(
+ interval: VariableInterval, *, team_name: str | None, session: Session
+ ) -> datetime.timedelta:
+ """
+ Resolve a ``VariableInterval`` to a concrete ``timedelta`` at DagRun
creation.
+
+ The Variable is resolved using the standard secrets lookup order. The
scheduler
+ session is passed to the metastore backend to avoid creating a new
session
+ during DagRun creation.
+
+ :param interval: The ``VariableInterval`` to resolve.
+ :param team_name: Team owning the DagRun, forwarded to scope the
Variable lookup.
+ :param session: Scheduler session used for metadata database lookups.
+ :return: The resolved ``timedelta``.
+ :raises ValueError: If the Variable cannot be resolved or converted to
a valid ``timedelta``.
+ """
+ for backend in ensure_secrets_loaded():
Review Comment:
The core version of this (Variable.get_variable_from_secrets) has more logic
like a tyr/except block and caching the returned value. Is there a reason we
don't need any of that in here?
##########
airflow-core/src/airflow/serialization/definitions/dag.py:
##########
@@ -741,52 +743,100 @@ def _process_dagrun_deadline_alerts(
if not deadline_alert:
continue
- deserialized_deadline_alert = decode_deadline_alert(
- {
- Encoding.TYPE: DAT.DEADLINE_ALERT,
- Encoding.VAR: {
- DeadlineAlertFields.REFERENCE:
deadline_alert.reference,
- DeadlineAlertFields.INTERVAL: deadline_alert.interval,
- DeadlineAlertFields.CALLBACK:
deadline_alert.callback_def,
- },
- }
- )
-
- interval = deserialized_deadline_alert.interval
+ # Deadline creation is best-effort. A failure here must not
prevent the DagRun
+ # itself from being created. Use a plain try/except rather than
+ # ``session.begin_nested()`` since ``create_dagrun`` runs under
+ # ``prohibit_commit`` and releasing a SAVEPOINT would trip that
guard.
+ try:
+ deserialized_deadline_alert = decode_deadline_alert(
+ {
+ Encoding.TYPE: DAT.DEADLINE_ALERT,
+ Encoding.VAR: {
+ DeadlineAlertFields.REFERENCE:
deadline_alert.reference,
+ DeadlineAlertFields.INTERVAL:
deadline_alert.interval,
+ DeadlineAlertFields.CALLBACK:
deadline_alert.callback_def,
+ },
+ }
+ )
- if isinstance(interval, VariableInterval):
- interval = interval.resolve()
+ interval = deserialized_deadline_alert.interval
- if isinstance(deserialized_deadline_alert.reference,
SerializedReferenceModels.TYPES.DAGRUN):
- deadline_time =
deserialized_deadline_alert.reference.evaluate_with(
- session=session,
- interval=interval,
- # TODO : Pretty sure we can drop these last two; verify
after testing is complete
- dag_id=self.dag_id,
- run_id=orm_dagrun.run_id,
+ # Resolve the DagRun's team once, so a team-scoped
VariableInterval is looked up
+ # against the right team (not the global scope) and the stats
tag is consistent.
+ team_name = (
+ DagModel.get_team_name(self.dag_id, session=session)
+ if airflow_conf.getboolean("core", "multi_team")
+ else None
)
- if deadline_time is not None:
- session.add(
- Deadline(
- deadline_time=deadline_time,
- callback=deserialized_deadline_alert.callback,
- dagrun_id=orm_dagrun.id,
- deadline_alert_id=deadline_alert.id,
- dag_id=orm_dagrun.dag_id,
- bundle_name=orm_dagrun.dag_model.bundle_name,
- )
- )
- team_name = (
- DagModel.get_team_name(self.dag_id, session=session)
- if airflow_conf.getboolean("core", "multi_team")
- else None
- )
- stats.incr(
- "deadline_alerts.deadline_created",
- tags=prune_dict({"dag_id": self.dag_id, "team_name":
team_name}),
+ if isinstance(interval, VariableInterval):
+ interval = self._resolve_variable_interval(interval,
team_name=team_name, session=session)
+
+ if isinstance(deserialized_deadline_alert.reference,
SerializedReferenceModels.TYPES.DAGRUN):
+ deadline_time =
deserialized_deadline_alert.reference.evaluate_with(
+ session=session,
+ interval=interval,
+ # TODO : Pretty sure we can drop these last two;
verify after testing is complete
+ dag_id=self.dag_id,
+ run_id=orm_dagrun.run_id,
)
+ if deadline_time is not None:
+ session.add(
+ Deadline(
+ deadline_time=deadline_time,
+ callback=deserialized_deadline_alert.callback,
+ dagrun_id=orm_dagrun.id,
+ deadline_alert_id=deadline_alert.id,
+ dag_id=orm_dagrun.dag_id,
+ bundle_name=orm_dagrun.dag_model.bundle_name,
+ )
+ )
+ stats.incr(
+ "deadline_alerts.deadline_created",
+ tags=prune_dict({"dag_id": self.dag_id,
"team_name": team_name}),
+ )
+ except Exception:
+ log.exception(
+ "Failed to create deadline for alert %s on DagRun %s
(dag_id=%s); "
+ "skipping this deadline, the DagRun is unaffected",
+ getattr(deadline_alert, "id", "<unknown>"),
+ orm_dagrun.run_id,
+ self.dag_id,
+ )
+ stats.incr("deadline_alerts.deadline_creation_failed",
tags={"dag_id": self.dag_id})
+
+ @staticmethod
+ def _resolve_variable_interval(
+ interval: VariableInterval, *, team_name: str | None, session: Session
+ ) -> datetime.timedelta:
+ """
+ Resolve a ``VariableInterval`` to a concrete ``timedelta`` at DagRun
creation.
+
+ The Variable is resolved using the standard secrets lookup order. The
scheduler
+ session is passed to the metastore backend to avoid creating a new
session
+ during DagRun creation.
+
+ :param interval: The ``VariableInterval`` to resolve.
+ :param team_name: Team owning the DagRun, forwarded to scope the
Variable lookup.
+ :param session: Scheduler session used for metadata database lookups.
+ :return: The resolved ``timedelta``.
+ :raises ValueError: If the Variable cannot be resolved or converted to
a valid ``timedelta``.
+ """
+ for backend in ensure_secrets_loaded():
+ value = call_secrets_backend_method(
+ backend.get_variable,
+ team_name=team_name,
+ key=interval.key,
+ **({"session": session} if isinstance(backend,
MetastoreBackend) else {}),
+ )
+ if value is not None:
+ return interval.coerce_to_timedelta(value)
Review Comment:
We're using the secrets backend, do we need to mask or unmask here? In
theory it's retrieving a number and that number itself may not be a secret, but
I'm not positive if everything in there gets masked by default, etc.
##########
airflow-core/tests/unit/models/test_dagrun.py:
##########
@@ -1509,71 +1511,121 @@ def
test_dagrun_success_handles_empty_deadline_list(self, mock_prune, dag_maker,
mock_prune.assert_not_called()
assert dag_run.state == DagRunState.SUCCESS
- @mock.patch.object(Variable, "get")
+ @pytest.mark.parametrize(
+ ("interval", "failure"),
+ [
+ pytest.param(VariableInterval("missing_key"), nullcontext(),
id="missing_variable"),
+ pytest.param(
+ datetime.timedelta(hours=1),
+ mock.patch(
+
"airflow.serialization.definitions.dag.decode_deadline_alert",
+ autospec=True,
+ side_effect=ValueError("corrupt deadline alert blob"),
+ ),
+ id="decode_failure",
+ ),
+ pytest.param(
+ datetime.timedelta(hours=1),
+ mock.patch.object(
+ SerializedReferenceModels.FixedDatetimeDeadline,
+ "evaluate_with",
+ autospec=True,
+ side_effect=RuntimeError("evaluate_with failed"),
+ ),
+ id="evaluate_with_failure",
+ ),
+ ],
+ )
@mock.patch.object(Deadline, "prune_deadlines")
- def test_dagrun_deadline_variable_interval_stable(self, _, mock_get,
session, deadline_test_dag):
- future_date = datetime.datetime.now() + datetime.timedelta(days=365)
+ def test_dagrun_deadline_failure_is_isolated(self, _, interval, failure,
session, deadline_test_dag):
+ """A failure while creating any single deadline must not abort DagRun
creation."""
+ future_date = datetime.datetime(2037, 1, 1,
tzinfo=datetime.timezone.utc)
+
+ scheduler_dag = deadline_test_dag(
+ deadline=DeadlineAlert(
+ reference=DeadlineReference.FIXED_DATETIME(future_date),
+ interval=interval,
+ callback=AsyncCallback(empty_callback_for_deadline),
+ ),
+ )
+
+ with failure:
+ dag_run = self.create_dag_run(
+ dag=scheduler_dag,
+ task_states={"task_1": TaskInstanceState.SUCCESS},
+ session=session,
+ )
- # First value used during resolution.
- mock_get.return_value = "60"
+ assert dag_run is not None
+ assert session.execute(select(Deadline)).scalars().one_or_none() is
None
+
+ @mock.patch.object(Deadline, "prune_deadlines")
+ def test_dagrun_deadline_variable_interval_resolves_from_env_var(
+ self, _, session, deadline_test_dag, monkeypatch
+ ):
+ """A VariableInterval backed by an ``AIRFLOW_VAR_*`` env var (no DB
row) must resolve."""
+ monkeypatch.setenv("AIRFLOW_VAR_ENV_INTERVAL_KEY", "7")
+ future_date = datetime.datetime(2037, 1, 1,
tzinfo=datetime.timezone.utc)
scheduler_dag = deadline_test_dag(
deadline=DeadlineAlert(
reference=DeadlineReference.FIXED_DATETIME(future_date),
- interval=VariableInterval("my_key"),
+ interval=VariableInterval("env_interval_key"),
callback=AsyncCallback(empty_callback_for_deadline),
),
)
dag_run = self.create_dag_run(
dag=scheduler_dag,
- task_states={"task_1": TaskInstanceState.SUCCESS, "task_2":
TaskInstanceState.SUCCESS},
+ task_states={"task_1": TaskInstanceState.SUCCESS},
session=session,
)
- dag_run.dag = scheduler_dag
-
- # First update resolve interval to "5".
- dag_run.update_state(session=session)
-
- deadline = session.execute(select(Deadline)).scalars().one_or_none()
- first_deadline_time = deadline.deadline_time
-
- # Change Variable value after resolution.
- mock_get.return_value = "120"
-
- # Run again (This should not change existing deadline).
- dag_run.update_state(session=session)
+ assert dag_run is not None
deadline = session.execute(select(Deadline)).scalars().one_or_none()
- assert deadline.deadline_time == first_deadline_time
+ assert deadline is not None
+ assert deadline.deadline_time == future_date +
datetime.timedelta(seconds=7)
+ @pytest.mark.parametrize(
+ ("multi_team", "team_name"),
+ [
+ pytest.param("true", "team_alpha", id="team_scoped"),
+ pytest.param("false", None, id="global"),
+ ],
+ )
@mock.patch.object(Deadline, "prune_deadlines")
- def test_dagrun_deadline_variable_interval_missing_variable_fails(self, _,
session, deadline_test_dag):
- mock_err = mock.Mock()
- mock_err.error.value = "MISSING_DEADLINE"
- mock_err.detail = "missing deadline"
+ def test_dagrun_deadline_variable_interval_scoped_to_team(
+ self, _, multi_team, team_name, session, deadline_test_dag
+ ):
+ """A VariableInterval must resolve against the DagRun's team, not the
global scope."""
+ future_date = datetime.datetime(2037, 1, 1,
tzinfo=datetime.timezone.utc)
- with mock.patch.object(
- Variable,
- "get",
- side_effect=AirflowRuntimeError(mock_err),
- ):
- future_date = datetime.datetime.now() +
datetime.timedelta(days=365)
+ scheduler_dag = deadline_test_dag(
+ deadline=DeadlineAlert(
+ reference=DeadlineReference.FIXED_DATETIME(future_date),
+ interval=VariableInterval("team_interval_key"),
+ callback=AsyncCallback(empty_callback_for_deadline),
+ ),
+ )
- scheduler_dag = deadline_test_dag(
- deadline=DeadlineAlert(
- reference=DeadlineReference.FIXED_DATETIME(future_date),
- interval=VariableInterval("missing_key"),
- callback=AsyncCallback(empty_callback_for_deadline),
- ),
+ with (
+ conf_vars({("core", "multi_team"): multi_team}),
+ mock.patch("airflow.models.dag.DagModel.get_team_name",
return_value=team_name),
+ mock.patch(
+
"airflow.serialization.definitions.dag.call_secrets_backend_method",
+ return_value="5",
+ ) as mock_call,
+ ):
+ dag_run = self.create_dag_run(
+ dag=scheduler_dag,
+ task_states={"task_1": TaskInstanceState.SUCCESS},
+ session=session,
)
- with pytest.raises(ValueError, match="not found"):
- self.create_dag_run(
- dag=scheduler_dag,
- task_states={"task_1": TaskInstanceState.SUCCESS},
- session=session,
- )
+ assert dag_run is not None
+ # The team the DagRun belongs to must be forwarded to the backend
lookup.
+ mock_call.assert_called_once()
+ assert mock_call.call_args.kwargs["team_name"] == team_name
Review Comment:
It's possible I'm missing something, but if the bug was that getting the
Variable from the db was running into the `prohibit_commit` wrapper, none of
these new tests are running into that.
The isolation half is covered, but I don't think the session-forwarding half
is. If you delete just the `session=session` forwarding to `MetastoreBackend`,
I believe every test here still passes: without a `prohibit_commit` guard in
the test, the commit just succeeds and the Variable is still found. Could we
get one test that creates the run inside the guard, the way
`test_mutation_hook_committing_session_crashes_under_prohibit_commit` below
does? That would actually test the thing the PR is fixing.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
