diff --git a/providers/tableau/src/airflow/providers/tableau/operators/tableau.py b/providers/tableau/src/airflow/providers/tableau/operators/tableau.py index 951ad44535197..ea435e47e6307 100644 --- a/providers/tableau/src/airflow/providers/tableau/operators/tableau.py +++ b/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: diff --git a/providers/tableau/src/airflow/providers/tableau/sensors/tableau.py b/providers/tableau/src/airflow/providers/tableau/sensors/tableau.py index ce823b2fdc023..2953039cd25ab 100644 --- a/providers/tableau/src/airflow/providers/tableau/sensors/tableau.py +++ b/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: diff --git a/providers/tableau/src/airflow/providers/tableau/version_compat.py b/providers/tableau/src/airflow/providers/tableau/version_compat.py index 48d122b669696..fde11c60d3655 100644 --- a/providers/tableau/src/airflow/providers/tableau/version_compat.py +++ b/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"] diff --git a/providers/telegram/src/airflow/providers/telegram/operators/telegram.py b/providers/telegram/src/airflow/providers/telegram/operators/telegram.py index 641a2760839b7..7531a740879f8 100644 --- a/providers/telegram/src/airflow/providers/telegram/operators/telegram.py +++ b/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: diff --git a/providers/telegram/src/airflow/providers/telegram/version_compat.py b/providers/telegram/src/airflow/providers/telegram/version_compat.py new file mode 100644 index 0000000000000..4f8d5e32bca4a --- /dev/null +++ b/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"] diff --git a/providers/teradata/src/airflow/providers/teradata/operators/bteq.py b/providers/teradata/src/airflow/providers/teradata/operators/bteq.py index c1373d51e9cd1..0496201fbae07 100644 --- a/providers/teradata/src/airflow/providers/teradata/operators/bteq.py +++ b/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): diff --git a/providers/teradata/src/airflow/providers/teradata/operators/teradata.py b/providers/teradata/src/airflow/providers/teradata/operators/teradata.py index beaca171106c6..da9d23ed3c699 100644 --- a/providers/teradata/src/airflow/providers/teradata/operators/teradata.py +++ b/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: diff --git a/providers/teradata/src/airflow/providers/teradata/operators/teradata_compute_cluster.py b/providers/teradata/src/airflow/providers/teradata/operators/teradata_compute_cluster.py index 81e4fc22b34cc..81ce3ff69e6fb 100644 --- a/providers/teradata/src/airflow/providers/teradata/operators/teradata_compute_cluster.py +++ b/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: diff --git a/providers/teradata/src/airflow/providers/teradata/transfers/azure_blob_to_teradata.py b/providers/teradata/src/airflow/providers/teradata/transfers/azure_blob_to_teradata.py index 125e63119f6bd..e046a4dcd9d72 100644 --- a/providers/teradata/src/airflow/providers/teradata/transfers/azure_blob_to_teradata.py +++ b/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 diff --git a/providers/teradata/src/airflow/providers/teradata/transfers/s3_to_teradata.py b/providers/teradata/src/airflow/providers/teradata/transfers/s3_to_teradata.py index 707de137ab063..85d6827c43faf 100644 --- a/providers/teradata/src/airflow/providers/teradata/transfers/s3_to_teradata.py +++ b/providers/teradata/src/airflow/providers/teradata/transfers/s3_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.amazon.aws.hooks.s3 import S3Hook except ModuleNotFoundError as e: @@ -30,6 +28,7 @@ raise AirflowOptionalProviderFeatureException(e) from airflow.providers.teradata.hooks.teradata import TeradataHook +from airflow.providers.teradata.version_compat import BaseOperator if TYPE_CHECKING: try: diff --git a/providers/teradata/src/airflow/providers/teradata/transfers/teradata_to_teradata.py b/providers/teradata/src/airflow/providers/teradata/transfers/teradata_to_teradata.py index d2465bd91f71e..825ffbc324d6b 100644 --- a/providers/teradata/src/airflow/providers/teradata/transfers/teradata_to_teradata.py +++ b/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: diff --git a/providers/teradata/src/airflow/providers/teradata/version_compat.py b/providers/teradata/src/airflow/providers/teradata/version_compat.py new file mode 100644 index 0000000000000..4f8d5e32bca4a --- /dev/null +++ b/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"] diff --git a/providers/teradata/tests/unit/teradata/operators/test_bteq.py b/providers/teradata/tests/unit/teradata/operators/test_bteq.py index 0e9a0ed0b150d..7afac54d311f8 100644 --- a/providers/teradata/tests/unit/teradata/operators/test_bteq.py +++ b/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;"