Make statement_timeout local for Postgres schemas (DataDog#23954)

* Make statement_timeout local for Postgres schemas

* Changelog

* Disable autocommit for schemas

* Manually set autocommit

* Manually commit

* Try/finally

* with transaction

* Disable in tests

Author: sethsamuel

postgres/changelog.d/23954.fixed

@@ -0,0 +1 @@
+Make statement_timeout local for Postgres schemas collection

postgres/datadog_checks/postgres/schemas.py

@@ -230,10 +230,12 @@ def _get_databases(self):
     def _get_cursor(self, database_name):
         with self._check.db_pool.get_connection(database_name) as conn:
             with conn.cursor(row_factory=dict_row) as cursor:
-                query, params = self.get_rows_query()
-                cursor.execute(f"SET statement_timeout = '{self._config.max_query_duration}s';")
-                cursor.execute(query, params)
-                yield cursor
+                # Explicitly wrap these queries in a transaction so we can set the statement timeout locally
+                with conn.transaction():
+                    query, params = self.get_rows_query()
+                    cursor.execute(f"SET LOCAL statement_timeout = '{self._config.max_query_duration}s';")
+                    cursor.execute(query, params)
+                    yield cursor
 
     def _get_schemas_query(self):
         query = SCHEMA_QUERY

postgres/tests/test_statements.py

@@ -1960,6 +1960,7 @@ def test_pg_stat_statements_max_warning(
 def test_pg_stat_statements_dealloc(aggregator, integration_check, dbm_instance_replica2):
     dbm_instance_replica2['query_samples'] = {'enabled': False}
     dbm_instance_replica2['query_activity'] = {'enabled': False}
+    dbm_instance_replica2['collect_schemas'] = {'enabled': False}
     with _get_superconn(dbm_instance_replica2) as superconn:
         with superconn.cursor() as cur:
             cur.execute("select pg_stat_statements_reset();")