[airflow] branch main updated: Add parameter to turn off SQL query logging (#24570)

[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 53284cf272 Add parameter to turn off SQL query logging (#24570)
53284cf272 is described below

commit 53284cf27260122ff0a56d397e677fb6ad667370
Author: Holly Evans <39...@users.noreply.github.com>
AuthorDate: Mon Jul 4 15:46:34 2022 -0500

    Add parameter to turn off SQL query logging (#24570)
---
 airflow/hooks/dbapi.py          | 8 ++++++--
 newsfragments/24570.feature.rst | 1 +
 tests/hooks/test_dbapi.py       | 6 ++++++
 tests/operators/test_sql.py     | 2 ++
 4 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/airflow/hooks/dbapi.py b/airflow/hooks/dbapi.py
index 0b9ce4377b..d3d3fa5fc8 100644
--- a/airflow/hooks/dbapi.py
+++ b/airflow/hooks/dbapi.py
@@ -56,6 +56,7 @@ class DbApiHook(BaseHook):
     :param schema: Optional DB schema that overrides the schema specified in the connection. Make sure that
         if you change the schema parameter value in the constructor of the derived Hook, such change
         should be done before calling the ``DBApiHook.__init__()``.
+    :param log_sql: Whether to log SQL query when it's executed. Defaults to *True*.
     """
 
     # Override to provide the connection name.
@@ -69,7 +70,7 @@ class DbApiHook(BaseHook):
     # Override with db-specific query to check connection
     _test_connection_sql = "select 1"
 
-    def __init__(self, *args, schema: Optional[str] = None, **kwargs):
+    def __init__(self, *args, schema: Optional[str] = None, log_sql: bool = True, **kwargs):
         super().__init__()
         if not self.conn_name_attr:
             raise AirflowException("conn_name_attr is not defined")
@@ -84,6 +85,7 @@ class DbApiHook(BaseHook):
         # from kwargs and store it on its own. We do not run "pop" here as we want to give the
         # Hook deriving from the DBApiHook to still have access to the field in it's constructor
         self.__schema = schema
+        self.log_sql = log_sql
 
     def get_conn(self):
         """Returns a connection object"""
@@ -228,7 +230,9 @@ class DbApiHook(BaseHook):
 
     def _run_command(self, cur, sql_statement, parameters):
         """Runs a statement using an already open cursor."""
-        self.log.info("Running statement: %s, parameters: %s", sql_statement, parameters)
+        if self.log_sql:
+            self.log.info("Running statement: %s, parameters: %s", sql_statement, parameters)
+
         if parameters:
             cur.execute(sql_statement, parameters)
         else:
diff --git a/newsfragments/24570.feature.rst b/newsfragments/24570.feature.rst
new file mode 100644
index 0000000000..382e0e8b2b
--- /dev/null
+++ b/newsfragments/24570.feature.rst
@@ -0,0 +1 @@
+DbApiHook accepts log_sql to turn off logging SQL queries.
diff --git a/tests/hooks/test_dbapi.py b/tests/hooks/test_dbapi.py
index a17c24aedb..ad5e7ba6af 100644
--- a/tests/hooks/test_dbapi.py
+++ b/tests/hooks/test_dbapi.py
@@ -44,6 +44,7 @@ class TestDbApiHook(unittest.TestCase):
                 return conn
 
         self.db_hook = UnitTestDbApiHook()
+        self.db_hook_no_log_sql = UnitTestDbApiHook(log_sql=False)
         self.db_hook_schema_override = UnitTestDbApiHook(schema='schema-override')
 
     def test_get_records(self):
@@ -346,6 +347,11 @@ class TestDbApiHook(unittest.TestCase):
         self.db_hook.run(statement)
         assert self.db_hook.log.info.call_count == 2
 
+    def test_run_no_log(self):
+        statement = 'SQL'
+        self.db_hook_no_log_sql.run(statement)
+        assert self.db_hook_no_log_sql.log.info.call_count == 1
+
     def test_run_with_handler(self):
         sql = 'SQL'
         param = ('p1', 'p2')
diff --git a/tests/operators/test_sql.py b/tests/operators/test_sql.py
index 2e73c3ac33..43d202819a 100644
--- a/tests/operators/test_sql.py
+++ b/tests/operators/test_sql.py
@@ -97,12 +97,14 @@ class TestSQLCheckOperatorDbHook:
             'database': 'database',
             'role': 'role',
             'schema': 'schema',
+            'log_sql': False,
         }
         assert self._operator._hook.conn_type == 'snowflake'
         assert self._operator._hook.warehouse == 'warehouse'
         assert self._operator._hook.database == 'database'
         assert self._operator._hook.role == 'role'
         assert self._operator._hook.schema == 'schema'
+        assert not self._operator._hook.log_sql
 
     def test_sql_operator_hook_params_biguery(self, mock_get_conn):
         mock_get_conn.return_value = Connection(