From 9c241b5cb76c952e2320c75fafbd9fa1869a4752 Mon Sep 17 00:00:00 2001 From: holly-evans Date: Mon, 20 Jun 2022 14:39:22 -0500 Subject: [PATCH 1/7] Add parameter to turn off SQL query logging --- airflow/hooks/dbapi.py | 5 ++++- tests/hooks/test_dbapi.py | 6 ++++++ tests/operators/test_sql.py | 2 ++ 3 files changed, 12 insertions(+), 1 deletion(-) diff --git a/airflow/hooks/dbapi.py b/airflow/hooks/dbapi.py index 0b9ce4377be23..52feb3883f192 100644 --- a/airflow/hooks/dbapi.py +++ b/airflow/hooks/dbapi.py @@ -84,6 +84,7 @@ def __init__(self, *args, schema: Optional[str] = None, **kwargs): # 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 = kwargs.get('log_sql', True) def get_conn(self): """Returns a connection object""" @@ -228,7 +229,9 @@ def run(self, sql, autocommit=False, parameters=None, handler=None): 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/tests/hooks/test_dbapi.py b/tests/hooks/test_dbapi.py index a17c24aedb384..23f9c0a215550 100644 --- a/tests/hooks/test_dbapi.py +++ b/tests/hooks/test_dbapi.py @@ -44,6 +44,7 @@ def get_conn(self): return conn self.db_hook = UnitTestDbApiHook() + self.db_hook_log_sql = UnitTestDbApiHook(log_sql=False) self.db_hook_schema_override = UnitTestDbApiHook(schema='schema-override') def test_get_records(self): @@ -346,6 +347,11 @@ def test_run_log(self): 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_log_sql.run(statement) + assert self.db_hook_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 2e73c3ac33487..43d202819afdc 100644 --- a/tests/operators/test_sql.py +++ b/tests/operators/test_sql.py @@ -97,12 +97,14 @@ def test_sql_operator_hook_params_snowflake(self, mock_get_conn): '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( From 12b3b54c5d90b73b72cc87739b4001e374221360 Mon Sep 17 00:00:00 2001 From: holly-evans Date: Mon, 20 Jun 2022 15:58:49 -0500 Subject: [PATCH 2/7] Add newsfragment --- newsfragments/24570.feature.rst | 1 + 1 file changed, 1 insertion(+) create mode 100644 newsfragments/24570.feature.rst diff --git a/newsfragments/24570.feature.rst b/newsfragments/24570.feature.rst new file mode 100644 index 0000000000000..382e0e8b2b1d9 --- /dev/null +++ b/newsfragments/24570.feature.rst @@ -0,0 +1 @@ +DbApiHook accepts log_sql to turn off logging SQL queries. From 61037928a8fe3c7b99d5a3cffbf3bf11aba850e4 Mon Sep 17 00:00:00 2001 From: holly-evans Date: Tue, 21 Jun 2022 11:51:41 -0500 Subject: [PATCH 3/7] Move log_sql to constructor and document --- airflow/hooks/dbapi.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/airflow/hooks/dbapi.py b/airflow/hooks/dbapi.py index 52feb3883f192..e7d89f476d85c 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 during `run`. 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: Optional[bool] = True, **kwargs): super().__init__() if not self.conn_name_attr: raise AirflowException("conn_name_attr is not defined") @@ -84,7 +85,7 @@ def __init__(self, *args, schema: Optional[str] = None, **kwargs): # 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 = kwargs.get('log_sql', True) + self.log_sql = log_sql def get_conn(self): """Returns a connection object""" From 886bb065630720b6639debf193210ef8eef0a13b Mon Sep 17 00:00:00 2001 From: Holly Evans <39742776+holly-evans@users.noreply.github.com> Date: Wed, 22 Jun 2022 08:11:26 -0500 Subject: [PATCH 4/7] Update docstring Co-authored-by: Tzu-ping Chung --- airflow/hooks/dbapi.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/hooks/dbapi.py b/airflow/hooks/dbapi.py index e7d89f476d85c..e9f5882b6eec1 100644 --- a/airflow/hooks/dbapi.py +++ b/airflow/hooks/dbapi.py @@ -56,7 +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 during `run`. Defaults to True. + :param log_sql: Whether to log SQL query when it's executed. Defaults to *True*. """ # Override to provide the connection name. From 5307cc40e52f81aa108c1e6ec5417130665d9b35 Mon Sep 17 00:00:00 2001 From: Holly Evans <39742776+holly-evans@users.noreply.github.com> Date: Wed, 22 Jun 2022 08:13:06 -0500 Subject: [PATCH 5/7] Remove Optional typing --- airflow/hooks/dbapi.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/hooks/dbapi.py b/airflow/hooks/dbapi.py index e9f5882b6eec1..d3d3fa5fc8972 100644 --- a/airflow/hooks/dbapi.py +++ b/airflow/hooks/dbapi.py @@ -70,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, log_sql: Optional[bool] = True, **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") From 51de59ed42b3e8c33bf881bf327b0e4f83bc2970 Mon Sep 17 00:00:00 2001 From: Holly Evans <39742776+holly-evans@users.noreply.github.com> Date: Thu, 23 Jun 2022 13:26:37 -0500 Subject: [PATCH 6/7] Rename test hook Co-authored-by: Tzu-ping Chung --- tests/hooks/test_dbapi.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/hooks/test_dbapi.py b/tests/hooks/test_dbapi.py index 23f9c0a215550..7370552c23c8f 100644 --- a/tests/hooks/test_dbapi.py +++ b/tests/hooks/test_dbapi.py @@ -44,7 +44,7 @@ def get_conn(self): return conn self.db_hook = UnitTestDbApiHook() - self.db_hook_log_sql = UnitTestDbApiHook(log_sql=False) + self.db_hook_no_log_sql = UnitTestDbApiHook(log_sql=False) self.db_hook_schema_override = UnitTestDbApiHook(schema='schema-override') def test_get_records(self): From 830bd9d0cd231540b24f1d84972117f63973524a Mon Sep 17 00:00:00 2001 From: holly-evans Date: Fri, 24 Jun 2022 11:23:24 -0500 Subject: [PATCH 7/7] Rename test hook --- tests/hooks/test_dbapi.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/hooks/test_dbapi.py b/tests/hooks/test_dbapi.py index 7370552c23c8f..ad5e7ba6af065 100644 --- a/tests/hooks/test_dbapi.py +++ b/tests/hooks/test_dbapi.py @@ -349,8 +349,8 @@ def test_run_log(self): def test_run_no_log(self): statement = 'SQL' - self.db_hook_log_sql.run(statement) - assert self.db_hook_log_sql.log.info.call_count == 1 + 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'