Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1#
2# Licensed to the Apache Software Foundation (ASF) under one
3# or more contributor license agreements. See the NOTICE file
4# distributed with this work for additional information
5# regarding copyright ownership. The ASF licenses this file
6# to you under the Apache License, Version 2.0 (the
7# "License"); you may not use this file except in compliance
8# with the License. You may obtain a copy of the License at
9#
10# http://www.apache.org/licenses/LICENSE-2.0
11#
12# Unless required by applicable law or agreed to in writing,
13# software distributed under the License is distributed on an
14# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15# KIND, either express or implied. See the License for the
16# specific language governing permissions and limitations
17# under the License.
18from __future__ import annotations
20import copy
21import functools
22import itertools
23import json
24import logging
25import os
26import sys
27import warnings
28import weakref
29from collections import abc, defaultdict, deque
30from collections.abc import Callable, Collection, Iterable, MutableSet
31from datetime import datetime, timedelta
32from inspect import signature
33from typing import TYPE_CHECKING, Any, ClassVar, TypeGuard, Union, cast, overload
34from urllib.parse import urlsplit
35from uuid import UUID
37import attrs
38import jinja2
39from dateutil.relativedelta import relativedelta
41from airflow import settings
42from airflow.sdk import TaskInstanceState, TriggerRule
43from airflow.sdk.bases.operator import BaseOperator
44from airflow.sdk.bases.timetable import BaseTimetable
45from airflow.sdk.definitions._internal.node import validate_key
46from airflow.sdk.definitions._internal.types import NOTSET, ArgNotSet, is_arg_set
47from airflow.sdk.definitions.asset import AssetAll, BaseAsset
48from airflow.sdk.definitions.context import Context
49from airflow.sdk.definitions.deadline import DeadlineAlert
50from airflow.sdk.definitions.param import DagParam, ParamsDict
51from airflow.sdk.definitions.timetables.assets import AssetTriggeredTimetable
52from airflow.sdk.definitions.timetables.simple import ContinuousTimetable, NullTimetable, OnceTimetable
53from airflow.sdk.exceptions import (
54 AirflowDagCycleException,
55 DuplicateTaskIdFound,
56 FailFastDagInvalidTriggerRule,
57 ParamValidationError,
58 RemovedInAirflow4Warning,
59 TaskNotFound,
60)
62if TYPE_CHECKING:
63 from re import Pattern
64 from typing import TypeAlias
66 from pendulum.tz.timezone import FixedTimezone, Timezone
67 from typing_extensions import Self, TypeIs
69 from airflow.models.taskinstance import TaskInstance as SchedulerTaskInstance
70 from airflow.sdk.definitions.decorators import TaskDecoratorCollection
71 from airflow.sdk.definitions.edges import EdgeInfoType
72 from airflow.sdk.definitions.mappedoperator import MappedOperator
73 from airflow.sdk.definitions.taskgroup import TaskGroup
74 from airflow.sdk.execution_time.supervisor import TaskRunResult
75 from airflow.timetables.base import DataInterval, Timetable as CoreTimetable
77 Operator: TypeAlias = BaseOperator | MappedOperator
79log = logging.getLogger(__name__)
81TAG_MAX_LEN = 100
83__all__ = [
84 "DAG",
85 "dag",
86]
88FINISHED_STATES = frozenset(
89 [
90 TaskInstanceState.SUCCESS,
91 TaskInstanceState.FAILED,
92 TaskInstanceState.SKIPPED,
93 TaskInstanceState.UPSTREAM_FAILED,
94 TaskInstanceState.REMOVED,
95 ]
96)
98DagStateChangeCallback = Callable[[Context], None]
99ScheduleInterval = None | str | timedelta | relativedelta
101ScheduleArg = Union[ScheduleInterval, BaseTimetable, "CoreTimetable", BaseAsset, Collection[BaseAsset]]
104_DAG_HASH_ATTRS = frozenset(
105 {
106 "dag_id",
107 "task_ids",
108 "start_date",
109 "end_date",
110 "fileloc",
111 "template_searchpath",
112 "last_loaded",
113 "schedule",
114 # TODO: Task-SDK: we should be hashing on timetable now, not schedule!
115 # "timetable",
116 }
117)
120def _is_core_timetable(schedule: ScheduleArg) -> TypeIs[CoreTimetable]:
121 try:
122 from airflow.timetables.base import Timetable
123 except ImportError:
124 return False
125 return isinstance(schedule, Timetable)
128def _create_timetable(interval: ScheduleInterval, timezone: Timezone | FixedTimezone) -> BaseTimetable:
129 """Create a Timetable instance from a plain ``schedule`` value."""
130 from airflow.sdk.configuration import conf as airflow_conf
131 from airflow.sdk.definitions.timetables.interval import (
132 CronDataIntervalTimetable,
133 DeltaDataIntervalTimetable,
134 )
135 from airflow.sdk.definitions.timetables.trigger import CronTriggerTimetable, DeltaTriggerTimetable
137 if interval is None:
138 return NullTimetable()
139 if interval == "@once":
140 return OnceTimetable()
141 if interval == "@continuous":
142 return ContinuousTimetable()
143 if isinstance(interval, timedelta | relativedelta):
144 if airflow_conf.getboolean("scheduler", "create_cron_data_intervals"):
145 return DeltaDataIntervalTimetable(interval)
146 return DeltaTriggerTimetable(interval)
147 if isinstance(interval, str):
148 if airflow_conf.getboolean("scheduler", "create_cron_data_intervals"):
149 return CronDataIntervalTimetable(interval, timezone)
150 return CronTriggerTimetable(interval, timezone=timezone)
151 raise ValueError(f"{interval!r} is not a valid schedule.")
154def _config_bool_factory(section: str, key: str) -> Callable[[], bool]:
155 from airflow.sdk.configuration import conf
157 return functools.partial(conf.getboolean, section, key)
160def _config_int_factory(section: str, key: str) -> Callable[[], int]:
161 from airflow.sdk.configuration import conf
163 return functools.partial(conf.getint, section, key)
166def _convert_params(val: abc.MutableMapping | None, self_: DAG) -> ParamsDict:
167 """
168 Convert the plain dict into a ParamsDict.
170 This will also merge in params from default_args
171 """
172 val = val or {}
174 # merging potentially conflicting default_args['params'] into params
175 if "params" in self_.default_args:
176 val.update(self_.default_args["params"])
177 del self_.default_args["params"]
179 params = ParamsDict(val)
180 object.__setattr__(self_, "params", params)
182 return params
185def _convert_str_to_tuple(val: str | Iterable[str] | None) -> Iterable[str] | None:
186 if isinstance(val, str):
187 return (val,)
188 return val
191def _convert_tags(tags: Collection[str] | None) -> MutableSet[str]:
192 return set(tags or [])
195def _convert_access_control(access_control):
196 if access_control is None:
197 return None
198 updated_access_control = {}
199 for role, perms in access_control.items():
200 updated_access_control[role] = updated_access_control.get(role, {})
201 if isinstance(perms, set | list):
202 # Support for old-style access_control where only the actions are specified
203 updated_access_control[role]["DAGs"] = set(perms)
204 else:
205 updated_access_control[role] = perms
206 return updated_access_control
209def _convert_deadline(deadline: list[DeadlineAlert] | DeadlineAlert | None) -> list[DeadlineAlert] | None:
210 """Convert deadline parameter to a list of DeadlineAlert objects."""
211 if deadline is None:
212 return None
213 if isinstance(deadline, DeadlineAlert):
214 return [deadline]
215 return list(deadline)
218def _convert_doc_md(doc_md: str | None) -> str | None:
219 if doc_md is None:
220 return doc_md
222 if doc_md.endswith(".md"):
223 try:
224 with open(doc_md) as fh:
225 return fh.read()
226 except FileNotFoundError:
227 return doc_md
229 return doc_md
232def _all_after_dag_id_to_kw_only(cls, fields: list[attrs.Attribute]):
233 i = iter(fields)
234 f = next(i)
235 if f.name != "dag_id":
236 raise RuntimeError("dag_id was not the first field")
237 yield f
239 for f in i:
240 yield f.evolve(kw_only=True)
243if TYPE_CHECKING:
244 # Given this attrs field:
245 #
246 # default_args: dict[str, Any] = attrs.field(factory=dict, converter=copy.copy)
247 #
248 # mypy ignores the type of the attrs and works out the type as the converter function. However it doesn't
249 # cope with generics properly and errors with 'incompatible type "dict[str, object]"; expected "_T"'
250 #
251 # https://github.com/python/mypy/issues/8625
252 def dict_copy(_: dict[str, Any]) -> dict[str, Any]: ...
253else:
254 dict_copy = copy.copy
257def _default_start_date(instance: DAG):
258 # Find start date inside default_args for compat with Airflow 2.
259 from airflow.sdk import timezone
261 if date := instance.default_args.get("start_date"):
262 if not isinstance(date, datetime):
263 date = timezone.parse(date)
264 instance.default_args["start_date"] = date
265 return date
266 return None
269def _default_dag_display_name(instance: DAG) -> str:
270 return instance.dag_id
273def _default_fileloc() -> str:
274 # Skip over this frame, and the 'attrs generated init'
275 back = sys._getframe().f_back
276 if not back or not (back := back.f_back):
277 # We expect two frames back, if not we don't know where we are
278 return ""
279 return back.f_code.co_filename if back else ""
282def _default_task_group(instance: DAG) -> TaskGroup:
283 from airflow.sdk.definitions.taskgroup import TaskGroup
285 return TaskGroup.create_root(dag=instance)
288# TODO: Task-SDK: look at re-enabling slots after we remove pickling
289@attrs.define(repr=False, field_transformer=_all_after_dag_id_to_kw_only, slots=False)
290class DAG:
291 """
292 A dag is a collection of tasks with directional dependencies.
294 A dag also has a schedule, a start date and an end date (optional). For each schedule,
295 (say daily or hourly), the DAG needs to run each individual tasks as their dependencies
296 are met. Certain tasks have the property of depending on their own past, meaning that
297 they can't run until their previous schedule (and upstream tasks) are completed.
299 Dags essentially act as namespaces for tasks. A task_id can only be
300 added once to a Dag.
302 Note that if you plan to use time zones all the dates provided should be pendulum
303 dates. See :ref:`timezone_aware_dags`.
305 .. versionadded:: 2.4
306 The *schedule* argument to specify either time-based scheduling logic
307 (timetable), or dataset-driven triggers.
309 .. versionchanged:: 3.0
310 The default value of *schedule* has been changed to *None* (no schedule).
311 The previous default was ``timedelta(days=1)``.
313 :param dag_id: The id of the DAG; must consist exclusively of alphanumeric
314 characters, dashes, dots and underscores (all ASCII)
315 :param description: The description for the DAG to e.g. be shown on the webserver
316 :param schedule: If provided, this defines the rules according to which DAG
317 runs are scheduled. Possible values include a cron expression string,
318 timedelta object, Timetable, or list of Asset objects.
319 See also :external:doc:`howto/timetable`.
320 :param start_date: The timestamp from which the scheduler will
321 attempt to backfill. If this is not provided, backfilling must be done
322 manually with an explicit time range.
323 :param end_date: A date beyond which your DAG won't run, leave to None
324 for open-ended scheduling.
325 :param template_searchpath: This list of folders (non-relative)
326 defines where jinja will look for your templates. Order matters.
327 Note that jinja/airflow includes the path of your DAG file by
328 default
329 :param template_undefined: Template undefined type.
330 :param user_defined_macros: a dictionary of macros that will be exposed
331 in your jinja templates. For example, passing ``dict(foo='bar')``
332 to this argument allows you to ``{{ foo }}`` in all jinja
333 templates related to this DAG. Note that you can pass any
334 type of object here.
335 :param user_defined_filters: a dictionary of filters that will be exposed
336 in your jinja templates. For example, passing
337 ``dict(hello=lambda name: 'Hello %s' % name)`` to this argument allows
338 you to ``{{ 'world' | hello }}`` in all jinja templates related to
339 this DAG.
340 :param default_args: A dictionary of default parameters to be used
341 as constructor keyword parameters when initialising operators.
342 Note that operators have the same hook, and precede those defined
343 here, meaning that if your dict contains `'depends_on_past': True`
344 here and `'depends_on_past': False` in the operator's call
345 `default_args`, the actual value will be `False`.
346 :param params: a dictionary of DAG level parameters that are made
347 accessible in templates, namespaced under `params`. These
348 params can be overridden at the task level.
349 :param max_active_tasks: the number of task instances allowed to run
350 concurrently
351 :param max_active_runs: maximum number of active DAG runs, beyond this
352 number of DAG runs in a running state, the scheduler won't create
353 new active DAG runs
354 :param max_consecutive_failed_dag_runs: (experimental) maximum number of consecutive failed DAG runs,
355 beyond this the scheduler will disable the DAG
356 :param dagrun_timeout: Specify the duration a DagRun should be allowed to run before it times out or
357 fails. Task instances that are running when a DagRun is timed out will be marked as skipped.
358 :param sla_miss_callback: DEPRECATED - The SLA feature is removed in Airflow 3.0, to be replaced with DeadlineAlerts in 3.1
359 :param deadline: An optional DeadlineAlert for the Dag.
360 :param catchup: Perform scheduler catchup (or only run latest)? Defaults to False
361 :param on_failure_callback: A function or list of functions to be called when a DagRun of this dag fails.
362 A context dictionary is passed as a single parameter to this function.
363 :param on_success_callback: Much like the ``on_failure_callback`` except
364 that it is executed when the dag succeeds.
365 :param access_control: Specify optional DAG-level actions, e.g.,
366 "{'role1': {'can_read'}, 'role2': {'can_read', 'can_edit', 'can_delete'}}"
367 or it can specify the resource name if there is a DAGs Run resource, e.g.,
368 "{'role1': {'DAG Runs': {'can_create'}}, 'role2': {'DAGs': {'can_read', 'can_edit', 'can_delete'}}"
369 :param is_paused_upon_creation: Specifies if the dag is paused when created for the first time.
370 If the dag exists already, this flag will be ignored. If this optional parameter
371 is not specified, the global config setting will be used.
372 :param jinja_environment_kwargs: additional configuration options to be passed to Jinja
373 ``Environment`` for template rendering
375 **Example**: to avoid Jinja from removing a trailing newline from template strings ::
377 DAG(
378 dag_id="my-dag",
379 jinja_environment_kwargs={
380 "keep_trailing_newline": True,
381 # some other jinja2 Environment options here
382 },
383 )
385 **See**: `Jinja Environment documentation
386 <https://jinja.palletsprojects.com/en/2.11.x/api/#jinja2.Environment>`_
388 :param render_template_as_native_obj: If True, uses a Jinja ``NativeEnvironment``
389 to render templates as native Python types. If False, a Jinja
390 ``Environment`` is used to render templates as string values.
391 :param tags: List of tags to help filtering Dags in the UI.
392 :param owner_links: Dict of owners and their links, that will be clickable on the Dags view UI.
393 Can be used as an HTTP link (for example the link to your Slack channel), or a mailto link.
394 e.g: ``{"dag_owner": "https://airflow.apache.org/"}``
395 :param auto_register: Automatically register this DAG when it is used in a ``with`` block
396 :param fail_fast: Fails currently running tasks when task in Dag fails.
397 **Warning**: A fail stop dag can only have tasks with the default trigger rule ("all_success").
398 An exception will be thrown if any task in a fail stop dag has a non default trigger rule.
399 :param dag_display_name: The display name of the Dag which appears on the UI.
400 """
402 __serialized_fields: ClassVar[frozenset[str]]
404 # Note: mypy gets very confused about the use of `@${attr}.default` for attrs without init=False -- and it
405 # doesn't correctly track/notice that they have default values (it gives errors about `Missing positional
406 # argument "description" in call to "DAG"`` etc), so for init=True args we use the `default=Factory()`
407 # style
409 def __rich_repr__(self):
410 yield "dag_id", self.dag_id
411 yield "schedule", self.schedule
412 yield "#tasks", len(self.tasks)
414 __rich_repr__.angular = True # type: ignore[attr-defined]
416 # NOTE: When updating arguments here, please also keep arguments in @dag()
417 # below in sync. (Search for 'def dag(' in this file.)
418 dag_id: str = attrs.field(kw_only=False, validator=lambda i, a, v: validate_key(v))
419 description: str | None = attrs.field(
420 default=None,
421 validator=attrs.validators.optional(attrs.validators.instance_of(str)),
422 )
423 default_args: dict[str, Any] = attrs.field(
424 factory=dict, validator=attrs.validators.instance_of(dict), converter=dict_copy
425 )
426 start_date: datetime | None = attrs.field(
427 default=attrs.Factory(_default_start_date, takes_self=True),
428 )
430 end_date: datetime | None = None
431 timezone: FixedTimezone | Timezone = attrs.field(init=False)
432 schedule: ScheduleArg = attrs.field(default=None, on_setattr=attrs.setters.frozen)
433 timetable: BaseTimetable | CoreTimetable = attrs.field(init=False)
434 template_searchpath: str | Iterable[str] | None = attrs.field(
435 default=None, converter=_convert_str_to_tuple
436 )
437 # TODO: Task-SDK: Work out how to not import jinj2 until we need it! It's expensive
438 template_undefined: type[jinja2.StrictUndefined] = jinja2.StrictUndefined
439 user_defined_macros: dict | None = None
440 user_defined_filters: dict | None = None
441 max_active_tasks: int = attrs.field(
442 factory=_config_int_factory("core", "max_active_tasks_per_dag"),
443 converter=attrs.converters.default_if_none( # type: ignore[misc]
444 # attrs only supports named callables or lambdas, but partial works
445 # OK here too. This is a false positive from attrs's Mypy plugin.
446 factory=_config_int_factory("core", "max_active_tasks_per_dag"),
447 ),
448 validator=attrs.validators.instance_of(int),
449 )
450 max_active_runs: int = attrs.field(
451 factory=_config_int_factory("core", "max_active_runs_per_dag"),
452 converter=attrs.converters.default_if_none( # type: ignore[misc]
453 # attrs only supports named callables or lambdas, but partial works
454 # OK here too. This is a false positive from attrs's Mypy plugin.
455 factory=_config_int_factory("core", "max_active_runs_per_dag"),
456 ),
457 validator=attrs.validators.instance_of(int),
458 )
459 max_consecutive_failed_dag_runs: int = attrs.field(
460 factory=_config_int_factory("core", "max_consecutive_failed_dag_runs_per_dag"),
461 converter=attrs.converters.default_if_none( # type: ignore[misc]
462 # attrs only supports named callables or lambdas, but partial works
463 # OK here too. This is a false positive from attrs's Mypy plugin.
464 factory=_config_int_factory("core", "max_consecutive_failed_dag_runs_per_dag"),
465 ),
466 validator=attrs.validators.instance_of(int),
467 )
468 dagrun_timeout: timedelta | None = attrs.field(
469 default=None,
470 validator=attrs.validators.optional(attrs.validators.instance_of(timedelta)),
471 )
472 deadline: list[DeadlineAlert] | DeadlineAlert | None = attrs.field(
473 default=None,
474 converter=_convert_deadline,
475 validator=attrs.validators.optional(
476 attrs.validators.deep_iterable(
477 member_validator=attrs.validators.instance_of(DeadlineAlert),
478 iterable_validator=attrs.validators.instance_of(list),
479 )
480 ),
481 )
483 sla_miss_callback: None = attrs.field(default=None)
484 catchup: bool = attrs.field(
485 factory=_config_bool_factory("scheduler", "catchup_by_default"),
486 )
487 on_success_callback: None | DagStateChangeCallback | list[DagStateChangeCallback] = None
488 on_failure_callback: None | DagStateChangeCallback | list[DagStateChangeCallback] = None
489 doc_md: str | None = attrs.field(default=None, converter=_convert_doc_md)
490 params: ParamsDict = attrs.field(
491 # mypy doesn't really like passing the Converter object
492 default=None,
493 converter=attrs.Converter(_convert_params, takes_self=True), # type: ignore[misc, call-overload]
494 )
495 access_control: dict[str, dict[str, Collection[str]]] | None = attrs.field(
496 default=None,
497 converter=attrs.Converter(_convert_access_control), # type: ignore[misc, call-overload]
498 )
499 is_paused_upon_creation: bool | None = None
500 jinja_environment_kwargs: dict | None = None
501 render_template_as_native_obj: bool = attrs.field(default=False, converter=bool)
502 tags: MutableSet[str] = attrs.field(factory=set, converter=_convert_tags)
503 owner_links: dict[str, str] = attrs.field(factory=dict)
504 auto_register: bool = attrs.field(default=True, converter=bool)
505 fail_fast: bool = attrs.field(default=False, converter=bool)
506 dag_display_name: str = attrs.field(
507 default=attrs.Factory(_default_dag_display_name, takes_self=True),
508 validator=attrs.validators.instance_of(str),
509 )
511 task_dict: dict[str, Operator] = attrs.field(factory=dict, init=False)
513 task_group: TaskGroup = attrs.field(
514 on_setattr=attrs.setters.frozen, default=attrs.Factory(_default_task_group, takes_self=True)
515 )
517 fileloc: str = attrs.field(init=False, factory=_default_fileloc)
518 relative_fileloc: str | None = attrs.field(init=False, default=None)
519 partial: bool = attrs.field(init=False, default=False)
521 edge_info: dict[str, dict[str, EdgeInfoType]] = attrs.field(init=False, factory=dict)
523 has_on_success_callback: bool = attrs.field(init=False)
524 has_on_failure_callback: bool = attrs.field(init=False)
525 disable_bundle_versioning: bool = attrs.field(
526 factory=_config_bool_factory("dag_processor", "disable_bundle_versioning")
527 )
529 # TODO (GH-52141): This is never used in the sdk dag (it only makes sense
530 # after this goes through the dag processor), but various parts of the code
531 # depends on its existence. We should remove this after completely splitting
532 # DAG classes in the SDK and scheduler.
533 last_loaded: datetime | None = attrs.field(init=False, default=None)
535 def __attrs_post_init__(self):
536 from airflow.sdk import timezone
538 # Apply the timezone we settled on to start_date, end_date if it wasn't supplied
539 if isinstance(_start_date := self.default_args.get("start_date"), str):
540 self.default_args["start_date"] = timezone.parse(_start_date, timezone=self.timezone)
541 if isinstance(_end_date := self.default_args.get("end_date"), str):
542 self.default_args["end_date"] = timezone.parse(_end_date, timezone=self.timezone)
544 self.start_date = timezone.convert_to_utc(self.start_date)
545 self.end_date = timezone.convert_to_utc(self.end_date)
546 if start_date := self.default_args.get("start_date", None):
547 self.default_args["start_date"] = timezone.convert_to_utc(start_date)
548 if end_date := self.default_args.get("end_date", None):
549 self.default_args["end_date"] = timezone.convert_to_utc(end_date)
550 if self.access_control is not None:
551 warnings.warn(
552 "The airflow.security.permissions module is deprecated; please see https://airflow.apache.org/docs/apache-airflow/stable/security/deprecated_permissions.html",
553 RemovedInAirflow4Warning,
554 stacklevel=2,
555 )
556 if (
557 active_runs_limit := self.timetable.active_runs_limit
558 ) is not None and active_runs_limit < self.max_active_runs:
559 raise ValueError(
560 f"Invalid max_active_runs: {type(self.timetable)} "
561 f"requires max_active_runs <= {active_runs_limit}"
562 )
564 @params.validator
565 def _validate_params(self, _, params: ParamsDict):
566 """
567 Validate Param values when the Dag has schedule defined.
569 Raise exception if there are any Params which can not be resolved by their schema definition.
570 """
571 if not self.timetable or not self.timetable.can_be_scheduled:
572 return
574 try:
575 params.validate()
576 except ParamValidationError as pverr:
577 raise ValueError(
578 f"Dag {self.dag_id!r} is not allowed to define a Schedule, "
579 "as there are required params without default values, or the default values are not valid."
580 ) from pverr
582 @catchup.validator
583 def _validate_catchup(self, _, catchup: bool):
584 requires_automatic_backfilling = self.timetable.can_be_scheduled and catchup
585 if requires_automatic_backfilling and not ("start_date" in self.default_args or self.start_date):
586 raise ValueError("start_date is required when catchup=True")
588 @tags.validator
589 def _validate_tags(self, _, tags: Collection[str]):
590 if tags and any(len(tag) > TAG_MAX_LEN for tag in tags):
591 raise ValueError(f"tag cannot be longer than {TAG_MAX_LEN} characters")
593 @max_active_runs.validator
594 def _validate_max_active_runs(self, _, max_active_runs):
595 if self.timetable.active_runs_limit is not None:
596 if self.timetable.active_runs_limit < self.max_active_runs:
597 raise ValueError(
598 f"Invalid max_active_runs: {type(self.timetable).__name__} "
599 f"requires max_active_runs <= {self.timetable.active_runs_limit}"
600 )
602 @timetable.default
603 def _default_timetable(instance: DAG) -> BaseTimetable | CoreTimetable:
604 schedule = instance.schedule
605 # TODO: Once
606 # delattr(self, "schedule")
607 if _is_core_timetable(schedule):
608 return schedule
609 if isinstance(schedule, BaseTimetable):
610 return schedule
611 if isinstance(schedule, BaseAsset):
612 return AssetTriggeredTimetable(schedule)
613 if isinstance(schedule, Collection) and not isinstance(schedule, str):
614 if not all(isinstance(x, BaseAsset) for x in schedule):
615 raise ValueError(
616 "All elements in 'schedule' should be either assets, asset references, or asset aliases"
617 )
618 return AssetTriggeredTimetable(AssetAll(*schedule))
619 return _create_timetable(schedule, instance.timezone)
621 @timezone.default
622 def _extract_tz(instance):
623 import pendulum
625 from airflow.sdk import timezone
627 start_date = instance.start_date or instance.default_args.get("start_date")
629 if start_date:
630 if not isinstance(start_date, datetime):
631 start_date = timezone.parse(start_date)
632 tzinfo = start_date.tzinfo or settings.TIMEZONE
633 tz = pendulum.instance(start_date, tz=tzinfo).timezone
634 else:
635 tz = settings.TIMEZONE
637 return tz
639 @has_on_success_callback.default
640 def _has_on_success_callback(self) -> bool:
641 return self.on_success_callback is not None
643 @has_on_failure_callback.default
644 def _has_on_failure_callback(self) -> bool:
645 return self.on_failure_callback is not None
647 @sla_miss_callback.validator
648 def _validate_sla_miss_callback(self, _, value):
649 if value is not None:
650 warnings.warn(
651 "The SLA feature is removed in Airflow 3.0, and replaced with a Deadline Alerts in >=3.1",
652 stacklevel=2,
653 )
654 return value
656 def __repr__(self):
657 return f"<DAG: {self.dag_id}>"
659 def __eq__(self, other: Self | Any):
660 # TODO: This subclassing behaviour seems wrong, but it's what Airflow has done for ~ever.
661 if type(self) is not type(other):
662 return False
663 return all(getattr(self, c, None) == getattr(other, c, None) for c in _DAG_HASH_ATTRS)
665 def __ne__(self, other: Any):
666 return not self == other
668 def __lt__(self, other):
669 return self.dag_id < other.dag_id
671 def __hash__(self):
672 hash_components: list[Any] = [type(self)]
673 for c in _DAG_HASH_ATTRS:
674 # If it is a list, convert to tuple because lists can't be hashed
675 if isinstance(getattr(self, c, None), list):
676 val = tuple(getattr(self, c))
677 else:
678 val = getattr(self, c, None)
679 try:
680 hash(val)
681 hash_components.append(val)
682 except TypeError:
683 hash_components.append(repr(val))
684 return hash(tuple(hash_components))
686 def __enter__(self) -> Self:
687 from airflow.sdk.definitions._internal.contextmanager import DagContext
689 DagContext.push(self)
690 return self
692 def __exit__(self, _type, _value, _tb):
693 from airflow.sdk.definitions._internal.contextmanager import DagContext
695 _ = DagContext.pop()
697 def validate(self):
698 """
699 Validate the Dag has a coherent setup.
701 This is called by the Dag bag before bagging the Dag.
702 """
703 self.timetable.validate()
704 self.validate_setup_teardown()
706 # We validate owner links on set, but since it's a dict it could be mutated without calling the
707 # setter. Validate again here
708 self._validate_owner_links(None, self.owner_links)
710 def validate_setup_teardown(self):
711 """
712 Validate that setup and teardown tasks are configured properly.
714 :meta private:
715 """
716 for task in self.tasks:
717 if task.is_setup:
718 for down_task in task.downstream_list:
719 if not down_task.is_teardown and down_task.trigger_rule != TriggerRule.ALL_SUCCESS:
720 # todo: we can relax this to allow out-of-scope tasks to have other trigger rules
721 # this is required to ensure consistent behavior of dag
722 # when clearing an indirect setup
723 raise ValueError("Setup tasks must be followed with trigger rule ALL_SUCCESS.")
725 def param(self, name: str, default: Any = NOTSET) -> DagParam:
726 """
727 Return a DagParam object for current dag.
729 :param name: dag parameter name.
730 :param default: fallback value for dag parameter.
731 :return: DagParam instance for specified name and current dag.
732 """
733 return DagParam(current_dag=self, name=name, default=default)
735 @property
736 def tasks(self) -> list[Operator]:
737 return list(self.task_dict.values())
739 @tasks.setter
740 def tasks(self, val):
741 raise AttributeError("DAG.tasks can not be modified. Use dag.add_task() instead.")
743 @property
744 def task_ids(self) -> list[str]:
745 return list(self.task_dict)
747 @property
748 def teardowns(self) -> list[Operator]:
749 return [task for task in self.tasks if getattr(task, "is_teardown", None)]
751 @property
752 def tasks_upstream_of_teardowns(self) -> list[Operator]:
753 upstream_tasks = [t.upstream_list for t in self.teardowns]
754 return [val for sublist in upstream_tasks for val in sublist if not getattr(val, "is_teardown", None)]
756 @property
757 def folder(self) -> str:
758 """Folder location of where the Dag object is instantiated."""
759 return os.path.dirname(self.fileloc)
761 @property
762 def owner(self) -> str:
763 """
764 Return list of all owners found in Dag tasks.
766 :return: Comma separated list of owners in Dag tasks
767 """
768 return ", ".join({t.owner for t in self.tasks})
770 def resolve_template_files(self):
771 for t in self.tasks:
772 # TODO: TaskSDK: move this on to BaseOperator and remove the check?
773 if hasattr(t, "resolve_template_files"):
774 t.resolve_template_files()
776 def get_template_env(self, *, force_sandboxed: bool = False) -> jinja2.Environment:
777 """Build a Jinja2 environment."""
778 from airflow.sdk.definitions._internal.templater import NativeEnvironment, SandboxedEnvironment
780 # Collect directories to search for template files
781 searchpath = [self.folder]
782 if self.template_searchpath:
783 searchpath += self.template_searchpath
785 # Default values (for backward compatibility)
786 jinja_env_options = {
787 "loader": jinja2.FileSystemLoader(searchpath),
788 "undefined": self.template_undefined,
789 "extensions": ["jinja2.ext.do"],
790 "cache_size": 0,
791 }
792 if self.jinja_environment_kwargs:
793 jinja_env_options.update(self.jinja_environment_kwargs)
794 env: jinja2.Environment
795 if self.render_template_as_native_obj and not force_sandboxed:
796 env = NativeEnvironment(**jinja_env_options)
797 else:
798 env = SandboxedEnvironment(**jinja_env_options)
800 # Add any user defined items. Safe to edit globals as long as no templates are rendered yet.
801 # http://jinja.pocoo.org/docs/2.10/api/#jinja2.Environment.globals
802 if self.user_defined_macros:
803 env.globals.update(self.user_defined_macros)
804 if self.user_defined_filters:
805 env.filters.update(self.user_defined_filters)
807 return env
809 def set_dependency(self, upstream_task_id, downstream_task_id):
810 """Set dependency between two tasks that already have been added to the Dag using add_task()."""
811 self.get_task(upstream_task_id).set_downstream(self.get_task(downstream_task_id))
813 @property
814 def roots(self) -> list[Operator]:
815 """Return nodes with no parents. These are first to execute and are called roots or root nodes."""
816 return [task for task in self.tasks if not task.upstream_list]
818 @property
819 def leaves(self) -> list[Operator]:
820 """Return nodes with no children. These are last to execute and are called leaves or leaf nodes."""
821 return [task for task in self.tasks if not task.downstream_list]
823 def topological_sort(self):
824 """
825 Sorts tasks in topographical order, such that a task comes after any of its upstream dependencies.
827 Deprecated in place of ``task_group.topological_sort``
828 """
829 from airflow.sdk.definitions.taskgroup import TaskGroup
831 # TODO: Remove in RemovedInAirflow3Warning
832 def nested_topo(group):
833 for node in group.topological_sort():
834 if isinstance(node, TaskGroup):
835 yield from nested_topo(node)
836 else:
837 yield node
839 return tuple(nested_topo(self.task_group))
841 def __deepcopy__(self, memo: dict[int, Any]):
842 # Switcharoo to go around deepcopying objects coming through the
843 # backdoor
844 cls = self.__class__
845 result = cls.__new__(cls)
846 memo[id(self)] = result
847 for k, v in self.__dict__.items():
848 if k not in ("user_defined_macros", "user_defined_filters", "_log"):
849 object.__setattr__(result, k, copy.deepcopy(v, memo))
851 result.user_defined_macros = self.user_defined_macros
852 result.user_defined_filters = self.user_defined_filters
853 if hasattr(self, "_log"):
854 result._log = self._log # type: ignore[attr-defined]
855 return result
857 def partial_subset(
858 self,
859 task_ids: str | Iterable[str],
860 include_downstream=False,
861 include_upstream=True,
862 include_direct_upstream=False,
863 ):
864 """
865 Return a subset of the current dag based on regex matching one or more tasks.
867 Returns a subset of the current dag as a deep copy of the current dag
868 based on a regex that should match one or many tasks, and includes
869 upstream and downstream neighbours based on the flag passed.
871 :param task_ids: Either a list of task_ids, or a string task_id
872 :param include_downstream: Include all downstream tasks of matched
873 tasks, in addition to matched tasks.
874 :param include_upstream: Include all upstream tasks of matched tasks,
875 in addition to matched tasks.
876 :param include_direct_upstream: Include all tasks directly upstream of matched
877 and downstream (if include_downstream = True) tasks
878 """
879 from airflow.sdk.definitions.mappedoperator import MappedOperator
881 def is_task(obj) -> TypeGuard[Operator]:
882 return isinstance(obj, BaseOperator | MappedOperator)
884 # deep-copying self.task_dict and self.task_group takes a long time, and we don't want all
885 # the tasks anyway, so we copy the tasks manually later
886 memo = {id(self.task_dict): None, id(self.task_group): None}
887 dag = copy.deepcopy(self, memo)
889 if isinstance(task_ids, str):
890 matched_tasks = [t for t in self.tasks if task_ids in t.task_id]
891 else:
892 matched_tasks = [t for t in self.tasks if t.task_id in task_ids]
894 also_include_ids: set[str] = set()
895 for t in matched_tasks:
896 if include_downstream:
897 for rel in t.get_flat_relatives(upstream=False):
898 also_include_ids.add(rel.task_id)
899 if rel not in matched_tasks: # if it's in there, we're already processing it
900 # need to include setups and teardowns for tasks that are in multiple
901 # non-collinear setup/teardown paths
902 if not rel.is_setup and not rel.is_teardown:
903 also_include_ids.update(
904 x.task_id for x in rel.get_upstreams_only_setups_and_teardowns()
905 )
906 if include_upstream:
907 also_include_ids.update(x.task_id for x in t.get_upstreams_follow_setups())
908 else:
909 if not t.is_setup and not t.is_teardown:
910 also_include_ids.update(x.task_id for x in t.get_upstreams_only_setups_and_teardowns())
911 if t.is_setup and not include_downstream:
912 also_include_ids.update(x.task_id for x in t.downstream_list if x.is_teardown)
914 also_include: list[Operator] = [self.task_dict[x] for x in also_include_ids]
915 direct_upstreams: list[Operator] = []
916 if include_direct_upstream:
917 for t in itertools.chain(matched_tasks, also_include):
918 upstream = (u for u in t.upstream_list if is_task(u))
919 direct_upstreams.extend(upstream)
921 # Make sure to not recursively deepcopy the dag or task_group while copying the task.
922 # task_group is reset later
923 def _deepcopy_task(t) -> Operator:
924 memo.setdefault(id(t.task_group), None)
925 return copy.deepcopy(t, memo)
927 # Compiling the unique list of tasks that made the cut
928 dag.task_dict = {
929 t.task_id: _deepcopy_task(t)
930 for t in itertools.chain(matched_tasks, also_include, direct_upstreams)
931 }
933 def filter_task_group(group, parent_group):
934 """Exclude tasks not included in the partial dag from the given TaskGroup."""
935 # We want to deepcopy _most but not all_ attributes of the task group, so we create a shallow copy
936 # and then manually deep copy the instances. (memo argument to deepcopy only works for instances
937 # of classes, not "native" properties of an instance)
938 copied = copy.copy(group)
940 memo[id(group.children)] = {}
941 if parent_group:
942 memo[id(group.parent_group)] = parent_group
943 for attr in type(group).__slots__:
944 value = getattr(group, attr)
945 value = copy.deepcopy(value, memo)
946 object.__setattr__(copied, attr, value)
948 proxy = weakref.proxy(copied)
950 for child in group.children.values():
951 if is_task(child):
952 if child.task_id in dag.task_dict:
953 task = copied.children[child.task_id] = dag.task_dict[child.task_id]
954 task.task_group = proxy
955 else:
956 copied.used_group_ids.discard(child.task_id)
957 else:
958 filtered_child = filter_task_group(child, proxy)
960 # Only include this child TaskGroup if it is non-empty.
961 if filtered_child.children:
962 copied.children[child.group_id] = filtered_child
964 return copied
966 object.__setattr__(dag, "task_group", filter_task_group(self.task_group, None))
968 # Removing upstream/downstream references to tasks and TaskGroups that did not make
969 # the cut.
970 groups = dag.task_group.get_task_group_dict()
971 for g in groups.values():
972 g.upstream_group_ids.intersection_update(groups)
973 g.downstream_group_ids.intersection_update(groups)
974 g.upstream_task_ids.intersection_update(dag.task_dict)
975 g.downstream_task_ids.intersection_update(dag.task_dict)
977 for t in dag.tasks:
978 # Removing upstream/downstream references to tasks that did not
979 # make the cut
980 t.upstream_task_ids.intersection_update(dag.task_dict)
981 t.downstream_task_ids.intersection_update(dag.task_dict)
983 dag.partial = len(dag.tasks) < len(self.tasks)
985 return dag
987 def has_task(self, task_id: str):
988 return task_id in self.task_dict
990 def has_task_group(self, task_group_id: str) -> bool:
991 return task_group_id in self.task_group_dict
993 @functools.cached_property
994 def task_group_dict(self):
995 return {k: v for k, v in self.task_group.get_task_group_dict().items() if k is not None}
997 def get_task(self, task_id: str) -> Operator:
998 if task_id in self.task_dict:
999 return self.task_dict[task_id]
1000 raise TaskNotFound(f"Task {task_id} not found")
1002 @property
1003 def task(self) -> TaskDecoratorCollection:
1004 from airflow.sdk.definitions.decorators import task
1006 return cast("TaskDecoratorCollection", functools.partial(task, dag=self))
1008 def add_task(self, task: Operator) -> None:
1009 """
1010 Add a task to the Dag.
1012 :param task: the task you want to add
1013 """
1014 # FailStopDagInvalidTriggerRule.check(dag=self, trigger_rule=task.trigger_rule)
1016 from airflow.sdk.definitions._internal.contextmanager import TaskGroupContext
1018 # if the task has no start date, assign it the same as the Dag
1019 if not task.start_date:
1020 task.start_date = self.start_date
1021 # otherwise, the task will start on the later of its own start date and
1022 # the Dag's start date
1023 elif self.start_date:
1024 task.start_date = max(task.start_date, self.start_date)
1026 # if the task has no end date, assign it the same as the dag
1027 if not task.end_date:
1028 task.end_date = self.end_date
1029 # otherwise, the task will end on the earlier of its own end date and
1030 # the Dag's end date
1031 elif task.end_date and self.end_date:
1032 task.end_date = min(task.end_date, self.end_date)
1034 task_id = task.node_id
1035 if not task.task_group:
1036 task_group = TaskGroupContext.get_current(self)
1037 if task_group:
1038 task_id = task_group.child_id(task_id)
1039 task_group.add(task)
1041 if (
1042 task_id in self.task_dict and self.task_dict[task_id] is not task
1043 ) or task_id in self.task_group.used_group_ids:
1044 raise DuplicateTaskIdFound(f"Task id '{task_id}' has already been added to the DAG")
1045 self.task_dict[task_id] = task
1047 task.dag = self
1048 # Add task_id to used_group_ids to prevent group_id and task_id collisions.
1049 self.task_group.used_group_ids.add(task_id)
1051 FailFastDagInvalidTriggerRule.check(fail_fast=self.fail_fast, trigger_rule=task.trigger_rule)
1053 def add_tasks(self, tasks: Iterable[Operator]) -> None:
1054 """
1055 Add a list of tasks to the Dag.
1057 :param tasks: a lit of tasks you want to add
1058 """
1059 for task in tasks:
1060 self.add_task(task)
1062 def _remove_task(self, task_id: str) -> None:
1063 # This is "private" as removing could leave a hole in dependencies if done incorrectly, and this
1064 # doesn't guard against that
1065 task = self.task_dict.pop(task_id)
1066 tg = getattr(task, "task_group", None)
1067 if tg:
1068 tg._remove(task)
1070 def check_cycle(self) -> None:
1071 """
1072 Check to see if there are any cycles in the Dag.
1074 :raises AirflowDagCycleException: If cycle is found in the Dag.
1075 """
1076 # default of int is 0 which corresponds to CYCLE_NEW
1077 CYCLE_NEW = 0
1078 CYCLE_IN_PROGRESS = 1
1079 CYCLE_DONE = 2
1081 visited: dict[str, int] = defaultdict(int)
1082 path_stack: deque[str] = deque()
1083 task_dict = self.task_dict
1085 def _check_adjacent_tasks(task_id, current_task):
1086 """Return first untraversed child task, else None if all tasks traversed."""
1087 for adjacent_task in current_task.get_direct_relative_ids():
1088 if visited[adjacent_task] == CYCLE_IN_PROGRESS:
1089 msg = f"Cycle detected in Dag: {self.dag_id}. Faulty task: {task_id}"
1090 raise AirflowDagCycleException(msg)
1091 if visited[adjacent_task] == CYCLE_NEW:
1092 return adjacent_task
1093 return None
1095 for dag_task_id in self.task_dict.keys():
1096 if visited[dag_task_id] == CYCLE_DONE:
1097 continue
1098 path_stack.append(dag_task_id)
1099 while path_stack:
1100 current_task_id = path_stack[-1]
1101 if visited[current_task_id] == CYCLE_NEW:
1102 visited[current_task_id] = CYCLE_IN_PROGRESS
1103 task = task_dict[current_task_id]
1104 child_to_check = _check_adjacent_tasks(current_task_id, task)
1105 if not child_to_check:
1106 visited[current_task_id] = CYCLE_DONE
1107 path_stack.pop()
1108 else:
1109 path_stack.append(child_to_check)
1111 def cli(self):
1112 """Exposes a CLI specific to this Dag."""
1113 self.check_cycle()
1115 from airflow.cli import cli_parser
1117 parser = cli_parser.get_parser(dag_parser=True)
1118 args = parser.parse_args()
1119 args.func(args, self)
1121 @classmethod
1122 def get_serialized_fields(cls):
1123 """Stringified Dags and operators contain exactly these fields."""
1124 return cls.__serialized_fields
1126 def get_edge_info(self, upstream_task_id: str, downstream_task_id: str) -> EdgeInfoType:
1127 """Return edge information for the given pair of tasks or an empty edge if there is no information."""
1128 empty = cast("EdgeInfoType", {})
1129 if self.edge_info:
1130 return self.edge_info.get(upstream_task_id, {}).get(downstream_task_id, empty)
1131 return empty
1133 def set_edge_info(self, upstream_task_id: str, downstream_task_id: str, info: EdgeInfoType):
1134 """
1135 Set the given edge information on the Dag.
1137 Note that this will overwrite, rather than merge with, existing info.
1138 """
1139 self.edge_info.setdefault(upstream_task_id, {})[downstream_task_id] = info
1141 @owner_links.validator
1142 def _validate_owner_links(self, _, owner_links):
1143 wrong_links = {}
1145 for owner, link in owner_links.items():
1146 result = urlsplit(link)
1147 if result.scheme == "mailto":
1148 # netloc is not existing for 'mailto' link, so we are checking that the path is parsed
1149 if not result.path:
1150 wrong_links[result.path] = link
1151 elif not result.scheme or not result.netloc:
1152 wrong_links[owner] = link
1153 if wrong_links:
1154 raise ValueError(
1155 "Wrong link format was used for the owner. Use a valid link \n"
1156 f"Bad formatted links are: {wrong_links}"
1157 )
1159 def test(
1160 self,
1161 run_after: datetime | None = None,
1162 logical_date: datetime | None | ArgNotSet = NOTSET,
1163 run_conf: dict[str, Any] | None = None,
1164 conn_file_path: str | None = None,
1165 variable_file_path: str | None = None,
1166 use_executor: bool = False,
1167 mark_success_pattern: Pattern | str | None = None,
1168 ):
1169 """
1170 Execute one single DagRun for a given Dag and logical date.
1172 :param run_after: the datetime before which to Dag cannot run.
1173 :param logical_date: logical date for the Dag run
1174 :param run_conf: configuration to pass to newly created dagrun
1175 :param conn_file_path: file path to a connection file in either yaml or json
1176 :param variable_file_path: file path to a variable file in either yaml or json
1177 :param use_executor: if set, uses an executor to test the Dag
1178 :param mark_success_pattern: regex of task_ids to mark as success instead of running
1179 """
1180 import re
1181 import time
1182 from contextlib import ExitStack
1183 from unittest.mock import patch
1185 from airflow import settings
1186 from airflow.models.dagrun import DagRun, get_or_create_dagrun
1187 from airflow.sdk import DagRunState, timezone
1188 from airflow.serialization.encoders import coerce_to_core_timetable
1189 from airflow.serialization.serialized_objects import SerializedDAG
1190 from airflow.utils.types import DagRunTriggeredByType, DagRunType
1192 exit_stack = ExitStack()
1194 if conn_file_path or variable_file_path:
1195 backend_kwargs = {}
1196 if conn_file_path:
1197 backend_kwargs["connections_file_path"] = conn_file_path
1198 if variable_file_path:
1199 backend_kwargs["variables_file_path"] = variable_file_path
1201 exit_stack.enter_context(
1202 patch.dict(
1203 os.environ,
1204 {
1205 "AIRFLOW__SECRETS__BACKEND": "airflow.secrets.local_filesystem.LocalFilesystemBackend",
1206 "AIRFLOW__SECRETS__BACKEND_KWARGS": json.dumps(backend_kwargs),
1207 },
1208 )
1209 )
1211 if settings.Session is None:
1212 raise RuntimeError("Session not configured. Call configure_orm() first.")
1213 session = settings.Session()
1215 with exit_stack:
1216 self.validate()
1218 # Allow users to explicitly pass None. If it isn't set, we default to current time.
1219 logical_date = logical_date if is_arg_set(logical_date) else timezone.utcnow()
1221 log.debug("Clearing existing task instances for logical date %s", logical_date)
1222 # TODO: Replace with calling client.dag_run.clear in Execution API at some point
1223 SerializedDAG.clear_dags(
1224 dags=[self],
1225 start_date=logical_date,
1226 end_date=logical_date,
1227 dag_run_state=False,
1228 )
1230 log.debug("Getting dagrun for dag %s", self.dag_id)
1231 logical_date = timezone.coerce_datetime(logical_date)
1232 run_after = timezone.coerce_datetime(run_after) or timezone.coerce_datetime(timezone.utcnow())
1233 if logical_date is None:
1234 data_interval: DataInterval | None = None
1235 else:
1236 timetable = coerce_to_core_timetable(self.timetable)
1237 data_interval = timetable.infer_manual_data_interval(run_after=logical_date)
1238 from airflow.models.dag_version import DagVersion
1240 version = DagVersion.get_version(self.dag_id)
1241 if not version:
1242 from airflow.dag_processing.bundles.manager import DagBundlesManager
1243 from airflow.dag_processing.dagbag import DagBag, sync_bag_to_db
1244 from airflow.sdk.definitions._internal.dag_parsing_context import (
1245 _airflow_parsing_context_manager,
1246 )
1248 manager = DagBundlesManager()
1249 manager.sync_bundles_to_db(session=session)
1250 session.commit()
1251 # sync all bundles? or use the dags-folder bundle?
1252 # What if the test dag is in a different bundle?
1253 for bundle in manager.get_all_dag_bundles():
1254 if not bundle.is_initialized:
1255 bundle.initialize()
1256 with _airflow_parsing_context_manager(dag_id=self.dag_id):
1257 dagbag = DagBag(
1258 dag_folder=bundle.path, bundle_path=bundle.path, include_examples=False
1259 )
1260 sync_bag_to_db(dagbag, bundle.name, bundle.version)
1261 version = DagVersion.get_version(self.dag_id)
1262 if version:
1263 break
1264 scheduler_dag = SerializedDAG.deserialize_dag(SerializedDAG.serialize_dag(self))
1265 # Preserve callback functions from original Dag since they're lost during serialization
1266 # and yes it is a hack for now! It is a tradeoff for code simplicity.
1267 # Without it, we need "Scheduler Dag" (Serialized dag) for the scheduler bits
1268 # -- dep check, scheduling tis
1269 # and need real dag to get and run callbacks without having to load the dag model
1271 # Scheduler DAG shouldn't have these attributes, but assigning them
1272 # here is an easy hack to get this test() thing working.
1273 scheduler_dag.on_success_callback = self.on_success_callback # type: ignore[attr-defined, union-attr]
1274 scheduler_dag.on_failure_callback = self.on_failure_callback # type: ignore[attr-defined, union-attr]
1276 dr: DagRun = get_or_create_dagrun(
1277 dag=scheduler_dag,
1278 start_date=logical_date or run_after,
1279 logical_date=logical_date,
1280 data_interval=data_interval,
1281 run_after=run_after,
1282 run_id=DagRun.generate_run_id(
1283 run_type=DagRunType.MANUAL,
1284 logical_date=logical_date,
1285 run_after=run_after,
1286 ),
1287 session=session,
1288 conf=run_conf,
1289 triggered_by=DagRunTriggeredByType.TEST,
1290 triggering_user_name="dag_test",
1291 )
1292 # Start a mock span so that one is present and not started downstream. We
1293 # don't care about otel in dag.test and starting the span during dagrun update
1294 # is not functioning properly in this context anyway.
1295 dr.start_dr_spans_if_needed(tis=[])
1297 log.debug("starting dagrun")
1298 # Instead of starting a scheduler, we run the minimal loop possible to check
1299 # for task readiness and dependency management.
1300 # Instead of starting a scheduler, we run the minimal loop possible to check
1301 # for task readiness and dependency management.
1303 # ``Dag.test()`` works in two different modes depending on ``use_executor``:
1304 # - if ``use_executor`` is False, runs the task locally with no executor using ``_run_task``
1305 # - if ``use_executor`` is True, sends workloads to the executor with
1306 # ``BaseExecutor.queue_workload``
1307 if use_executor:
1308 from airflow.executors.base_executor import ExecutorLoader
1310 executor = ExecutorLoader.get_default_executor()
1311 executor.start()
1313 while dr.state == DagRunState.RUNNING:
1314 session.expire_all()
1315 schedulable_tis, _ = dr.update_state(session=session)
1316 for s in schedulable_tis:
1317 if s.state != TaskInstanceState.UP_FOR_RESCHEDULE:
1318 s.try_number += 1
1319 s.state = TaskInstanceState.SCHEDULED
1320 s.scheduled_dttm = timezone.utcnow()
1321 session.commit()
1322 # triggerer may mark tasks scheduled so we read from DB
1323 all_tis = set(dr.get_task_instances(session=session))
1324 scheduled_tis = {x for x in all_tis if x.state == TaskInstanceState.SCHEDULED}
1325 ids_unrunnable = {x for x in all_tis if x.state not in FINISHED_STATES} - scheduled_tis
1326 if not scheduled_tis and ids_unrunnable:
1327 log.warning("No tasks to run. unrunnable tasks: %s", ids_unrunnable)
1328 time.sleep(1)
1330 for ti in scheduled_tis:
1331 task = self.task_dict[ti.task_id]
1333 mark_success = (
1334 re.compile(mark_success_pattern).fullmatch(ti.task_id) is not None
1335 if mark_success_pattern is not None
1336 else False
1337 )
1339 if use_executor:
1340 if executor.has_task(ti):
1341 continue
1343 from pathlib import Path
1345 from airflow.executors import workloads
1346 from airflow.executors.base_executor import ExecutorLoader
1347 from airflow.executors.workloads import BundleInfo
1349 workload = workloads.ExecuteTask.make(
1350 ti,
1351 dag_rel_path=Path(self.fileloc),
1352 generator=executor.jwt_generator,
1353 sentry_integration=executor.sentry_integration,
1354 # For the system test/debug purpose, we use the default bundle which uses
1355 # local file system. If it turns out to be a feature people want, we could
1356 # plumb the Bundle to use as a parameter to dag.test
1357 bundle_info=BundleInfo(name="dags-folder"),
1358 )
1359 executor.queue_workload(workload, session=session)
1360 ti.state = TaskInstanceState.QUEUED
1361 session.commit()
1362 else:
1363 # Run the task locally
1364 try:
1365 if mark_success:
1366 ti.set_state(TaskInstanceState.SUCCESS)
1367 log.info("[DAG TEST] Marking success for %s on %s", task, ti.logical_date)
1368 else:
1369 _run_task(ti=ti, task=task, run_triggerer=True)
1370 except Exception:
1371 log.exception("Task failed; ti=%s", ti)
1372 if use_executor:
1373 executor.heartbeat()
1374 from airflow.jobs.scheduler_job_runner import SchedulerJobRunner
1375 from airflow.models.dagbag import DBDagBag
1377 SchedulerJobRunner.process_executor_events(
1378 executor=executor, job_id=None, scheduler_dag_bag=DBDagBag(), session=session
1379 )
1380 if use_executor:
1381 executor.end()
1382 return dr
1385def _run_task(
1386 *,
1387 ti: SchedulerTaskInstance,
1388 task: Operator,
1389 run_triggerer: bool = False,
1390) -> TaskRunResult | None:
1391 """
1392 Run a single task instance, and push result to Xcom for downstream tasks.
1394 Bypasses a lot of extra steps used in `task.run` to keep our local running as fast as
1395 possible. This function is only meant for the `dag.test` function as a helper function.
1396 """
1397 from airflow.sdk._shared.module_loading import import_string
1399 taskrun_result: TaskRunResult | None
1400 log.info("[DAG TEST] starting task_id=%s map_index=%s", ti.task_id, ti.map_index)
1401 while True:
1402 try:
1403 log.info("[DAG TEST] running task %s", ti)
1405 from airflow.sdk.api.datamodels._generated import TaskInstance as TaskInstanceSDK
1406 from airflow.sdk.execution_time.comms import DeferTask
1407 from airflow.sdk.execution_time.supervisor import run_task_in_process
1408 from airflow.serialization.serialized_objects import create_scheduler_operator
1410 # The API Server expects the task instance to be in QUEUED state before
1411 # it is run.
1412 ti.set_state(TaskInstanceState.QUEUED)
1413 task_sdk_ti = TaskInstanceSDK(
1414 id=UUID(str(ti.id)),
1415 task_id=ti.task_id,
1416 dag_id=ti.dag_id,
1417 run_id=ti.run_id,
1418 try_number=ti.try_number,
1419 map_index=ti.map_index,
1420 dag_version_id=UUID(str(ti.dag_version_id)),
1421 )
1423 taskrun_result = run_task_in_process(ti=task_sdk_ti, task=task)
1424 msg = taskrun_result.msg
1425 ti.set_state(taskrun_result.ti.state)
1426 ti.task = create_scheduler_operator(taskrun_result.ti.task)
1428 if ti.state == TaskInstanceState.DEFERRED and isinstance(msg, DeferTask) and run_triggerer:
1429 from airflow.utils.session import create_session
1431 # API Server expects the task instance to be in QUEUED state before
1432 # resuming from deferral.
1433 ti.set_state(TaskInstanceState.QUEUED)
1435 log.info("[DAG TEST] running trigger in line")
1436 trigger = import_string(msg.classpath)(**msg.trigger_kwargs)
1437 event = _run_inline_trigger(trigger, task_sdk_ti)
1438 ti.next_method = msg.next_method
1439 ti.next_kwargs = {"event": event.payload} if event else msg.next_kwargs
1440 log.info("[DAG TEST] Trigger completed")
1442 # Set the state to SCHEDULED so that the task can be resumed.
1443 with create_session() as session:
1444 ti.state = TaskInstanceState.SCHEDULED
1445 session.add(ti)
1446 continue
1448 break
1449 except Exception:
1450 log.exception("[DAG TEST] Error running task %s", ti)
1451 if ti.state not in FINISHED_STATES:
1452 ti.set_state(TaskInstanceState.FAILED)
1453 taskrun_result = None
1454 break
1455 raise
1457 log.info("[DAG TEST] end task task_id=%s map_index=%s", ti.task_id, ti.map_index)
1458 return taskrun_result
1461def _run_inline_trigger(trigger, task_sdk_ti):
1462 from airflow.sdk.execution_time.supervisor import InProcessTestSupervisor
1464 return InProcessTestSupervisor.run_trigger_in_process(trigger=trigger, ti=task_sdk_ti)
1467# Since we define all the attributes of the class with attrs, we can compute this statically at parse time
1468DAG._DAG__serialized_fields = frozenset(a.name for a in attrs.fields(DAG)) - { # type: ignore[attr-defined]
1469 "schedule_asset_references",
1470 "schedule_asset_alias_references",
1471 "task_outlet_asset_references",
1472 "_old_context_manager_dags",
1473 "safe_dag_id",
1474 "last_loaded",
1475 "user_defined_filters",
1476 "user_defined_macros",
1477 "partial",
1478 "params",
1479 "_log",
1480 "task_dict",
1481 "template_searchpath",
1482 "sla_miss_callback",
1483 "on_success_callback",
1484 "on_failure_callback",
1485 "template_undefined",
1486 "jinja_environment_kwargs",
1487 # has_on_*_callback are only stored if the value is True, as the default is False
1488 "has_on_success_callback",
1489 "has_on_failure_callback",
1490 "auto_register",
1491 "schedule",
1492}
1494if TYPE_CHECKING:
1495 # NOTE: Please keep the list of arguments in sync with DAG.__init__.
1496 # Only exception: dag_id here should have a default value, but not in DAG.
1497 @overload
1498 def dag(
1499 dag_id: str = "",
1500 *,
1501 description: str | None = None,
1502 schedule: ScheduleArg = None,
1503 start_date: datetime | None = None,
1504 end_date: datetime | None = None,
1505 template_searchpath: str | Iterable[str] | None = None,
1506 template_undefined: type[jinja2.StrictUndefined] = jinja2.StrictUndefined,
1507 user_defined_macros: dict | None = None,
1508 user_defined_filters: dict | None = None,
1509 default_args: dict[str, Any] | None = None,
1510 max_active_tasks: int = ...,
1511 max_active_runs: int = ...,
1512 max_consecutive_failed_dag_runs: int = ...,
1513 dagrun_timeout: timedelta | None = None,
1514 catchup: bool = ...,
1515 on_success_callback: None | DagStateChangeCallback | list[DagStateChangeCallback] = None,
1516 on_failure_callback: None | DagStateChangeCallback | list[DagStateChangeCallback] = None,
1517 deadline: list[DeadlineAlert] | DeadlineAlert | None = None,
1518 doc_md: str | None = None,
1519 params: ParamsDict | dict[str, Any] | None = None,
1520 access_control: dict[str, dict[str, Collection[str]]] | dict[str, Collection[str]] | None = None,
1521 is_paused_upon_creation: bool | None = None,
1522 jinja_environment_kwargs: dict | None = None,
1523 render_template_as_native_obj: bool = False,
1524 tags: Collection[str] | None = None,
1525 owner_links: dict[str, str] | None = None,
1526 auto_register: bool = True,
1527 fail_fast: bool = False,
1528 dag_display_name: str | None = None,
1529 disable_bundle_versioning: bool = False,
1530 ) -> Callable[[Callable], Callable[..., DAG]]:
1531 """
1532 Python dag decorator which wraps a function into an Airflow Dag.
1534 Accepts kwargs for operator kwarg. Can be used to parameterize Dags.
1536 :param dag_args: Arguments for DAG object
1537 :param dag_kwargs: Kwargs for DAG object.
1538 """
1540 @overload
1541 def dag(func: Callable[..., DAG]) -> Callable[..., DAG]:
1542 """Python dag decorator to use without any arguments."""
1545def dag(dag_id_or_func=None, __DAG_class=DAG, __warnings_stacklevel_delta=2, **decorator_kwargs):
1546 from airflow.sdk.definitions._internal.decorators import fixup_decorator_warning_stack
1548 # TODO: Task-SDK: remove __DAG_class
1549 # __DAG_class is a temporary hack to allow the dag decorator in airflow.models.dag to continue to
1550 # return SchedulerDag objects
1551 DAG = __DAG_class
1553 def wrapper(f: Callable) -> Callable[..., DAG]:
1554 # Determine dag_id: prioritize keyword arg, then positional string, fallback to function name
1555 if "dag_id" in decorator_kwargs:
1556 dag_id = decorator_kwargs.pop("dag_id", "")
1557 elif isinstance(dag_id_or_func, str) and dag_id_or_func.strip():
1558 dag_id = dag_id_or_func
1559 else:
1560 dag_id = f.__name__
1562 @functools.wraps(f)
1563 def factory(*args, **kwargs):
1564 # Generate signature for decorated function and bind the arguments when called
1565 # we do this to extract parameters, so we can annotate them on the DAG object.
1566 # In addition, this fails if we are missing any args/kwargs with TypeError as expected.
1567 f_sig = signature(f).bind(*args, **kwargs)
1568 # Apply defaults to capture default values if set.
1569 f_sig.apply_defaults()
1571 # Initialize Dag with bound arguments
1572 with DAG(dag_id, **decorator_kwargs) as dag_obj:
1573 # Set Dag documentation from function documentation if it exists and doc_md is not set.
1574 if f.__doc__ and not dag_obj.doc_md:
1575 dag_obj.doc_md = f.__doc__
1577 # Generate DAGParam for each function arg/kwarg and replace it for calling the function.
1578 # All args/kwargs for function will be DAGParam object and replaced on execution time.
1579 f_kwargs = {}
1580 for name, value in f_sig.arguments.items():
1581 f_kwargs[name] = dag_obj.param(name, value)
1583 # set file location to caller source path
1584 back = sys._getframe().f_back
1585 dag_obj.fileloc = back.f_code.co_filename if back else ""
1587 # Invoke function to create operators in the Dag scope.
1588 f(**f_kwargs)
1590 # Return dag object such that it's accessible in Globals.
1591 return dag_obj
1593 # Ensure that warnings from inside DAG() are emitted from the caller, not here
1594 fixup_decorator_warning_stack(factory)
1595 return factory
1597 if callable(dag_id_or_func) and not isinstance(dag_id_or_func, str):
1598 return wrapper(dag_id_or_func)
1600 return wrapper