This is an automated email from the ASF dual-hosted git repository.
dabla pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 8afc36c1a24 Added get_async_hook in common.compat provider (#69208)
8afc36c1a24 is described below
commit 8afc36c1a246320eb214480b88aa5cebafd0cc48
Author: David Blain <[email protected]>
AuthorDate: Mon Jul 6 17:26:52 2026 +0200
Added get_async_hook in common.compat provider (#69208)
* refactor: Added get_async_hook in common.compat provider
Co-authored-by: Tzu-ping Chung <[email protected]>
---------
Co-authored-by: David Blain <[email protected]>
Co-authored-by: Tzu-ping Chung <[email protected]>
---
.../providers/common/compat/hook/__init__.py | 46 ++++++++++++++++++
.../tests/unit/common/compat/hook/__init__.py | 16 +++++++
.../tests/unit/common/compat/hook/test_hook.py | 54 ++++++++++++++++++++++
3 files changed, 116 insertions(+)
diff --git
a/providers/common/compat/src/airflow/providers/common/compat/hook/__init__.py
b/providers/common/compat/src/airflow/providers/common/compat/hook/__init__.py
new file mode 100644
index 00000000000..87a5f4d12ea
--- /dev/null
+++
b/providers/common/compat/src/airflow/providers/common/compat/hook/__init__.py
@@ -0,0 +1,46 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+import logging
+
+from airflow.providers.common.compat.sdk import BaseHook
+
+log = logging.getLogger(__name__)
+
+
+async def get_async_hook(conn_id: str, hook_params: dict | None = None) ->
BaseHook:
+ """
+ Get an asynchronous Airflow hook that is backwards compatible.
+
+ :param conn_id: The provided connection ID.
+ :param hook_params: Additional hook params.
+ :returns: BaseHook
+ """
+ from asgiref.sync import sync_to_async
+
+ if hasattr(BaseHook, "aget_hook"):
+ log.debug("Get hook using `BaseHook.aget_hook()`.")
+ return await BaseHook.aget_hook(conn_id=conn_id,
hook_params=hook_params)
+ log.debug("Get hook using `BaseHook.get_hook()`.")
+ return await sync_to_async(BaseHook.get_hook)(conn_id=conn_id,
hook_params=hook_params)
+
+
+__all__ = [
+ "get_async_hook",
+]
diff --git a/providers/common/compat/tests/unit/common/compat/hook/__init__.py
b/providers/common/compat/tests/unit/common/compat/hook/__init__.py
new file mode 100644
index 00000000000..13a83393a91
--- /dev/null
+++ b/providers/common/compat/tests/unit/common/compat/hook/__init__.py
@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
diff --git a/providers/common/compat/tests/unit/common/compat/hook/test_hook.py
b/providers/common/compat/tests/unit/common/compat/hook/test_hook.py
new file mode 100644
index 00000000000..da71a555630
--- /dev/null
+++ b/providers/common/compat/tests/unit/common/compat/hook/test_hook.py
@@ -0,0 +1,54 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from unittest import mock
+
+import pytest
+
+from airflow.providers.common.compat.hook import get_async_hook
+from airflow.providers.common.compat.sdk import BaseHook
+
+_MOCK_HOOK = mock.MagicMock(spec=BaseHook)
+
+
+class MockAgetBaseHook:
+ def __init__(self, *args, **kwargs):
+ self.aget_hook = mock.AsyncMock(return_value=_MOCK_HOOK)
+
+
+class MockBaseHook:
+ def __init__(self, *args, **kwargs):
+ self.get_hook = mock.MagicMock(return_value=_MOCK_HOOK)
+
+
+class TestGetAsyncHook:
+ @pytest.mark.parametrize("hook_params", [None, {"key": "value"}])
+ @mock.patch("airflow.providers.common.compat.hook.BaseHook",
new_callable=MockAgetBaseHook)
+ @pytest.mark.asyncio
+ async def test_get_async_hook_uses_aget_hook_when_available(self,
mock_hook_class, hook_params):
+ result = await get_async_hook("test_conn", hook_params=hook_params)
+ assert result is _MOCK_HOOK
+ mock_hook_class.aget_hook.assert_called_once_with(conn_id="test_conn",
hook_params=hook_params)
+
+ @pytest.mark.parametrize("hook_params", [None, {"key": "value"}])
+ @mock.patch("airflow.providers.common.compat.hook.BaseHook",
new_callable=MockBaseHook)
+ @pytest.mark.asyncio
+ async def test_get_async_hook_falls_back_to_get_hook(self,
mock_hook_class, hook_params):
+ result = await get_async_hook("test_conn", hook_params=hook_params)
+ assert result is _MOCK_HOOK
+ mock_hook_class.get_hook.assert_called_once_with(conn_id="test_conn",
hook_params=hook_params)