Replace models.BaseOperator to Task SDK one for Tableau, Telegram, and Teradata

providers/tableau/src/airflow/providers/tableau/operators/tableau.py

@@ -20,12 +20,12 @@
 from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException
-from airflow.models import BaseOperator
 from airflow.providers.tableau.hooks.tableau import (
     TableauHook,
     TableauJobFailedException,
     TableauJobFinishCode,
 )
+from airflow.providers.tableau.version_compat import BaseOperator
 
 if TYPE_CHECKING:
     try:

providers/tableau/src/airflow/providers/tableau/sensors/tableau.py

@@ -24,12 +24,7 @@
     TableauJobFailedException,
     TableauJobFinishCode,
 )
-from airflow.providers.tableau.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.tableau.version_compat import BaseSensorOperator
 
 if TYPE_CHECKING:
     try:

providers/tableau/src/airflow/providers/tableau/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  # type: ignore[no-redef]
+    from airflow.sensors.base import BaseSensorOperator  # type: ignore[no-redef]
+
+__all__ = ["AIRFLOW_V_3_0_PLUS", "BaseOperator", "BaseSensorOperator"]

providers/telegram/src/airflow/providers/telegram/operators/telegram.py

@@ -23,8 +23,8 @@
 from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException
-from airflow.models import BaseOperator
 from airflow.providers.telegram.hooks.telegram import TelegramHook
+from airflow.providers.telegram.version_compat import BaseOperator
 
 if TYPE_CHECKING:
     try:

providers/telegram/src/airflow/providers/telegram/version_compat.py

@@ -0,0 +1,42 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+# NOTE! THIS FILE IS COPIED MANUALLY IN OTHER PROVIDERS DELIBERATELY TO AVOID ADDING UNNECESSARY
+# DEPENDENCIES BETWEEN PROVIDERS. IF YOU WANT TO ADD CONDITIONAL CODE IN YOUR PROVIDER THAT DEPENDS
+# ON AIRFLOW VERSION, PLEASE COPY THIS FILE TO THE ROOT PACKAGE OF YOUR PROVIDER AND IMPORT
+# THOSE CONSTANTS FROM IT RATHER THAN IMPORTING THEM FROM ANOTHER PROVIDER OR TEST CODE
+#
+from __future__ import annotations
+
+
+def get_base_airflow_version_tuple() -> tuple[int, int, int]:
+    from packaging.version import Version
+
+    from airflow import __version__
+
+    airflow_version = Version(__version__)
+    return airflow_version.major, airflow_version.minor, airflow_version.micro
+
+
+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  # type: ignore[no-redef]
+
+__all__ = ["AIRFLOW_V_3_0_PLUS", "BaseOperator"]

providers/teradata/src/airflow/providers/teradata/operators/bteq.py

@@ -36,10 +36,10 @@
     except ImportError:
         from airflow.utils.context import Context
 
-from airflow.models import BaseOperator
 from airflow.providers.ssh.hooks.ssh import SSHHook
 from airflow.providers.teradata.hooks.bteq import BteqHook
 from airflow.providers.teradata.hooks.teradata import TeradataHook
+from airflow.providers.teradata.version_compat import BaseOperator
 
 
 def contains_template(parameter_value):

providers/teradata/src/airflow/providers/teradata/operators/teradata.py

@@ -20,9 +20,9 @@
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, ClassVar
 
-from airflow.models import BaseOperator
 from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator
 from airflow.providers.teradata.hooks.teradata import TeradataHook
+from airflow.providers.teradata.version_compat import BaseOperator
 
 if TYPE_CHECKING:
     try:

providers/teradata/src/airflow/providers/teradata/operators/teradata_compute_cluster.py

@@ -23,9 +23,9 @@
 from functools import cached_property
 from typing import TYPE_CHECKING
 
-from airflow.models import BaseOperator
 from airflow.providers.teradata.hooks.teradata import TeradataHook
 from airflow.providers.teradata.utils.constants import Constants
+from airflow.providers.teradata.version_compat import BaseOperator
 
 if TYPE_CHECKING:
     try:

providers/teradata/src/airflow/providers/teradata/transfers/azure_blob_to_teradata.py

@@ -21,8 +21,6 @@
 from textwrap import dedent
 from typing import TYPE_CHECKING
 
-from airflow.models import BaseOperator
-
 try:
     from airflow.providers.microsoft.azure.hooks.wasb import WasbHook
 except ModuleNotFoundError as e:
@@ -31,6 +29,7 @@
     raise AirflowOptionalProviderFeatureException(e)
 
 from airflow.providers.teradata.hooks.teradata import TeradataHook
+from airflow.providers.teradata.version_compat import BaseOperator
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

providers/teradata/src/airflow/providers/teradata/transfers/s3_to_teradata.py

@@ -21,15 +21,14 @@
 from textwrap import dedent
 from typing import TYPE_CHECKING
 
-from airflow.models import BaseOperator
-
 try:
     from airflow.providers.amazon.aws.hooks.s3 import S3Hook
 except ModuleNotFoundError as e:
     from airflow.exceptions import AirflowOptionalProviderFeatureException
 
     raise AirflowOptionalProviderFeatureException(e)
 from airflow.providers.teradata.hooks.teradata import TeradataHook
+from airflow.providers.teradata.version_compat import BaseOperator
 
 if TYPE_CHECKING:
     try:

providers/teradata/src/airflow/providers/teradata/transfers/teradata_to_teradata.py

@@ -21,8 +21,8 @@
 from functools import cached_property
 from typing import TYPE_CHECKING
 
-from airflow.models import BaseOperator
 from airflow.providers.teradata.hooks.teradata import TeradataHook
+from airflow.providers.teradata.version_compat import BaseOperator
 
 if TYPE_CHECKING:
     try:

providers/teradata/src/airflow/providers/teradata/version_compat.py

@@ -0,0 +1,42 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+# NOTE! THIS FILE IS COPIED MANUALLY IN OTHER PROVIDERS DELIBERATELY TO AVOID ADDING UNNECESSARY
+# DEPENDENCIES BETWEEN PROVIDERS. IF YOU WANT TO ADD CONDITIONAL CODE IN YOUR PROVIDER THAT DEPENDS
+# ON AIRFLOW VERSION, PLEASE COPY THIS FILE TO THE ROOT PACKAGE OF YOUR PROVIDER AND IMPORT
+# THOSE CONSTANTS FROM IT RATHER THAN IMPORTING THEM FROM ANOTHER PROVIDER OR TEST CODE
+#
+from __future__ import annotations
+
+
+def get_base_airflow_version_tuple() -> tuple[int, int, int]:
+    from packaging.version import Version
+
+    from airflow import __version__
+
+    airflow_version = Version(__version__)
+    return airflow_version.major, airflow_version.minor, airflow_version.micro
+
+
+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  # type: ignore[no-redef]
+
+__all__ = ["AIRFLOW_V_3_0_PLUS", "BaseOperator"]

providers/teradata/tests/unit/teradata/operators/test_bteq.py

@@ -252,7 +252,7 @@ def test_remote_execution_with_sql(
         mock_execute_bteq_script.assert_called_once()
         assert result == 0
 
-    @mock.patch("airflow.models.BaseOperator.render_template")
+    @mock.patch("airflow.providers.teradata.version_compat.BaseOperator.render_template")
     def test_render_template_in_sql(self, mock_render):
         op = BteqOperator(task_id="render_test", sql="SELECT * FROM {{ params.table }};")
         mock_render.return_value = "SELECT * FROM my_table;"