From 1e8789c9c2c1cdf2ee28a4dec62185fafe640043 Mon Sep 17 00:00:00 2001 From: geonwoo Date: Sun, 29 Jun 2025 19:05:11 +0900 Subject: [PATCH] add: version_compat --- .../providers/apache/spark/operators/spark_sql.py | 2 +- .../providers/apache/spark/operators/spark_submit.py | 2 +- .../airflow/providers/apache/spark/version_compat.py | 10 ++++++++++ 3 files changed, 12 insertions(+), 2 deletions(-) diff --git a/providers/apache/spark/src/airflow/providers/apache/spark/operators/spark_sql.py b/providers/apache/spark/src/airflow/providers/apache/spark/operators/spark_sql.py index b2de5d0cb237c..e4fc07e6ec4f5 100644 --- a/providers/apache/spark/src/airflow/providers/apache/spark/operators/spark_sql.py +++ b/providers/apache/spark/src/airflow/providers/apache/spark/operators/spark_sql.py @@ -20,8 +20,8 @@ from collections.abc import Sequence from typing import TYPE_CHECKING, Any -from airflow.models import BaseOperator from airflow.providers.apache.spark.hooks.spark_sql import SparkSqlHook +from airflow.providers.apache.spark.version_compat import BaseOperator if TYPE_CHECKING: from airflow.utils.context import Context diff --git a/providers/apache/spark/src/airflow/providers/apache/spark/operators/spark_submit.py b/providers/apache/spark/src/airflow/providers/apache/spark/operators/spark_submit.py index 9c1b38e5eee40..44fd5f3f43bf7 100644 --- a/providers/apache/spark/src/airflow/providers/apache/spark/operators/spark_submit.py +++ b/providers/apache/spark/src/airflow/providers/apache/spark/operators/spark_submit.py @@ -21,8 +21,8 @@ from typing import TYPE_CHECKING, Any from airflow.configuration import conf -from airflow.models import BaseOperator from airflow.providers.apache.spark.hooks.spark_submit import SparkSubmitHook +from airflow.providers.apache.spark.version_compat import BaseOperator from airflow.providers.common.compat.openlineage.utils.spark import ( inject_parent_job_information_into_spark_properties, inject_transport_information_into_spark_properties, diff --git a/providers/apache/spark/src/airflow/providers/apache/spark/version_compat.py b/providers/apache/spark/src/airflow/providers/apache/spark/version_compat.py index 48d122b669696..e7a259afb357c 100644 --- a/providers/apache/spark/src/airflow/providers/apache/spark/version_compat.py +++ b/providers/apache/spark/src/airflow/providers/apache/spark/version_compat.py @@ -33,3 +33,13 @@ 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 +else: + from airflow.models import BaseOperator + +__all__ = [ + "AIRFLOW_V_3_0_PLUS", + "BaseOperator", +]