Fix get_hook() got an unexpected keyword argument 'hook_params' in SQLExecuteQueryTrigger

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:

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)