Move BaseHook class to task SDK

airflow-core/docs/public-airflow-interface.rst

@@ -138,7 +138,7 @@ Hooks
 
 Hooks are interfaces to external platforms and databases, implementing a common
 interface when possible and acting as building blocks for operators. All hooks
-are derived from :class:`~airflow.hooks.base.BaseHook`.
+are derived from :class:`~airflow.sdk.bases.hook.BaseHook`.
 
 Airflow has a set of Hooks that are considered public. You are free to extend their functionality
 by extending them:

airflow-core/src/airflow/cli/commands/connection_command.py

@@ -33,9 +33,9 @@
 from airflow.cli.utils import is_stdout, print_export_output
 from airflow.configuration import conf
 from airflow.exceptions import AirflowNotFoundException
-from airflow.hooks.base import BaseHook
 from airflow.models import Connection
 from airflow.providers_manager import ProvidersManager
+from airflow.sdk import BaseHook
 from airflow.secrets.local_filesystem import load_connections_dict
 from airflow.utils import cli as cli_utils, helpers, yaml
 from airflow.utils.cli import suppress_logs_and_warning
@@ -375,7 +375,8 @@ def connections_test(args) -> None:
         raise SystemExit(1)
 
     print("\nTesting...")
-    status, message = conn.test_connection()
+    # TODO: Revisit this
+    status, message = conn.test_connection()  # type: ignore
     if status is True:
         console.print("[bold green]\nConnection success!\n")
     else:

airflow-core/src/airflow/hooks/__init__.py

@@ -31,5 +31,8 @@
     "subprocess": {
         "SubprocessHook": "airflow.providers.standard.hooks.subprocess.SubprocessHook",
     },
+    "base": {
+        "BaseHook": "airflow.sdk.bases.hook.BaseHook",
+    },
 }
 add_deprecated_classes(__deprecated_classes, __name__)

airflow-core/src/airflow/hooks/base.py

@@ -19,75 +19,7 @@
 
 from __future__ import annotations
 
-import logging
-from typing import TYPE_CHECKING, Any, Protocol
-
-from airflow.utils.log.logging_mixin import LoggingMixin
-
-if TYPE_CHECKING:
-    from airflow.models.connection import Connection  # Avoid circular imports.
-
-log = logging.getLogger(__name__)
-
-
-class BaseHook(LoggingMixin):
-    """
-    Abstract base class for hooks.
-
-    Hooks are meant as an interface to
-    interact with external systems. MySqlHook, HiveHook, PigHook return
-    object that can handle the connection and interaction to specific
-    instances of these systems, and expose consistent methods to interact
-    with them.
-
-    :param logger_name: Name of the logger used by the Hook to emit logs.
-        If set to `None` (default), the logger name will fall back to
-        `airflow.task.hooks.{class.__module__}.{class.__name__}` (e.g. DbApiHook will have
-        *airflow.task.hooks.airflow.providers.common.sql.hooks.sql.DbApiHook* as logger).
-    """
-
-    def __init__(self, logger_name: str | None = None):
-        super().__init__()
-        self._log_config_logger_name = "airflow.task.hooks"
-        self._logger_name = logger_name
-
-    @classmethod
-    def get_connection(cls, conn_id: str) -> Connection:
-        """
-        Get connection, given connection id.
-
-        :param conn_id: connection id
-        :return: connection
-        """
-        from airflow.models.connection import Connection
-
-        conn = Connection.get_connection_from_secrets(conn_id)
-        log.info("Connection Retrieved '%s'", conn.conn_id)
-        return conn
-
-    @classmethod
-    def get_hook(cls, conn_id: str, hook_params: dict | None = None) -> BaseHook:
-        """
-        Return default hook for this connection id.
-
-        :param conn_id: connection id
-        :param hook_params: hook parameters
-        :return: default hook for this connection
-        """
-        connection = cls.get_connection(conn_id)
-        return connection.get_hook(hook_params=hook_params)
-
-    def get_conn(self) -> Any:
-        """Return connection for the hook."""
-        raise NotImplementedError()
-
-    @classmethod
-    def get_connection_form_widgets(cls) -> dict[str, Any]:
-        return {}
-
-    @classmethod
-    def get_ui_field_behaviour(cls) -> dict[str, Any]:
-        return {}
+from typing import Any, Protocol
 
 
 class DiscoverableHook(Protocol):

