diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/fs/adls.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/fs/adls.py index ef9fe01d4872c..299d9f775df76 100644 --- a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/fs/adls.py +++ b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/fs/adls.py @@ -21,11 +21,7 @@ from azure.identity import ClientSecretCredential from airflow.providers.microsoft.azure.utils import get_field, parse_blob_account_url - -try: - from airflow.sdk import BaseHook -except ImportError: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] +from airflow.providers.microsoft.azure.version_compat import BaseHook if TYPE_CHECKING: from fsspec import AbstractFileSystem diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/adx.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/adx.py index 819725ca2c394..d83cab1c05b9e 100644 --- a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/adx.py +++ b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/adx.py @@ -38,11 +38,7 @@ add_managed_identity_connection_widgets, get_sync_default_azure_credential, ) - -try: - from airflow.sdk import BaseHook -except ImportError: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] +from airflow.providers.microsoft.azure.version_compat import BaseHook if TYPE_CHECKING: from azure.kusto.data.response import KustoResponseDataSet diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/asb.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/asb.py index a49c9598d4936..cac10ca0b3e8a 100644 --- a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/asb.py +++ b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/asb.py @@ -42,11 +42,7 @@ get_field, get_sync_default_azure_credential, ) - -try: - from airflow.sdk import BaseHook -except ImportError: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] +from airflow.providers.microsoft.azure.version_compat import BaseHook if TYPE_CHECKING: import datetime diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/base_azure.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/base_azure.py index 2a625becba763..2a59234e0a4c7 100644 --- a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/base_azure.py +++ b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/base_azure.py @@ -26,11 +26,7 @@ AzureIdentityCredentialAdapter, add_managed_identity_connection_widgets, ) - -try: - from airflow.sdk import BaseHook -except ImportError: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] +from airflow.providers.microsoft.azure.version_compat import BaseHook class AzureBaseHook(BaseHook): diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/batch.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/batch.py index e60ff63597865..96494e56f5e48 100644 --- a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/batch.py +++ b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/batch.py @@ -30,11 +30,7 @@ add_managed_identity_connection_widgets, get_field, ) - -try: - from airflow.sdk import BaseHook -except ImportError: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] +from airflow.providers.microsoft.azure.version_compat import BaseHook from airflow.utils import timezone if TYPE_CHECKING: diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_registry.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_registry.py index d53f1dd31766b..521665caf12d0 100644 --- a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_registry.py +++ b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_registry.py @@ -30,11 +30,7 @@ get_field, get_sync_default_azure_credential, ) - -try: - from airflow.sdk import BaseHook -except ImportError: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] +from airflow.providers.microsoft.azure.version_compat import BaseHook class AzureContainerRegistryHook(BaseHook): diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_volume.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_volume.py index c652940cb5200..ed6a063011bd6 100644 --- a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_volume.py +++ b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_volume.py @@ -26,11 +26,7 @@ get_field, get_sync_default_azure_credential, ) - -try: - from airflow.sdk import BaseHook -except ImportError: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] +from airflow.providers.microsoft.azure.version_compat import BaseHook class AzureContainerVolumeHook(BaseHook): diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/cosmos.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/cosmos.py index 296c49297ad35..8c695e1cd686b 100644 --- a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/cosmos.py +++ b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/cosmos.py @@ -41,11 +41,7 @@ get_field, get_sync_default_azure_credential, ) - -try: - from airflow.sdk import BaseHook -except ImportError: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] +from airflow.providers.microsoft.azure.version_compat import BaseHook if TYPE_CHECKING: PartitionKeyType = str | list[str] diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/data_factory.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/data_factory.py index fc33473f5f957..c49b066ea2662 100644 --- a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/data_factory.py +++ b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/data_factory.py @@ -54,11 +54,7 @@ get_async_default_azure_credential, get_sync_default_azure_credential, ) - -try: - from airflow.sdk import BaseHook -except ImportError: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] +from airflow.providers.microsoft.azure.version_compat import BaseHook if TYPE_CHECKING: from azure.core.polling import LROPoller diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/data_lake.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/data_lake.py index c132a68f532d6..43fad25cf57a6 100644 --- a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/data_lake.py +++ b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/data_lake.py @@ -39,11 +39,7 @@ get_field, get_sync_default_azure_credential, ) - -try: - from airflow.sdk import BaseHook -except ImportError: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] +from airflow.providers.microsoft.azure.version_compat import BaseHook Credentials = ClientSecretCredential | AzureIdentityCredentialAdapter | DefaultAzureCredential diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/fileshare.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/fileshare.py index f45c9766e7d72..439be31bb0317 100644 --- a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/fileshare.py +++ b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/fileshare.py @@ -25,11 +25,7 @@ add_managed_identity_connection_widgets, get_sync_default_azure_credential, ) - -try: - from airflow.sdk import BaseHook -except ImportError: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] +from airflow.providers.microsoft.azure.version_compat import BaseHook class AzureFileShareHook(BaseHook): diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/msgraph.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/msgraph.py index 71f286fb641ab..2674b304f7acc 100644 --- a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/msgraph.py +++ b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/msgraph.py @@ -50,11 +50,7 @@ AirflowException, AirflowNotFoundException, ) - -try: - from airflow.sdk import BaseHook -except ImportError: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] +from airflow.providers.microsoft.azure.version_compat import BaseHook if TYPE_CHECKING: from azure.identity._internal.client_credential_base import ClientCredentialBase diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/synapse.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/synapse.py index f470e4e375e94..ceb86142aecb9 100644 --- a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/synapse.py +++ b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/synapse.py @@ -30,11 +30,7 @@ get_field, get_sync_default_azure_credential, ) - -try: - from airflow.sdk import BaseHook -except ImportError: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] +from airflow.providers.microsoft.azure.version_compat import BaseHook if TYPE_CHECKING: from azure.synapse.artifacts.models import CreateRunResponse, PipelineRun diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/wasb.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/wasb.py index b981e1bd049e9..9db0130ce77d0 100644 --- a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/wasb.py +++ b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/wasb.py @@ -52,11 +52,7 @@ get_sync_default_azure_credential, parse_blob_account_url, ) - -try: - from airflow.sdk import BaseHook -except ImportError: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] +from airflow.providers.microsoft.azure.version_compat import BaseHook if TYPE_CHECKING: from azure.core.credentials import TokenCredential diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/data_factory.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/data_factory.py index b733e74f39104..8ad795e5a467e 100644 --- a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/data_factory.py +++ b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/data_factory.py @@ -31,12 +31,10 @@ get_field, ) from airflow.providers.microsoft.azure.triggers.data_factory import AzureDataFactoryTrigger - -try: - from airflow.sdk import BaseHook -except ImportError: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] -from airflow.providers.microsoft.azure.version_compat import BaseOperator +from airflow.providers.microsoft.azure.version_compat import ( + BaseHook, + BaseOperator, +) from airflow.utils.log.logging_mixin import LoggingMixin if TYPE_CHECKING: diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/synapse.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/synapse.py index 85f40863526b2..17c80b8bd45e6 100644 --- a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/synapse.py +++ b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/synapse.py @@ -29,12 +29,10 @@ AzureSynapsePipelineRunStatus, AzureSynapseSparkBatchRunStatus, ) -from airflow.providers.microsoft.azure.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.microsoft.azure.version_compat import ( + BaseHook, + BaseOperator, +) if TYPE_CHECKING: from azure.synapse.spark.models import SparkBatchJobOptions diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/version_compat.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/version_compat.py index 8a5eeb817bfea..3825100866897 100644 --- a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/version_compat.py +++ b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/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 = 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 ( @@ -46,6 +52,7 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: __all__ = [ "AIRFLOW_V_3_0_PLUS", + "BaseHook", "BaseOperator", "BaseOperatorLink", "BaseSensorOperator",