Revert "Update Postgres to psycopg3 (DataDog#20617)" (DataDog#20885)

This reverts commit fbb0579.

Author: Kyle-Neale

agent_requirements.in

@@ -33,7 +33,6 @@ prometheus-client==0.22.1
 protobuf==6.31.1
 psutil==6.0.0
 psycopg2-binary==2.9.9
-psycopg[binary]==3.2.9
 pyasn1==0.4.8
 pycryptodomex==3.23.0
 pydantic==2.11.7

postgres/datadog_checks/postgres/connections.py

@@ -8,7 +8,7 @@
 import time
 from typing import Callable, Dict
 
-import psycopg
+import psycopg2
 
 
 class ConnectionPoolFullError(Exception):
@@ -23,7 +23,7 @@ def __str__(self):
 class ConnectionInfo:
     def __init__(
         self,
-        connection: psycopg.Connection,
+        connection: psycopg2.extensions.connection,
         deadline: int,
         active: bool,
         last_accessed: int,
@@ -86,9 +86,9 @@ def _get_connection_raw(
         dbname: str,
         ttl_ms: int,
         timeout: int = None,
-        startup_fn: Callable[[psycopg.Connection], None] = None,
+        startup_fn: Callable[[psycopg2.extensions.connection], None] = None,
         persistent: bool = False,
-    ) -> psycopg.Connection:
+    ) -> psycopg2.extensions.connection:
         """
         Return a connection from the pool.
         Pass a function to startup_func if there is an action needed with the connection
@@ -117,7 +117,7 @@ def _get_connection_raw(
                 # if already in pool, retain persistence status
                 persistent = conn.persistent
 
-            if db.info.status != psycopg.pq.ConnStatus.OK:
+            if db.status != psycopg2.extensions.STATUS_READY:
                 # Some transaction went wrong and the connection is in an unhealthy state. Let's fix that
                 db.rollback()
 
@@ -130,7 +130,6 @@ def _get_connection_raw(
                 thread=threading.current_thread(),
                 persistent=persistent,
             )
-
             return db
 
     @contextlib.contextmanager
@@ -139,7 +138,7 @@ def get_connection(
         dbname: str,
         ttl_ms: int,
         timeout: int = None,
-        startup_fn: Callable[[psycopg.Connection], None] = None,
+        startup_fn: Callable[[psycopg2.extensions.connection], None] = None,
         persistent: bool = False,
     ):
         """
@@ -148,14 +147,12 @@ def get_connection(
         make a new connection if the max_conn limit hasn't been reached.
         Blocks until a connection can be added to the pool,
         and optionally takes a timeout in seconds.
-        Note that leaving a connection context here does NOT close the connection in psycopg;
+        Note that leaving a connection context here does NOT close the connection in psycopg2;
         connections must be manually closed by `close_all_connections()`.
         """
         try:
             with self._mu:
-                db = self._get_connection_raw(
-                    dbname=dbname, ttl_ms=ttl_ms, timeout=timeout, startup_fn=startup_fn, persistent=persistent
-                )
+                db = self._get_connection_raw(dbname, ttl_ms, timeout, startup_fn, persistent)
             yield db
         finally:
             with self._mu:

postgres/datadog_checks/postgres/cursor.py

@@ -2,10 +2,10 @@
 # All rights reserved
 # Licensed under a 3-clause BSD style license (see LICENSE)
 
-import psycopg
+import psycopg2.extensions
+import psycopg2.extras
 
 from datadog_checks.base.utils.db.sql_commenter import add_sql_comment
-from datadog_checks.postgres.encoding import decode_with_encodings
 
 DD_QUERY_ATTRIBUTES = {
     'service': 'datadog-agent',
@@ -17,41 +17,20 @@ def __init__(self, *args, **kwargs):
         self.__attributes = DD_QUERY_ATTRIBUTES
         super().__init__(*args, **kwargs)
 
-    def execute(self, query, params=None, ignore_query_metric=False, binary=False, prepare=None):
+    def execute(self, query, vars=None, ignore_query_metric=False):
         '''
         When ignore is True, a /* DDIGNORE */ comment will be added to the query.
         This comment indicates that the query should be ignored in query metrics.
         '''
         query = add_sql_comment(query, prepand=True, **self.__attributes)
         if ignore_query_metric:
             query = '{} {}'.format('/* DDIGNORE */', query)
-        return super().execute(query, params, binary=binary, prepare=prepare)
+        return super().execute(query, vars)
 
 
-class CommenterCursor(BaseCommenterCursor, psycopg.ClientCursor):
+class CommenterCursor(BaseCommenterCursor, psycopg2.extensions.cursor):
     pass
 
 
-class CommenterDictCursor(BaseCommenterCursor, psycopg.ClientCursor):
+class CommenterDictCursor(BaseCommenterCursor, psycopg2.extras.DictCursor):
     pass
-
-
-class SQLASCIITextLoader(psycopg.adapt.Loader):
-    """
-    Custom loader for SQLASCII encoding.
-    """
-
-    encodings = ['utf-8']
-    format = psycopg.pq.Format.TEXT
-
-    def load(self, data):
-        if isinstance(data, memoryview):
-            # Convert memoryview to bytes
-            data = data.tobytes()
-        if not isinstance(data, bytes) or data is None:
-            return data
-        try:
-            return decode_with_encodings(data, self.encodings)
-        except:
-            # Fallback to utf8 with replacement
-            return data.decode('utf-8', errors='backslashreplace')

postgres/datadog_checks/postgres/discovery.py

@@ -6,6 +6,7 @@
 
 from datadog_checks.base import AgentCheck
 from datadog_checks.base.utils.discovery import Discovery
+from datadog_checks.postgres.cursor import CommenterCursor
 from datadog_checks.postgres.util import DatabaseConfigurationError, warning_with_tags
 
 AUTODISCOVERY_QUERY: str = """select datname from pg_catalog.pg_database where datistemplate = false;"""
@@ -71,7 +72,7 @@ def get_items(self) -> List[str]:
 
     def _get_databases(self) -> List[str]:
         with self.db_pool.get_connection(self._db, self._default_ttl) as conn:
-            with conn.cursor() as cursor:
+            with conn.cursor(cursor_factory=CommenterCursor) as cursor:
                 cursor.execute(AUTODISCOVERY_QUERY)
                 databases = list(cursor.fetchall())
                 databases = [

postgres/datadog_checks/postgres/explain_parameterized_queries.py

@@ -5,9 +5,10 @@
 import logging
 import re
 
-import psycopg
+import psycopg2
 
 from datadog_checks.base.utils.tracking import tracked_method
+from datadog_checks.postgres.cursor import CommenterDictCursor
 
 from .util import DBExplainError
 from .version_utils import V12
@@ -75,7 +76,7 @@ def explain_statement(self, dbname, statement, obfuscated_statement, query_signa
         if self._check.version < V12:
             # if pg version < 12, skip explaining parameterized queries because
             # plan_cache_mode is not supported
-            e = psycopg.errors.UndefinedParameter("Unable to explain parameterized query")
+            e = psycopg2.errors.UndefinedParameter("Unable to explain parameterized query")
             logger.debug(
                 "Unable to explain parameterized query. Postgres version %s does not support plan_cache_mode",
                 self._check.version,
@@ -85,9 +86,9 @@ def explain_statement(self, dbname, statement, obfuscated_statement, query_signa
 
         try:
             self._create_prepared_statement(dbname, statement, obfuscated_statement, query_signature)
-        except psycopg.errors.IndeterminateDatatype as e:
+        except psycopg2.errors.IndeterminateDatatype as e:
             return None, DBExplainError.indeterminate_datatype, '{}'.format(type(e))
-        except psycopg.errors.UndefinedFunction as e:
+        except psycopg2.errors.UndefinedFunction as e:
             return None, DBExplainError.undefined_function, '{}'.format(type(e))
         except Exception as e:
             # if we fail to create a prepared statement, we cannot explain the query
@@ -187,14 +188,16 @@ def _deallocate_prepared_statement(self, dbname, query_signature):
             )
 
     def _execute_query(self, dbname, query):
+        # Psycopg2 connections do not get closed when context ends;
+        # leaving context will just mark the connection as inactive in MultiDatabaseConnectionPool
         with self._check.db_pool.get_connection(dbname, self._check._config.idle_connection_timeout) as conn:
-            with conn.cursor() as cursor:
+            with conn.cursor(cursor_factory=CommenterDictCursor) as cursor:
                 logger.debug('Executing query=[%s]', query)
                 cursor.execute(query, ignore_query_metric=True)
 
     def _execute_query_and_fetch_rows(self, dbname, query):
         with self._check.db_pool.get_connection(dbname, self._check._config.idle_connection_timeout) as conn:
-            with conn.cursor() as cursor:
+            with conn.cursor(cursor_factory=CommenterDictCursor) as cursor:
                 cursor.execute(query, ignore_query_metric=True)
                 return cursor.fetchall()
 

postgres/datadog_checks/postgres/metadata.py

@@ -6,8 +6,9 @@
 import time
 from typing import Dict, List, Optional, Tuple, Union  # noqa: F401
 
-import psycopg
-from psycopg.rows import dict_row
+import psycopg2
+
+from datadog_checks.postgres.cursor import CommenterDictCursor
 
 try:
     import datadog_agent
@@ -60,8 +61,7 @@
 """
 
 PG_EXTENSIONS_QUERY = """
-SELECT extname, nspname schemaname
-FROM pg_extension left join pg_namespace on extnamespace = pg_namespace.oid;
+SELECT extname, nspname schemaname FROM pg_extension left join pg_namespace on extnamespace = pg_namespace.oid;
 """
 
 PG_EXTENSION_LOADER_QUERY = {
@@ -256,7 +256,7 @@ def __init__(self, check, config, shutdown_callback):
             enabled=is_affirmative(config.resources_metadata_config.get("enabled", True)),
             dbms="postgres",
             min_collection_interval=config.min_collection_interval,
-            expected_db_exceptions=(psycopg.errors.DatabaseError,),
+            expected_db_exceptions=(psycopg2.errors.DatabaseError,),
             job_name="database-metadata",
             shutdown_callback=shutdown_callback,
         )
@@ -319,7 +319,7 @@ def report_postgres_extensions(self):
     @tracked_method(agent_check_getter=agent_check_getter)
     def _collect_postgres_extensions(self):
         with self._check._get_main_db() as conn:
-            with conn.cursor(row_factory=dict_row) as cursor:
+            with conn.cursor(cursor_factory=CommenterDictCursor) as cursor:
                 self._time_since_last_extension_query = time.time()
 
                 # Get loaded extensions
@@ -394,7 +394,7 @@ def _collect_postgres_schemas(self):
                     continue
 
                 with self.db_pool.get_connection(dbname, self._config.idle_connection_timeout) as conn:
-                    with conn.cursor(row_factory=dict_row) as cursor:
+                    with conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cursor:
                         for schema in database["schemas"]:
                             if not self._should_collect_metadata(schema["name"], "schema"):
                                 continue
@@ -510,7 +510,9 @@ def _collect_schema_info(self):
         self._last_schemas_query_time = time.time()
         return metadata
 
-    def _query_database_information(self, cursor: psycopg.Cursor, dbname: str) -> Dict[str, Union[str, int]]:
+    def _query_database_information(
+        self, cursor: psycopg2.extensions.cursor, dbname: str
+    ) -> Dict[str, Union[str, int]]:
         """
         Collect database info. Returns
             description: str
@@ -523,7 +525,7 @@ def _query_database_information(self, cursor: psycopg.Cursor, dbname: str) -> Di
         row = cursor.fetchone()
         return row
 
-    def _query_schema_information(self, cursor: psycopg.Cursor, dbname: str) -> Dict[str, str]:
+    def _query_schema_information(self, cursor: psycopg2.extensions.cursor, dbname: str) -> Dict[str, str]:
         """
         Collect user schemas. Returns
             id: str
@@ -639,7 +641,7 @@ def sort_tables(info):
         return table_info[:limit]
 
     def _query_tables_for_schema(
-        self, cursor: psycopg.Cursor, schema_id: str, dbname: str
+        self, cursor: psycopg2.extensions.cursor, schema_id: str, dbname: str
     ) -> List[Dict[str, Union[str, Dict]]]:
         """
         Collect list of tables for a schema. Returns a list of dictionaries
@@ -670,7 +672,7 @@ def _query_tables_for_schema(
         return table_payloads
 
     def _query_table_information(
-        self, cursor: psycopg.Cursor, schema_name: str, table_info: List[Dict[str, Union[str, bool]]]
+        self, cursor: psycopg2.extensions.cursor, schema_name: str, table_info: List[Dict[str, Union[str, bool]]]
     ) -> List[Dict[str, Union[str, Dict]]]:
         """
         Collect table information . Returns a dictionary
@@ -750,7 +752,7 @@ def _query_table_information(
     def _collect_metadata_for_database(self, dbname):
         metadata = {}
         with self.db_pool.get_connection(dbname, self._config.idle_connection_timeout) as conn:
-            with conn.cursor(row_factory=dict_row) as cursor:
+            with conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cursor:
                 database_info = self._query_database_information(cursor, dbname)
                 metadata.update(
                     {
@@ -771,7 +773,7 @@ def _collect_metadata_for_database(self, dbname):
     @tracked_method(agent_check_getter=agent_check_getter)
     def _collect_postgres_settings(self):
         with self._check._get_main_db() as conn:
-            with conn.cursor(row_factory=dict_row) as cursor:
+            with conn.cursor(cursor_factory=CommenterDictCursor) as cursor:
                 # Get loaded extensions
                 cursor.execute(PG_EXTENSIONS_QUERY)
                 rows = cursor.fetchall()
@@ -798,14 +800,6 @@ def _collect_postgres_settings(self):
                 )
                 self._time_since_last_settings_query = time.time()
                 cursor.execute(query, (self.pg_settings_ignored_patterns,))
-                # pg3 returns a set of results for each statement in the multiple statement query
-                # We want to retrieve the last one that actually has the settings results
-                rows = []
-                has_more_results = True
-                while has_more_results:
-                    if cursor.pgresult.status == psycopg.pq.ExecStatus.TUPLES_OK:
-                        rows = cursor.fetchall()
-                    has_more_results = cursor.nextset()
-                self._log.debug("Loaded %s rows from pg_settings", rows)
+                rows = cursor.fetchall()
                 self._log.debug("Loaded %s rows from pg_settings", len(rows))
-                return rows
+                return [dict(row) for row in rows]