From 7d72b4136a905ed56c9c81c1bf6dd1796461b5dd Mon Sep 17 00:00:00 2001
From: Dov Benyomin Sohacheski <b@kloud.email>
Date: Mon, 30 Jun 2025 01:55:33 +0300
Subject: [PATCH] Replace models.BaseOperator to Task SDK for apache/livy

---
 .../livy/src/airflow/providers/apache/livy/hooks/livy.py | 2 --
 .../src/airflow/providers/apache/livy/operators/livy.py  | 4 +---
 .../src/airflow/providers/apache/livy/sensors/livy.py    | 9 +--------
 .../src/airflow/providers/apache/livy/triggers/livy.py   | 3 ---
 .../src/airflow/providers/apache/livy/version_compat.py  | 8 ++++++++
 5 files changed, 10 insertions(+), 16 deletions(-)

diff --git a/providers/apache/livy/src/airflow/providers/apache/livy/hooks/livy.py b/providers/apache/livy/src/airflow/providers/apache/livy/hooks/livy.py
index 2cf8957b1120b..68912c30fe9d3 100644
--- a/providers/apache/livy/src/airflow/providers/apache/livy/hooks/livy.py
+++ b/providers/apache/livy/src/airflow/providers/apache/livy/hooks/livy.py
@@ -14,8 +14,6 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-"""This module contains the Apache Livy hook."""
-
 from __future__ import annotations
 
 import asyncio
diff --git a/providers/apache/livy/src/airflow/providers/apache/livy/operators/livy.py b/providers/apache/livy/src/airflow/providers/apache/livy/operators/livy.py
index 645d81542c6c2..57f90832bd27a 100644
--- a/providers/apache/livy/src/airflow/providers/apache/livy/operators/livy.py
+++ b/providers/apache/livy/src/airflow/providers/apache/livy/operators/livy.py
@@ -14,8 +14,6 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-"""This module contains the Apache Livy operator."""
-
 from __future__ import annotations
 
 import time
@@ -25,9 +23,9 @@
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
-from airflow.models import BaseOperator
 from airflow.providers.apache.livy.hooks.livy import BatchState, LivyHook
 from airflow.providers.apache.livy.triggers.livy import LivyTrigger
+from airflow.providers.apache.livy.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/livy/src/airflow/providers/apache/livy/sensors/livy.py b/providers/apache/livy/src/airflow/providers/apache/livy/sensors/livy.py
index cf2f4aa15f017..ca0b352d7ee1a 100644
--- a/providers/apache/livy/src/airflow/providers/apache/livy/sensors/livy.py
+++ b/providers/apache/livy/src/airflow/providers/apache/livy/sensors/livy.py
@@ -14,20 +14,13 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-"""This module contains the Apache Livy sensor."""
-
 from __future__ import annotations
 
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
 from airflow.providers.apache.livy.hooks.livy import LivyHook
-from airflow.providers.apache.livy.version_compat import AIRFLOW_V_3_0_PLUS
-
-if AIRFLOW_V_3_0_PLUS:
-    from airflow.sdk import BaseSensorOperator
-else:
-    from airflow.sensors.base import BaseSensorOperator  # type: ignore[no-redef]
+from airflow.providers.apache.livy.version_compat import BaseSensorOperator
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
diff --git a/providers/apache/livy/src/airflow/providers/apache/livy/triggers/livy.py b/providers/apache/livy/src/airflow/providers/apache/livy/triggers/livy.py
index 9c47706b73718..20a8e0382d7d9 100644
--- a/providers/apache/livy/src/airflow/providers/apache/livy/triggers/livy.py
+++ b/providers/apache/livy/src/airflow/providers/apache/livy/triggers/livy.py
@@ -14,9 +14,6 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-
-"""This module contains the Apache Livy Trigger."""
-
 from __future__ import annotations
 
 import asyncio
diff --git a/providers/apache/livy/src/airflow/providers/apache/livy/version_compat.py b/providers/apache/livy/src/airflow/providers/apache/livy/version_compat.py
index 48d122b669696..126f95c848abd 100644
--- a/providers/apache/livy/src/airflow/providers/apache/livy/version_compat.py
+++ b/providers/apache/livy/src/airflow/providers/apache/livy/version_compat.py
@@ -33,3 +33,11 @@ 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
+else:
+    from airflow.models import BaseOperator
+    from airflow.sensors.base import BaseSensorOperator  # type: ignore[no-redef]
+
+__all__ = ["AIRFLOW_V_3_0_PLUS", "BaseOperator", "BaseSensorOperator"]