Upgrade PGBouncer to psycopg3 (#19325)

* WIP

* WIP

* v3

* Changelog

* Use source

* Comments

* Binary

* Build psycopg (#19326)

* WIP

* Merge fix

* Dep

* Fix license

* WIP

* WIP

* Changelog

* test

* test

* Version

* Changelog

* Always add bin path to $PATH on macos builders

* Add a dummy change to test that a macos build with a cache hit works

* Revert "Add a dummy change to test that a macos build with a cache hit works"

This reverts commit 4df3963.

---------

Co-authored-by: Alex Lopez <alex.lopez.zorzano@gmail.com>
Co-authored-by: Alex Lopez <alex.lopez@datadoghq.com>

Author: sethsamuel

.builders/images/macos-x86_64/builder_setup.sh

@@ -96,8 +96,8 @@ rm "${DD_PREFIX_PATH}/bin/curl"
 
 # libpq and pg_config as needed by psycopg
 DOWNLOAD_URL="https://ftp.postgresql.org/pub/source/v{{version}}/postgresql-{{version}}.tar.bz2" \
-VERSION="16.0" \
-SHA256="df9e823eb22330444e1d48e52cc65135a652a6fdb3ce325e3f08549339f51b99" \
+VERSION="16.6" \
+SHA256="23369cdaccd45270ac5dcc30fa9da205d5be33fa505e1f17a0418d2caeca477b" \
 RELATIVE_PATH="postgresql-{{version}}" \
   install-from-source --without-readline --with-openssl --without-icu
 

.ddev/ci/scripts/pgbouncer/linux/50_install_postgres.sh

@@ -0,0 +1,6 @@
+#!/bin/bash
+
+set -ex
+
+sudo apt update
+sudo apt install -y --no-install-recommends build-essential python3-dev libpq-dev

.ddev/config.toml

@@ -85,9 +85,6 @@ mmh3 = ['CC0-1.0']
 paramiko = ['LGPL-2.1-only']
 # https://github.com/psycopg/psycopg/blob/master/LICENSE.txt
 psycopg = ['LGPL-3.0-only']
-# https://github.com/psycopg/psycopg2/blob/master/LICENSE
-# https://github.com/psycopg/psycopg2/blob/master/doc/COPYING.LESSER
-psycopg2-binary = ['LGPL-3.0-only', 'BSD-3-Clause']
 # https://github.com/Legrandin/pycryptodome/blob/master/LICENSE.rst
 pycryptodomex = ['Unlicense', 'BSD-2-Clause']
 # https://github.com/mongodb/mongo-python-driver/blob/master/LICENSE
@@ -118,7 +115,6 @@ lxml = 'https://github.com/lxml/lxml'
 packaging = 'https://github.com/pypa/packaging'
 paramiko = 'https://github.com/paramiko/paramiko'
 protobuf = 'https://github.com/protocolbuffers/protobuf'
-psycopg2-binary = 'https://github.com/psycopg/psycopg2'
 pycryptodomex = 'https://github.com/Legrandin/pycryptodome'
 redis = 'https://github.com/redis/redis-py'
 requests = 'https://github.com/psf/requests'
@@ -170,9 +166,6 @@ exclude = [
     'aerospike',  # v8+ breaks agent build.
     # https://github.com/DataDog/integrations-core/pull/16080
     'lxml',
-    # We're not ready to switch to v3 of the postgres library, see:
-    # https://github.com/DataDog/integrations-core/pull/15859
-    'psycopg2-binary',
     'psutil',
     'pymongo[srv]', # Upgrade from 4.8.0 to 4.10.1 causes "AttributeError: module 'pymongo' has no attribute 'mongo_client'"
 ]

LICENSE-3rdparty.csv

@@ -46,8 +46,6 @@ prometheus-client,PyPI,Apache-2.0,Copyright 2015 The Prometheus Authors
 protobuf,PyPI,BSD-3-Clause,Copyright 2008 Google Inc.  All rights reserved.
 psutil,PyPI,BSD-3-Clause,"Copyright (c) 2009, Jay Loden, Dave Daeschler, Giampaolo Rodola"
 psycopg,PyPI,LGPL-3.0-only,Copyright (C) 2020 The Psycopg Team
-psycopg2-binary,PyPI,BSD-3-Clause,Copyright 2013 Federico Di Gregorio
-psycopg2-binary,PyPI,LGPL-3.0-only,Copyright (C) 2013 Federico Di Gregorio
 pyOpenSSL,PyPI,Apache-2.0,Copyright The pyOpenSSL developers
 pyasn1,PyPI,BSD-3-Clause,"Copyright (c) 2005-2019, Ilya Etingof <etingof@gmail.com>"
 pycryptodomex,PyPI,BSD-2-Clause,Copyright 2014 Helder Eijs

agent_requirements.in

@@ -31,7 +31,6 @@ ply==3.11
 prometheus-client==0.21.1
 protobuf==5.29.3
 psutil==6.0.0
-psycopg2-binary==2.9.9
 psycopg[c]==3.2.3
 pyasn1==0.4.8
 pycryptodomex==3.21.0

datadog_checks_dev/changelog.d/19325.security

@@ -0,0 +1 @@
+Removed dependency on psycopg2

datadog_checks_dev/datadog_checks/dev/tooling/commands/validate/licenses.py

@@ -45,9 +45,6 @@
     'mmh3': ['CC0-1.0'],
     # https://github.com/paramiko/paramiko/blob/master/LICENSE
     'paramiko': ['LGPL-2.1-only'],
-    # https://github.com/psycopg/psycopg2/blob/master/LICENSE
-    # https://github.com/psycopg/psycopg2/blob/master/doc/COPYING.LESSER
-    'psycopg2-binary': ['LGPL-3.0-only', 'BSD-3-Clause'],
     # https://github.com/psycopg/psycopg/blob/master/LICENSE.txt
     'psycopg': ['LGPL-3.0-only'],
     # https://github.com/psycopg/psycopg/blob/master/psycopg_pool/LICENSE.txt
@@ -161,7 +158,6 @@
     'packaging': 'https://github.com/pypa/packaging',
     'paramiko': 'https://github.com/paramiko/paramiko',
     'protobuf': 'https://github.com/protocolbuffers/protobuf',
-    'psycopg2-binary': 'https://github.com/psycopg/psycopg2',
     'psycopg': 'https://github.com/psycopg/psycopg',
     'pycryptodomex': 'https://github.com/Legrandin/pycryptodome',
     'redis': 'https://github.com/redis/redis-py',

pgbouncer/changelog.d/19325.added

@@ -0,0 +1 @@
+Upgrade pgbouncer to psycopg3

pgbouncer/changelog.d/19325.security

@@ -0,0 +1 @@
+Removed dependency on psycopg2

pgbouncer/datadog_checks/pgbouncer/pgbouncer.py

@@ -5,8 +5,9 @@
 import time
 from urllib.parse import urlparse
 
-import psycopg2 as pg
-from psycopg2 import extras as pgextras
+import psycopg as pg
+from psycopg import ClientCursor
+from psycopg.rows import dict_row
 
 from datadog_checks.base import AgentCheck, ConfigurationError, is_affirmative
 from datadog_checks.pgbouncer.metrics import (
@@ -73,47 +74,46 @@ def _collect_stats(self, db):
             metric_scope.append(SERVERS_METRICS)
 
         try:
-            with db.cursor(cursor_factory=pgextras.DictCursor) as cursor:
-                for scope in metric_scope:
-                    descriptors = scope['descriptors']
-                    metrics = scope['metrics']
-                    query = scope['query']
-
-                    try:
-                        self.log.debug("Running query: %s", query)
-                        cursor.execute(query)
-                        rows = self.iter_rows(cursor)
-
-                    except Exception as e:
-                        self.log.exception("Not all metrics may be available: %s", str(e))
-
-                    else:
-                        for row in rows:
-                            if 'key' in row:  # We are processing "config metrics"
-                                # Make a copy of the row to allow mutation
-                                # (a `psycopg2.lib.extras.DictRow` object doesn't accept a new key)
-                                row = row.copy()
-                                # We flip/rotate the row: row value becomes the column name
-                                row[row['key']] = row['value']
-                            # Skip the "pgbouncer" database
-                            elif row.get('database') == self.DB_NAME:
-                                continue
-
-                            tags = list(self.tags)
-                            tags += ["%s:%s" % (tag, row[column]) for (column, tag) in descriptors if column in row]
-                            for column, (name, reporter) in metrics:
-                                if column in row:
-                                    value = row[column]
-                                    if column in ['connect_time', 'request_time']:
-                                        self.log.debug("Parsing timestamp; original value: %s", value)
-                                        # First get rid of any UTC suffix.
-                                        value = re.findall(r'^[^ ]+ [^ ]+', value)[0]
-                                        value = time.strptime(value, '%Y-%m-%d %H:%M:%S')
-                                        value = time.mktime(value)
-                                    reporter(self, name, value, tags)
-
-                        if not rows:
-                            self.log.warning("No results were found for query: %s", query)
+            for scope in metric_scope:
+                descriptors = scope['descriptors']
+                metrics = scope['metrics']
+                query = scope['query']
+
+                try:
+                    cursor = db.cursor(row_factory=dict_row)
+                    self.log.debug("Running query: %s", query)
+                    cursor.execute(query)
+                    rows = self.iter_rows(cursor)
+
+                except Exception as e:
+                    self.log.exception("Not all metrics may be available: %s", str(e))
+
+                else:
+                    for row in rows:
+                        if 'key' in row:  # We are processing "config metrics"
+                            # Make a copy of the row to allow mutation
+                            row = row.copy()
+                            # We flip/rotate the row: row value becomes the column name
+                            row[row['key']] = row['value']
+                        # Skip the "pgbouncer" database
+                        elif row.get('database') == self.DB_NAME:
+                            continue
+
+                        tags = list(self.tags)
+                        tags += ["%s:%s" % (tag, row[column]) for (column, tag) in descriptors if column in row]
+                        for column, (name, reporter) in metrics:
+                            if column in row:
+                                value = row[column]
+                                if column in ['connect_time', 'request_time']:
+                                    self.log.debug("Parsing timestamp; original value: %s", value)
+                                    # First get rid of any UTC suffix.
+                                    value = re.findall(r'^[^ ]+ [^ ]+', value)[0]
+                                    value = time.strptime(value, '%Y-%m-%d %H:%M:%S')
+                                    value = time.mktime(value)
+                                reporter(self, name, value, tags)
+
+                    if not rows:
+                        self.log.warning("No results were found for query: %s", query)
 
         except pg.Error:
             self.log.exception("Connection error")
@@ -138,21 +138,25 @@ def iter_rows(self, cursor):
 
     def _get_connect_kwargs(self):
         """
-        Get the params to pass to psycopg2.connect() based on passed-in vals
+        Get the params to pass to psycopg.connect() based on passed-in vals
         from yaml settings file
         """
+        # It's important to set the client_encoding to utf-8
+        # PGBouncer defaults to an encoding of 'UNICODE`, which will cause psycopg to error out
         if self.database_url:
-            return {'dsn': self.database_url}
+            return {'conninfo': self.database_url, 'client_encoding': 'utf-8'}
 
         if self.host in ('localhost', '127.0.0.1') and self.password == '':
             # Use ident method
-            return {'dsn': "user={} dbname={}".format(self.user, self.DB_NAME)}
+            return {'conninfo': "user={} dbname={} client_encoding=utf-8".format(self.user, self.DB_NAME)}
 
         args = {
             'host': self.host,
             'user': self.user,
             'password': self.password,
-            'database': self.DB_NAME,
+            'dbname': self.DB_NAME,
+            'cursor_factory': ClientCursor,
+            'client_encoding': 'utf-8',
         }
         if self.port:
             args['port'] = self.port
@@ -166,8 +170,9 @@ def _get_connection(self, use_cached=None):
             return self.connection
         try:
             connect_kwargs = self._get_connect_kwargs()
-            connection = pg.connect(**connect_kwargs)
-            connection.set_isolation_level(pg.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
+            # Somewhat counterintuitively, we need to set autocommit to True to avoid a BEGIN/COMMIT block
+            # https://www.psycopg.org/psycopg3/docs/basic/transactions.html#autocommit-transactions
+            connection = pg.connect(**connect_kwargs, autocommit=True)
         except Exception:
             redacted_url = self._get_redacted_dsn()
             message = u'Cannot establish connection to {}'.format(redacted_url)
@@ -180,6 +185,10 @@ def _get_connection(self, use_cached=None):
         self.connection = connection
         return connection
 
+    def _close_connection(self):
+        self.connection.close()
+        self.connection = None
+
     def _get_redacted_dsn(self):
         if not self.database_url:
             return u'pgbouncer://%s:******@%s:%s/%s' % (self.user, self.host, self.port, self.DB_NAME)
@@ -200,6 +209,8 @@ def check(self, instance):
 
         self.service_check(self.SERVICE_CHECK_NAME, AgentCheck.OK, tags=self._get_service_checks_tags())
         self._set_metadata()
+        # Avoid holding an open connection
+        self._close_connection()
 
     def _set_metadata(self):
         if self.is_metadata_collection_enabled():
@@ -209,14 +220,5 @@ def _set_metadata(self):
 
     def get_version(self):
         db = self._get_connection()
-        regex = r'\d+\.\d+\.\d+'
-        with db.cursor(cursor_factory=pgextras.DictCursor) as cursor:
-            cursor.execute('SHOW VERSION;')
-            if db.notices:
-                data = db.notices[0]
-            else:
-                data = cursor.fetchone()[0]
-            res = re.findall(regex, data)
-            if res:
-                return res[0]
-            self.log.debug("Couldn't detect version from %s", data)
+        version = pg.pq.version_pretty(db.connection.info.server_version)
+        return version

pgbouncer/pyproject.toml

@@ -37,7 +37,7 @@ license = "BSD-3-Clause"
 
 [project.optional-dependencies]
 deps = [
-    "psycopg2-binary==2.9.9",
+    "psycopg[c]==3.2.3",
 ]
 
 [project.urls]

pgbouncer/tests/common.py

@@ -20,6 +20,13 @@
     'tags': ['optional:tag1'],
 }
 
+E2E_METADATA = {
+    'start_commands': [
+        'apt update',
+        'apt install -y --no-install-recommends build-essential python3-dev libpq-dev',
+    ],
+}
+
 
 def get_version_from_env():
     return version.parse(os.environ.get('PGBOUNCER_VERSION'))

pgbouncer/tests/conftest.py

@@ -5,7 +5,7 @@
 import os
 from copy import deepcopy
 
-import psycopg2
+import psycopg
 import pytest
 from packaging import version
 
@@ -20,8 +20,8 @@ def container_up(service_name, port):
     """
     Try to connect to postgres/pgbouncer
     """
-    psycopg2.connect(
-        host=common.HOST, port=port, user=common.USER, password=common.PASS, database=common.DB, connect_timeout=2
+    psycopg.connect(
+        host=common.HOST, port=port, user=common.USER, password=common.PASS, dbname=common.DB, connect_timeout=2
     )
 
 
@@ -45,7 +45,7 @@ def dd_environment():
         ],
     ):
 
-        yield common.DEFAULT_INSTANCE
+        yield common.DEFAULT_INSTANCE, common.E2E_METADATA
 
 
 @pytest.fixture

pgbouncer/tests/test_pgbouncer_integration_e2e.py

@@ -2,7 +2,7 @@
 # All rights reserved
 # Licensed under Simplified BSD License (see LICENSE)
 
-import psycopg2
+import psycopg
 import pytest
 from packaging import version
 
@@ -15,15 +15,14 @@
 @pytest.mark.usefixtures("dd_environment")
 def test_check(instance, aggregator, datadog_agent, dd_run_check):
     # add some stats
-    connection = psycopg2.connect(
+    connection = psycopg.connect(
         host=common.HOST,
         port=common.PORT,
         user=common.USER,
         password=common.PASS,
-        database=common.DB,
+        dbname=common.DB,
         connect_timeout=1,
     )
-    connection.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
     cur = connection.cursor()
     cur.execute('SELECT * FROM persons;')
 
@@ -49,15 +48,14 @@ def test_check(instance, aggregator, datadog_agent, dd_run_check):
 @pytest.mark.usefixtures("dd_environment")
 def test_check_with_clients(instance, aggregator, datadog_agent, dd_run_check):
     # add some stats
-    connection = psycopg2.connect(
+    connection = psycopg.connect(
         host=common.HOST,
         port=common.PORT,
         user=common.USER,
         password=common.PASS,
-        database=common.DB,
+        dbname=common.DB,
         connect_timeout=1,
     )
-    connection.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
     cur = connection.cursor()
     cur.execute('SELECT * FROM persons;')
 
@@ -80,15 +78,14 @@ def test_check_with_clients(instance, aggregator, datadog_agent, dd_run_check):
 @pytest.mark.usefixtures("dd_environment")
 def test_check_with_servers(instance, aggregator, datadog_agent, dd_run_check):
     # add some stats
-    connection = psycopg2.connect(
+    connection = psycopg.connect(
         host=common.HOST,
         port=common.PORT,
         user=common.USER,
         password=common.PASS,
-        database=common.DB,
+        dbname=common.DB,
         connect_timeout=1,
     )
-    connection.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
     cur = connection.cursor()
     cur.execute('SELECT * FROM persons;')