From fc655ae28a9e64f3d6a0bf2672e3f6ffedf85415 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Thu, 11 Sep 2025 14:57:45 -0600 Subject: [PATCH 1/2] Remove SDK dependency from SerializedDAG --- .../core_api/routes/public/extra_links.py | 9 +- .../execution_api/routes/task_instances.py | 9 +- .../src/airflow/cli/commands/task_command.py | 10 +- .../src/airflow/jobs/scheduler_job_runner.py | 9 +- airflow-core/src/airflow/models/dagrun.py | 14 +- .../src/airflow/models/mappedoperator.py | 4 + .../src/airflow/models/taskinstance.py | 13 +- airflow-core/src/airflow/models/xcom_arg.py | 5 +- .../serialization/definitions/taskgroup.py | 3 +- .../serialization/serialized_objects.py | 124 +++++++++++++++--- .../src/airflow/ti_deps/dep_context.py | 8 +- airflow-core/src/airflow/utils/dag_edges.py | 4 +- .../src/airflow/utils/dot_renderer.py | 4 +- airflow-core/tests/unit/models/test_dag.py | 7 +- .../serialization/test_dag_serialization.py | 4 +- .../auth_manager/security_manager/override.py | 5 +- .../openlineage/utils/selective_enable.py | 9 +- .../providers/openlineage/utils/utils.py | 33 +++-- .../sensors/test_external_task_sensor.py | 3 +- 19 files changed, 180 insertions(+), 97 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/extra_links.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/extra_links.py index 1dbbf0db8b70c..83d17b404ad82 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/extra_links.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/extra_links.py @@ -17,8 +17,6 @@ from __future__ import annotations -from typing import TYPE_CHECKING, cast - from fastapi import Depends, HTTPException, status from sqlalchemy.sql import select @@ -31,10 +29,6 @@ from airflow.exceptions import TaskNotFound from airflow.models import DagRun -if TYPE_CHECKING: - from airflow.models.mappedoperator import MappedOperator - from airflow.serialization.serialized_objects import SerializedBaseOperator - extra_links_router = AirflowRouter( tags=["Extra Links"], prefix="/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/links" ) @@ -62,8 +56,7 @@ def get_extra_links( dag = get_dag_for_run_or_latest_version(dag_bag, dag_run, dag_id, session) try: - # TODO (GH-52141): Make dag a db-backed object so it only returns db-backed tasks. - task = cast("MappedOperator | SerializedBaseOperator", dag.get_task(task_id)) + task = dag.get_task(task_id) except TaskNotFound: raise HTTPException(status.HTTP_404_NOT_FOUND, f"Task with ID = {task_id} not found") diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py index c30cf1e9c4c4b..f2f8fc8e11565 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py @@ -254,14 +254,7 @@ def ti_run( if dag := dag_bag.get_dag_for_run(dag_run=dr, session=session): upstream_map_indexes = dict( - _get_upstream_map_indexes( - # TODO (GH-52141): This get_task should return scheduler - # types instead, but currently it inherits SDK's DAG. - cast("MappedOperator | SerializedBaseOperator", dag.get_task(ti.task_id)), - ti.map_index, - ti.run_id, - session=session, - ) + _get_upstream_map_indexes(dag.get_task(ti.task_id), ti.map_index, ti.run_id, session=session) ) else: upstream_map_indexes = None diff --git a/airflow-core/src/airflow/cli/commands/task_command.py b/airflow-core/src/airflow/cli/commands/task_command.py index 5b54d76100a50..b19ea6161a20a 100644 --- a/airflow-core/src/airflow/cli/commands/task_command.py +++ b/airflow-core/src/airflow/cli/commands/task_command.py @@ -82,7 +82,7 @@ def _generate_temporary_run_id() -> str: def _get_dag_run( *, - dag: DAG, + dag: DAG | SerializedDAG, create_if_necessary: CreateIfNecessary, logical_date_or_run_id: str | None = None, session: Session | None = None, @@ -274,9 +274,7 @@ def task_state(args) -> None: """ if not (dag := SerializedDagModel.get_dag(args.dag_id)): raise SystemExit(f"Can not find dag {args.dag_id!r}") - # TODO (GH-52141): get_task in scheduler needs to return scheduler types - # instead, but currently it inherits SDK's DAG. - task = cast("Operator", dag.get_task(task_id=args.task_id)) + task = dag.get_task(task_id=args.task_id) ti, _ = _get_ti(task, args.map_index, logical_date_or_run_id=args.logical_date_or_run_id) print(ti.state) @@ -434,9 +432,7 @@ def task_render(args, dag: DAG | None = None) -> None: dag = get_bagged_dag(args.bundle_name, args.dag_id) serialized_dag = SerializedDAG.deserialize_dag(SerializedDAG.serialize_dag(dag)) ti, _ = _get_ti( - # TODO (GH-52141): get_task in scheduler needs to return scheduler types - # instead, but currently it inherits SDK's DAG. - cast("Operator", serialized_dag.get_task(task_id=args.task_id)), + serialized_dag.get_task(task_id=args.task_id), args.map_index, logical_date_or_run_id=args.logical_date_or_run_id, create_if_necessary="memory", diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 5d91870473648..aac4544478a63 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -30,7 +30,7 @@ from datetime import date, datetime, timedelta from functools import lru_cache, partial from itertools import groupby -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any from sqlalchemy import and_, delete, desc, exists, func, inspect, or_, select, text, tuple_, update from sqlalchemy.exc import OperationalError @@ -97,9 +97,8 @@ from airflow._shared.logging.types import Logger from airflow.executors.base_executor import BaseExecutor from airflow.executors.executor_utils import ExecutorName - from airflow.models.mappedoperator import MappedOperator from airflow.models.taskinstance import TaskInstanceKey - from airflow.serialization.serialized_objects import SerializedBaseOperator, SerializedDAG + from airflow.serialization.serialized_objects import SerializedDAG from airflow.utils.sqlalchemy import CommitProhibitorGuard TI = TaskInstance @@ -917,9 +916,7 @@ def process_executor_events( ) if TYPE_CHECKING: assert dag - # TODO (GH-52141): get_task in scheduler needs to return scheduler types - # instead, but currently it inherits SDK's DAG. - task = cast("MappedOperator | SerializedBaseOperator", dag.get_task(ti.task_id)) + task = dag.get_task(ti.task_id) except Exception: cls.logger().exception("Marking task instance %s as %s", ti, state) ti.set_state(state) diff --git a/airflow-core/src/airflow/models/dagrun.py b/airflow-core/src/airflow/models/dagrun.py index 516ef64b30cb5..ec7dff2ba71df 100644 --- a/airflow-core/src/airflow/models/dagrun.py +++ b/airflow-core/src/airflow/models/dagrun.py @@ -1173,7 +1173,7 @@ def recalculate(self) -> _UnfinishedStates: self.notify_dagrun_state_changed(msg="task_failure") if execute_callbacks and dag.has_on_failure_callback: - self.handle_dag_callback(dag=dag, success=False, reason="task_failure") + self.handle_dag_callback(dag=cast("SDKDAG", dag), success=False, reason="task_failure") elif dag.has_on_failure_callback: callback = DagCallbackRequest( filepath=self.dag_model.relative_fileloc, @@ -1206,7 +1206,7 @@ def recalculate(self) -> _UnfinishedStates: self.notify_dagrun_state_changed(msg="success") if execute_callbacks and dag.has_on_success_callback: - self.handle_dag_callback(dag=dag, success=True, reason="success") + self.handle_dag_callback(dag=cast("SDKDAG", dag), success=True, reason="success") elif dag.has_on_success_callback: callback = DagCallbackRequest( filepath=self.dag_model.relative_fileloc, @@ -1237,7 +1237,11 @@ def recalculate(self) -> _UnfinishedStates: self.notify_dagrun_state_changed(msg="all_tasks_deadlocked") if execute_callbacks and dag.has_on_failure_callback: - self.handle_dag_callback(dag=dag, success=False, reason="all_tasks_deadlocked") + self.handle_dag_callback( + dag=cast("SDKDAG", dag), + success=False, + reason="all_tasks_deadlocked", + ) elif dag.has_on_failure_callback: callback = DagCallbackRequest( filepath=self.dag_model.relative_fileloc, @@ -1306,9 +1310,7 @@ def _filter_tis_and_exclude_removed(dag: SerializedDAG, tis: list[TI]) -> Iterab """Populate ``ti.task`` while excluding those missing one, marking them as REMOVED.""" for ti in tis: try: - # TODO (GH-52141): get_task in scheduler needs to return scheduler types - # instead, but currently it inherits SDK's DAG. - ti.task = cast("Operator", dag.get_task(ti.task_id)) + ti.task = dag.get_task(ti.task_id) except TaskNotFound: if ti.state != TaskInstanceState.REMOVED: self.log.error("Failed to get task for ti %s. Marking it as removed.", ti) diff --git a/airflow-core/src/airflow/models/mappedoperator.py b/airflow-core/src/airflow/models/mappedoperator.py index 7c7988d3c9bd3..33d77215500f8 100644 --- a/airflow-core/src/airflow/models/mappedoperator.py +++ b/airflow-core/src/airflow/models/mappedoperator.py @@ -114,6 +114,10 @@ class MappedOperator(DAGNode): dag: SerializedDAG = attrs.field(init=False) # type: ignore[assignment] task_group: SerializedTaskGroup = attrs.field(init=False) # type: ignore[assignment] + doc: str | None = attrs.field(init=False) + doc_json: str | None = attrs.field(init=False) + doc_rst: str | None = attrs.field(init=False) + doc_yaml: str | None = attrs.field(init=False) start_date: pendulum.DateTime | None = attrs.field(init=False, default=None) end_date: pendulum.DateTime | None = attrs.field(init=False, default=None) upstream_task_ids: set[str] = attrs.field(factory=set, init=False) diff --git a/airflow-core/src/airflow/models/taskinstance.py b/airflow-core/src/airflow/models/taskinstance.py index 4d440bd38a042..36193d077f36a 100644 --- a/airflow-core/src/airflow/models/taskinstance.py +++ b/airflow-core/src/airflow/models/taskinstance.py @@ -27,7 +27,7 @@ from collections.abc import Collection, Iterable from datetime import timedelta from functools import cache -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any from urllib.parse import quote import attrs @@ -237,8 +237,7 @@ def clear_task_instances( log.warning("No serialized dag found for dag '%s'", dr.dag_id) task_id = ti.task_id if ti_dag and ti_dag.has_task(task_id): - # TODO (GH-52141): Make dag a db-backed object so it only returns db-backed tasks. - task = cast("Operator", ti_dag.get_task(task_id)) + task = ti_dag.get_task(task_id) ti.refresh_from_task(task) if TYPE_CHECKING: assert ti.task @@ -1455,9 +1454,11 @@ def run( assert original_task is not None assert original_task.dag is not None - self.task = SerializedDAG.deserialize_dag(SerializedDAG.serialize_dag(original_task.dag)).task_dict[ - original_task.task_id - ] + # We don't set up all tests well... + if not isinstance(original_task.dag, SerializedDAG): + serialized_dag = SerializedDAG.deserialize_dag(SerializedDAG.serialize_dag(original_task.dag)) + self.task = serialized_dag.get_task(original_task.task_id) + res = self.check_and_change_state_before_execution( verbose=verbose, ignore_all_deps=ignore_all_deps, diff --git a/airflow-core/src/airflow/models/xcom_arg.py b/airflow-core/src/airflow/models/xcom_arg.py index 75cccba50334d..bc9326b2b2f52 100644 --- a/airflow-core/src/airflow/models/xcom_arg.py +++ b/airflow-core/src/airflow/models/xcom_arg.py @@ -19,7 +19,7 @@ from collections.abc import Iterator, Sequence from functools import singledispatch -from typing import TYPE_CHECKING, Any, TypeAlias, cast +from typing import TYPE_CHECKING, Any, TypeAlias import attrs from sqlalchemy import func, or_, select @@ -92,8 +92,7 @@ class SchedulerPlainXComArg(SchedulerXComArg): @classmethod def _deserialize(cls, data: dict[str, Any], dag: SerializedDAG) -> Self: - # TODO (GH-52141): SerializedDAG should return scheduler operator instead. - return cls(cast("Operator", dag.get_task(data["task_id"])), data["key"]) + return cls(dag.get_task(data["task_id"]), data["key"]) def iter_references(self) -> Iterator[tuple[Operator, str]]: yield self.operator, self.key diff --git a/airflow-core/src/airflow/serialization/definitions/taskgroup.py b/airflow-core/src/airflow/serialization/definitions/taskgroup.py index e26c6cfb4aeea..fef9dad303a69 100644 --- a/airflow-core/src/airflow/serialization/definitions/taskgroup.py +++ b/airflow-core/src/airflow/serialization/definitions/taskgroup.py @@ -45,7 +45,8 @@ class SerializedTaskGroup(DAGNode): group_display_name: str | None = attrs.field() prefix_group_id: bool = attrs.field() parent_group: SerializedTaskGroup | None = attrs.field() - dag: SerializedDAG = attrs.field() + # TODO (GH-52141): Replace DAGNode dependency. + dag: SerializedDAG = attrs.field() # type: ignore[assignment] tooltip: str = attrs.field() default_args: dict[str, Any] = attrs.field(factory=dict) diff --git a/airflow-core/src/airflow/serialization/serialized_objects.py b/airflow-core/src/airflow/serialization/serialized_objects.py index 280b5907a89b4..ca9b6ba5f9a80 100644 --- a/airflow-core/src/airflow/serialization/serialized_objects.py +++ b/airflow-core/src/airflow/serialization/serialized_objects.py @@ -58,7 +58,13 @@ from airflow._shared.timezones.timezone import coerce_datetime, from_timestamp, parse_timezone, utcnow from airflow.callbacks.callback_requests import DagCallbackRequest, TaskCallbackRequest from airflow.configuration import conf as airflow_conf -from airflow.exceptions import AirflowException, DeserializationError, SerializationError, TaskDeferred +from airflow.exceptions import ( + AirflowException, + DeserializationError, + SerializationError, + TaskDeferred, + TaskNotFound, +) from airflow.models.connection import Connection from airflow.models.dag import DagModel from airflow.models.dag_version import DagVersion @@ -125,6 +131,7 @@ from airflow.models.mappedoperator import MappedOperator as SerializedMappedOperator from airflow.models.taskinstance import TaskInstance from airflow.sdk import BaseOperatorLink + from airflow.sdk.definitions.edges import EdgeInfoType from airflow.serialization.json_schema import Validator from airflow.task.trigger_rule import TriggerRule from airflow.ti_deps.deps.base_ti_dep import BaseTIDep @@ -1354,7 +1361,8 @@ def __eq__(self, other: Any) -> bool: def node_id(self) -> str: return self.task_id - def get_dag(self) -> DAG | None: + # TODO (GH-52141): Replace DAGNode with a scheduler type. + def get_dag(self) -> SerializedDAG | None: # type: ignore[override] return self.dag @property @@ -2331,7 +2339,7 @@ def _create_orm_dagrun( return run -class SerializedDAG(DAG, BaseSerialization): +class SerializedDAG(BaseSerialization): """ A JSON serializable representation of DAG. @@ -2342,10 +2350,31 @@ class SerializedDAG(DAG, BaseSerialization): _decorated_fields: ClassVar[set[str]] = {"default_args", "access_control"} - # TODO (GH-52141): These should contain serialized containers, but currently - # this class inherits from an SDK one. - task_group: SerializedTaskGroup # type: ignore[assignment] - task_dict: dict[str, SerializedBaseOperator | SerializedMappedOperator] # type: ignore[assignment] + access_control: dict[str, dict[str, Collection[str]]] | None + catchup: bool + dag_id: str + dagrun_timeout: datetime.timedelta | None + deadline: list[DeadlineAlert] | DeadlineAlert | None + default_args: dict[str, Any] + description: str | None + disable_bundle_versioning: bool + doc_md: str | None + edge_info: dict[str, dict[str, EdgeInfoType]] + end_date: datetime.datetime | None + fail_fast: bool + has_on_failure_callback: bool + has_on_success_callback: bool + max_active_runs: int + max_active_tasks: int + max_consecutive_failed_dag_runs: int + params: ParamsDict # TODO (GH-52141): These should use scheduler-specific types. + partial: bool + start_date: datetime.datetime | None + tags: set[str] + task_dict: dict[str, SerializedBaseOperator | SerializedMappedOperator] + task_group: SerializedTaskGroup + timetable: Timetable + timezone: FixedTimezone | Timezone last_loaded: datetime.datetime # this will only be set at serialization time @@ -2368,6 +2397,39 @@ def __get_constructor_defaults(): _json_schema: ClassVar[Validator] = lazy_object_proxy.Proxy(load_dag_schema) + @classmethod + def get_serialized_fields(cls) -> frozenset[str]: + return frozenset( + { + "access_control", + "catchup", + "dag_display_name", + "dag_id", + "dagrun_timeout", + "deadline", + "default_args", + "description", + "disable_bundle_versioning", + "doc_md", + "edge_info", + "end_date", + "fail_fast", + "fileloc", + "is_paused_upon_creation", + "max_active_runs", + "max_active_tasks", + "max_consecutive_failed_dag_runs", + "owner_links", + "relative_fileloc", + "render_template_as_native_obj", + "start_date", + "tags", + "task_group", + "timetable", + "timezone", + } + ) + @classmethod def serialize_dag(cls, dag: DAG) -> dict: """Serialize a DAG into a JSON object.""" @@ -2413,7 +2475,8 @@ def deserialize_dag( """Deserializes a DAG from a JSON object.""" if "dag_id" not in encoded_dag: raise DeserializationError( - message="Encoded dag object has no dag_id key. You may need to run `airflow dags reserialize`." + message="Encoded dag object has no dag_id key. " + "You may need to run `airflow dags reserialize`." ) dag_id = encoded_dag["dag_id"] @@ -2432,7 +2495,8 @@ def _deserialize_dag_internal( cls, encoded_dag: dict[str, Any], client_defaults: dict[str, Any] | None = None ) -> SerializedDAG: """Handle the main Dag deserialization logic.""" - dag = SerializedDAG(dag_id=encoded_dag["dag_id"], schedule=None) + dag = SerializedDAG() + dag.dag_id = encoded_dag["dag_id"] dag.last_loaded = utcnow() # Note: Context is passed explicitly through method parameters, no class attributes needed @@ -2747,12 +2811,34 @@ def bulk_write_to_db( dag_op.update_dag_asset_expression(orm_dags=orm_dags, orm_assets=orm_assets) session.flush() - # TODO (GH-52141): This needs to take scheduler types, but currently it inherits SDK's DAG. - # TODO (GH-52141): This shouldn't need to be writable, but SDK's DAG defines it as such. - @property # type: ignore[misc] - def tasks(self) -> Sequence[SerializedOperator]: # type: ignore[override] + @property + def tasks(self) -> Sequence[SerializedOperator]: return list(self.task_dict.values()) + @property + def task_ids(self) -> list[str]: + return list(self.task_dict) + + @property + def roots(self) -> list[SerializedOperator]: + return [task for task in self.tasks if not task.upstream_list] + + @property + def owner(self) -> str: + return ", ".join({t.owner for t in self.tasks}) + + def has_task(self, task_id: str) -> bool: + return task_id in self.task_dict + + def get_task(self, task_id: str) -> SerializedOperator: + if task_id in self.task_dict: + return self.task_dict[task_id] + raise TaskNotFound(f"Task {task_id} not found") + + @property + def task_group_dict(self): + return {k: v for k, v in self.task_group.get_task_group_dict().items() if k is not None} + def partial_subset( self, task_ids: str | Iterable[str], @@ -3156,9 +3242,7 @@ def set_task_instance_state( """ from airflow.api.common.mark_tasks import set_state - # TODO (GH-52141): get_task in scheduler needs to return scheduler types - # instead, but currently it inherits SDK's DAG. - task = cast("SerializedOperator", self.get_task(task_id)) + task = self.get_task(task_id) task.dag = self tasks_to_set_state: list[SerializedOperator | tuple[SerializedOperator, int]] @@ -3536,6 +3620,14 @@ def _coerce_dag(dag): for dag in dags ) + def get_edge_info(self, upstream_task_id: str, downstream_task_id: str) -> EdgeInfoType: + """Return edge information for the given pair of tasks or an empty edge if there is no information.""" + # Note - older serialized Dags may not have edge_info being a dict at all + empty = cast("EdgeInfoType", {}) + if self.edge_info: + return self.edge_info.get(upstream_task_id, {}).get(downstream_task_id, empty) + return empty + class TaskGroupSerialization(BaseSerialization): """JSON serializable representation of a task group.""" diff --git a/airflow-core/src/airflow/ti_deps/dep_context.py b/airflow-core/src/airflow/ti_deps/dep_context.py index 056b633f36122..1feafdd041ae1 100644 --- a/airflow-core/src/airflow/ti_deps/dep_context.py +++ b/airflow-core/src/airflow/ti_deps/dep_context.py @@ -18,7 +18,7 @@ from __future__ import annotations import contextlib -from typing import TYPE_CHECKING, cast +from typing import TYPE_CHECKING import attr @@ -29,9 +29,7 @@ from sqlalchemy.orm.session import Session from airflow.models.dagrun import DagRun - from airflow.models.mappedoperator import MappedOperator from airflow.models.taskinstance import TaskInstance - from airflow.serialization.serialized_objects import SerializedBaseOperator @attr.define @@ -100,9 +98,7 @@ def ensure_finished_tis(self, dag_run: DagRun, session: Session) -> list[TaskIns if getattr(ti, "task", None) is not None or (dag := dag_run.dag) is None: continue with contextlib.suppress(TaskNotFound): - # TODO (GH-52141): get_task in scheduler should contain scheduler - # types instead, but currently it inherits SDK's DAG. - ti.task = cast("MappedOperator | SerializedBaseOperator", dag.get_task(ti.task_id)) + ti.task = dag.get_task(ti.task_id) self.finished_tis = finished_tis else: finished_tis = self.finished_tis diff --git a/airflow-core/src/airflow/utils/dag_edges.py b/airflow-core/src/airflow/utils/dag_edges.py index f8ed846f25532..94c6069f91b02 100644 --- a/airflow-core/src/airflow/utils/dag_edges.py +++ b/airflow-core/src/airflow/utils/dag_edges.py @@ -20,7 +20,7 @@ from airflow.models.mappedoperator import MappedOperator from airflow.sdk.definitions._internal.abstractoperator import AbstractOperator -from airflow.serialization.serialized_objects import SerializedBaseOperator +from airflow.serialization.serialized_objects import SerializedBaseOperator, SerializedDAG if TYPE_CHECKING: from collections.abc import Iterable @@ -30,7 +30,7 @@ Operator: TypeAlias = MappedOperator | SerializedBaseOperator -def dag_edges(dag: DAG): +def dag_edges(dag: DAG | SerializedDAG): """ Create the list of edges needed to construct the Graph view. diff --git a/airflow-core/src/airflow/utils/dot_renderer.py b/airflow-core/src/airflow/utils/dot_renderer.py index 259b4ced25204..caa5e27a2ffc1 100644 --- a/airflow-core/src/airflow/utils/dot_renderer.py +++ b/airflow-core/src/airflow/utils/dot_renderer.py @@ -28,7 +28,7 @@ from airflow.sdk import DAG, BaseOperator, TaskGroup from airflow.sdk.definitions.mappedoperator import MappedOperator from airflow.serialization.definitions.taskgroup import SerializedTaskGroup -from airflow.serialization.serialized_objects import SerializedBaseOperator +from airflow.serialization.serialized_objects import SerializedBaseOperator, SerializedDAG from airflow.utils.dag_edges import dag_edges from airflow.utils.state import State @@ -194,7 +194,7 @@ def render_dag_dependencies(deps: dict[str, list[DagDependency]]) -> graphviz.Di return dot -def render_dag(dag: DAG, tis: list[TaskInstance] | None = None) -> graphviz.Digraph: +def render_dag(dag: DAG | SerializedDAG, tis: list[TaskInstance] | None = None) -> graphviz.Digraph: """ Render the DAG object to the DOT object. diff --git a/airflow-core/tests/unit/models/test_dag.py b/airflow-core/tests/unit/models/test_dag.py index ba22ed2aaddc7..22953cb7cee9a 100644 --- a/airflow-core/tests/unit/models/test_dag.py +++ b/airflow-core/tests/unit/models/test_dag.py @@ -2242,12 +2242,7 @@ def test_asset_expression(self, session: Session, testing_dag_bundle) -> None: ), start_date=datetime.datetime.min, ) - SerializedDAG.bulk_write_to_db( - "testing", - None, - [SerializedDAG.deserialize_dag(SerializedDAG.serialize_dag(dag))], - session=session, - ) + SerializedDAG.bulk_write_to_db("testing", None, [dag], session=session) expression = session.scalars(select(DagModel.asset_expression).filter_by(dag_id=dag.dag_id)).one() assert expression == { diff --git a/airflow-core/tests/unit/serialization/test_dag_serialization.py b/airflow-core/tests/unit/serialization/test_dag_serialization.py index 82f9323f1807c..7b2d43f8a4d44 100644 --- a/airflow-core/tests/unit/serialization/test_dag_serialization.py +++ b/airflow-core/tests/unit/serialization/test_dag_serialization.py @@ -1097,8 +1097,8 @@ def test_full_param_roundtrip(self, param: Param): """ Test to make sure that only native Param objects are being passed as dag or task params """ - dag = DAG(dag_id="simple_dag", schedule=None, params={"my_param": param}) - serialized_json = SerializedDAG.to_json(dag) + sdk_dag = DAG(dag_id="simple_dag", schedule=None, params={"my_param": param}) + serialized_json = SerializedDAG.to_json(sdk_dag) serialized = json.loads(serialized_json) SerializedDAG.validate_schema(serialized) dag = SerializedDAG.from_dict(serialized) diff --git a/providers/fab/src/airflow/providers/fab/auth_manager/security_manager/override.py b/providers/fab/src/airflow/providers/fab/auth_manager/security_manager/override.py index 2fa5c0c5b4b1f..488d57b66c4e8 100644 --- a/providers/fab/src/airflow/providers/fab/auth_manager/security_manager/override.py +++ b/providers/fab/src/airflow/providers/fab/auth_manager/security_manager/override.py @@ -112,6 +112,7 @@ RESOURCE_ASSET_ALIAS, ) from airflow.sdk import DAG + from airflow.serialization.serialized_objects import SerializedDAG else: from airflow.providers.common.compat.security.permissions import ( RESOURCE_ASSET, @@ -122,13 +123,13 @@ from airflow.models.dagbag import DBDagBag from airflow.utils.session import create_session - def _iter_dags() -> Iterable[DAG]: + def _iter_dags() -> Iterable[DAG | SerializedDAG]: with create_session() as session: yield from DBDagBag().iter_all_latest_version_dags(session=session) else: from airflow.models.dagbag import DagBag - def _iter_dags() -> Iterable[DAG]: + def _iter_dags() -> Iterable[DAG | SerializedDAG]: dagbag = DagBag(read_dags_from_db=True) # type: ignore[call-arg] dagbag.collect_dags_from_db() # type: ignore[attr-defined] return dagbag.dags.values() diff --git a/providers/openlineage/src/airflow/providers/openlineage/utils/selective_enable.py b/providers/openlineage/src/airflow/providers/openlineage/utils/selective_enable.py index 30dc825d4394f..6f7c08f6a4ae5 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/utils/selective_enable.py +++ b/providers/openlineage/src/airflow/providers/openlineage/utils/selective_enable.py @@ -24,9 +24,14 @@ from airflow.models.xcom_arg import XComArg if TYPE_CHECKING: + from typing import TypeAlias + + from airflow.models.mappedoperator import MappedOperator as SerializedMappedOperator from airflow.sdk import DAG, BaseOperator from airflow.sdk.definitions.mappedoperator import MappedOperator + from airflow.serialization.serialized_objects import SerializedBaseOperator, SerializedDAG + AnyOperator: TypeAlias = BaseOperator | MappedOperator | SerializedBaseOperator | SerializedMappedOperator T = TypeVar("T", bound=DAG | BaseOperator | MappedOperator) else: try: @@ -76,7 +81,7 @@ def disable_lineage(obj: T) -> T: return obj -def is_task_lineage_enabled(task: BaseOperator | MappedOperator) -> bool: +def is_task_lineage_enabled(task: AnyOperator) -> bool: """Check if selective enable OpenLineage parameter is set to True on task level.""" if task.params.get(ENABLE_OL_PARAM_NAME) is False: log.debug( @@ -85,7 +90,7 @@ def is_task_lineage_enabled(task: BaseOperator | MappedOperator) -> bool: return task.params.get(ENABLE_OL_PARAM_NAME) is True -def is_dag_lineage_enabled(dag: DAG) -> bool: +def is_dag_lineage_enabled(dag: DAG | SerializedDAG) -> bool: """ Check if DAG is selectively enabled to emit OpenLineage events. diff --git a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py index b9b66b5f42953..2a425266870f7 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py +++ b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py @@ -34,6 +34,7 @@ # TODO: move this maybe to Airflow's logic? from airflow.models import DagRun, TaskReschedule +from airflow.models.mappedoperator import MappedOperator as SerializedMappedOperator from airflow.providers.openlineage import ( __version__ as OPENLINEAGE_PROVIDER_VERSION, conf, @@ -53,7 +54,7 @@ is_task_lineage_enabled, ) from airflow.providers.openlineage.version_compat import AIRFLOW_V_3_0_PLUS, get_base_airflow_version_tuple -from airflow.serialization.serialized_objects import SerializedBaseOperator +from airflow.serialization.serialized_objects import SerializedBaseOperator, SerializedDAG from airflow.utils.module_loading import import_string if AIRFLOW_V_3_0_PLUS: @@ -65,6 +66,8 @@ from airflow.utils.session import NEW_SESSION, provide_session if TYPE_CHECKING: + from typing import TypeAlias + from openlineage.client.event_v2 import Dataset as OpenLineageDataset from openlineage.client.facet_v2 import RunFacet, processing_engine_run @@ -79,6 +82,8 @@ ) from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance from airflow.utils.state import DagRunState, TaskInstanceState + + AnyOperator: TypeAlias = BaseOperator | MappedOperator | SerializedBaseOperator | SerializedMappedOperator else: try: from airflow.sdk import DAG, BaseOperator @@ -128,13 +133,13 @@ def try_import_from_string(string: str) -> Any: return import_string(string) -def get_operator_class(task: BaseOperator) -> type: +def get_operator_class(task: BaseOperator | MappedOperator) -> type: if task.__class__.__name__ in ("DecoratedMappedOperator", "MappedOperator"): return task.operator_class return task.__class__ -def get_operator_provider_version(operator: BaseOperator | MappedOperator) -> str | None: +def get_operator_provider_version(operator: AnyOperator) -> str | None: """Get the provider package version for the given operator.""" try: class_path = get_fully_qualified_class_name(operator) @@ -214,7 +219,7 @@ def _truncate_string_to_byte_size(s: str, max_size: int = _MAX_DOC_BYTES) -> str return truncated.decode("utf-8", errors="ignore") -def get_task_documentation(operator: BaseOperator | MappedOperator | None) -> tuple[str | None, str | None]: +def get_task_documentation(operator: AnyOperator | None) -> tuple[str | None, str | None]: """Get task documentation and mime type, truncated to _MAX_DOC_BYTES bytes length, if present.""" if not operator: return None, None @@ -241,7 +246,7 @@ def get_task_documentation(operator: BaseOperator | MappedOperator | None) -> tu return None, None -def get_dag_documentation(dag: DAG | None) -> tuple[str | None, str | None]: +def get_dag_documentation(dag: DAG | SerializedDAG | None) -> tuple[str | None, str | None]: """Get dag documentation and mime type, truncated to _MAX_DOC_BYTES bytes length, if present.""" if not dag: return None, None @@ -312,23 +317,23 @@ def get_user_provided_run_facets(ti: TaskInstance, ti_state: TaskInstanceState) return custom_facets -def get_fully_qualified_class_name(operator: BaseOperator | MappedOperator) -> str: - if isinstance(operator, (MappedOperator, SerializedBaseOperator)): +def get_fully_qualified_class_name(operator: AnyOperator) -> str: + if isinstance(operator, (SerializedMappedOperator, SerializedBaseOperator)): # as in airflow.api_connexion.schemas.common_schema.ClassReferenceSchema return operator._task_module + "." + operator.task_type op_class = get_operator_class(operator) return op_class.__module__ + "." + op_class.__name__ -def is_operator_disabled(operator: BaseOperator | MappedOperator) -> bool: +def is_operator_disabled(operator: AnyOperator) -> bool: return get_fully_qualified_class_name(operator) in conf.disabled_operators() -def is_selective_lineage_enabled(obj: DAG | BaseOperator | MappedOperator) -> bool: +def is_selective_lineage_enabled(obj: DAG | SerializedDAG | AnyOperator) -> bool: """If selective enable is active check if DAG or Task is enabled to emit events.""" if not conf.selective_enable(): return True - if isinstance(obj, DAG): + if isinstance(obj, (DAG, SerializedDAG)): return is_dag_lineage_enabled(obj) if isinstance(obj, (BaseOperator, MappedOperator)): return is_task_lineage_enabled(obj) @@ -749,7 +754,7 @@ def get_airflow_state_run_facet( } -def _get_tasks_details(dag: DAG) -> dict: +def _get_tasks_details(dag: DAG | SerializedDAG) -> dict: tasks = { single_task.task_id: { "operator": get_fully_qualified_class_name(single_task), @@ -768,7 +773,7 @@ def _get_tasks_details(dag: DAG) -> dict: return tasks -def _get_task_groups_details(dag: DAG) -> dict: +def _get_task_groups_details(dag: DAG | SerializedDAG) -> dict: return { tg_id: { "parent_group": tg.parent_group.group_id, @@ -780,7 +785,9 @@ def _get_task_groups_details(dag: DAG) -> dict: } -def _emits_ol_events(task: BaseOperator | MappedOperator) -> bool: +def _emits_ol_events( + task: BaseOperator | MappedOperator | SerializedBaseOperator | SerializedMappedOperator, +) -> bool: config_selective_enabled = is_selective_lineage_enabled(task) config_disabled_for_operators = is_operator_disabled(task) # empty operators without callbacks/outlets are skipped for optimization by Airflow diff --git a/providers/standard/tests/unit/standard/sensors/test_external_task_sensor.py b/providers/standard/tests/unit/standard/sensors/test_external_task_sensor.py index a5a068f088f7c..2c6d230e1740a 100644 --- a/providers/standard/tests/unit/standard/sensors/test_external_task_sensor.py +++ b/providers/standard/tests/unit/standard/sensors/test_external_task_sensor.py @@ -1917,7 +1917,8 @@ def _factory(depth: int) -> DagBag: for dag in dags: if AIRFLOW_V_3_0_PLUS: - dag_bag.bag_dag(dag=sync_dag_to_db(dag)) + sync_dag_to_db(dag) + dag_bag.bag_dag(dag=dag) else: dag_bag.bag_dag(dag=dag, root_dag=dag) # type: ignore[call-arg] From ce491c5f542a75ae5996d19c00fd060bad05b30b Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Thu, 11 Sep 2025 15:25:28 -0600 Subject: [PATCH 2/2] Remove SDK dependency from SerializedDAG --- .../api_fastapi/core_api/datamodels/dags.py | 7 ++- .../serialization/serialized_objects.py | 44 ++++++++++++++++--- airflow-core/tests/unit/models/test_dag.py | 4 +- airflow-core/tests/unit/models/test_dagrun.py | 15 ++++--- .../serialization/test_dag_serialization.py | 2 +- .../openlineage/utils/selective_enable.py | 7 +-- .../providers/openlineage/utils/utils.py | 4 +- .../unit/standard/operators/test_datetime.py | 10 +++-- task-sdk/src/airflow/sdk/bases/operator.py | 6 +-- task-sdk/src/airflow/sdk/definitions/dag.py | 7 +-- 10 files changed, 71 insertions(+), 35 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/dags.py b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/dags.py index b14b8b18cbdc3..23aa4ec300f47 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/dags.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/dags.py @@ -19,7 +19,6 @@ import inspect from collections import abc -from collections.abc import Iterable from datetime import datetime, timedelta from typing import Any @@ -151,9 +150,9 @@ class DAGDetailsResponse(DAGResponse): start_date: datetime | None end_date: datetime | None is_paused_upon_creation: bool | None - params: abc.MutableMapping | None + params: abc.Mapping | None render_template_as_native_obj: bool - template_search_path: Iterable[str] | None + template_search_path: list[str] | None timezone: str | None last_parsed: datetime | None default_args: abc.Mapping | None @@ -177,7 +176,7 @@ def get_doc_md(cls, doc_md: str | None) -> str | None: @field_validator("params", mode="before") @classmethod - def get_params(cls, params: abc.MutableMapping | None) -> dict | None: + def get_params(cls, params: abc.Mapping | None) -> dict | None: """Convert params attribute to dict representation.""" if params is None: return None diff --git a/airflow-core/src/airflow/serialization/serialized_objects.py b/airflow-core/src/airflow/serialization/serialized_objects.py index ca9b6ba5f9a80..37c09b4a9254e 100644 --- a/airflow-core/src/airflow/serialization/serialized_objects.py +++ b/airflow-core/src/airflow/serialization/serialized_objects.py @@ -2350,9 +2350,10 @@ class SerializedDAG(BaseSerialization): _decorated_fields: ClassVar[set[str]] = {"default_args", "access_control"} - access_control: dict[str, dict[str, Collection[str]]] | None + access_control: dict[str, dict[str, Collection[str]]] | None = None catchup: bool dag_id: str + dag_display_name: str dagrun_timeout: datetime.timedelta | None deadline: list[DeadlineAlert] | DeadlineAlert | None default_args: dict[str, Any] @@ -2364,15 +2365,19 @@ class SerializedDAG(BaseSerialization): fail_fast: bool has_on_failure_callback: bool has_on_success_callback: bool + is_paused_upon_creation: bool | None max_active_runs: int max_active_tasks: int max_consecutive_failed_dag_runs: int - params: ParamsDict # TODO (GH-52141): These should use scheduler-specific types. + owner_links: dict[str, str] + params: ParamsDict # TODO (GH-52141): Should use a scheduler-specific type. partial: bool + render_template_as_native_obj: bool start_date: datetime.datetime | None tags: set[str] - task_dict: dict[str, SerializedBaseOperator | SerializedMappedOperator] + task_dict: dict[str, SerializedOperator] task_group: SerializedTaskGroup + template_searchpath: tuple[str, ...] | None timetable: Timetable timezone: FixedTimezone | Timezone @@ -2381,6 +2386,34 @@ class SerializedDAG(BaseSerialization): # it's only use is for determining the relative fileloc based only on the serialize dag _processor_dags_folder: str + def __init__(self, *, dag_id: str) -> None: + self.catchup = airflow_conf.getboolean("scheduler", "catchup_by_default") + self.dag_id = self.dag_display_name = dag_id + self.dagrun_timeout = None + self.deadline = None + self.default_args = {} + self.description = None + self.disable_bundle_versioning = airflow_conf.getboolean("dag_processor", "disable_bundle_versioning") + self.doc_md = None + self.edge_info = {} + self.end_date = None + self.fail_fast = False + self.has_on_failure_callback = False + self.has_on_success_callback = False + self.is_paused_upon_creation = None + self.max_active_runs = airflow_conf.getint("core", "max_active_runs_per_dag") + self.max_active_tasks = airflow_conf.getint("core", "max_active_tasks_per_dag") + self.max_consecutive_failed_dag_runs = airflow_conf.getint( + "core", "max_consecutive_failed_dag_runs_per_dag" + ) + self.owner_links = {} + self.params = ParamsDict() + self.partial = False + self.render_template_as_native_obj = False + self.start_date = None + self.tags = set() + self.template_searchpath = None + @staticmethod def __get_constructor_defaults(): param_to_attr = { @@ -2495,8 +2528,7 @@ def _deserialize_dag_internal( cls, encoded_dag: dict[str, Any], client_defaults: dict[str, Any] | None = None ) -> SerializedDAG: """Handle the main Dag deserialization logic.""" - dag = SerializedDAG() - dag.dag_id = encoded_dag["dag_id"] + dag = SerializedDAG(dag_id=encoded_dag["dag_id"]) dag.last_loaded = utcnow() # Note: Context is passed explicitly through method parameters, no class attributes needed @@ -3622,7 +3654,7 @@ def _coerce_dag(dag): def get_edge_info(self, upstream_task_id: str, downstream_task_id: str) -> EdgeInfoType: """Return edge information for the given pair of tasks or an empty edge if there is no information.""" - # Note - older serialized Dags may not have edge_info being a dict at all + # Note - older serialized dags may not have edge_info being a dict at all empty = cast("EdgeInfoType", {}) if self.edge_info: return self.edge_info.get(upstream_task_id, {}).get(downstream_task_id, empty) diff --git a/airflow-core/tests/unit/models/test_dag.py b/airflow-core/tests/unit/models/test_dag.py index 22953cb7cee9a..53cafab354334 100644 --- a/airflow-core/tests/unit/models/test_dag.py +++ b/airflow-core/tests/unit/models/test_dag.py @@ -925,8 +925,8 @@ def test_dag_handle_callback_with_removed_task(self, dag_maker, session, testing assert dag_run.get_task_instance(task_removed.task_id).state == TaskInstanceState.REMOVED # should not raise any exception - dag_run.handle_dag_callback(dag=scheduler_dag, success=False) - dag_run.handle_dag_callback(dag=scheduler_dag, success=True) + dag_run.handle_dag_callback(dag=dag, success=False) + dag_run.handle_dag_callback(dag=dag, success=True) @pytest.mark.parametrize("catchup,expected_next_dagrun", [(True, DEFAULT_DATE), (False, None)]) def test_next_dagrun_after_fake_scheduled_previous( diff --git a/airflow-core/tests/unit/models/test_dagrun.py b/airflow-core/tests/unit/models/test_dagrun.py index 369f530077b4a..73284f673563a 100644 --- a/airflow-core/tests/unit/models/test_dagrun.py +++ b/airflow-core/tests/unit/models/test_dagrun.py @@ -401,7 +401,9 @@ def on_success_callable(context): } dag_run = self.create_dag_run(dag=dag, task_states=initial_task_states, session=session) - _, callback = dag_run.update_state() + with mock.patch.object(dag_run, "handle_dag_callback") as handle_dag_callback: + _, callback = dag_run.update_state() + assert handle_dag_callback.mock_calls == [mock.call(dag=dag, success=True, reason="success")] assert dag_run.state == DagRunState.SUCCESS # Callbacks are not added until handle_callback = False is passed to dag_run.update_state() assert callback is None @@ -426,7 +428,9 @@ def on_failure_callable(context): dag_task1.set_downstream(dag_task2) dag_run = self.create_dag_run(dag=dag, task_states=initial_task_states, session=session) - _, callback = dag_run.update_state() + with mock.patch.object(dag_run, "handle_dag_callback") as handle_dag_callback: + _, callback = dag_run.update_state() + assert handle_dag_callback.mock_calls == [mock.call(dag=dag, success=False, reason="task_failure")] assert dag_run.state == DagRunState.FAILED # Callbacks are not added until handle_callback = False is passed to dag_run.update_state() assert callback is None @@ -790,8 +794,7 @@ def test_get_task_instance_on_empty_dagrun(self, dag_maker, session): schedule=datetime.timedelta(days=1), start_date=timezone.datetime(2017, 1, 1), ) as dag: - ... - ShortCircuitOperator(task_id="test_short_circuit_false", dag=dag, python_callable=lambda: False) + ShortCircuitOperator(task_id="test_short_circuit_false", python_callable=lambda: False) now = timezone.utcnow() @@ -1282,7 +1285,9 @@ def on_success_callable(context): dag_run = session.merge(dag_run) dag_run.dag = dag - _, callback = dag_run.update_state() + with mock.patch.object(dag_run, "handle_dag_callback") as handle_dag_callback: + _, callback = dag_run.update_state() + assert handle_dag_callback.mock_calls == [mock.call(dag=dag, success=True, reason="success")] assert dag_run.state == DagRunState.SUCCESS # Callbacks are not added until handle_callback = False is passed to dag_run.update_state() assert callback is None diff --git a/airflow-core/tests/unit/serialization/test_dag_serialization.py b/airflow-core/tests/unit/serialization/test_dag_serialization.py index 7b2d43f8a4d44..5ff0a42df4953 100644 --- a/airflow-core/tests/unit/serialization/test_dag_serialization.py +++ b/airflow-core/tests/unit/serialization/test_dag_serialization.py @@ -623,7 +623,7 @@ def test_deserialization_across_process(self): if v is None: break dag = SerializedDAG.from_json(v) - assert isinstance(dag, DAG) + assert isinstance(dag, SerializedDAG) stringified_dags[dag.dag_id] = dag dags, _ = collect_dags("airflow/example_dags") diff --git a/providers/openlineage/src/airflow/providers/openlineage/utils/selective_enable.py b/providers/openlineage/src/airflow/providers/openlineage/utils/selective_enable.py index 6f7c08f6a4ae5..edf6f9c858f84 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/utils/selective_enable.py +++ b/providers/openlineage/src/airflow/providers/openlineage/utils/selective_enable.py @@ -24,14 +24,11 @@ from airflow.models.xcom_arg import XComArg if TYPE_CHECKING: - from typing import TypeAlias - - from airflow.models.mappedoperator import MappedOperator as SerializedMappedOperator + from airflow.providers.openlineage.utils.utils import AnyOperator from airflow.sdk import DAG, BaseOperator from airflow.sdk.definitions.mappedoperator import MappedOperator - from airflow.serialization.serialized_objects import SerializedBaseOperator, SerializedDAG + from airflow.serialization.serialized_objects import SerializedDAG - AnyOperator: TypeAlias = BaseOperator | MappedOperator | SerializedBaseOperator | SerializedMappedOperator T = TypeVar("T", bound=DAG | BaseOperator | MappedOperator) else: try: diff --git a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py index 2a425266870f7..d08868961fe5c 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py +++ b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py @@ -785,9 +785,7 @@ def _get_task_groups_details(dag: DAG | SerializedDAG) -> dict: } -def _emits_ol_events( - task: BaseOperator | MappedOperator | SerializedBaseOperator | SerializedMappedOperator, -) -> bool: +def _emits_ol_events(task: AnyOperator) -> bool: config_selective_enabled = is_selective_lineage_enabled(task) config_disabled_for_operators = is_operator_disabled(task) # empty operators without callbacks/outlets are skipped for optimization by Airflow diff --git a/providers/standard/tests/unit/standard/operators/test_datetime.py b/providers/standard/tests/unit/standard/operators/test_datetime.py index afc1d8ff012e2..0630a7e187912 100644 --- a/providers/standard/tests/unit/standard/operators/test_datetime.py +++ b/providers/standard/tests/unit/standard/operators/test_datetime.py @@ -28,11 +28,15 @@ from airflow.providers.standard.operators.datetime import BranchDateTimeOperator from airflow.providers.standard.operators.empty import EmptyOperator from airflow.timetables.base import DataInterval -from airflow.utils import timezone from airflow.utils.session import create_session from airflow.utils.state import State -from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_1, AIRFLOW_V_3_0_PLUS +from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_1, AIRFLOW_V_3_0_PLUS, AIRFLOW_V_3_1_PLUS + +if AIRFLOW_V_3_1_PLUS: + from airflow.sdk import timezone +else: + from airflow.utils import timezone # type: ignore[attr-defined,no-redef] pytestmark = pytest.mark.db_test @@ -113,7 +117,7 @@ def test_no_target_time(self): follow_task_ids_if_false="branch_2", target_upper=None, target_lower=None, - dag=self.dag, + dag=self.dag_maker.dag, ) @pytest.mark.parametrize( diff --git a/task-sdk/src/airflow/sdk/bases/operator.py b/task-sdk/src/airflow/sdk/bases/operator.py index 7a8e80ec4d98a..2cc891d106b7a 100644 --- a/task-sdk/src/airflow/sdk/bases/operator.py +++ b/task-sdk/src/airflow/sdk/bases/operator.py @@ -1309,12 +1309,12 @@ def _convert__dag(self, dag: DAG | None) -> DAG | None: return dag if not isinstance(dag, DAG): - raise TypeError(f"Expected Dag; received {dag.__class__.__name__}") + raise TypeError(f"Expected dag; received {dag.__class__.__name__}") if self._dag is not None and self._dag is not dag: - raise ValueError(f"The Dag assigned to {self} can not be changed.") + raise ValueError(f"The dag assigned to {self} can not be changed.") if self.__from_mapped: - pass # Don't add to Dag -- the mapped task takes the place. + pass # Don't add to dag -- the mapped task takes the place. elif dag.task_dict.get(self.task_id) is not self: dag.add_task(self) return dag diff --git a/task-sdk/src/airflow/sdk/definitions/dag.py b/task-sdk/src/airflow/sdk/definitions/dag.py index ed14b7dd5878d..7448b321389cb 100644 --- a/task-sdk/src/airflow/sdk/definitions/dag.py +++ b/task-sdk/src/airflow/sdk/definitions/dag.py @@ -1087,7 +1087,6 @@ def get_serialized_fields(cls): def get_edge_info(self, upstream_task_id: str, downstream_task_id: str) -> EdgeInfoType: """Return edge information for the given pair of tasks or an empty edge if there is no information.""" - # Note - older serialized Dags may not have edge_info being a dict at all empty = cast("EdgeInfoType", {}) if self.edge_info: return self.edge_info.get(upstream_task_id, {}).get(downstream_task_id, empty) @@ -1192,8 +1191,10 @@ def test( # -- dep check, scheduling tis # and need real dag to get and run callbacks without having to load the dag model - scheduler_dag.on_success_callback = self.on_success_callback - scheduler_dag.on_failure_callback = self.on_failure_callback + # Scheduler DAG shouldn't have these attributes, but assigning them + # here is an easy hack to get this test() thing working. + scheduler_dag.on_success_callback = self.on_success_callback # type: ignore[attr-defined] + scheduler_dag.on_failure_callback = self.on_failure_callback # type: ignore[attr-defined] dr: DagRun = get_or_create_dagrun( dag=scheduler_dag,