Create DAGNode in Core

airflow-core/pyproject.toml

@@ -226,8 +226,9 @@ exclude = [
 
 [tool.hatch.build.targets.sdist.force-include]
 "../shared/configuration/src/airflow_shared/configuration" = "src/airflow/_shared/configuration"
-"../shared/module_loading/src/airflow_shared/module_loading" = "src/airflow/_shared/module_loading"
+"../shared/dagnode/src/airflow_shared/dagnode" = "src/airflow/_shared/dagnode"
 "../shared/logging/src/airflow_shared/logging" = "src/airflow/_shared/logging"
+"../shared/module_loading/src/airflow_shared/module_loading" = "src/airflow/_shared/module_loading"
 "../shared/observability/src/airflow_shared/observability" = "src/airflow/_shared/observability"
 "../shared/secrets_backend/src/airflow_shared/secrets_backend" = "src/airflow/_shared/secrets_backend"
 "../shared/secrets_masker/src/airflow_shared/secrets_masker" = "src/airflow/_shared/secrets_masker"
@@ -300,10 +301,11 @@ apache-airflow-devel-common = { workspace = true }
 [tool.airflow]
 shared_distributions = [
     "apache-airflow-shared-configuration",
+    "apache-airflow-shared-dagnode",
     "apache-airflow-shared-logging",
     "apache-airflow-shared-module-loading",
+    "apache-airflow-shared-observability",
     "apache-airflow-shared-secrets-backend",
     "apache-airflow-shared-secrets-masker",
     "apache-airflow-shared-timezones",
-    "apache-airflow-shared-observability",
 ]

airflow-core/src/airflow/_shared/dagnode

@@ -0,0 +1 @@
+../../../../shared/dagnode/src/airflow_shared/dagnode

airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py

@@ -77,8 +77,6 @@
 if TYPE_CHECKING:
     from sqlalchemy.sql.dml import Update
 
-    from airflow.models.expandinput import SchedulerExpandInput
-
 router = VersionedAPIRouter()
 
 ti_id_router = VersionedAPIRouter(
@@ -318,7 +316,7 @@ def _get_upstream_map_indexes(
             except NotFullyPopulated:
                 # Second try: resolve XCom for correct count
                 try:
-                    expand_input = cast("SchedulerExpandInput", upstream_mapped_group._expand_input)
+                    expand_input = upstream_mapped_group._expand_input
                     mapped_ti_count = expand_input.get_total_map_length(ti.run_id, session=session)
                 except NotFullyPopulated:
                     # For these trigger rules, unresolved map indexes are acceptable.

airflow-core/src/airflow/models/mappedoperator.py

@@ -31,8 +31,8 @@
 from airflow.exceptions import AirflowException, NotMapped
 from airflow.sdk import BaseOperator as TaskSDKBaseOperator
 from airflow.sdk.definitions._internal.abstractoperator import DEFAULT_RETRY_DELAY_MULTIPLIER
-from airflow.sdk.definitions._internal.node import DAGNode
 from airflow.sdk.definitions.mappedoperator import MappedOperator as TaskSDKMappedOperator
+from airflow.serialization.definitions.node import DAGNode
 from airflow.serialization.definitions.param import SerializedParamsDict
 from airflow.serialization.definitions.taskgroup import SerializedMappedTaskGroup, SerializedTaskGroup
 from airflow.serialization.enums import DagAttributeTypes
@@ -48,6 +48,7 @@
     from airflow.models import TaskInstance
     from airflow.models.expandinput import SchedulerExpandInput
     from airflow.sdk import BaseOperatorLink, Context
+    from airflow.sdk.definitions._internal.node import DAGNode as TaskSDKDAGNode
     from airflow.sdk.definitions.operator_resources import Resources
     from airflow.serialization.serialized_objects import SerializedDAG
     from airflow.task.trigger_rule import TriggerRule
@@ -503,7 +504,7 @@ def expand_start_trigger_args(self, *, context: Context) -> StartTriggerArgs | N
 
 
 @functools.singledispatch
-def get_mapped_ti_count(task: DAGNode, run_id: str, *, session: Session) -> int:
+def get_mapped_ti_count(task: DAGNode | TaskSDKDAGNode, run_id: str, *, session: Session) -> int:
     raise NotImplementedError(f"Not implemented for {type(task)}")
 
 

airflow-core/src/airflow/serialization/definitions/node.py

@@ -0,0 +1,39 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from airflow._shared.dagnode.node import GenericDAGNode
+
+if TYPE_CHECKING:
+    from typing import TypeAlias
+
+    from airflow.models.mappedoperator import MappedOperator
+    from airflow.serialization.definitions.taskgroup import SerializedTaskGroup  # noqa: F401
+    from airflow.serialization.serialized_objects import SerializedBaseOperator, SerializedDAG  # noqa: F401
+
+    Operator: TypeAlias = SerializedBaseOperator | MappedOperator
+
+
+class DAGNode(GenericDAGNode["SerializedDAG", "Operator", "SerializedTaskGroup"]):
+    """
+    Base class for a node in the graph of a workflow.
+
+    A node may be an operator or task group, either mapped or unmapped.
+    """

airflow-core/src/airflow/serialization/definitions/taskgroup.py

@@ -27,7 +27,7 @@
 import attrs
 import methodtools
 
-from airflow.sdk.definitions._internal.node import DAGNode
+from airflow.serialization.definitions.node import DAGNode
 
 if TYPE_CHECKING:
     from collections.abc import Generator, Iterator

airflow-core/src/airflow/serialization/serialized_objects.py

@@ -79,7 +79,6 @@
 from airflow.observability.stats import Stats
 from airflow.sdk import DAG, Asset, AssetAlias, BaseOperator, XComArg
 from airflow.sdk.bases.operator import OPERATOR_DEFAULTS  # TODO: Copy this into the scheduler?
-from airflow.sdk.definitions._internal.node import DAGNode
 from airflow.sdk.definitions.asset import (
     AssetAliasEvent,
     AssetAliasUniqueKey,
@@ -101,6 +100,7 @@
     SerializedAssetBase,
     SerializedAssetUniqueKey,
 )
+from airflow.serialization.definitions.node import DAGNode
 from airflow.serialization.definitions.param import SerializedParam, SerializedParamsDict
 from airflow.serialization.definitions.taskgroup import SerializedMappedTaskGroup, SerializedTaskGroup
 from airflow.serialization.encoders import (
@@ -139,26 +139,22 @@
 if TYPE_CHECKING:
     from inspect import Parameter
 
+    from kubernetes.client import models as k8s  # noqa: TC004
     from pydantic import NonNegativeInt
     from sqlalchemy.orm import Session
 
     from airflow.models.expandinput import SchedulerExpandInput
     from airflow.models.mappedoperator import MappedOperator as SerializedMappedOperator
     from airflow.models.taskinstance import TaskInstance
+    from airflow.providers.cncf.kubernetes.pod_generator import PodGenerator  # noqa: TC004
     from airflow.sdk import BaseOperatorLink
+    from airflow.sdk.definitions._internal.node import DAGNode as SDKDAGNode
     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
     from airflow.timetables.simple import PartitionMapper
 
-    try:
-        from kubernetes.client import models as k8s  # noqa: TC004
-
-        from airflow.providers.cncf.kubernetes.pod_generator import PodGenerator  # noqa: TC004
-    except ImportError:
-        pass
-
     SerializedOperator: TypeAlias = "SerializedMappedOperator | SerializedBaseOperator"
     SdkOperator: TypeAlias = BaseOperator | MappedOperator
 
@@ -1057,7 +1053,6 @@ def detect_dag_dependencies(dag: DAG | None) -> Iterable[DagDependency]:
         yield from tt.asset_condition.iter_dag_dependencies(source="", target=dag.dag_id)
 
 
-# TODO (GH-52141): Duplicate DAGNode in the scheduler.
 class SerializedBaseOperator(DAGNode, BaseSerialization):
     """
     A JSON serializable representation of operator.
@@ -1194,8 +1189,7 @@ def __repr__(self) -> str:
     def node_id(self) -> str:
         return self.task_id
 
-    # TODO (GH-52141): Replace DAGNode with a scheduler type.
-    def get_dag(self) -> SerializedDAG | None:  # type: ignore[override]
+    def get_dag(self) -> SerializedDAG | None:
         return self.dag
 
     @property
@@ -1715,7 +1709,7 @@ def _matches_client_defaults(cls, var: Any, attrname: str) -> bool:
         return False
 
     @classmethod
-    def _is_excluded(cls, var: Any, attrname: str, op: DAGNode):
+    def _is_excluded(cls, var: Any, attrname: str, op: SDKDAGNode) -> bool:
         """
         Determine if a variable is excluded from the serialized object.
 

airflow-core/src/airflow/utils/dag_edges.py

@@ -23,8 +23,6 @@
 from airflow.serialization.serialized_objects import SerializedBaseOperator, SerializedDAG
 
 if TYPE_CHECKING:
-    from collections.abc import Iterable
-
     from airflow.sdk import DAG
 
     Operator: TypeAlias = MappedOperator | SerializedBaseOperator
@@ -118,11 +116,7 @@ def collect_edges(task_group):
     while tasks_to_trace:
         tasks_to_trace_next: list[Operator] = []
         for task in tasks_to_trace:
-            # TODO (GH-52141): downstream_list on DAGNode needs to be able to
-            # return scheduler types when used in scheduler, but SDK types when
-            # used at runtime. This means DAGNode needs to be rewritten as a
-            # generic class.
-            for child in cast("Iterable[Operator]", task.downstream_list):
+            for child in task.downstream_list:
                 edge = (task.task_id, child.task_id)
                 if task.is_setup and child.is_teardown:
                     setup_teardown_edges.add(edge)

pyproject.toml

@@ -1346,6 +1346,7 @@ apache-airflow-kubernetes-tests = { workspace = true }
 apache-airflow-providers = { workspace = true }
 apache-aurflow-docker-stack = { workspace = true }
 apache-airflow-shared-configuration = { workspace = true }
+apache-airflow-shared-dagnode = { workspace = true }
 apache-airflow-shared-logging = { workspace = true }
 apache-airflow-shared-module-loading = { workspace = true }
 apache-airflow-shared-secrets-backend = { workspace = true }

shared/dagnode/pyproject.toml

@@ -0,0 +1,48 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+[project]
+name = "apache-airflow-shared-dagnode"
+description = "Shared DAGNode logic for Airflow distributions"
+version = "0.0"
+classifiers = [
+    "Private :: Do Not Upload",
+]
+
+dependencies = [
+    "structlog>=25.4.0",
+]
+
+[dependency-groups]
+dev = [
+    "apache-airflow-devel-common",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/airflow_shared"]
+
+[tool.ruff]
+extend = "../../pyproject.toml"
+src = ["src"]
+
+[tool.ruff.lint.per-file-ignores]
+# Ignore Doc rules et al for anything outside of tests
+"!src/*" = ["D", "S101", "TRY002"]

airflow-core/src/airflow/models/taskmixin.py → shared/dagnode/src/airflow_shared/dagnode/__init__.py

@@ -14,15 +14,3 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-from __future__ import annotations
-
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
-    from typing import TypeAlias
-
-import airflow.sdk.definitions._internal.mixins
-import airflow.sdk.definitions._internal.node
-
-DependencyMixin: TypeAlias = airflow.sdk.definitions._internal.mixins.DependencyMixin
-DAGNode: TypeAlias = airflow.sdk.definitions._internal.node.DAGNode