-
Notifications
You must be signed in to change notification settings - Fork 17.6k
Resolve VariableInterval deadlines safely at DagRun creation #68917
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
0cf3e1b
c6d243e
e1a907b
3cb692a
e9d840f
e65ee16
f27bb1f
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1 @@ | ||
| Deadline alerts using a ``VariableInterval`` no longer risk aborting DagRun creation. The interval is now resolved through the full secrets chain (environment variables, configured secrets backends, then the metadata database) on the scheduler's own session, so ``AIRFLOW_VAR_*`` and secrets-backend-backed Variables resolve correctly and the read does not commit inside the scheduler's ``prohibit_commit`` guard. Each deadline alert is also isolated: a single unresolvable or undecodable alert is logged and skipped instead of preventing the DagRun from being created. |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -31,8 +31,9 @@ | |
| from sqlalchemy import func, or_, select, tuple_ | ||
|
|
||
| from airflow._shared.observability.metrics import stats | ||
| from airflow._shared.secrets_backend.base import call_secrets_backend_method | ||
| from airflow._shared.timezones.timezone import coerce_datetime | ||
| from airflow.configuration import conf as airflow_conf | ||
| from airflow.configuration import conf as airflow_conf, ensure_secrets_loaded | ||
| from airflow.exceptions import ( | ||
| AirflowException, | ||
| DagNotPartitionedError, | ||
|
|
@@ -50,7 +51,9 @@ | |
| from airflow.models.deadline_alert import DeadlineAlert as DeadlineAlertModel | ||
| from airflow.models.taskinstancekey import TaskInstanceKey | ||
| from airflow.models.tasklog import LogTemplate | ||
| from airflow.models.variable import AirflowSecretsBackendAccessDenied | ||
| from airflow.sdk.definitions.deadline import VariableInterval | ||
| from airflow.secrets.metastore import MetastoreBackend | ||
| from airflow.serialization.decoders import decode_deadline_alert | ||
| from airflow.serialization.definitions.deadline import DeadlineAlertFields, SerializedReferenceModels | ||
| from airflow.serialization.definitions.param import SerializedParamsDict | ||
|
|
@@ -741,52 +744,112 @@ 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``. | ||
| :raises AirflowSecretsBackendAccessDenied: If a backend authoritatively denies access. | ||
| """ | ||
| for backend in ensure_secrets_loaded(): | ||
| try: | ||
| value = call_secrets_backend_method( | ||
| backend.get_variable, | ||
| team_name=team_name, | ||
| key=interval.key, | ||
| **({"session": session} if isinstance(backend, MetastoreBackend) else {}), | ||
| ) | ||
| except AirflowSecretsBackendAccessDenied: | ||
| # Authoritative deny — must NOT fall through to a less-restrictive backend. | ||
| raise | ||
| except Exception: | ||
| log.exception( | ||
| "Unable to retrieve variable from secrets backend (%s). " | ||
| "Checking subsequent secrets backend.", | ||
| type(backend).__name__, | ||
| ) | ||
| continue | ||
| if value is not None: | ||
| return interval.coerce_to_timedelta(value) | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 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.
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Nothing fetched here would be masked, so I don't believe there's any need to consider that here. |
||
| raise ValueError( | ||
| f"VariableInterval '{interval.key}' could not be resolved from any " | ||
| f"secrets backend, environment variable, or the metadata database" | ||
| ) | ||
|
|
||
| @provide_session | ||
| def set_task_instance_state( | ||
| self, | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I believe there's no notion of caching in the scheduler. The analogous try/except blocks, I agree, we can add that. Added in the latest commit.