Skip to content

Commit b13f5e0

Browse files
sethsamuelsethsamuel
authored
Update Posgres to psycopg3 (#19322)
* WIP * WIP * Changelog * Changelog * Dep * License * License * Lint * Fix * Build deps * E2E * Lint * +x
1 parent 6c5362e commit b13f5e0

34 files changed

+443
-375
lines changed

‎.builders/images/helpers.ps1

+8
Original file line numberDiff line numberDiff line change
@@ -38,3 +38,11 @@ function Add-ToPath() {
3838
$target="$oldPath;$Append"
3939
[Environment]::SetEnvironmentVariable("Path", $target, [System.EnvironmentVariableTarget]::User)
4040
}
41+
42+
function RunOnVSConsole() {
43+
param(
44+
[Parameter(Mandatory = $true)][string] $Command
45+
)
46+
Write-Host "Running $Command"
47+
Start-Process -Wait -NoNewWindow "cmd.exe" -ArgumentList "/c ""$Env:VCVARSALL_BAT"" $Env:DD_TARGET_ARCH && $Command"
48+
}

‎.builders/images/linux-aarch64/Dockerfile

+1-1
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,7 @@ RUN \
9191
--without-debugger \
9292
--disable-static
9393

94-
# libpq and pg_config as needed by psycopg2
94+
# libpq and pg_config as needed by psycopg
9595
RUN \
9696
DOWNLOAD_URL="https://ftp.postgresql.org/pub/source/v{{version}}/postgresql-{{version}}.tar.bz2" \
9797
VERSION="16.0" \

‎.builders/images/linux-x86_64/Dockerfile

+1-1
Original file line numberDiff line numberDiff line change
@@ -95,7 +95,7 @@ RUN \
9595
--without-debugger \
9696
--disable-static
9797

98-
# libpq and pg_config as needed by psycopg2
98+
# libpq and pg_config as needed by psycopg
9999
RUN \
100100
DOWNLOAD_URL="https://ftp.postgresql.org/pub/source/v{{version}}/postgresql-{{version}}.tar.bz2" \
101101
VERSION="16.0" \

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

+9
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,15 @@ RELATIVE_PATH="curl-{{version}}" \
9494
# Remove the binary installed so that we consistenly use the same original `curl` binary
9595
rm "${DD_PREFIX_PATH}/bin/curl"
9696

97+
# libpq and pg_config as needed by psycopg
98+
DOWNLOAD_URL="https://ftp.postgresql.org/pub/source/v{{version}}/postgresql-{{version}}.tar.bz2" \
99+
VERSION="16.0" \
100+
SHA256="df9e823eb22330444e1d48e52cc65135a652a6fdb3ce325e3f08549339f51b99" \
101+
RELATIVE_PATH="postgresql-{{version}}" \
102+
install-from-source --without-readline --with-openssl --without-icu
103+
# Add paths to pg_config and to the library
104+
echo PATH="${DD_PREFIX_PATH}/bin:${PATH:-}" >> $DD_ENV_FILE
105+
97106
# Dependencies needed to build librdkafka (and thus, confluent-kafka) with kerberos support
98107
DOWNLOAD_URL="https://github.com/LMDB/lmdb/archive/LMDB_{{version}}.tar.gz" \
99108
VERSION="0.9.29" \

‎.builders/images/windows-x86_64/Dockerfile

+36
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ RUN curl -SL --output vs_buildtools.exe https://aka.ms/vs/17/release/vs_buildtoo
2525
--add Microsoft.VisualStudio.Workload.VCTools `
2626
|| IF "%ERRORLEVEL%"=="3010" EXIT 0) `
2727
&& del /q vs_buildtools.exe
28+
ENV VCVARSALL_BAT="C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools\VC\Auxiliary\Build\vcvarsall.bat"
2829

2930
# Upgrade PowerShell
3031
ENV POWERSHELL_VERSION="7.4.0"
@@ -114,7 +115,42 @@ RUN Get-RemoteFile `
114115
Add-ToPath -Append "C:\perl\perl\bin" && `
115116
Remove-Item "strawberry-perl-$Env:PERL_VERSION-64bit.zip"
116117

118+
# Nasm
119+
ENV NASM_VERSION="2.16.03"
120+
RUN Get-RemoteFile `
121+
-Uri https://www.nasm.us/pub/nasm/releasebuilds/$Env:NASM_VERSION/win64/nasm-$Env:NASM_VERSION-win64.zip `
122+
-Path "nasm-$Env:NASM_VERSION-win64.zip" `
123+
-Hash '3ee4782247bcb874378d02f7eab4e294a84d3d15f3f6ee2de2f47a46aa7226e6' && `
124+
7z x "nasm-$Env:NASM_VERSION-win64.zip" -o"C:\nasm" && `
125+
Add-ToPath -Append "C:\nasm\nasm-$Env:NASM_VERSION" && `
126+
Remove-Item "nasm-$Env:NASM_VERSION-win64.zip"
127+
# openssl
117128
ENV OPENSSL_VERSION="3.3.2"
129+
RUN Get-RemoteFile `
130+
-Uri https://www.openssl.org/source/openssl-$Env:OPENSSL_VERSION.tar.gz `
131+
-Path openssl-$Env:OPENSSL_VERSION.tar.gz `
132+
-Hash '2e8a40b01979afe8be0bbfb3de5dc1c6709fedb46d6c89c10da114ab5fc3d281'; `
133+
7z x openssl-$Env:OPENSSL_VERSION.tar.gz -r -y && `
134+
7z x openssl-$Env:OPENSSL_VERSION.tar -oC:\openssl_3 && `
135+
cd C:\openssl_3\openssl-$Env:OPENSSL_VERSION && `
136+
RunOnVSConsole -Command `
137+
'C:\perl\perl\bin\perl.exe Configure && `
138+
nmake && `
139+
nmake install_sw'
140+
# libpq and pg_config as needed by psycopg
141+
ENV PG_VERSION="16.0"
142+
COPY pg_config.pl C:\pg_config.pl
143+
RUN Get-RemoteFile `
144+
-Uri https://ftp.postgresql.org/pub/source/v$Env:PG_VERSION/postgresql-$Env:PG_VERSION.tar.bz2 `
145+
-Path postgresql-$Env:PG_VERSION.tar.bz2 `
146+
-Hash 'df9e823eb22330444e1d48e52cc65135a652a6fdb3ce325e3f08549339f51b99'; `
147+
7z x postgresql-$Env:PG_VERSION.tar.bz2 -r -y && `
148+
7z x postgresql-$Env:PG_VERSION.tar -oC:\postgresql_src && `
149+
cd C:\postgresql_src\postgresql-$Env:PG_VERSION\src\tools\msvc && `
150+
Copy-Item C:\pg_config.pl -Destination .\config.pl && `
151+
RunOnVSConsole -Command 'C:\perl\perl\bin\perl.exe build.pl' && `
152+
RunOnVSConsole -Command 'C:\perl\perl\bin\perl.exe install.pl C:\postgresql' && `
153+
Add-ToPath -Append "C:\postgresql\bin"
118154

119155
ENV CURL_VERSION="8.11.1"
120156

‎.builders/images/windows-x86_64/pg_config.pl

+4
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
# Overrides over default postgres config
2+
# https://github.com/postgres/postgres/blob/REL_16_0/src/tools/msvc/config_default.pl
3+
4+
$config->{openssl} = 'c:\Program Files\OpenSSL';

‎.ddev/ci/scripts/postgres/linux/50_install_postgres.sh

+6
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
#!/bin/bash
2+
3+
set -ex
4+
5+
sudo apt update
6+
sudo apt install -y --no-install-recommends build-essential python3-dev libpq-dev

‎.ddev/config.toml

+2
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,8 @@ oauthlib = ['BSD-3-Clause']
8383
mmh3 = ['CC0-1.0']
8484
# https://github.com/paramiko/paramiko/blob/master/LICENSE
8585
paramiko = ['LGPL-2.1-only']
86+
# https://github.com/psycopg/psycopg/blob/master/LICENSE.txt
87+
psycopg = ['LGPL-3.0-only']
8688
# https://github.com/psycopg/psycopg2/blob/master/LICENSE
8789
# https://github.com/psycopg/psycopg2/blob/master/doc/COPYING.LESSER
8890
psycopg2-binary = ['LGPL-3.0-only', 'BSD-3-Clause']

‎LICENSE-3rdparty.csv

+1
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ ply,PyPI,BSD-3-Clause,Copyright (C) 2001-2018
4545
prometheus-client,PyPI,Apache-2.0,Copyright 2015 The Prometheus Authors
4646
protobuf,PyPI,BSD-3-Clause,Copyright 2008 Google Inc. All rights reserved.
4747
psutil,PyPI,BSD-3-Clause,"Copyright (c) 2009, Jay Loden, Dave Daeschler, Giampaolo Rodola"
48+
psycopg,PyPI,LGPL-3.0-only,Copyright (C) 2020 The Psycopg Team
4849
psycopg2-binary,PyPI,BSD-3-Clause,Copyright 2013 Federico Di Gregorio
4950
psycopg2-binary,PyPI,LGPL-3.0-only,Copyright (C) 2013 Federico Di Gregorio
5051
pyOpenSSL,PyPI,Apache-2.0,Copyright The pyOpenSSL developers

‎agent_requirements.in

+1
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ prometheus-client==0.20.0
3232
protobuf==5.27.3
3333
psutil==6.0.0
3434
psycopg2-binary==2.9.9
35+
psycopg[c]==3.2.3
3536
pyasn1==0.4.8
3637
pycryptodomex==3.20.0
3738
pydantic==2.8.2

‎postgres/changelog.d/19322.added

+1
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Upgrade psycopg to version 3 for Postgres integration

‎postgres/datadog_checks/postgres/connections.py

+10-8
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
import time
99
from typing import Callable, Dict
1010

11-
import psycopg2
11+
import psycopg
1212

1313

1414
class ConnectionPoolFullError(Exception):
@@ -23,7 +23,7 @@ def __str__(self):
2323
class ConnectionInfo:
2424
def __init__(
2525
self,
26-
connection: psycopg2.extensions.connection,
26+
connection: psycopg.Connection,
2727
deadline: int,
2828
active: bool,
2929
last_accessed: int,
@@ -86,9 +86,9 @@ def _get_connection_raw(
8686
dbname: str,
8787
ttl_ms: int,
8888
timeout: int = None,
89-
startup_fn: Callable[[psycopg2.extensions.connection], None] = None,
89+
startup_fn: Callable[[psycopg.Connection], None] = None,
9090
persistent: bool = False,
91-
) -> psycopg2.extensions.connection:
91+
) -> psycopg.Connection:
9292
"""
9393
Return a connection from the pool.
9494
Pass a function to startup_func if there is an action needed with the connection
@@ -117,7 +117,7 @@ def _get_connection_raw(
117117
# if already in pool, retain persistence status
118118
persistent = conn.persistent
119119

120-
if db.status != psycopg2.extensions.STATUS_READY:
120+
if db.info.status != psycopg.pq.ConnStatus.OK:
121121
# Some transaction went wrong and the connection is in an unhealthy state. Let's fix that
122122
db.rollback()
123123

@@ -138,7 +138,7 @@ def get_connection(
138138
dbname: str,
139139
ttl_ms: int,
140140
timeout: int = None,
141-
startup_fn: Callable[[psycopg2.extensions.connection], None] = None,
141+
startup_fn: Callable[[psycopg.Connection], None] = None,
142142
persistent: bool = False,
143143
):
144144
"""
@@ -147,12 +147,14 @@ def get_connection(
147147
make a new connection if the max_conn limit hasn't been reached.
148148
Blocks until a connection can be added to the pool,
149149
and optionally takes a timeout in seconds.
150-
Note that leaving a connection context here does NOT close the connection in psycopg2;
150+
Note that leaving a connection context here does NOT close the connection in psycopg;
151151
connections must be manually closed by `close_all_connections()`.
152152
"""
153153
try:
154154
with self._mu:
155-
db = self._get_connection_raw(dbname, ttl_ms, timeout, startup_fn, persistent)
155+
db = self._get_connection_raw(
156+
dbname=dbname, ttl_ms=ttl_ms, timeout=timeout, startup_fn=startup_fn, persistent=persistent
157+
)
156158
yield db
157159
finally:
158160
with self._mu:

‎postgres/datadog_checks/postgres/cursor.py

+3-4
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,7 @@
22
# All rights reserved
33
# Licensed under a 3-clause BSD style license (see LICENSE)
44

5-
import psycopg2.extensions
6-
import psycopg2.extras
5+
import psycopg
76

87
from datadog_checks.base.utils.db.sql_commenter import add_sql_comment
98

@@ -28,9 +27,9 @@ def execute(self, query, vars=None, ignore_query_metric=False):
2827
return super().execute(query, vars)
2928

3029

31-
class CommenterCursor(BaseCommenterCursor, psycopg2.extensions.cursor):
30+
class CommenterCursor(BaseCommenterCursor, psycopg.ClientCursor):
3231
pass
3332

3433

35-
class CommenterDictCursor(BaseCommenterCursor, psycopg2.extras.DictCursor):
34+
class CommenterDictCursor(BaseCommenterCursor, psycopg.ClientCursor):
3635
pass

‎postgres/datadog_checks/postgres/discovery.py

+1-2
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@
66

77
from datadog_checks.base import AgentCheck
88
from datadog_checks.base.utils.discovery import Discovery
9-
from datadog_checks.postgres.cursor import CommenterCursor
109
from datadog_checks.postgres.util import DatabaseConfigurationError, warning_with_tags
1110

1211
AUTODISCOVERY_QUERY: str = """select datname from pg_catalog.pg_database where datistemplate = false;"""
@@ -72,7 +71,7 @@ def get_items(self) -> List[str]:
7271

7372
def _get_databases(self) -> List[str]:
7473
with self.db_pool.get_connection(self._db, self._default_ttl) as conn:
75-
with conn.cursor(cursor_factory=CommenterCursor) as cursor:
74+
with conn.cursor() as cursor:
7675
cursor.execute(AUTODISCOVERY_QUERY)
7776
databases = list(cursor.fetchall())
7877
databases = [

‎postgres/datadog_checks/postgres/explain_parameterized_queries.py

+4-7
Original file line numberDiff line numberDiff line change
@@ -5,11 +5,10 @@
55
import logging
66
import re
77

8-
import psycopg2
8+
import psycopg
99

1010
from datadog_checks.base.utils.db.sql import compute_sql_signature
1111
from datadog_checks.base.utils.tracking import tracked_method
12-
from datadog_checks.postgres.cursor import CommenterDictCursor
1312

1413
from .util import DBExplainError
1514
from .version_utils import V12
@@ -77,7 +76,7 @@ def explain_statement(self, dbname, statement, obfuscated_statement):
7776
if self._check.version < V12:
7877
# if pg version < 12, skip explaining parameterized queries because
7978
# plan_cache_mode is not supported
80-
e = psycopg2.errors.UndefinedParameter("Unable to explain parameterized query")
79+
e = psycopg.errors.UndefinedParameter("Unable to explain parameterized query")
8180
logger.debug(
8281
"Unable to explain parameterized query. Postgres version %s does not support plan_cache_mode",
8382
self._check.version,
@@ -180,16 +179,14 @@ def _deallocate_prepared_statement(self, dbname, query_signature):
180179
)
181180

182181
def _execute_query(self, dbname, query):
183-
# Psycopg2 connections do not get closed when context ends;
184-
# leaving context will just mark the connection as inactive in MultiDatabaseConnectionPool
185182
with self._check.db_pool.get_connection(dbname, self._check._config.idle_connection_timeout) as conn:
186-
with conn.cursor(cursor_factory=CommenterDictCursor) as cursor:
183+
with conn.cursor() as cursor:
187184
logger.debug('Executing query=[%s]', query)
188185
cursor.execute(query, ignore_query_metric=True)
189186

190187
def _execute_query_and_fetch_rows(self, dbname, query):
191188
with self._check.db_pool.get_connection(dbname, self._check._config.idle_connection_timeout) as conn:
192-
with conn.cursor(cursor_factory=CommenterDictCursor) as cursor:
189+
with conn.cursor() as cursor:
193190
cursor.execute(query, ignore_query_metric=True)
194191
return cursor.fetchall()
195192

0 commit comments

Comments
 (0)