Fix type import to AbstractOperator (#51773)

Author: uranusjr

providers/openlineage/src/airflow/providers/openlineage/utils/utils.py

@@ -32,7 +32,7 @@
 from airflow import __version__ as AIRFLOW_VERSION
 
 # TODO: move this maybe to Airflow's logic?
-from airflow.models import BaseOperator, DagRun, TaskReschedule
+from airflow.models import DagRun, TaskReschedule
 from airflow.providers.openlineage import (
     __version__ as OPENLINEAGE_PROVIDER_VERSION,
     conf,
@@ -59,18 +59,14 @@
 if not AIRFLOW_V_3_0_PLUS:
     from airflow.utils.session import NEW_SESSION, provide_session
 
-try:
-    from airflow.sdk import BaseOperator as SdkBaseOperator
-except ImportError:
-    SdkBaseOperator = BaseOperator  # type: ignore[misc]
-
 if TYPE_CHECKING:
     from openlineage.client.event_v2 import Dataset as OpenLineageDataset
     from openlineage.client.facet_v2 import RunFacet, processing_engine_run
 
     from airflow.models import TaskInstance
     from airflow.providers.common.compat.assets import Asset
     from airflow.sdk import DAG
+    from airflow.sdk.bases.operator import BaseOperator
     from airflow.sdk.definitions.mappedoperator import MappedOperator
     from airflow.sdk.execution_time.secrets_masker import (
         Redactable,
@@ -82,9 +78,10 @@
 else:
     try:
         from airflow.sdk import DAG
+        from airflow.sdk.bases.operator import BaseOperator
         from airflow.sdk.definitions.mappedoperator import MappedOperator
     except ImportError:
-        from airflow.models import DAG, MappedOperator
+        from airflow.models import DAG, BaseOperator, MappedOperator
 
     try:
         from airflow.providers.common.compat.assets import Asset
@@ -119,7 +116,7 @@ def try_import_from_string(string: str) -> Any:
         return import_string(string)
 
 
-def get_operator_class(task: BaseOperator | SdkBaseOperator) -> type:
+def get_operator_class(task: BaseOperator) -> type:
     if task.__class__.__name__ in ("DecoratedMappedOperator", "MappedOperator"):
         return task.operator_class
     return task.__class__
@@ -203,25 +200,25 @@ def get_user_provided_run_facets(ti: TaskInstance, ti_state: TaskInstanceState)
     return custom_facets
 
 
-def get_fully_qualified_class_name(operator: BaseOperator | MappedOperator | SdkBaseOperator) -> str:
+def get_fully_qualified_class_name(operator: BaseOperator | MappedOperator) -> str:
     if isinstance(operator, (MappedOperator, SerializedBaseOperator)):
         # as in airflow.api_connexion.schemas.common_schema.ClassReferenceSchema
         return operator._task_module + "." + operator._task_type  # type: ignore
     op_class = get_operator_class(operator)
     return op_class.__module__ + "." + op_class.__name__
 
 
-def is_operator_disabled(operator: BaseOperator | MappedOperator | SdkBaseOperator) -> bool:
+def is_operator_disabled(operator: BaseOperator | MappedOperator) -> bool:
     return get_fully_qualified_class_name(operator) in conf.disabled_operators()
 
 
-def is_selective_lineage_enabled(obj: DAG | BaseOperator | MappedOperator | SdkBaseOperator) -> bool:
+def is_selective_lineage_enabled(obj: DAG | BaseOperator | MappedOperator) -> 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):
         return is_dag_lineage_enabled(obj)
-    if isinstance(obj, (BaseOperator, MappedOperator, SdkBaseOperator)):
+    if isinstance(obj, (BaseOperator, MappedOperator)):
         return is_task_lineage_enabled(obj)
     raise TypeError("is_selective_lineage_enabled can only be used on DAG or Operator objects")
 

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

@@ -27,7 +27,7 @@
 from airflow.sdk.definitions._internal.mixins import DependencyMixin
 
 if TYPE_CHECKING:
-    from airflow.sdk.definitions.abstractoperator import Operator
+    from airflow.sdk.definitions._internal.abstractoperator import Operator
     from airflow.sdk.definitions.dag import DAG
     from airflow.sdk.definitions.edges import EdgeModifier
     from airflow.sdk.definitions.taskgroup import TaskGroup

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

@@ -74,7 +74,7 @@
 
     from pendulum.tz.timezone import FixedTimezone, Timezone
 
-    from airflow.sdk.definitions.abstractoperator import Operator
+    from airflow.sdk.definitions._internal.abstractoperator import Operator
     from airflow.sdk.definitions.decorators import TaskDecoratorCollection
     from airflow.sdk.definitions.taskgroup import TaskGroup
     from airflow.typing_compat import Self

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

@@ -663,7 +663,7 @@ def iter_mapped_dependencies(self) -> Iterator[Operator]:
 
 def task_group_to_dict(task_item_or_group):
     """Create a nested dict representation of this TaskGroup and its children used to construct the Graph."""
-    from airflow.sdk.definitions.abstractoperator import AbstractOperator
+    from airflow.sdk.definitions._internal.abstractoperator import AbstractOperator
     from airflow.sdk.definitions.mappedoperator import MappedOperator
     from airflow.sensors.base import BaseSensorOperator