diff --git a/providers/apache/flink/src/airflow/providers/apache/flink/operators/flink_kubernetes.py b/providers/apache/flink/src/airflow/providers/apache/flink/operators/flink_kubernetes.py index ed05397ec98a2..2c90ce9cc0e42 100644 --- a/providers/apache/flink/src/airflow/providers/apache/flink/operators/flink_kubernetes.py +++ b/providers/apache/flink/src/airflow/providers/apache/flink/operators/flink_kubernetes.py @@ -21,13 +21,13 @@ from functools import cached_property from typing import TYPE_CHECKING -from airflow.models import BaseOperator +from airflow.providers.apache.flink.version_compat import BaseOperator from airflow.providers.cncf.kubernetes.hooks.kubernetes import KubernetesHook if TYPE_CHECKING: from kubernetes.client import CoreV1Api - from airflow.utils.context import Context + from airflow.providers.apache.flink.version_compat import Context class FlinkKubernetesOperator(BaseOperator): diff --git a/providers/apache/flink/src/airflow/providers/apache/flink/sensors/flink_kubernetes.py b/providers/apache/flink/src/airflow/providers/apache/flink/sensors/flink_kubernetes.py index 4a83956989f7c..ad0cc1cfa54fc 100644 --- a/providers/apache/flink/src/airflow/providers/apache/flink/sensors/flink_kubernetes.py +++ b/providers/apache/flink/src/airflow/providers/apache/flink/sensors/flink_kubernetes.py @@ -23,16 +23,11 @@ from kubernetes import client from airflow.exceptions import AirflowException -from airflow.providers.apache.flink.version_compat import AIRFLOW_V_3_0_PLUS +from airflow.providers.apache.flink.version_compat import BaseSensorOperator from airflow.providers.cncf.kubernetes.hooks.kubernetes import KubernetesHook -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk import BaseSensorOperator -else: - from airflow.sensors.base import BaseSensorOperator # type: ignore[no-redef] - if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.apache.flink.version_compat import Context class FlinkKubernetesSensor(BaseSensorOperator): diff --git a/providers/apache/flink/src/airflow/providers/apache/flink/version_compat.py b/providers/apache/flink/src/airflow/providers/apache/flink/version_compat.py index 48d122b669696..5faab587ffcef 100644 --- a/providers/apache/flink/src/airflow/providers/apache/flink/version_compat.py +++ b/providers/apache/flink/src/airflow/providers/apache/flink/version_compat.py @@ -33,3 +33,18 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 0, 0) + +if AIRFLOW_V_3_0_PLUS: + from airflow.sdk import BaseOperator, BaseSensorOperator + from airflow.sdk.definitions.context import Context +else: + from airflow.models import BaseOperator + from airflow.sensors.base import BaseSensorOperator # type: ignore[no-redef] + from airflow.utils.context import Context + +__all__ = [ + "AIRFLOW_V_3_0_PLUS", + "BaseOperator", + "BaseSensorOperator", + "Context", +]