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 0000000000000..87a5f4d12eae2 --- /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 0000000000000..13a83393a9124 --- /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 0000000000000..da71a55563063 --- /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)