diff --git a/airflow/providers/common/sql/hooks/sql.py b/airflow/providers/common/sql/hooks/sql.py index 06d91f8f4103e..73dfb754a9c4e 100644 --- a/airflow/providers/common/sql/hooks/sql.py +++ b/airflow/providers/common/sql/hooks/sql.py @@ -447,6 +447,8 @@ def run( # If autocommit was set to False or db does not support autocommit, we do a manual commit. if not self.get_autocommit(conn): conn.commit() + # Logs all database messages or errors sent to the client + self.get_db_log_messages(conn) if handler is None: return None @@ -718,3 +720,10 @@ def get_openlineage_authority_part(connection, default_port: int | None = None) else: authority = parsed.hostname return authority + + def get_db_log_messages(self, conn) -> None: + """ + Log all database messages sent to the client during the session. + + :param conn: Connection object + """ diff --git a/airflow/providers/common/sql/hooks/sql.pyi b/airflow/providers/common/sql/hooks/sql.pyi index 16c3d6592a341..0f22202bd6ba7 100644 --- a/airflow/providers/common/sql/hooks/sql.pyi +++ b/airflow/providers/common/sql/hooks/sql.pyi @@ -131,3 +131,4 @@ class DbApiHook(BaseHook): def get_openlineage_database_specific_lineage(self, task_instance) -> OperatorLineage | None: ... @staticmethod def get_openlineage_authority_part(connection, default_port: int | None = None) -> str: ... + def get_db_log_messages(self, conn) -> None: ... diff --git a/airflow/providers/postgres/hooks/postgres.py b/airflow/providers/postgres/hooks/postgres.py index 4b7074ec75df4..5dcc2bd8230ab 100644 --- a/airflow/providers/postgres/hooks/postgres.py +++ b/airflow/providers/postgres/hooks/postgres.py @@ -69,6 +69,10 @@ class PostgresHook(DbApiHook): :param options: Optional. Specifies command-line options to send to the server at connection start. For example, setting this to ``-c search_path=myschema`` sets the session's value of the ``search_path`` to ``myschema``. + :param enable_log_db_messages: Optional. If enabled logs database messages sent to the client + during the session. To avoid a memory leak psycopg2 only saves the last 50 messages. + For details, see: `PostgreSQL logging configuration parameters + `__ """ conn_name_attr = "postgres_conn_id" @@ -78,7 +82,9 @@ class PostgresHook(DbApiHook): supports_autocommit = True supports_executemany = True - def __init__(self, *args, options: str | None = None, **kwargs) -> None: + def __init__( + self, *args, options: str | None = None, enable_log_db_messages: bool = False, **kwargs + ) -> None: if "schema" in kwargs: warnings.warn( 'The "schema" arg has been renamed to "database" as it contained the database name.' @@ -92,6 +98,7 @@ def __init__(self, *args, options: str | None = None, **kwargs) -> None: self.conn: connection = None self.database: str | None = kwargs.pop("database", None) self.options = options + self.enable_log_db_messages = enable_log_db_messages @property @deprecated( @@ -396,3 +403,13 @@ def get_ui_field_behaviour(cls) -> dict[str, Any]: "schema": "Database", }, } + + def get_db_log_messages(self, conn) -> None: + """ + Log all database messages sent to the client during the session. + + :param conn: Connection object + """ + if self.enable_log_db_messages: + for output in conn.notices: + self.log.info(output) diff --git a/docs/apache-airflow-providers-postgres/operators/postgres_operator_howto_guide.rst b/docs/apache-airflow-providers-postgres/operators/postgres_operator_howto_guide.rst index 5bc8c627f1cb4..f9dafe34196b1 100644 --- a/docs/apache-airflow-providers-postgres/operators/postgres_operator_howto_guide.rst +++ b/docs/apache-airflow-providers-postgres/operators/postgres_operator_howto_guide.rst @@ -157,6 +157,23 @@ class. params={"begin_date": "2020-01-01", "end_date": "2020-12-31"}, ) + +Enable logging of database messages sent to the client +------------------------------------------------------------- + +SQLExecuteQueryOperator provides ``hook_params`` attribute that allows you to pass add parameters to DbApiHook. +You can use ``enable_log_db_messages`` to log database messages or errors emitted by the ``RAISE`` statement. + +.. code-block:: python + + call_proc = SQLExecuteQueryOperator( + task_id="call_proc", + conn_id="postgres_default", + sql="call proc();", + hook_params={"enable_log_db_messages": True}, + ) + + Passing Server Configuration Parameters into PostgresOperator ------------------------------------------------------------- diff --git a/tests/providers/common/sql/hooks/test_dbapi.py b/tests/providers/common/sql/hooks/test_dbapi.py index 090ec80e682b1..3e91e5cd25e45 100644 --- a/tests/providers/common/sql/hooks/test_dbapi.py +++ b/tests/providers/common/sql/hooks/test_dbapi.py @@ -65,6 +65,9 @@ def get_connection(cls, conn_id: str) -> Connection: def get_conn(self): return conn + def get_db_log_messages(self, conn) -> None: + return conn.get_messages() + self.db_hook = DbApiHookMock(**kwargs) self.db_hook_no_log_sql = DbApiHookMock(log_sql=False) self.db_hook_schema_override = DbApiHookMock(schema="schema-override") @@ -531,6 +534,11 @@ def test_run_no_queries(self): self.db_hook.run(sql=[]) assert err.value.args[0] == "List of SQL statements is empty" + def test_run_and_log_db_messages(self): + statement = "SQL" + self.db_hook.run(statement) + self.conn.get_messages.assert_called() + def test_instance_check_works_for_provider_derived_hook(self): assert isinstance(DbApiHookInProvider(), DbApiHook) diff --git a/tests/providers/postgres/hooks/test_postgres.py b/tests/providers/postgres/hooks/test_postgres.py index d73311427fc85..740f957643f8f 100644 --- a/tests/providers/postgres/hooks/test_postgres.py +++ b/tests/providers/postgres/hooks/test_postgres.py @@ -18,6 +18,7 @@ from __future__ import annotations import json +import logging import os from unittest import mock @@ -513,3 +514,36 @@ def test_rowcount(self): cur.execute(f"INSERT INTO {self.table} VALUES {values}") conn.commit() assert cur.rowcount == len(input_data) + + @pytest.mark.usefixtures("reset_logging_config") + def test_get_all_db_log_messages(self, caplog): + messages = ["a", "b", "c"] + + class FakeLogger: + notices = messages + + with caplog.at_level(logging.INFO): + hook = PostgresHook(enable_log_db_messages=True) + hook.get_db_log_messages(FakeLogger) + for msg in messages: + assert msg in caplog.text + + @pytest.mark.usefixtures("reset_logging_config") + def test_log_db_messages_by_db_proc(self, caplog): + proc_name = "raise_notice" + notice_proc = f""" + CREATE PROCEDURE {proc_name} (s text) LANGUAGE PLPGSQL AS + $$ + BEGIN + raise notice 'Message from db: %', s; + END; + $$; + """ + with caplog.at_level(logging.INFO): + hook = PostgresHook(enable_log_db_messages=True) + try: + hook.run(sql=notice_proc) + hook.run(sql=f"call {proc_name}('42')") + assert "NOTICE: Message from db: 42" in caplog.text + finally: + hook.run(sql=f"DROP PROCEDURE {proc_name} (s text)")