diff --git a/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py b/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py index 49fa489b1a286..f4a41192acda3 100644 --- a/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py +++ b/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py @@ -61,25 +61,35 @@ def serialize(self) -> tuple[str, dict[str, Any]]: }, ) + def get_hook(self) -> DbApiHook: + """ + Return DbApiHook. + + :return: DbApiHook for this connection + """ + connection = BaseHook.get_connection(self.conn_id) + hook = connection.get_hook(hook_params=self.hook_params) + if not isinstance(hook, DbApiHook): + raise AirflowException( + f"You are trying to use `common-sql` with {hook.__class__.__name__}," + " but its provider does not support it. Please upgrade the provider to a version that" + " supports `common-sql`. The hook class should be a subclass of" + f" `{hook.__class__.__module__}.{hook.__class__.__name__}`." + f" Got {hook.__class__.__name__} hook with class hierarchy: {hook.__class__.mro()}" + ) + return hook + async def run(self) -> AsyncIterator[TriggerEvent]: try: - hook = BaseHook.get_hook(self.conn_id, hook_params=self.hook_params) - - if not isinstance(hook, DbApiHook): - raise AirflowException( - f"You are trying to use `common-sql` with {hook.__class__.__name__}," - " but its provider does not support it. Please upgrade the provider to a version that" - " supports `common-sql`. The hook class should be a subclass of" - f" `{hook.__class__.__module__}.{hook.__class__.__name__}`." - f" Got {hook.__class__.__name__} hook with class hierarchy: {hook.__class__.mro()}" - ) + hook = self.get_hook() self.log.info("Extracting data from %s", self.conn_id) self.log.info("Executing: \n %s", self.sql) self.log.info("Reading records from %s", self.conn_id) + results = hook.get_records(self.sql) - self.log.info("Reading records from %s done!", self.conn_id) + self.log.info("Reading records from %s done!", self.conn_id) self.log.debug("results: %s", results) yield TriggerEvent({"status": "success", "results": results}) except Exception as e: diff --git a/providers/common/sql/tests/unit/common/sql/triggers/test_sql.py b/providers/common/sql/tests/unit/common/sql/triggers/test_sql.py index aad3953f7f4f5..fbcb5da6fa6fd 100644 --- a/providers/common/sql/tests/unit/common/sql/triggers/test_sql.py +++ b/providers/common/sql/tests/unit/common/sql/triggers/test_sql.py @@ -18,6 +18,7 @@ from unittest import mock +from airflow.models.connection import Connection from airflow.providers.common.sql.hooks.sql import DbApiHook from airflow.providers.common.sql.triggers.sql import SQLExecuteQueryTrigger from airflow.triggers.base import TriggerEvent @@ -26,12 +27,14 @@ class TestSQLExecuteQueryTrigger: - @mock.patch("airflow.hooks.base.BaseHook.get_hook") - def test_run(self, mock_get_hook): + @mock.patch("airflow.hooks.base.BaseHook.get_connection") + def test_run(self, mock_get_connection): data = [(1, "Alice"), (2, "Bob")] + mock_connection = mock.MagicMock(spec=Connection) mock_hook = mock.MagicMock(spec=DbApiHook) mock_hook.get_records.side_effect = lambda sql: data - mock_get_hook.return_value = mock_hook + mock_get_connection.return_value = mock_connection + mock_connection.get_hook.side_effect = lambda hook_params: mock_hook trigger = SQLExecuteQueryTrigger(sql="SELECT * FROM users;", conn_id="test_conn_id") actual = run_trigger(trigger)