Set statement_timeout at session level & use reset_val as setting whe…

…n source is session (#17264)

* set statement_timeout at session level

* add changelog

postgres/changelog.d/17264.fixed

@@ -0,0 +1,3 @@
+Fixed bug where `statement_timeout` setting incorrectly reflected integration connection value instead of database level
+- Adjusted `statement_timeout` to apply at the session level post-database connection.
+- Modified the `pg_settings` query to select `reset_val` when sourced from 'session', guaranteeing the retrieval of the accurate server-level setting.

postgres/datadog_checks/postgres/metadata.py

@@ -27,8 +27,15 @@
 DEFAULT_RESOURCES_COLLECTION_INTERVAL = 300
 DEFAULT_SETTINGS_IGNORED_PATTERNS = ["plpgsql%"]
 
+# PG_SETTINGS_QURERY is used to collect all the settings from the pg_settings table
+# Edge case: If source is 'session', it uses reset_val
+# (which represents the value that the setting would revert to on session end or reset),
+# otherwise, it uses the current setting value.
 PG_SETTINGS_QUERY = """
-SELECT name, setting FROM pg_settings
+SELECT
+name,
+case when source = 'session' then reset_val else setting end as setting
+FROM pg_settings
 """
 
 DATABASE_INFORMATION_QUERY = """

postgres/datadog_checks/postgres/postgres.py

@@ -753,8 +753,6 @@ def _new_connection(self, dbname):
                 dbname,
                 self._config.application_name,
             )
-            if self._config.query_timeout:
-                connection_string += " options='-c statement_timeout=%s'" % self._config.query_timeout
             conn = psycopg2.connect(connection_string)
         else:
             password = self._config.password
@@ -787,8 +785,6 @@ def _new_connection(self, dbname):
             }
             if self._config.port:
                 args['port'] = self._config.port
-            if self._config.query_timeout:
-                args['options'] = '-c statement_timeout=%s' % self._config.query_timeout
             if self._config.ssl_cert:
                 args['sslcert'] = self._config.ssl_cert
             if self._config.ssl_root_cert:
@@ -800,6 +796,10 @@ def _new_connection(self, dbname):
             conn = psycopg2.connect(**args)
         # Autocommit is enabled by default for safety for all new connections (to prevent long-lived transactions).
         conn.set_session(autocommit=True, readonly=True)
+        if self._config.query_timeout:
+            # Set the statement_timeout for the session
+            with conn.cursor() as cursor:
+                cursor.execute("SET statement_timeout TO %d" % self._config.query_timeout)
         return conn
 
     def _connect(self):

postgres/tests/compose/etc/postgresql/postgresql.conf

@@ -4,6 +4,7 @@ pg_stat_statements.max=10000
 pg_stat_statements.track=all
 track_io_timing=on
 track_functions=pl
+statement_timeout = 10000
 
 wal_level = logical
 max_wal_senders = 10

postgres/tests/compose/etc/postgresql_logical_replica/postgresql.conf

@@ -5,3 +5,4 @@ pg_stat_statements.max=100
 pg_stat_statements.track=all
 track_io_timing=on
 track_functions=pl
+statement_timeout = 10000

postgres/tests/compose/etc/postgresql_replica/postgresql.conf

@@ -5,6 +5,7 @@ pg_stat_statements.max=10000
 pg_stat_statements.track=all
 track_io_timing=on
 track_functions=pl
+statement_timeout = 10000
 
 hot_standby = on
 hot_standby_feedback = on

postgres/tests/compose/etc/postgresql_replica2/postgresql.conf

@@ -5,6 +5,7 @@ pg_stat_statements.max=100
 pg_stat_statements.track=all
 track_io_timing=on
 track_functions=pl
+statement_timeout = 10000
 
 hot_standby = on
 hot_standby_feedback = off

postgres/tests/test_connections.py

@@ -355,3 +355,18 @@ def _terminate_connection(conn, dbname):
 
     # new check run will re-open connection
     check.check(pg_instance)
+
+
+@pytest.mark.integration
+@pytest.mark.usefixtures('dd_environment')
+def test_conn_statement_timeout(pg_instance):
+    """
+    Test db connection statement timeout is set at the session level
+    """
+    pg_instance["query_timeout"] = 500
+    check = PostgreSql('postgres', {}, [pg_instance])
+    check._connect()
+    with pytest.raises(psycopg2.errors.QueryCanceled):
+        with check.db() as conn:
+            with conn.cursor() as cursor:
+                cursor.execute("SELECT pg_sleep(1)")

postgres/tests/test_metadata.py

@@ -45,6 +45,10 @@ def test_collect_metadata(integration_check, dbm_instance, aggregator):
     assert event['kind'] == "pg_settings"
     assert len(event["metadata"]) > 0
     assert all(not k['name'].startswith('max_wal') for k in event['metadata'])
+    statement_timeout_setting = next((k for k in event['metadata'] if k['name'] == 'statement_timeout'), None)
+    assert statement_timeout_setting is not None
+    # statement_timeout should be server level setting not session level
+    assert statement_timeout_setting['setting'] == '10000'
 
 
 def test_collect_schemas(integration_check, dbm_instance, aggregator):