[airflow] branch main updated: Add `sql_hook_params` parameter to `S3ToSqlOperator` (#33427)

[po...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

potiuk 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 1407e277ae Add `sql_hook_params` parameter to `S3ToSqlOperator` (#33427)
1407e277ae is described below

commit 1407e277aeb059cbfd1bb96fb3f43c4bf4f15cea
Author: Alex Begg <45...@users.noreply.github.com>
AuthorDate: Sat Aug 19 00:52:40 2023 -0700

    Add `sql_hook_params` parameter to `S3ToSqlOperator` (#33427)
    
    Adding `sql_hook_params` parameter to `SqlToS3Operator`. This will allow you to pass extra config params to the underlying SQL hook.
    
    This uses the same "sql_hook_params" parameter name as already used in `SqlToSlackOperator`.
---
 airflow/providers/amazon/aws/transfers/s3_to_sql.py    |  7 ++++++-
 tests/providers/amazon/aws/transfers/test_s3_to_sql.py | 17 ++++++++++++++---
 2 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/airflow/providers/amazon/aws/transfers/s3_to_sql.py b/airflow/providers/amazon/aws/transfers/s3_to_sql.py
index 8e0613ea6d..667e3c174b 100644
--- a/airflow/providers/amazon/aws/transfers/s3_to_sql.py
+++ b/airflow/providers/amazon/aws/transfers/s3_to_sql.py
@@ -44,6 +44,8 @@ class S3ToSqlOperator(BaseOperator):
     :param s3_bucket: reference to a specific S3 bucket
     :param s3_key: reference to a specific S3 key
     :param sql_conn_id: reference to a specific SQL database. Must be of type DBApiHook
+    :param sql_hook_params: Extra config params to be passed to the underlying hook.
+        Should match the desired hook constructor params.
     :param aws_conn_id: reference to a specific S3 / AWS connection
     :param column_list: list of column names to use in the insert SQL.
     :param commit_every: The maximum number of rows to insert in one
@@ -83,6 +85,7 @@ class S3ToSqlOperator(BaseOperator):
         commit_every: int = 1000,
         schema: str | None = None,
         sql_conn_id: str = "sql_default",
+        sql_hook_params: dict | None = None,
         aws_conn_id: str = "aws_default",
         **kwargs,
     ) -> None:
@@ -96,6 +99,7 @@ class S3ToSqlOperator(BaseOperator):
         self.column_list = column_list
         self.commit_every = commit_every
         self.parser = parser
+        self.sql_hook_params = sql_hook_params
 
     def execute(self, context: Context) -> None:
         self.log.info("Loading %s to SQL table %s...", self.s3_key, self.table)
@@ -120,7 +124,8 @@ class S3ToSqlOperator(BaseOperator):
     @cached_property
     def db_hook(self):
         self.log.debug("Get connection for %s", self.sql_conn_id)
-        hook = BaseHook.get_hook(self.sql_conn_id)
+        conn = BaseHook.get_connection(self.sql_conn_id)
+        hook = conn.get_hook(hook_params=self.sql_hook_params)
         if not callable(getattr(hook, "insert_rows", None)):
             raise AirflowException(
                 "This hook is not supported. The hook class must have an `insert_rows` method."
diff --git a/tests/providers/amazon/aws/transfers/test_s3_to_sql.py b/tests/providers/amazon/aws/transfers/test_s3_to_sql.py
index 0239624034..3e6e6913e7 100644
--- a/tests/providers/amazon/aws/transfers/test_s3_to_sql.py
+++ b/tests/providers/amazon/aws/transfers/test_s3_to_sql.py
@@ -76,7 +76,7 @@ class TestS3ToSqlTransfer:
         return bad_hook
 
     @patch("airflow.providers.amazon.aws.transfers.s3_to_sql.NamedTemporaryFile")
-    @patch("airflow.providers.amazon.aws.transfers.s3_to_sql.BaseHook")
+    @patch("airflow.models.connection.Connection.get_hook")
     @patch("airflow.providers.amazon.aws.transfers.s3_to_sql.S3Hook.get_key")
     def test_execute(self, mock_get_key, mock_hook, mock_tempfile, mock_parser):
 
@@ -93,7 +93,7 @@ class TestS3ToSqlTransfer:
 
         mock_parser.assert_called_once_with(mock_tempfile.return_value.__enter__.return_value.name)
 
-        mock_hook.get_hook.return_value.insert_rows.assert_called_once_with(
+        mock_hook.return_value.insert_rows.assert_called_once_with(
             table=self.s3_to_sql_transfer_kwargs["table"],
             schema=self.s3_to_sql_transfer_kwargs["schema"],
             target_fields=self.s3_to_sql_transfer_kwargs["column_list"],
@@ -102,13 +102,24 @@ class TestS3ToSqlTransfer:
         )
 
     @patch("airflow.providers.amazon.aws.transfers.s3_to_sql.NamedTemporaryFile")
-    @patch("airflow.providers.amazon.aws.transfers.s3_to_sql.BaseHook.get_hook", return_value=mock_bad_hook)
+    @patch("airflow.models.connection.Connection.get_hook", return_value=mock_bad_hook)
     @patch("airflow.providers.amazon.aws.transfers.s3_to_sql.S3Hook.get_key")
     def test_execute_with_bad_hook(self, mock_get_key, mock_bad_hook, mock_tempfile, mock_parser):
 
         with pytest.raises(AirflowException):
             S3ToSqlOperator(parser=mock_parser, **self.s3_to_sql_transfer_kwargs).execute({})
 
+    def test_hook_params(self, mock_parser):
+        op = S3ToSqlOperator(
+            parser=mock_parser,
+            sql_hook_params={
+                "log_sql": False,
+            },
+            **self.s3_to_sql_transfer_kwargs,
+        )
+        hook = op.db_hook
+        assert hook.log_sql == op.sql_hook_params["log_sql"]
+
     def teardown_method(self):
         with create_session() as session:
             (