airflow-core/src/airflow/lineage/hook.py

@@ -29,8 +29,7 @@
 from airflow.utils.log.logging_mixin import LoggingMixin
 
 if TYPE_CHECKING:
-    from airflow.hooks.base import BaseHook
-    from airflow.sdk import ObjectStoragePath
+    from airflow.sdk import BaseHook, ObjectStoragePath
 
     # Store context what sent lineage.
     LineageContext = Union[BaseHook, ObjectStoragePath]

airflow-core/src/airflow/providers_manager.py

@@ -85,7 +85,7 @@ def ensure_prefix(field):
 if TYPE_CHECKING:
     from urllib.parse import SplitResult
 
-    from airflow.hooks.base import BaseHook
+    from airflow.sdk import BaseHook
     from airflow.sdk.bases.decorator import TaskDecorator
     from airflow.sdk.definitions.asset import Asset
 

airflow-core/src/airflow/utils/email.py

@@ -246,7 +246,7 @@ def send_mime_email(
 
     if conn_id is not None:
         try:
-            from airflow.hooks.base import BaseHook
+            from airflow.sdk import BaseHook
 
             airflow_conn = BaseHook.get_connection(conn_id)
             smtp_user = airflow_conn.login

airflow-core/tests/unit/always/test_connection.py

@@ -29,8 +29,8 @@
 from cryptography.fernet import Fernet
 
 from airflow.exceptions import AirflowException
-from airflow.hooks.base import BaseHook
 from airflow.models import Connection, crypto
+from airflow.sdk import BaseHook
 
 sqlite = pytest.importorskip("airflow.providers.sqlite.hooks.sqlite")
 

airflow-core/tests/unit/always/test_example_dags.py

@@ -27,8 +27,8 @@
 from packaging.specifiers import SpecifierSet
 from packaging.version import Version
 
-from airflow.hooks.base import BaseHook
 from airflow.models import Connection, DagBag
+from airflow.sdk import BaseHook
 from airflow.utils import yaml
 
 from tests_common.test_utils.asserts import assert_queries_count

airflow-core/tests/unit/cli/commands/test_rotate_fernet_key_command.py

@@ -23,8 +23,8 @@
 
 from airflow.cli import cli_parser
 from airflow.cli.commands import rotate_fernet_key_command
-from airflow.hooks.base import BaseHook
 from airflow.models import Connection, Variable
+from airflow.sdk import BaseHook
 from airflow.utils.session import provide_session
 
 from tests_common.test_utils.config import conf_vars

airflow-core/tests/unit/dag_processing/test_processor.py

@@ -281,8 +281,7 @@ def test_top_level_connection_access(
         logger_filehandle = MagicMock()
 
         def dag_in_a_fn():
-            from airflow.hooks.base import BaseHook
-            from airflow.sdk import DAG
+            from airflow.sdk import DAG, BaseHook
 
             with DAG(f"test_{BaseHook.get_connection(conn_id='my_conn').conn_id}"):
                 ...
@@ -311,8 +310,7 @@ def test_top_level_connection_access_not_found(self, tmp_path: pathlib.Path, inp
         logger_filehandle = MagicMock()
 
         def dag_in_a_fn():
-            from airflow.hooks.base import BaseHook
-            from airflow.sdk import DAG
+            from airflow.sdk import DAG, BaseHook
 
             with DAG(f"test_{BaseHook.get_connection(conn_id='my_conn').conn_id}"):
                 ...

airflow-core/tests/unit/jobs/test_triggerer_job.py

@@ -31,7 +31,6 @@
 from asgiref.sync import sync_to_async
 
 from airflow.executors import workloads
-from airflow.hooks.base import BaseHook
 from airflow.jobs.job import Job
 from airflow.jobs.triggerer_job_runner import (
     TriggerCommsDecoder,
@@ -49,6 +48,7 @@
 from airflow.providers.standard.operators.empty import EmptyOperator
 from airflow.providers.standard.operators.python import PythonOperator
 from airflow.providers.standard.triggers.temporal import DateTimeTrigger, TimeDeltaTrigger
+from airflow.sdk import BaseHook
 from airflow.triggers.base import BaseTrigger, TriggerEvent
 from airflow.triggers.testing import FailureTrigger, SuccessTrigger
 from airflow.utils import timezone

airflow-core/tests/unit/lineage/test_hook.py

@@ -22,7 +22,6 @@
 import pytest
 
 from airflow import plugins_manager
-from airflow.hooks.base import BaseHook
 from airflow.lineage import hook
 from airflow.lineage.hook import (
     AssetLineageInfo,
@@ -32,6 +31,7 @@
     NoOpCollector,
     get_hook_lineage_collector,
 )
+from airflow.sdk import BaseHook
 from airflow.sdk.definitions.asset import Asset
 
 from tests_common.test_utils.mock_plugins import mock_plugin_manager

airflow-core/tests/unit/serialization/test_dag_serialization.py

@@ -53,7 +53,6 @@
     ParamValidationError,
     SerializationError,
 )
-from airflow.hooks.base import BaseHook
 from airflow.models.asset import AssetModel
 from airflow.models.baseoperator import BaseOperator
 from airflow.models.connection import Connection
@@ -64,7 +63,7 @@
 from airflow.providers.cncf.kubernetes.pod_generator import PodGenerator
 from airflow.providers.standard.operators.bash import BashOperator
 from airflow.providers.standard.sensors.bash import BashSensor
-from airflow.sdk import AssetAlias, teardown
+from airflow.sdk import AssetAlias, BaseHook, teardown
 from airflow.sdk.bases.decorator import DecoratedOperator
 from airflow.sdk.definitions._internal.expandinput import EXPAND_INPUT_EMPTY
 from airflow.sdk.definitions.asset import Asset, AssetUniqueKey

devel-common/src/tests_common/test_utils/common_sql.py

@@ -23,7 +23,7 @@
 from airflow.models import Connection
 
 if TYPE_CHECKING:
-    from airflow.hooks.base import BaseHook
+    from airflow.sdk import BaseHook
 
 
 def mock_db_hook(hook_class: type[BaseHook], hook_params=None, conn_params=None):

providers/airbyte/src/airflow/providers/airbyte/hooks/airbyte.py

@@ -26,7 +26,11 @@
 from requests import Session
 
 from airflow.exceptions import AirflowException
-from airflow.hooks.base import BaseHook
+
+try:
+    from airflow.sdk import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook  # type: ignore
 
 T = TypeVar("T", bound=Any)
 

providers/alibaba/src/airflow/providers/alibaba/cloud/hooks/analyticdb_spark.py

@@ -34,7 +34,11 @@
 from alibabacloud_tea_openapi.models import Config
 
 from airflow.exceptions import AirflowException
-from airflow.hooks.base import BaseHook
+
+try:
+    from airflow.sdk import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook  # type: ignore
 from airflow.utils.log.logging_mixin import LoggingMixin
 
 

providers/alibaba/src/airflow/providers/alibaba/cloud/hooks/base_alibaba.py

@@ -18,7 +18,10 @@
 
 from typing import Any, NamedTuple
 
-from airflow.hooks.base import BaseHook
+try:
+    from airflow.sdk import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook  # type: ignore
 
 
 class AccessKeyCredentials(NamedTuple):

providers/alibaba/src/airflow/providers/alibaba/cloud/hooks/oss.py

@@ -26,7 +26,11 @@
 from oss2.exceptions import ClientError
 
 from airflow.exceptions import AirflowException
-from airflow.hooks.base import BaseHook
+
+try:
+    from airflow.sdk import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook  # type: ignore
 
 if TYPE_CHECKING:
     from airflow.models.connection import Connection

providers/amazon/src/airflow/providers/amazon/aws/hooks/base_aws.py

@@ -54,12 +54,16 @@
     AirflowNotFoundException,
     AirflowProviderDeprecationWarning,
 )
-from airflow.hooks.base import BaseHook
 from airflow.providers.amazon.aws.utils.connection_wrapper import AwsConnectionWrapper
 from airflow.providers.amazon.aws.utils.identifiers import generate_uuid
 from airflow.providers.amazon.aws.utils.suppress import return_on_error
 from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_0_PLUS
 from airflow.providers_manager import ProvidersManager
+
+try:
+    from airflow.sdk import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook  # type: ignore
 from airflow.utils.helpers import exactly_one
 from airflow.utils.log.logging_mixin import LoggingMixin
 

providers/amazon/src/airflow/providers/amazon/aws/hooks/sagemaker_unified_studio.py

@@ -25,9 +25,13 @@
 from sagemaker_studio.sagemaker_studio_api import SageMakerStudioAPI
 
 from airflow.exceptions import AirflowException
-from airflow.hooks.base import BaseHook
 from airflow.providers.amazon.aws.utils.sagemaker_unified_studio import is_local_runner
 
+try:
+    from airflow.sdk import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook  # type: ignore
+
 
 class SageMakerNotebookHook(BaseHook):
     """

providers/amazon/src/airflow/providers/amazon/aws/transfers/s3_to_sql.py

@@ -22,10 +22,14 @@
 from typing import TYPE_CHECKING, Callable
 
 from airflow.exceptions import AirflowException
-from airflow.hooks.base import BaseHook
 from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
 
+try:
+    from airflow.sdk import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook  # type: ignore
+
 if TYPE_CHECKING:
     from airflow.utils.context import Context
 

providers/amazon/src/airflow/providers/amazon/aws/transfers/sql_to_s3.py

@@ -27,10 +27,14 @@
 from typing_extensions import Literal
 
 from airflow.exceptions import AirflowException
-from airflow.hooks.base import BaseHook
 from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
 
+try:
+    from airflow.sdk import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook  # type: ignore
+
 if TYPE_CHECKING:
     import pandas as pd
 

providers/amazon/tests/unit/amazon/aws/hooks/test_s3.py

@@ -1748,7 +1748,7 @@ def test_delete_bucket_tagging_with_no_tags(self):
         ("rel_key", "with_conn", "with_bucket", "provide", ["kwargs_bucket", "key.txt"]),
     ],
 )
-@patch("airflow.hooks.base.BaseHook.get_connection")
+@patch("airflow.sdk.bases.hook.BaseHook.get_connection")
 def test_unify_and_provide_bucket_name_combination(
     mock_base, key_kind, has_conn, has_bucket, precedence, expected, caplog
 ):
@@ -1811,7 +1811,7 @@ def do_something(self, bucket_name=None, key=None):
         ("rel_key", "with_conn", "with_bucket", ["kwargs_bucket", "key.txt"]),
     ],
 )
-@patch("airflow.hooks.base.BaseHook.get_connection")
+@patch("airflow.sdk.bases.hook.BaseHook.get_connection")
 def test_s3_head_object_decorated_behavior(mock_conn, has_conn, has_bucket, key_kind, expected):
     if has_conn == "with_conn":
         c = Connection(extra={"service_config": {"s3": {"bucket_name": "conn_bucket"}}})