diff --git a/providers/common/compat/src/airflow/providers/common/compat/connection/__init__.py b/providers/common/compat/src/airflow/providers/common/compat/connection/__init__.py index ff0aa346ae88e..f37c09199d9bc 100644 --- a/providers/common/compat/src/airflow/providers/common/compat/connection/__init__.py +++ b/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__ = [ diff --git a/providers/common/compat/tests/unit/common/compat/connection/test_connection.py b/providers/common/compat/tests/unit/common/compat/connection/test_connection.py index dcbf9b58212ef..079608e765fef 100644 --- a/providers/common/compat/tests/unit/common/compat/connection/test_connection.py +++ b/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" diff --git a/providers/http/pyproject.toml b/providers/http/pyproject.toml index 38dd5dc776c82..cc028bd928429 100644 --- a/providers/http/pyproject.toml +++ b/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", diff --git a/providers/http/src/airflow/providers/http/hooks/http.py b/providers/http/src/airflow/providers/http/hooks/http.py index 3d22867f265d1..0c5067743f090 100644 --- a/providers/http/src/airflow/providers/http/hooks/http.py +++ b/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