diff --git a/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py b/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py index 28cefecf8021f..86c01d396da21 100644 --- a/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py +++ b/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py @@ -41,11 +41,7 @@ ) from airflow.providers.common.sql.dialects.dialect import Dialect from airflow.providers.common.sql.hooks import handlers - -try: - from airflow.sdk import BaseHook -except ImportError: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] +from airflow.providers.common.sql.version_compat import BaseHook from airflow.utils.module_loading import import_string if TYPE_CHECKING: diff --git a/providers/common/sql/src/airflow/providers/common/sql/operators/generic_transfer.py b/providers/common/sql/src/airflow/providers/common/sql/operators/generic_transfer.py index 748a146e0d060..414ef8bad8462 100644 --- a/providers/common/sql/src/airflow/providers/common/sql/operators/generic_transfer.py +++ b/providers/common/sql/src/airflow/providers/common/sql/operators/generic_transfer.py @@ -24,12 +24,7 @@ from airflow.exceptions import AirflowException from airflow.providers.common.sql.hooks.sql import DbApiHook from airflow.providers.common.sql.triggers.sql import SQLExecuteQueryTrigger -from airflow.providers.common.sql.version_compat import BaseOperator - -try: - from airflow.sdk import BaseHook -except ImportError: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] +from airflow.providers.common.sql.version_compat import BaseHook, BaseOperator if TYPE_CHECKING: import jinja2 diff --git a/providers/common/sql/src/airflow/providers/common/sql/operators/sql.py b/providers/common/sql/src/airflow/providers/common/sql/operators/sql.py index ce0cbe5214bbb..fa6963c86e235 100644 --- a/providers/common/sql/src/airflow/providers/common/sql/operators/sql.py +++ b/providers/common/sql/src/airflow/providers/common/sql/operators/sql.py @@ -27,14 +27,9 @@ from airflow.models import SkipMixin from airflow.providers.common.sql.hooks.handlers import fetch_all_handler, return_single_query_results from airflow.providers.common.sql.hooks.sql import DbApiHook -from airflow.providers.common.sql.version_compat import BaseOperator +from airflow.providers.common.sql.version_compat import BaseHook, BaseOperator from airflow.utils.helpers import merge_dicts -try: - from airflow.sdk import BaseHook -except ImportError: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] - if TYPE_CHECKING: from airflow.providers.openlineage.extractors import OperatorLineage from airflow.utils.context import Context diff --git a/providers/common/sql/src/airflow/providers/common/sql/sensors/sql.py b/providers/common/sql/src/airflow/providers/common/sql/sensors/sql.py index 687b39f63c0d1..4ef79769ef329 100644 --- a/providers/common/sql/src/airflow/providers/common/sql/sensors/sql.py +++ b/providers/common/sql/src/airflow/providers/common/sql/sensors/sql.py @@ -22,18 +22,7 @@ from airflow.exceptions import AirflowException from airflow.providers.common.sql.hooks.sql import DbApiHook - -try: - from airflow.sdk import BaseHook -except ImportError: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] - -from airflow.providers.common.sql.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.common.sql.version_compat import BaseHook, BaseSensorOperator if TYPE_CHECKING: from airflow.utils.context import Context diff --git a/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py b/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py index 64618fc8f5f38..f137293617e30 100644 --- a/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py +++ b/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py @@ -21,11 +21,7 @@ from airflow.exceptions import AirflowException from airflow.providers.common.sql.hooks.sql import DbApiHook - -try: - from airflow.sdk import BaseHook -except ImportError: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] +from airflow.providers.common.sql.version_compat import BaseHook from airflow.triggers.base import BaseTrigger, TriggerEvent if TYPE_CHECKING: diff --git a/providers/common/sql/src/airflow/providers/common/sql/version_compat.py b/providers/common/sql/src/airflow/providers/common/sql/version_compat.py index b326387fea2a1..03bf6973a704b 100644 --- a/providers/common/sql/src/airflow/providers/common/sql/version_compat.py +++ b/providers/common/sql/src/airflow/providers/common/sql/version_compat.py @@ -33,6 +33,12 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 0, 0) +AIRFLOW_V_3_1_PLUS: bool = get_base_airflow_version_tuple() >= (3, 1, 0) + +if AIRFLOW_V_3_1_PLUS: + from airflow.sdk import BaseHook +else: + from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] if AIRFLOW_V_3_0_PLUS: from airflow.sdk import BaseOperator, BaseSensorOperator @@ -42,6 +48,8 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: __all__ = [ "AIRFLOW_V_3_0_PLUS", + "AIRFLOW_V_3_1_PLUS", + "BaseHook", "BaseOperator", "BaseSensorOperator", ] diff --git a/providers/common/sql/tests/unit/common/sql/hooks/test_dbapi.py b/providers/common/sql/tests/unit/common/sql/hooks/test_dbapi.py index 29a9f8350e8fe..69c3147a40e19 100644 --- a/providers/common/sql/tests/unit/common/sql/hooks/test_dbapi.py +++ b/providers/common/sql/tests/unit/common/sql/hooks/test_dbapi.py @@ -30,11 +30,7 @@ from airflow.providers.common.sql.dialects.dialect import Dialect from airflow.providers.common.sql.hooks.handlers import fetch_all_handler, fetch_one_handler from airflow.providers.common.sql.hooks.sql import DbApiHook - -try: - from airflow.sdk import BaseHook -except ImportError: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] +from airflow.providers.common.sql.version_compat import BaseHook class DbApiHookInProvider(DbApiHook):