Provider Migration: Replace models.BaseOperator to Task SDK for apache/livy

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

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,

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

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

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"]