Honor async hook subclass overrides in get_async_connection

The async connection helper resolved connections through the literal
BaseHook class, so a hook that overrides aget_connection/get_connection
was ignored in the triggerer even though the sync path honors it on the
worker. Accept the calling hook (instance or class) and dispatch through
it, defaulting to BaseHook, and pass the hook instance from HttpAsyncHook
so a custom subclass behaves the same in the triggerer as on the worker.
The debug log now names the hook that actually resolved the connection.

Signed-off-by: Anas Khan <83116240+anxkhn@users.noreply.github.com>

Author: anxkhn

providers/common/compat/src/airflow/providers/common/compat/connection/__init__.py

@@ -29,20 +29,25 @@
 log = logging.getLogger(__name__)
 
 
-async def get_async_connection(conn_id: str) -> Connection:
+async def get_async_connection(conn_id: str, hook: BaseHook | type[BaseHook] | None = None) -> Connection:
     """
     Get an asynchronous Airflow connection that is backwards compatible.
 
     :param conn_id: The provided connection ID.
+    :param hook: Hook (class or instance) to resolve the connection through, so a
+        subclass override of ``aget_connection``/``get_connection`` is honored.
+        Defaults to ``BaseHook``.
     :returns: Connection
     """
     from asgiref.sync import sync_to_async
 
-    if hasattr(BaseHook, "aget_connection"):
-        log.debug("Get connection using `BaseHook.aget_connection().")
-        return await BaseHook.aget_connection(conn_id=conn_id)
-    log.debug("Get connection using `BaseHook.get_connection().")
-    return await sync_to_async(BaseHook.get_connection)(conn_id=conn_id)
+    hook = hook or BaseHook
+    hook_name = hook.__name__ if isinstance(hook, type) else type(hook).__name__
+    if hasattr(hook, "aget_connection"):
+        log.debug("Get connection using `%s.aget_connection()`.", hook_name)
+        return await hook.aget_connection(conn_id=conn_id)
+    log.debug("Get connection using `%s.get_connection()`.", hook_name)
+    return await sync_to_async(hook.get_connection)(conn_id=conn_id)
 
 
 __all__ = [

providers/common/compat/tests/unit/common/compat/connection/test_connection.py

@@ -57,7 +57,7 @@ async def test_get_async_connection_with_aget(self, _, caplog):
             conn = await get_async_connection("test_conn")
         assert conn.password == "secret_token_aget"
         assert conn.conn_type == "http"
-        assert "Get connection using `BaseHook.aget_connection()." in caplog.text
+        assert "Get connection using `MockAgetBaseHook.aget_connection()`." in caplog.text
 
     @mock.patch("airflow.providers.common.compat.connection.BaseHook", new_callable=MockBaseHook)
     @pytest.mark.asyncio
@@ -66,4 +66,26 @@ async def test_get_async_connection_with_get_connection(self, _, caplog):
             conn = await get_async_connection("test_conn")
         assert conn.password == "secret_token"
         assert conn.conn_type == "http"
-        assert "Get connection using `BaseHook.get_connection()." in caplog.text
+        assert "Get connection using `MockBaseHook.get_connection()`." in caplog.text
+
+    @mock.patch("airflow.providers.common.compat.connection.BaseHook", new_callable=MockAgetBaseHook)
+    @pytest.mark.asyncio
+    async def test_get_async_connection_honors_passed_hook(self, _):
+        class OverrideHook:
+            @classmethod
+            async def aget_connection(cls, conn_id: str):
+                return Connection(conn_id="override", conn_type="http", password="override_token")
+
+        conn = await get_async_connection("test_conn", hook=OverrideHook)
+        assert conn.password == "override_token"
+
+    @mock.patch("airflow.providers.common.compat.connection.BaseHook", new_callable=MockBaseHook)
+    @pytest.mark.asyncio
+    async def test_get_async_connection_honors_passed_hook_get_connection(self, _):
+        class OverrideHook:
+            @classmethod
+            def get_connection(cls, conn_id: str):
+                return Connection(conn_id="override", conn_type="http", password="override_token")
+
+        conn = await get_async_connection("test_conn", hook=OverrideHook)
+        assert conn.password == "override_token"

providers/http/pyproject.toml

@@ -60,7 +60,7 @@ requires-python = ">=3.10"
 # After you modify the dependencies, and rebuild your Breeze CI image with ``breeze ci-image build``
 dependencies = [
     "apache-airflow>=2.11.0",
-    "apache-airflow-providers-common-compat>=1.12.0",
+    "apache-airflow-providers-common-compat>=1.12.0",  # use next version
     # The 2.26.0 release of requests got rid of the chardet LGPL mandatory dependency, allowing us to
     # release it as a requirement for airflow
     "requests>=2.32.0,<3",

providers/http/src/airflow/providers/http/hooks/http.py

@@ -615,7 +615,7 @@ async def config(self) -> SessionConfig:
             extra_options: dict[str, Any] = {}
 
             if self.http_conn_id:
-                conn = await get_async_connection(conn_id=self.http_conn_id)
+                conn = await get_async_connection(conn_id=self.http_conn_id, hook=self)
 
                 if conn.host and "://" in conn.host:
                     base_url = conn.host