Exploring possibility of moving DagAttributeTypes to execution API spec (#61251)

Author: amoghrajesh

airflow-core/src/airflow/api_fastapi/execution_api/app.py

@@ -265,6 +265,7 @@ def get_extra_schemas() -> dict[str, dict]:
     """Get all the extra schemas that are not part of the main FastAPI app."""
     from airflow.api_fastapi.execution_api.datamodels.taskinstance import TaskInstance
     from airflow.executors.workloads import BundleInfo
+    from airflow.serialization.enums import DagAttributeTypes
     from airflow.task.trigger_rule import TriggerRule
     from airflow.task.weight_rule import WeightRule
     from airflow.utils.state import TaskInstanceState, TerminalTIState
@@ -278,6 +279,11 @@ def get_extra_schemas() -> dict[str, dict]:
         "TaskInstanceState": {"type": "string", "enum": list(TaskInstanceState)},
         "WeightRule": {"type": "string", "enum": list(WeightRule)},
         "TriggerRule": {"type": "string", "enum": list(TriggerRule)},
+        "DagAttributeTypes": {
+            "type": "string",
+            "enum": [DagAttributeTypes.OP.value, DagAttributeTypes.TASK_GROUP.value],
+            "x-enum-varnames": [DagAttributeTypes.OP.name, DagAttributeTypes.TASK_GROUP.name],
+        },
     }
 
 

task-sdk/src/airflow/sdk/api/datamodels/_generated.py

@@ -485,6 +485,11 @@ class TriggerRule(str, Enum):
     ALL_SKIPPED = "all_skipped"
 
 
+class DagAttributeTypes(str, Enum):
+    OP = "operator"
+    TASK_GROUP = "taskgroup"
+
+
 class AssetReferenceAssetEventDagRun(BaseModel):
     """
     Schema for AssetModel used in AssetEventDagRunReference.

task-sdk/src/airflow/sdk/bases/operator.py

@@ -89,13 +89,13 @@ def db_safe_priority(priority_weight: int) -> int:
     import jinja2
     from typing_extensions import Self
 
+    from airflow.sdk.api.datamodels._generated import DagAttributeTypes
     from airflow.sdk.bases.operatorlink import BaseOperatorLink
     from airflow.sdk.definitions.context import Context
     from airflow.sdk.definitions.dag import DAG
     from airflow.sdk.definitions.operator_resources import Resources
     from airflow.sdk.definitions.taskgroup import TaskGroup
     from airflow.sdk.definitions.xcom_arg import XComArg
-    from airflow.serialization.enums import DagAttributeTypes
     from airflow.task.priority_strategy import PriorityWeightStrategy
     from airflow.triggers.base import BaseTrigger, StartTriggerArgs
 
@@ -1554,7 +1554,7 @@ def prepare_for_execution(self) -> Self:
 
     def serialize_for_task_group(self) -> tuple[DagAttributeTypes, Any]:
         """Serialize; required by DAGNode."""
-        from airflow.serialization.enums import DagAttributeTypes
+        from airflow.sdk.api.datamodels._generated import DagAttributeTypes
 
         return DagAttributeTypes.OP, self.task_id
 

task-sdk/src/airflow/sdk/definitions/_internal/node.py

@@ -27,11 +27,11 @@
 from airflow.sdk.definitions._internal.mixins import DependencyMixin
 
 if TYPE_CHECKING:
+    from airflow.sdk.api.datamodels._generated import DagAttributeTypes
     from airflow.sdk.definitions.dag import DAG
     from airflow.sdk.definitions.edges import EdgeModifier
     from airflow.sdk.definitions.taskgroup import TaskGroup  # noqa: F401
     from airflow.sdk.types import Operator  # noqa: F401
-    from airflow.serialization.enums import DagAttributeTypes
 
 
 KEY_REGEX = re.compile(r"^[\w.-]+$")

task-sdk/src/airflow/sdk/definitions/mappedoperator.py

@@ -27,6 +27,7 @@
 import methodtools
 from lazy_object_proxy import Proxy
 
+from airflow.sdk.api.datamodels._generated import DagAttributeTypes
 from airflow.sdk.bases.xcom import BaseXCom
 from airflow.sdk.definitions._internal.abstractoperator import (
     DEFAULT_EXECUTOR,
@@ -50,7 +51,6 @@
     is_mappable,
 )
 from airflow.sdk.definitions._internal.types import NOTSET
-from airflow.serialization.enums import DagAttributeTypes
 
 if TYPE_CHECKING:
     import datetime

task-sdk/src/airflow/sdk/definitions/taskgroup.py

@@ -36,14 +36,14 @@
 )
 
 if TYPE_CHECKING:
+    from airflow.sdk.api.datamodels._generated import DagAttributeTypes
     from airflow.sdk.bases.operator import BaseOperator
     from airflow.sdk.definitions._internal.abstractoperator import AbstractOperator
     from airflow.sdk.definitions._internal.expandinput import DictOfListsExpandInput, ListOfDictsExpandInput
     from airflow.sdk.definitions._internal.mixins import DependencyMixin
     from airflow.sdk.definitions.dag import DAG
     from airflow.sdk.definitions.edges import EdgeModifier
     from airflow.sdk.types import Operator
-    from airflow.serialization.enums import DagAttributeTypes
 
 
 def _default_parent_group() -> TaskGroup | None:
@@ -493,7 +493,7 @@ def get_child_by_label(self, label: str) -> DAGNode:
 
     def serialize_for_task_group(self) -> tuple[DagAttributeTypes, Any]:
         """Serialize task group; required by DagNode."""
-        from airflow.serialization.enums import DagAttributeTypes
+        from airflow.sdk.api.datamodels._generated import DagAttributeTypes
         from airflow.serialization.serialized_objects import TaskGroupSerialization
 
         return (