diff --git a/airflow/hooks/dbapi.py b/airflow/hooks/dbapi.py index 0b9ce4377be23..d3d3fa5fc8972 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 @@ 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 = log_sql def get_conn(self): """Returns a connection object""" @@ -228,7 +230,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/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. diff --git a/tests/hooks/test_dbapi.py b/tests/hooks/test_dbapi.py index a17c24aedb384..ad5e7ba6af065 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_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 @@ 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_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 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(