[GitHub] [airflow] uranusjr commented on a change in pull request #17592: Support for passing arguments to SqlSensor underlying hooks

[GitBox <gi...@apache.org>]

uranusjr commented on a change in pull request #17592:
URL: https://github.com/apache/airflow/pull/17592#discussion_r690955205



##########
File path: airflow/models/connection.py
##########
@@ -286,8 +286,11 @@ def rotate_fernet_key(self):
         if self._extra and self.is_extra_encrypted:
             self._extra = fernet.rotate(self._extra.encode('utf-8')).decode()
 
-    def get_hook(self):
-        """Return hook based on conn_type."""
+    def get_hook(self, hook_kwargs: Dict):

Review comment:
       Should this not be `**hook_kwargs`

##########
File path: airflow/sensors/sql.py
##########
@@ -90,7 +105,7 @@ def _get_hook(self):
                 f"Connection type ({conn.conn_type}) is not supported by SqlSensor. "
                 + f"Supported connection types: {list(allowed_conn_type)}"
             )
-        return conn.get_hook()
+        return conn.get_hook(**self.hook_kwargs)

Review comment:
       … or this should be `conn.get_hook(self.hook_kwargs)`?

##########
File path: tests/sensors/test_sql_sensor.py
##########
@@ -242,6 +242,22 @@ def test_sql_sensor_postgres_poke_invalid_success(self, mock_hook):
             op.poke(None)
         assert "self.success is present, but not callable -> [1]" == str(ctx.value)
 
+    @mock.patch('airflow.sensors.sql.BaseHook')
+    def test_sql_sensor_bigquery_hook_kwargs(self, mock_hook):
+        op = SqlSensor(
+            task_id='sql_sensor_check',
+            conn_id='postgres_default',
+            sql="SELECT 1",
+            hook_kwargs={
+                'use_legacy_sql': False,
+                'location': 'test_location',
+            },
+        )
+
+        mock_hook.get_connection('google_cloud_default').conn_type = "google_cloud_platform"
+        assert op._get_hook().use_legacy_sql
+        assert op._get_hook().location == 'test_location'

Review comment:
       And this test should ensure the kwargs are correctly passed onto the hook.




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org