feat: log client db messages for provider postgres

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
+        """

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: ...

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
+        <https://www.postgresql.org/docs/current/runtime-config-logging.html>`__
     """
 
     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)

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
 -------------------------------------------------------------
 

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)
 

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)")