Move MappedOperator to serialization (#59628)

Author: uranusjr

airflow-core/src/airflow/api/common/mark_tasks.py

@@ -20,7 +20,7 @@
 from __future__ import annotations
 
 from collections.abc import Collection, Iterable, Iterator
-from typing import TYPE_CHECKING, TypeAlias
+from typing import TYPE_CHECKING
 
 from sqlalchemy import and_, or_, select
 from sqlalchemy.orm import lazyload
@@ -33,11 +33,8 @@
 if TYPE_CHECKING:
     from sqlalchemy.orm import Session as SASession
 
-    from airflow.models.mappedoperator import MappedOperator
-    from airflow.serialization.definitions.baseoperator import SerializedBaseOperator
     from airflow.serialization.definitions.dag import SerializedDAG
-
-    Operator: TypeAlias = MappedOperator | SerializedBaseOperator
+    from airflow.serialization.definitions.mappedoperator import Operator
 
 
 @provide_session

airflow-core/src/airflow/api_fastapi/core_api/services/ui/grid.py

@@ -24,9 +24,9 @@
 
 from airflow.api_fastapi.common.parameters import state_priority
 from airflow.api_fastapi.core_api.services.ui.task_group import get_task_group_children_getter
-from airflow.models.mappedoperator import MappedOperator
 from airflow.models.taskmap import TaskMap
 from airflow.serialization.definitions.baseoperator import SerializedBaseOperator
+from airflow.serialization.definitions.mappedoperator import SerializedMappedOperator
 from airflow.serialization.definitions.taskgroup import SerializedTaskGroup
 
 log = structlog.get_logger(logger_name=__name__)
@@ -90,7 +90,7 @@ def _find_aggregates(
 
     if node is None:
         return
-    if isinstance(node, MappedOperator):
+    if isinstance(node, SerializedMappedOperator):
         # For unmapped tasks, reflect a single None state so UI shows one square
         mapped_details = details or [{"state": None, "start_date": None, "end_date": None}]
         yield {

airflow-core/src/airflow/api_fastapi/core_api/services/ui/task_group.py

@@ -24,8 +24,8 @@
 from operator import methodcaller
 
 from airflow.configuration import conf
-from airflow.models.mappedoperator import MappedOperator, is_mapped
 from airflow.serialization.definitions.baseoperator import SerializedBaseOperator
+from airflow.serialization.definitions.mappedoperator import SerializedMappedOperator, is_mapped
 
 
 @cache
@@ -39,7 +39,7 @@ def get_task_group_children_getter() -> Callable:
 
 def task_group_to_dict(task_item_or_group, parent_group_is_mapped=False):
     """Create a nested dict representation of this TaskGroup and its children used to construct the Graph."""
-    if isinstance(task := task_item_or_group, (SerializedBaseOperator, MappedOperator)):
+    if isinstance(task := task_item_or_group, (SerializedBaseOperator, SerializedMappedOperator)):
         # we explicitly want the short task ID here, not the full doted notation if in a group
         task_display_name = task.task_display_name if task.task_display_name != task.task_id else task.label
         node_operator = {
@@ -83,7 +83,7 @@ def task_group_to_dict(task_item_or_group, parent_group_is_mapped=False):
 
 def task_group_to_dict_grid(task_item_or_group, parent_group_is_mapped=False):
     """Create a nested dict representation of this TaskGroup and its children used to construct the Grid."""
-    if isinstance(task := task_item_or_group, (MappedOperator, SerializedBaseOperator)):
+    if isinstance(task := task_item_or_group, (SerializedMappedOperator, SerializedBaseOperator)):
         mapped = None
         if parent_group_is_mapped or is_mapped(task):
             mapped = True

airflow-core/src/airflow/cli/commands/task_command.py

@@ -62,11 +62,9 @@
 
     from sqlalchemy.orm.session import Session
 
-    from airflow.models.mappedoperator import MappedOperator
-    from airflow.serialization.definitions.baseoperator import SerializedBaseOperator
+    from airflow.serialization.definitions.mappedoperator import Operator
 
     CreateIfNecessary = Literal[False, "db", "memory"]
-    Operator = MappedOperator | SerializedBaseOperator
 
 log = logging.getLogger(__name__)
 

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

@@ -109,7 +109,7 @@ def __getattr__(name):
     "Deadline": "airflow.models.deadline",
     "Log": "airflow.models.log",
     "HITLDetail": "airflow.models.hitl",
-    "MappedOperator": "airflow.models.mappedoperator",
+    "MappedOperator": "airflow.sdk.definitions.mappedoperator",
     "Param": "airflow.sdk.definitions.param",
     "Pool": "airflow.models.pool",
     "RenderedTaskInstanceFields": "airflow.models.renderedtifields",
@@ -136,7 +136,6 @@ def __getattr__(name):
     from airflow.models.db_callback_request import DbCallbackRequest
     from airflow.models.deadline import Deadline
     from airflow.models.log import Log
-    from airflow.models.mappedoperator import MappedOperator
     from airflow.models.pool import Pool
     from airflow.models.renderedtifields import RenderedTaskInstanceFields
     from airflow.models.skipmixin import SkipMixin
@@ -147,6 +146,7 @@ def __getattr__(name):
     from airflow.models.variable import Variable
     from airflow.sdk import DAG, BaseOperator, BaseOperatorLink, Param
     from airflow.sdk.bases.xcom import BaseXCom
+    from airflow.sdk.definitions.mappedoperator import MappedOperator
     from airflow.sdk.execution_time.xcom import XCom
 
 

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

@@ -70,16 +70,13 @@
 if TYPE_CHECKING:
     from typing import TypeAlias
 
-    from airflow.models.mappedoperator import MappedOperator
     from airflow.serialization.definitions.assets import (
         SerializedAsset,
         SerializedAssetAlias,
         SerializedAssetBase,
     )
-    from airflow.serialization.definitions.baseoperator import SerializedBaseOperator
     from airflow.serialization.definitions.dag import SerializedDAG
 
-    Operator: TypeAlias = MappedOperator | SerializedBaseOperator
     UKey: TypeAlias = SerializedAssetUniqueKey
 
 log = logging.getLogger(__name__)

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

@@ -104,17 +104,15 @@
 
     from airflow._shared.observability.traces.base_tracer import EmptySpan
     from airflow.models.dag_version import DagVersion
-    from airflow.models.mappedoperator import MappedOperator
     from airflow.models.taskinstancekey import TaskInstanceKey
     from airflow.sdk import DAG as SDKDAG
-    from airflow.serialization.definitions.baseoperator import SerializedBaseOperator
     from airflow.serialization.definitions.dag import SerializedDAG
+    from airflow.serialization.definitions.mappedoperator import Operator
 
     CreatedTasks = TypeVar("CreatedTasks", Iterator["dict[str, Any]"], Iterator[TI])
     AttributeValueType: TypeAlias = (
         str | bool | int | float | Sequence[str] | Sequence[bool] | Sequence[int] | Sequence[float]
     )
-    Operator: TypeAlias = MappedOperator | SerializedBaseOperator
 
 RUN_ID_REGEX = r"^(?:manual|scheduled|asset_triggered)__(?:\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}\+00:00)$"
 
@@ -1529,7 +1527,7 @@ def _expand_mapped_task_if_needed(ti: TI) -> Iterable[TI] | None:
             If the ti does not need expansion, either because the task is not
             mapped, or has already been expanded, *None* is returned.
             """
-            from airflow.models.mappedoperator import is_mapped
+            from airflow.serialization.definitions.mappedoperator import is_mapped
 
             if TYPE_CHECKING:
                 assert ti.task
@@ -1749,7 +1747,7 @@ def _check_for_removed_or_restored_tasks(
 
         """
         from airflow.models.expandinput import NotFullyPopulated
-        from airflow.models.mappedoperator import get_mapped_ti_count
+        from airflow.serialization.definitions.mappedoperator import get_mapped_ti_count
 
         tis = self.get_task_instances(session=session)
 
@@ -1891,7 +1889,7 @@ def _create_tasks(
         :param task_creator: Function to create task instances
         """
         from airflow.models.expandinput import NotFullyPopulated
-        from airflow.models.mappedoperator import get_mapped_ti_count
+        from airflow.serialization.definitions.mappedoperator import get_mapped_ti_count
 
         map_indexes: Iterable[int]
         for task in tasks:
@@ -1965,7 +1963,7 @@ def _revise_map_indexes_if_mapped(
         for more details.
         """
         from airflow.models.expandinput import NotFullyPopulated
-        from airflow.models.mappedoperator import get_mapped_ti_count
+        from airflow.serialization.definitions.mappedoperator import get_mapped_ti_count
         from airflow.settings import task_instance_mutation_hook
 
         try:

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

@@ -35,16 +35,13 @@
 )
 
 if TYPE_CHECKING:
-    from typing import TypeAlias, TypeGuard
+    from typing import TypeGuard
 
     from sqlalchemy.orm import Session
 
-    from airflow.models.mappedoperator import MappedOperator
-    from airflow.serialization.definitions.baseoperator import SerializedBaseOperator
+    from airflow.serialization.definitions.mappedoperator import Operator
     from airflow.serialization.definitions.xcom_arg import SchedulerXComArg
 
-    Operator: TypeAlias = MappedOperator | SerializedBaseOperator
-
 
 __all__ = [
     "DictOfListsExpandInput",

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

@@ -21,12 +21,8 @@
 
 if TYPE_CHECKING:
     from collections.abc import Iterable
-    from typing import TypeAlias
 
-    from airflow.models.mappedoperator import MappedOperator
-    from airflow.serialization.definitions.baseoperator import SerializedBaseOperator
-
-    Operator: TypeAlias = MappedOperator | SerializedBaseOperator
+    from airflow.serialization.definitions.mappedoperator import Operator
 
 
 @runtime_checkable

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

@@ -103,7 +103,7 @@
 
 if TYPE_CHECKING:
     from datetime import datetime
-    from typing import Literal, TypeAlias
+    from typing import Literal
 
     import pendulum
     from sqlalchemy.engine import Connection as SAConnection, Engine
@@ -114,14 +114,11 @@
     from airflow.api_fastapi.execution_api.datamodels.asset import AssetProfile
     from airflow.models.dag import DagModel
     from airflow.models.dagrun import DagRun
-    from airflow.models.mappedoperator import MappedOperator
-    from airflow.serialization.definitions.baseoperator import SerializedBaseOperator
     from airflow.serialization.definitions.dag import SerializedDAG
+    from airflow.serialization.definitions.mappedoperator import Operator
     from airflow.serialization.definitions.taskgroup import SerializedTaskGroup
     from airflow.utils.context import Context
 
-    Operator: TypeAlias = MappedOperator | SerializedBaseOperator
-
 
 PAST_DEPENDS_MET = "past_depends_met"
 
@@ -1660,7 +1657,6 @@ def get_template_context(
             session = settings.get_session()()
 
         from airflow.exceptions import NotMapped
-        from airflow.models.mappedoperator import get_mapped_ti_count
         from airflow.sdk.api.datamodels._generated import (
             DagRun as DagRunSDK,
             PrevSuccessfulDagRunResponse,
@@ -1669,6 +1665,7 @@ def get_template_context(
         from airflow.sdk.definitions.param import process_params
         from airflow.sdk.execution_time.context import InletEventsAccessors
         from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
+        from airflow.serialization.definitions.mappedoperator import get_mapped_ti_count
         from airflow.utils.context import (
             ConnectionAccessor,
             OutletEventAccessors,
@@ -2193,7 +2190,7 @@ def _find_common_ancestor_mapped_group(node1: Operator, node2: Operator) -> Seri
 
 def _is_further_mapped_inside(operator: Operator, container: SerializedTaskGroup) -> bool:
     """Whether given operator is *further* mapped inside a task group."""
-    from airflow.models.mappedoperator import is_mapped
+    from airflow.serialization.definitions.mappedoperator import is_mapped
 
     if is_mapped(operator):
         return True
@@ -2260,7 +2257,7 @@ def tg2(inp):
     :return: Specific map index or map indexes to pull, or ``None`` if we
         want to "whole" return value (i.e. no mapped task groups involved).
     """
-    from airflow.models.mappedoperator import get_mapped_ti_count
+    from airflow.serialization.definitions.mappedoperator import get_mapped_ti_count
 
     # This value should never be None since we already know the current task
     # is in a mapped task group, and should have been expanded, despite that,
@@ -2307,7 +2304,7 @@ def find_relevant_relatives(
     run_id: str,
     session: Session,
 ) -> Collection[str | tuple[str, int]]:
-    from airflow.models.mappedoperator import get_mapped_ti_count
+    from airflow.serialization.definitions.mappedoperator import get_mapped_ti_count
 
     visited: set[str | tuple[str, int]] = set()