From 2ce4840482202d24189bfe5158fe6f765c84c902 Mon Sep 17 00:00:00 2001 From: Jelte Fennema-Nio Date: Tue, 16 Dec 2025 09:25:48 +0100 Subject: [PATCH v1 05/12] Add pytest infrastructure to interact with PostgreSQL servers This adds a Python framework to write tests against PostgreSQL servers similar to the framework we have for Perl tests. It's designed such that tests are both easy to read and write, by making common things (like running queries) easy to do. One of the primary differences with the Perl framework is that this uses libpq directly to run queries, as opposed to using the psql binary for that. Two Python packages are added under src/test/pytest: 1. libpq/ is a small ctypes wrapper libpq that was just built. This is used instead of a third party driver such as psycopg so that the tests have no dependency beyond pytest itself, and so that unreleased libpq functionality can be tested. 2. pypg/ builds the test framework on top of that: A PostgresServer class to query/configure/backup/replicate/etc a server. And pytest fixtures[1] to easily create instances of that class. The guiding principle is that a test should contain only the logic of the behaviour it tests, and no boilerplate. So the query helpers do the obvious thing by default: values are converted to and from their Python counterparts, errors are raised as Python exceptions that can be caught with `pytest.raises`, and a query returning a single row or a single cell returns that row or cell instead of a nested list. No tests are included with this infrastructure commit yet. That will be done in follow-up commits. --- doc/src/sgml/regress.sgml | 21 +- meson.build | 25 +- pyproject.toml | 3 + src/test/pytest/README | 196 +++- src/test/pytest/libpq/__init__.py | 51 + src/test/pytest/libpq/_bindings.py | 260 +++++ src/test/pytest/libpq/_conversions.py | 201 ++++ src/test/pytest/libpq/_core.py | 556 +++++++++++ src/test/pytest/libpq/errors.py | 96 ++ src/test/pytest/pypg/__init__.py | 30 + src/test/pytest/pypg/_env.py | 182 ++++ src/test/pytest/pypg/bins.py | 36 + src/test/pytest/pypg/fixtures.py | 348 +++++++ src/test/pytest/pypg/paths.py | 48 + src/test/pytest/pypg/portlock.py | 299 ++++++ src/test/pytest/pypg/proc.py | 176 ++++ src/test/pytest/pypg/server.py | 1287 +++++++++++++++++++++++++ src/test/pytest/pypg/util.py | 126 +++ src/test/pytest/pypg/wait.py | 48 + 19 files changed, 3977 insertions(+), 12 deletions(-) create mode 100644 src/test/pytest/libpq/__init__.py create mode 100644 src/test/pytest/libpq/_bindings.py create mode 100644 src/test/pytest/libpq/_conversions.py create mode 100644 src/test/pytest/libpq/_core.py create mode 100644 src/test/pytest/libpq/errors.py create mode 100644 src/test/pytest/pypg/__init__.py create mode 100644 src/test/pytest/pypg/_env.py create mode 100644 src/test/pytest/pypg/bins.py create mode 100644 src/test/pytest/pypg/fixtures.py create mode 100644 src/test/pytest/pypg/paths.py create mode 100644 src/test/pytest/pypg/portlock.py create mode 100644 src/test/pytest/pypg/proc.py create mode 100644 src/test/pytest/pypg/server.py create mode 100644 src/test/pytest/pypg/util.py create mode 100644 src/test/pytest/pypg/wait.py diff --git a/doc/src/sgml/regress.sgml b/doc/src/sgml/regress.sgml index 7d44cfff4d5..0eab3f66e97 100644 --- a/doc/src/sgml/regress.sgml +++ b/doc/src/sgml/regress.sgml @@ -1022,7 +1022,9 @@ PG_TEST_NOCLEAN=1 make -C src/bin/pg_dump check Tests in pyt directories use the Python - pytest framework. + pytest framework. These tests provide a + convenient way to test libpq client functionality and scenarios requiring + multiple PostgreSQL server instances. @@ -1059,6 +1061,23 @@ make -C src/test/recovery check + + You can also run pytest directly, to run every pytest test in the tree, or + to select tests the ways pytest itself supports: + +pytest +pytest src/test/recovery/pyt/test_049_wait_for_lsn.py +pytest -k "wait_for_lsn" + + + + + Many operations in the test suites use a 180-second timeout, which on slow + hosts may lead to load-induced timeouts. Setting the environment variable + PG_TEST_TIMEOUT_DEFAULT to a higher number will change + the default to avoid this. + + For more information on writing pytest tests, see the src/test/pytest/README file. diff --git a/meson.build b/meson.build index bcd4e6dc486..21fffde804b 100644 --- a/meson.build +++ b/meson.build @@ -1834,13 +1834,20 @@ endif pytest_enabled = false pytest_version = '' pytest_cmd = ['pytest'] # dummy, overwritten when pytest is found -# The mtest runner passes our pgtap plugin with -p, and that plugin lives in -# the directory below. pyproject.toml puts the same directory on pytest's -# pythonpath, but versions below 8.4 only apply that setting after loading the -# plugins passed with -p, so on those the plugin is only importable if -# PYTHONPATH is set here too. This won't help people manually running pytest -# outside of meson/make, but we expect those to use a recent enough version of -# pytest anyway (and if not they can manually configure PYTHONPATH too). +# We also configure the same PYTHONPATH in the pytest settings in +# pyproject.toml, but pytest versions below 8.4 only apply that setting after +# loading the plugins listed in its addopts. So on lower versions, any pytest +# invocation from the source root (even just 'pytest --version') fails with an +# import error for the 'pypg.fixtures' plugin, because it's only findable +# through that not-yet-applied PYTHONPATH. So we need to configure PYTHONPATH +# here too. This won't help people manually running pytest outside of +# meson/make, but we expect those to use a recent enough version of pytest +# anyway (and if not they can manually configure PYTHONPATH too). +# +# This PYTHONPATH requirement is also why we cannot use pytest.version() +# below, instead of parsing the 'pytest --version' output ourselves: meson +# provides no way to set env for its --version probe. And pytest.version() +# would need meson >= 0.62 anyway which we don't require yet. pytest_env = {'PYTHONPATH': meson.project_source_root() / 'src' / 'test' / 'pytest'} pytestopt = get_option('pytest') @@ -1849,11 +1856,11 @@ if not pytestopt.disabled() if pytest.found() pytest_enabled = true - pytest_version = run_command(pytest, '--version', check: false).stdout().strip().split(' ')[-1] + pytest_version = run_command(pytest, '--version', env: pytest_env, check: false).stdout().strip().split(' ')[-1] pytest_cmd = [pytest.full_path()] else # Try python -m pytest as a fallback - pytest_check = run_command(python, '-m', 'pytest', '--version', check: false) + pytest_check = run_command(python, '-m', 'pytest', '--version', env: pytest_env, check: false) if pytest_check.returncode() == 0 pytest_enabled = true pytest_version = pytest_check.stdout().strip().split(' ')[-1] diff --git a/pyproject.toml b/pyproject.toml index aec4e7e9804..72e654d7d9a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,3 +22,6 @@ minversion = "7.2" # Common test code can be found here. pythonpath = ["src/test/pytest"] + +# Load the shared fixtures plugin +addopts = ["-p", "pypg.fixtures"] diff --git a/src/test/pytest/README b/src/test/pytest/README index c15e0556dd2..fa37a777821 100644 --- a/src/test/pytest/README +++ b/src/test/pytest/README @@ -3,8 +3,8 @@ src/test/pytest/README Pytest-based tests ================== -This directory contains the infrastructure for Python-based tests using -pytest. +This directory contains the infrastructure and helper functions and classes for +Python-based tests using pytest. The tests themselves live in pyt/ directories next to the code they cover, the same way the Perl tests live in t/ directories, in files named test_.py. @@ -57,10 +57,202 @@ You can run specific test files and/or use pytest's -k option to select tests: Directory structure =================== +pypg/ + Python library providing common functions and pytest fixtures that can be + used in tests. + +libpq/ + A simple but user-friendly python wrapper around libpq + pgtap.py A pytest plugin to output results in TAP format +Writing tests +============= + +The simplest test you can write is: + + def test_addition(pg): + assert pg.sql("SELECT $1 + $2", 40, 2) == 42 + +Taking the "pg" fixture[2] provisions a default server. Each test file gets +its own server: it is started the first time a test in that file uses "pg", +and stopped once the file is done. The tests in a file therefore share one +server, but not with any other file. + +A best effort cleanup happens between tests: connections are closed and config +changes are rolled back (see "Server configuration"). It stops there, though. +Anything a test creates with SQL -- tables, roles, databases -- is still there +for the tests that follow, so drop it yourself, or pick names that don't +collide with the other tests in the file. + +It's also possible to create a dedicated Postgres server for the test using +"create_pg". This can be used a non-default configured server is needed +or more than one server is needed in a single test: + + def test_two_servers(create_pg): + node1 = create_pg("node1", conf={"work_mem": "1MB"}) + node2 = create_pg("node2") + + assert node1.sql("SHOW work_mem") == "1MB" + assert node1.port != node2.port + +When several tests in a file can share one setup, build it once in a +module-scoped fixture with create_pg_module. Those servers live as long as the +file does, and are rolled back between tests just like "pg": + + @pytest.fixture(scope="module") + def node(create_pg_module): + return create_pg_module("node", conf={"wal_level": "logical"}) + + def test_wal_level(node): + assert node.sql("SHOW wal_level") == "logical" + + def test_slot_can_be_created(node): + node.sql("SELECT pg_create_logical_replication_slot('s', 'pgoutput')") + assert node.sql("SELECT count(*) FROM pg_replication_slots") == 1 + node.sql("SELECT pg_drop_replication_slot('s')") + + +Running queries +=============== + +All the examples above use the pg.sql() method to run queries, it's the one to +use by default. It runs a single query and returns its result, with any +arguments bound to the query's $1, $2, ... placeholders. Building queries using +string formatting should usually not be needed. Results are converted to Python +and simplified to the shape of the query, and errors from the server are raised +as LibpqError: + + def test_error(pg): + with pytest.raises(LibpqError, match="division by zero"): + pg.sql("SELECT 1/0") + +Every pg.sql() call reuses one cached connection, so its session lasts as long +as the test: a SET, a temporary table or an open transaction is still in +effect in the next call. The connection is dropped between tests, so none of +that reaches the next one. There are quite a few alternative ways of running +queries, each of which are useful in its own scenario (an example using all of +them is below): + +pg.sql_batch(...) + Pure syntactic sugar for consecutive pg.sql() calls on the same connection, + returning a list with each statement's result. Convenient to run a series + of statements in a row, e.g. for setup or cleanup. + +pg.sql_oneshot(...) + Runs a query on a fresh connection and closes it right away. Use it over + pg.sql() when the query requires a temporary clean connection. It can also + take connection options as keyword arguments, e.g. + pg.sql_oneshot("SELECT current_user", user="alice"). + +pg.background_sql(...) + Runs a query in the background on its own thread and returns a Future + instead of a result. Use it over sql() for a query that is expected to + block, which sql() would simply hang on. You can use result() on the Future + to wait for the query once it should be unblocked. + +pg.poll_query_until(...) + Opens a new connection and runs a query over on that connection until it + returns the expected result, or the timeout expires. Use it over sql() + whenever the answer depends on something that happens in the background, + e.g. a standby catching up or a checkpoint finishing. + +pg.connect(...) + Returns a connection of its own and takes any connection options. Use it + over sql() when the test needs to keep two sessions going at once, or wants + to hold on to session state while running other queries through pg.sql(). + The sql(), sql_batch() and background_sql() methods work on the returned + connection exactly the same way as on pg itself. + + +These go together whenever a test has to observe blocking: + + def test_lock_waits(pg): + pg.sql_batch("CREATE TABLE t (i int)", "BEGIN", "LOCK TABLE t") + + other = pg.connect(application_name="locker") + other.sql("BEGIN") + waiting = other.background_sql("LOCK TABLE t") + pg.poll_query_until( + "SELECT count(*) FROM pg_stat_activity " + "WHERE application_name = 'locker' AND wait_event_type = 'Lock'", + expected=1, + ) + + pg.sql("COMMIT") + waiting.result() + +These are all PostgresServer methods; see pypg/server.py for the rest of them, +including the log helpers, the other wait helpers, and the backup and recovery +ones. + + +Server configuration +==================== + +Server configuration is changed by editing the config files and then +reloading (or restarting) the server: + + pg.append_conf(log_connections="all") # append a setting + pg.adjust_conf(work_mem="4MB") # replace any existing line(s) + pg.pg_ctl("reload") # or pg.pg_ctl("restart") + +Both helpers quote and escape values for postgresql.conf automatically, and +boolean GUCs accept Python True/False. The pg_hba.conf and pg_ident.conf files +can be reset to a single rule with pg.reset_hba(database, role, method) and +pg.reset_ident(map, system_user, pg_user); both reload the server themselves. + +A test does not have to put any of this back: the config files are restored +when it ends, and the server is reloaded or restarted to match. Config given +to a server when it is created is part of the server rather than of a test, so +that is kept. + + +Testing libpq itself +==================== + +It's also possible to write tests for libpq that don't need or even want a +Postgres server at all. For example tests of libpq itself that connect to a +mock server. Those tests can use the connect fixture instead to connect +anywhere: + + def test_must_connect_errors(connect): + with pytest.raises(LibpqError, match="invalid connection option"): + connect(some_unknown_keyword="whatever") + + +Timeouts +======== + +Tests inherit the PG_TEST_TIMEOUT_DEFAULT environment variable (defaulting +to 180 seconds). It is used as the default connect_timeout for connections +and as the default timeout for the polling/wait helpers (poll_query_until, +wait_for_log, wait_for_catchup, and similar). + + +Environment variables +===================== + +PG_TEST_TIMEOUT_DEFAULT + Default timeout in seconds for connections and for the wait helpers, not + for a test as a whole (default: 180) + +PG_CONFIG + Path to pg_config (default: uses PATH) + +TESTDATADIR + Directory for test data (default: pytest temp directory) + +PG_TEST_EXTRA + Space-separated list of optional test categories to run (e.g., "ssl") + +PG_TEST_NOCLEAN + Keep the data directories of servers whose tests passed, which are + otherwise removed at the end of each test file + + References ========== diff --git a/src/test/pytest/libpq/__init__.py b/src/test/pytest/libpq/__init__.py new file mode 100644 index 00000000000..787e728130a --- /dev/null +++ b/src/test/pytest/libpq/__init__.py @@ -0,0 +1,51 @@ +# Copyright (c) 2025, PostgreSQL Global Development Group + +""" +libpq testing utilities - ctypes bindings and helpers for PostgreSQL's libpq library. + +This module provides Python wrappers around libpq for use in pytest tests. +""" + +from __future__ import annotations + +from . import errors +from ._bindings import ( + ConnectionStatus, + DiagField, + ExecStatus, + load_libpq_handle, +) +from ._conversions import register_type_info +from ._core import ( + Notify, + PGconn, + PGresult, + PreparedStatement, + connect, + connstr, +) +from .errors import ( + LibpqError, + PostgresMessage, + PostgresNotice, + PostgresWarning, +) + +__all__ = [ + "ConnectionStatus", + "DiagField", + "ExecStatus", + "LibpqError", + "Notify", + "PGconn", + "PGresult", + "PostgresMessage", + "PostgresNotice", + "PostgresWarning", + "PreparedStatement", + "connect", + "connstr", + "errors", + "load_libpq_handle", + "register_type_info", +] diff --git a/src/test/pytest/libpq/_bindings.py b/src/test/pytest/libpq/_bindings.py new file mode 100644 index 00000000000..1563fa4897e --- /dev/null +++ b/src/test/pytest/libpq/_bindings.py @@ -0,0 +1,260 @@ +# Copyright (c) 2025, PostgreSQL Global Development Group + +""" +Low-level ctypes bindings for libpq. + +This is the FFI layer: the opaque struct/pointer types, the enums mirroring +libpq's status codes, the loader that opens the shared library and declares the +function prototypes, and the helper that reads diagnostic fields off a raw +result handle. Nothing here knows about the friendly PGconn/PGresult wrappers +in ``_core``; it only deals in raw libpq handles. +""" + +from __future__ import annotations + +import ctypes +import enum +import os +import platform +from typing import Any + + +# PG_DIAG field identifiers from postgres_ext.h +class DiagField(enum.IntEnum): + SEVERITY = ord("S") + SEVERITY_NONLOCALIZED = ord("V") + SQLSTATE = ord("C") + MESSAGE_PRIMARY = ord("M") + MESSAGE_DETAIL = ord("D") + MESSAGE_HINT = ord("H") + STATEMENT_POSITION = ord("P") + INTERNAL_POSITION = ord("p") + INTERNAL_QUERY = ord("q") + CONTEXT = ord("W") + SCHEMA_NAME = ord("s") + TABLE_NAME = ord("t") + COLUMN_NAME = ord("c") + DATATYPE_NAME = ord("d") + CONSTRAINT_NAME = ord("n") + SOURCE_FILE = ord("F") + SOURCE_LINE = ord("L") + SOURCE_FUNCTION = ord("R") + + +class ConnectionStatus(enum.IntEnum): + """PostgreSQL connection status codes from libpq.""" + + CONNECTION_OK = 0 + CONNECTION_BAD = 1 + + +class ExecStatus(enum.IntEnum): + """PostgreSQL result status codes from PQresultStatus.""" + + PGRES_EMPTY_QUERY = 0 + PGRES_COMMAND_OK = 1 + PGRES_TUPLES_OK = 2 + PGRES_COPY_OUT = 3 + PGRES_COPY_IN = 4 + PGRES_BAD_RESPONSE = 5 + PGRES_NONFATAL_ERROR = 6 + PGRES_FATAL_ERROR = 7 + PGRES_COPY_BOTH = 8 + PGRES_SINGLE_TUPLE = 9 + PGRES_PIPELINE_SYNC = 10 + PGRES_PIPELINE_ABORTED = 11 + + +class _PGconn(ctypes.Structure): + pass + + +class _PGresult(ctypes.Structure): + pass + + +class _PGnotify(ctypes.Structure): + """Mirror of libpq's PGnotify (postgres_ext.h). Only the public fields are + used; ``next`` is libpq-internal and kept opaque.""" + + _fields_ = [ + ("relname", ctypes.c_char_p), + ("be_pid", ctypes.c_int), + ("extra", ctypes.c_char_p), + ("next", ctypes.c_void_p), + ] + + +_PGconn_p = ctypes.POINTER(_PGconn) +_PGresult_p = ctypes.POINTER(_PGresult) +_PGnotify_p = ctypes.POINTER(_PGnotify) + +# Signature of a libpq notice receiver: void (*)(void *arg, const PGresult *res). +_NOTICE_RECEIVER = ctypes.CFUNCTYPE(None, ctypes.c_void_p, _PGresult_p) + + +def _extract_diag_fields( + lib: ctypes.CDLL, res: ctypes._Pointer[_PGresult] +) -> dict[str, Any]: + """Pull the PostgreSQL diagnostic fields off a raw result handle into the + keyword arguments shared by LibpqError and PostgresMessage. Takes a bare + _PGresult_p so the notice receiver callback can use it too.""" + + def field(diag: DiagField) -> str | None: + val = lib.PQresultErrorField(res, int(diag)) + return val.decode() if val else None + + position_str = field(DiagField.STATEMENT_POSITION) + return dict( + sqlstate=field(DiagField.SQLSTATE), + severity=field(DiagField.SEVERITY), + primary=field(DiagField.MESSAGE_PRIMARY), + detail=field(DiagField.MESSAGE_DETAIL), + hint=field(DiagField.MESSAGE_HINT), + schema_name=field(DiagField.SCHEMA_NAME), + table_name=field(DiagField.TABLE_NAME), + column_name=field(DiagField.COLUMN_NAME), + datatype_name=field(DiagField.DATATYPE_NAME), + constraint_name=field(DiagField.CONSTRAINT_NAME), + context=field(DiagField.CONTEXT), + position=int(position_str) if position_str else None, + ) + + +def load_libpq_handle( + bindir: str | os.PathLike[str], libdir: str | os.PathLike[str] +) -> ctypes.CDLL: + """ + Loads a ctypes handle for libpq and declares common function prototypes. + + ``bindir`` and ``libdir`` are pg_config's values, passed in by the caller + so this module needs no install-discovery dependency of its own. + """ + system = platform.system() + + if system in ("Linux", "FreeBSD", "NetBSD", "OpenBSD"): + name = "libpq.so.5" + elif system == "Darwin": + name = "libpq.5.dylib" + elif system == "Windows": + name = "libpq.dll" + else: + assert False, f"the libpq fixture must be updated for {system}" + + if system == "Windows": + # On Windows, libpq.dll is confusingly in bindir, not libdir. Its + # dependent DLLs (OpenSSL, zstd, ...) resolve via PATH, which ctypes' + # default (LOAD_LIBRARY_SEARCH_DEFAULT_DIRS) does not search; winmode=0 + # selects the standard PATH-inclusive search instead. + libpq_path = os.path.join(bindir, name) + lib = ctypes.CDLL(libpq_path, winmode=0) + else: + libpq_path = os.path.join(libdir, name) + lib = ctypes.CDLL(libpq_path) + + # + # Function Prototypes + # + + lib.PQconnectdb.restype = _PGconn_p + lib.PQconnectdb.argtypes = [ctypes.c_char_p] + + lib.PQstatus.restype = ctypes.c_int + lib.PQstatus.argtypes = [_PGconn_p] + + lib.PQclosePortal.restype = _PGresult_p + lib.PQclosePortal.argtypes = [_PGconn_p, ctypes.c_char_p] + + lib.PQclosePrepared.restype = _PGresult_p + lib.PQclosePrepared.argtypes = [_PGconn_p, ctypes.c_char_p] + + lib.PQresultStatus.restype = ctypes.c_int + lib.PQresultStatus.argtypes = [_PGresult_p] + + lib.PQclear.restype = None + lib.PQclear.argtypes = [_PGresult_p] + + lib.PQerrorMessage.restype = ctypes.c_char_p + lib.PQerrorMessage.argtypes = [_PGconn_p] + + lib.PQfinish.restype = None + lib.PQfinish.argtypes = [_PGconn_p] + + lib.PQresultErrorMessage.restype = ctypes.c_char_p + lib.PQresultErrorMessage.argtypes = [_PGresult_p] + + lib.PQntuples.restype = ctypes.c_int + lib.PQntuples.argtypes = [_PGresult_p] + + lib.PQnfields.restype = ctypes.c_int + lib.PQnfields.argtypes = [_PGresult_p] + + lib.PQgetvalue.restype = ctypes.c_char_p + lib.PQgetvalue.argtypes = [_PGresult_p, ctypes.c_int, ctypes.c_int] + + lib.PQgetisnull.restype = ctypes.c_int + lib.PQgetisnull.argtypes = [_PGresult_p, ctypes.c_int, ctypes.c_int] + + lib.PQftype.restype = ctypes.c_uint + lib.PQftype.argtypes = [_PGresult_p, ctypes.c_int] + + lib.PQresultErrorField.restype = ctypes.c_char_p + lib.PQresultErrorField.argtypes = [_PGresult_p, ctypes.c_int] + + _char_pp = ctypes.POINTER(ctypes.c_char_p) + + lib.PQexecParams.restype = _PGresult_p + lib.PQexecParams.argtypes = [ + _PGconn_p, + ctypes.c_char_p, # command + ctypes.c_int, # nParams + ctypes.POINTER(ctypes.c_uint), # paramTypes (Oid *) + _char_pp, # paramValues + ctypes.c_void_p, # paramLengths + ctypes.c_void_p, # paramFormats + ctypes.c_int, # resultFormat + ] + + lib.PQprepare.restype = _PGresult_p + lib.PQprepare.argtypes = [ + _PGconn_p, + ctypes.c_char_p, # stmtName + ctypes.c_char_p, # query + ctypes.c_int, # nParams + ctypes.c_void_p, # paramTypes + ] + + lib.PQexecPrepared.restype = _PGresult_p + lib.PQexecPrepared.argtypes = [ + _PGconn_p, + ctypes.c_char_p, # stmtName + ctypes.c_int, # nParams + _char_pp, # paramValues + ctypes.c_void_p, # paramLengths + ctypes.c_void_p, # paramFormats + ctypes.c_int, # resultFormat + ] + + lib.PQgetResult.restype = _PGresult_p + lib.PQgetResult.argtypes = [_PGconn_p] + + lib.PQgetCopyData.restype = ctypes.c_int + lib.PQgetCopyData.argtypes = [ + _PGconn_p, + ctypes.POINTER(ctypes.c_char_p), + ctypes.c_int, + ] + + lib.PQfreemem.restype = None + lib.PQfreemem.argtypes = [ctypes.c_void_p] + + lib.PQconsumeInput.restype = ctypes.c_int + lib.PQconsumeInput.argtypes = [_PGconn_p] + + lib.PQnotifies.restype = _PGnotify_p + lib.PQnotifies.argtypes = [_PGconn_p] + + lib.PQsetNoticeReceiver.restype = ctypes.c_void_p + lib.PQsetNoticeReceiver.argtypes = [_PGconn_p, _NOTICE_RECEIVER, ctypes.c_void_p] + + return lib diff --git a/src/test/pytest/libpq/_conversions.py b/src/test/pytest/libpq/_conversions.py new file mode 100644 index 00000000000..b0c05d5a1ab --- /dev/null +++ b/src/test/pytest/libpq/_conversions.py @@ -0,0 +1,201 @@ +# Copyright (c) 2025, PostgreSQL Global Development Group + +""" +Conversion between PostgreSQL text values and Python objects. + +Two directions live here: the OID-keyed registry that turns a column's text +value into a typed Python object (``register_type_info`` / ``_convert_pg_value`` +/ ``_parse_array``), and ``_build_params``, which encodes Python parameters into +the C array libpq's extended-protocol functions expect. ``simplify_query_results`` +is the small ergonomic helper that unwraps single-row/single-column results. +""" + +from __future__ import annotations + +import ctypes +import datetime +import decimal +import json +import uuid +from collections.abc import Callable +from typing import Any + +# PostgreSQL type OIDs and conversion system +# Type registry - maps OID to converter function +_type_converters: dict[int, Callable[[str], Any]] = {} +_array_to_elem_map: dict[int, int] = {} + + +def register_type_info( + name: str, oid: int, array_oid: int, converter: Callable[[str], Any] +): + """ + Register a PostgreSQL type with its OID, array OID, and conversion function. + + Usage: + register_type_info("bool", 16, 1000, lambda v: v == "t") + """ + _type_converters[oid] = converter + if array_oid is not None: + _array_to_elem_map[array_oid] = oid + + +def _parse_array(value: str, elem_oid: int) -> list[Any]: + """Parse PostgreSQL array syntax into nested Python lists.""" + stack: list[list[Any]] = [] + current_element: list[str] = [] + in_quotes = False + was_quoted = False + pos = 0 + + while pos < len(value): + char = value[pos] + + if in_quotes: + if char == "\\": + next_char = value[pos + 1] + if next_char not in '"\\': + raise NotImplementedError('Only \\" and \\\\ escapes are supported') + current_element.append(next_char) + pos += 2 + continue + elif char == '"': + in_quotes = False + else: + current_element.append(char) + elif char == '"': + in_quotes = True + was_quoted = True + elif char == "{": + stack.append([]) + elif char in ",}": + if current_element or was_quoted: + elem = "".join(current_element) + if not was_quoted and elem == "NULL": + stack[-1].append(None) + else: + stack[-1].append(_convert_pg_value(elem, elem_oid)) + current_element = [] + was_quoted = False + if char == "}": + completed = stack.pop() + if not stack: + return completed + stack[-1].append(completed) + elif char != " ": + current_element.append(char) + pos += 1 + + raise ValueError(f"Malformed array literal: {value}") + + +# Register standard PostgreSQL types that we'll likely encounter in tests +register_type_info("bool", 16, 1000, lambda v: v == "t") +register_type_info("int2", 21, 1005, int) +register_type_info("int4", 23, 1007, int) +register_type_info("int8", 20, 1016, int) +register_type_info("float4", 700, 1021, float) +register_type_info("float8", 701, 1022, float) +register_type_info("numeric", 1700, 1231, decimal.Decimal) +register_type_info("text", 25, 1009, str) +register_type_info("varchar", 1043, 1015, str) +register_type_info("date", 1082, 1182, datetime.date.fromisoformat) +register_type_info("time", 1083, 1183, datetime.time.fromisoformat) +register_type_info("timestamp", 1114, 1115, datetime.datetime.fromisoformat) +register_type_info("timestamptz", 1184, 1185, datetime.datetime.fromisoformat) +register_type_info("uuid", 2950, 2951, uuid.UUID) +register_type_info("json", 114, 199, json.loads) +register_type_info("jsonb", 3802, 3807, json.loads) + + +def _convert_pg_value(value: str, type_oid: int) -> Any: + """ + Convert PostgreSQL string value to appropriate Python type based on OID. + Uses the registered type converters from register_type_info(). + """ + # Check if it's an array type + if type_oid in _array_to_elem_map: + elem_oid = _array_to_elem_map[type_oid] + return _parse_array(value, elem_oid) + + # Use registered converter if available + converter = _type_converters.get(type_oid) + if converter: + return converter(value) + + # Unknown types - return as string + return value + + +def simplify_query_results(results: list[tuple[Any, ...]]) -> Any: + """ + Simplify the results of a query so that the caller doesn't have to unpack + lists and tuples of length 1. + """ + if len(results) == 1: + row = results[0] + if len(row) == 1: + # If there's only a single cell, just return the value + return row[0] + # If there's only a single row, just return that row + return row + + if len(results) != 0 and len(results[0]) == 1: + # If there's only a single column, return an array of values + return [row[0] for row in results] + + # if there are multiple rows and columns, return the results as is + return results + + +# Type OIDs for the Python types whose SQL type is unambiguous. Sending these +# saves the caller writing casts into the query: without a type the server sees +# an "unknown" parameter, which it cannot resolve when a function is overloaded +# (generate_series($1, $2) being the common example). +_BOOL_OID = 16 +_INT8_OID = 20 +_INT4_OID = 23 +_FLOAT8_OID = 701 +_NUMERIC_OID = 1700 + +# The ranges of the integer types we can declare. Python has no equivalent of +# INT32_MAX to borrow -- sys.maxsize is the pointer width, not int4 -- so name +# them here rather than spelling powers of two inline. +_INT4_MIN, _INT4_MAX = -(2**31), 2**31 - 1 +_INT8_MIN, _INT8_MAX = -(2**63), 2**63 - 1 + +# Strings are deliberately absent: leaving them unknown is what lets a Python +# str stand in for whatever the context wants -- a regclass argument, an enum, +# a pg_lsn -- without the test spelling out the type. + + +def _param_oid(value: Any) -> int: + """The OID to declare for ``value``, or 0 to leave it to the server.""" + if isinstance(value, bool): # before int: bool is a subclass of it + return _BOOL_OID + if isinstance(value, int): + # The narrowest type that holds it, so a parameter can be passed to a + # function or compared against a column expecting a plain integer. + # Python ints are unbounded, so fall back to numeric beyond int8. + if _INT4_MIN <= value <= _INT4_MAX: + return _INT4_OID + if _INT8_MIN <= value <= _INT8_MAX: + return _INT8_OID + return _NUMERIC_OID + if isinstance(value, float): + return _FLOAT8_OID + return 0 + + +def _build_params(params: tuple[Any, ...]) -> tuple[int, Any, Any]: + """Build the (nParams, paramTypes, paramValues) triple libpq's + extended-protocol functions expect from a tuple of Python parameter values. + Values are passed in text format; ``None`` becomes a SQL NULL.""" + if not params: + return 0, None, None + values = (ctypes.c_char_p * len(params))() + types = (ctypes.c_uint * len(params))() + for i, p in enumerate(params): + values[i] = None if p is None else str(p).encode() + types[i] = _param_oid(p) + return len(params), types, values diff --git a/src/test/pytest/libpq/_core.py b/src/test/pytest/libpq/_core.py new file mode 100644 index 00000000000..04201cba11e --- /dev/null +++ b/src/test/pytest/libpq/_core.py @@ -0,0 +1,556 @@ +# Copyright (c) 2025, PostgreSQL Global Development Group + +""" +Friendly connection and result wrappers over libpq. + +PGconn and PGresult wrap the raw libpq handles from the _bindings module and +turn them into an easy to use API. The ctypes bindings live in _bindings, +value conversion in _conversions, and the server-message warning categories +in messages; this module ties them together. +""" + +from __future__ import annotations + +import contextlib +import ctypes +import warnings +from concurrent.futures import Future, ThreadPoolExecutor +from typing import Any, NamedTuple, NoReturn + +from ._bindings import ( + _NOTICE_RECEIVER, + ConnectionStatus, + DiagField, + ExecStatus, + _extract_diag_fields, + _PGconn, + _PGresult, +) +from ._conversions import _build_params, _convert_pg_value, simplify_query_results +from .errors import ( + LibpqError, + PostgresMessage, + PostgresNotice, + PostgresWarning, +) + + +# A LISTEN/NOTIFY notification, as returned by PGconn.notifies(). +class Notify(NamedTuple): + channel: str + pid: int + payload: str + + +class PGresult(contextlib.AbstractContextManager): + """Wraps a raw _PGresult_p with a more friendly interface.""" + + def __init__(self, lib: ctypes.CDLL, res: ctypes._Pointer[_PGresult]): + self._lib = lib + # Cleared to None by __exit__ once the result has been freed. + self._res: ctypes._Pointer[_PGresult] | None = res + + def __exit__(self, *exc: object) -> None: + self._lib.PQclear(self._res) + self._res = None + + def status(self) -> ExecStatus: + return ExecStatus(self._lib.PQresultStatus(self._res)) + + def error_message(self) -> str: + """Returns the error message associated with this result.""" + msg = self._lib.PQresultErrorMessage(self._res) + return msg.decode() if msg else "" + + def raise_error(self) -> NoReturn: + """ + Raises LibpqError with diagnostic information from the result. + """ + if not self._res: + raise LibpqError("query failed: out of memory or connection lost") + + fields = _extract_diag_fields(self._lib, self._res) + raise LibpqError(fields["primary"] or self.error_message(), **fields) + + def fetch_all(self) -> list[tuple[Any, ...]]: + """ + Fetch all rows and convert to Python types. + Returns a list of tuples, with values converted based on their PostgreSQL type. + """ + nrows = self._lib.PQntuples(self._res) + ncols = self._lib.PQnfields(self._res) + + # Get type OIDs for each column + type_oids = [self._lib.PQftype(self._res, col) for col in range(ncols)] + + results = [] + for row in range(nrows): + row_data = [] + for col in range(ncols): + if self._lib.PQgetisnull(self._res, row, col): + row_data.append(None) + else: + value = self._lib.PQgetvalue(self._res, row, col).decode() + row_data.append(_convert_pg_value(value, type_oids[col])) + results.append(tuple(row_data)) + + return results + + +class _MustConsumeFuture(Future): + """The Future returned by ``PGconn.background_sql()``. + + A plain ``concurrent.futures.Future`` silently drops an exception that + nobody retrieves, which would let a background query fail unnoticed. This + subclass records whether ``result()``/``exception()`` was called, so that + ``PGconn.close()`` can turn a forgotten ``.result()`` into a visible test + failure instead of a silent pass. + """ + + def __init__(self) -> None: + super().__init__() + self.consumed = False + + def result(self, timeout: float | None = None) -> Any: + self.consumed = True + return super().result(timeout) + + def exception(self, timeout: float | None = None) -> BaseException | None: + self.consumed = True + return super().exception(timeout) + + +class PGconn(contextlib.AbstractContextManager): + """ + Wraps a raw _PGconn_p with a more friendly interface. This is just a + stub; it's expected to grow. + """ + + def __init__( + self, + lib: ctypes.CDLL, + handle: ctypes._Pointer[_PGconn], + stack: contextlib.ExitStack, + ): + self._lib = lib + # Cleared to None by close() once the connection has been finished. + self._handle: ctypes._Pointer[_PGconn] | None = handle + self._stack = stack + + # background_sql() machinery. A single libpq connection must never be + # driven by two threads at once, so background queries run on one + # worker thread (created lazily on first use) and only one may be in + # flight at a time. ``_pending`` is that query's future, if any. + self._executor: ThreadPoolExecutor | None = None + self._pending: _MustConsumeFuture | None = None + + # Sequence for prepare()'s generated __p{n} statement names. + self._prepared_seq = 0 + + # Surface NOTICE/WARNING messages as Python warnings so tests can + # assert on them with pytest.warns(...). The callback object must stay + # alive as long as the connection, or ctypes frees it and libpq calls + # into freed memory. + self._notice_cb = _NOTICE_RECEIVER(self._receive_notice) + self._lib.PQsetNoticeReceiver(self._handle, self._notice_cb, None) + + def _receive_notice( + self, _arg: int | None, res: ctypes._Pointer[_PGresult] + ) -> None: + severity = self._lib.PQresultErrorField( + res, int(DiagField.SEVERITY_NONLOCALIZED) + ) + message = self._lib.PQresultErrorMessage(res) + # WARNING and NOTICE get their own categories; anything else (INFO, LOG, + # DEBUG, ...) falls back to the PostgresMessage base. + category = { + b"WARNING": PostgresWarning, + b"NOTICE": PostgresNotice, + }.get(severity, PostgresMessage) + # Passing a constructed warning instance (rather than string + + # category) makes warnings.warn use its type as the category, and lets + # us attach the same diagnostic fields a LibpqError carries. + fields = _extract_diag_fields(self._lib, res) + warnings.warn(category(message.decode().rstrip("\n"), **fields)) + + def __exit__(self, *exc: object) -> None: + # When another exception is already propagating, that is the real + # failure: abandon any pending background query rather than raising + # close()'s own "result never consumed" error on top of it. + if exc[0] is not None: + self._pending = None + self.close() + + def close(self) -> None: + """Close the connection (PQfinish). Idempotent, so it is safe to close + early even though the owning ExitStack will also close it at teardown. + + Closing with an unconsumed background_sql() future raises a + RuntimeError.""" + + self._check_pending() + + if self._executor is not None: + self._executor.shutdown(wait=True) + self._executor = None + + self._close_impl() + + def _close_impl(self) -> None: + """Release the libpq handle (PQfinish), without close()'s pending + guard or executor shutdown. Idempotent; also called directly by the + background_sql(close_when_done=True) worker, which cannot go through + close() (see background_sql).""" + if self._handle is not None: + self._lib.PQfinish(self._handle) + self._handle = None + + def close_portal(self, name: str) -> None: + """ + Close a portal, releasing its snapshot. Pass ``""`` for the unnamed + portal — the one sql()/sql_batch() statements bind. + + Inside an open transaction the unnamed portal survives until the next + statement, keeping its snapshot registered and the backend's xmin set. + Call this when a test needs an idle-in-transaction session holding an + XID but no snapshot — e.g. one that CREATE INDEX CONCURRENTLY's + WaitForOlderSnapshots must NOT wait for. + """ + self._check_pending() + res = self._lib.PQclosePortal(self._handle, name.encode()) + self._result_or_raise(self._stack.enter_context(PGresult(self._lib, res))) + + def sql(self, query: str, *params: Any, simplify_result: bool = True) -> Any: + """ + Runs ``query`` through the extended query protocol (an unnamed Parse/ + Bind/Execute), the same path real client drivers use, and raises an + exception if it fails. Any ``params`` are bound to the query's + ``$1, $2, ...`` placeholders in text format, the libpq equivalent of + psql's `` \\bind \\g``. + + Returns the query results with automatic type conversion and simplification. + For commands that don't return data (INSERT, UPDATE, etc.), returns None. + + Examples: + - SELECT 1 -> 1 + - SELECT 1, 2 -> (1, 2) + - SELECT * FROM generate_series(1, 3) -> [1, 2, 3] + - SELECT * FROM (VALUES (1, 'a'), (2, 'b')) t -> [(1, 'a'), (2, 'b')] + - CREATE TABLE ... -> None + - INSERT INTO ... -> None + + Pass ``simplify_result=False`` to always get a list of row tuples, with + no unwrapping and an empty list rather than None. That is what you want + when the row count is part of what the test is checking, or varies: + otherwise the caller has to undo the simplification to tell one row from + one column from no rows at all. + + - SELECT 1 -> [(1,)] + - SELECT * FROM generate_series(1, 3) -> [(1,), (2,), (3,)] + - SELECT WHERE false -> [] + - CREATE TABLE ... -> [] (still nothing to return) + """ + self._check_pending() + return self._sql_impl(query, *params, simplify_result=simplify_result) + + def _sql_impl(self, query: str, *params: Any, simplify_result: bool = True) -> Any: + """The actual PQexecParams call behind sql()/background_sql().""" + nparams, types, values = _build_params(params) + res = self._lib.PQexecParams( + self._handle, query.encode(), nparams, types, values, None, None, 0 + ) + return self._result_or_raise( + self._stack.enter_context(PGresult(self._lib, res)), + simplify_result=simplify_result, + ) + + def _check_pending(self) -> None: + """Guard run before anything else touches the connection: raises until + a background_sql() future has been consumed via result()/exception(). + + While the query runs the connection is genuinely busy (a second libpq + call would race the worker thread); once it finishes, raising forces + the caller to deal with the outcome — including any error — instead of + letting it silently leak past the next query.""" + if self._pending is None: + return + + if self._pending.consumed: + self._pending = None + return + + if self._pending.done(): + raise RuntimeError( + "the previous background_sql() result was never " + "consumed; call .result() on its future before " + "issuing another query" + ) + raise RuntimeError( + "connection is busy with an unresolved background_sql(); " + "call .result() on its future before issuing another query" + ) + + def background_sql( + self, + query: str, + *params: Any, + close_when_done: bool = False, + simplify_result: bool = True, + ) -> Future[Any]: + """Dispatch a query that is expected to *block* (on a lock, an + injection point, ...) and return an already-running Future. The test + can carry on (e.g. confirm the wait with + ``PostgresServer.wait_for_event()``, then release it) and collect the + outcome with ``.result()``, which re-raises any LibpqError. + + The query runs on a worker thread over this same connection rather + than in a separate process, so the session state it builds up + (transactions, locks, session-local settings) is visible to later + calls on this connection. Only one background query at a time: its + future must be consumed before the connection is used again, or + _check_pending() raises. + + With ``close_when_done=True`` the worker thread finishes the + connection (PQfinish) as soon as the query completes, for callers that + dispatch on a throwaway connection (see + ``PostgresServer.background_sql_oneshot``). The worker cannot go + through close() for this — close() rejects an unconsumed future and + would join the worker's own thread — so it releases the handle + directly; the must-consume guard still fires at teardown. + + ``simplify_result`` works as it does for ``sql()``: pass False when the + future's result is a row set whose length matters, so it arrives as a + list of tuples rather than being unwrapped.""" + return self._dispatch_background( + lambda: self._sql_impl(query, *params, simplify_result=simplify_result), + close_when_done=close_when_done, + ) + + def _dispatch_background( + self, work: Any, *, close_when_done: bool = False + ) -> Future[Any]: + """Run ``work`` (a no-argument callable issuing exactly one query on + this connection) on the worker thread, with the single-pending-query + bookkeeping shared by background_sql() and + PreparedStatement.background_exec().""" + self._check_pending() + if self._executor is None: + self._executor = ThreadPoolExecutor(max_workers=1) + fut = _MustConsumeFuture() + + def run(): + if not fut.set_running_or_notify_cancel(): + return + try: + fut.set_result(work()) + except BaseException as e: + fut.set_exception(e) + finally: + if close_when_done: + self._close_impl() + + self._executor.submit(run) + self._pending = fut + return fut + + def sql_batch(self, *queries: str) -> list[Any]: + """ + Runs each of ``queries`` through the extended query protocol, exactly + like consecutive sql() calls, and returns a list with every statement's + simplified result. + + Raises on the first failing statement; earlier statements stay + executed (and committed). + """ + self._check_pending() + return [self._sql_impl(query) for query in queries] + + def notifies(self) -> list[Notify]: + """ + Return and consume all pending LISTEN/NOTIFY notifications, each a + ``Notify(channel, pid, payload)``. + + Input is consumed first (``PQconsumeInput``) so notifications already + waiting on the socket are picked up. A LISTENing session only receives + notifications once its transaction ends, so call this after the + relevant command — and poll, since they may arrive slightly after the + command's own result. + """ + self._check_pending() + self._lib.PQconsumeInput(self._handle) + out = [] + while True: + n = self._lib.PQnotifies(self._handle) + if not n: + break + c = n.contents + out.append(Notify(c.relname.decode(), c.be_pid, c.extra.decode())) + self._lib.PQfreemem(n) + return out + + def _result_or_raise(self, res: PGresult, *, simplify_result: bool = True) -> Any: + """Turn a PGresult into a Python value, raising LibpqError on any error + status. Shared by sql() and the extended-protocol helpers. With + ``simplify_result=False`` the rows are returned as-is, so callers that + care about the row count don't have to guess whether a bare value means + one row, one column, or none.""" + status = res.status() + if status == ExecStatus.PGRES_COMMAND_OK: + return None if simplify_result else [] + if status == ExecStatus.PGRES_TUPLES_OK: + rows = res.fetch_all() + return simplify_query_results(rows) if simplify_result else rows + if status == ExecStatus.PGRES_COPY_OUT: + # Drain the COPY OUT stream. An error raised mid-copy surfaces as + # PQgetCopyData returning -2 with the real result coming from + # PQgetResult, so resolve that result (raising on error) before + # handing back the copied bytes — undecoded, since COPY can stream + # binary data. + chunks = [] + buf = ctypes.c_char_p() + while True: + # Read exactly the returned length; COPY data is not + # guaranteed to be NUL-terminated. + n = self._lib.PQgetCopyData(self._handle, ctypes.byref(buf), 0) + if n <= 0: + break + chunks.append(ctypes.string_at(buf, n)) + self._lib.PQfreemem(buf) + buf = ctypes.c_char_p() + final = self._lib.PQgetResult(self._handle) + self._result_or_raise(self._stack.enter_context(PGresult(self._lib, final))) + return b"".join(chunks) + res.raise_error() + + def prepare(self, query: str, *, name: str | None = None) -> PreparedStatement: + """ + Parse ``query`` into a named prepared statement and return a + ``PreparedStatement`` to run it with. This is the libpq equivalent of + psql's `` \\parse ``. + + Pass ``name=`` only when the test cares about the statement's name + (e.g. it appears in a log line or pg_prepared_statements); otherwise a + connection-unique ``__p{n}`` name is generated. + """ + self._check_pending() + if name is None: + self._prepared_seq += 1 + name = f"__p{self._prepared_seq}" + res = self._lib.PQprepare(self._handle, name.encode(), query.encode(), 0, None) + self._result_or_raise(self._stack.enter_context(PGresult(self._lib, res))) + return PreparedStatement(self, name) + + def _exec_prepared_impl(self, name: str, *params: Any) -> Any: + """The PQexecPrepared call behind PreparedStatement.exec()/ + background_exec(); like _sql_impl it runs without the pending guard so + the worker thread can use it.""" + # PQexecPrepared takes no paramTypes: the types were fixed at Parse. + nparams, _types, values = _build_params(params) + res = self._lib.PQexecPrepared( + self._handle, name.encode(), nparams, values, None, None, 0 + ) + return self._result_or_raise( + self._stack.enter_context(PGresult(self._lib, res)) + ) + + +class PreparedStatement(contextlib.AbstractContextManager): + """A named server-side prepared statement, created by ``PGconn.prepare()``. + + Execute it with ``exec(*params)`` (or ``background_exec(*params)`` when + the execution is expected to block), and release it with ``close()`` — or + use it as a context manager to do that automatically. The statement lives + on the connection that prepared it, so all methods are subject to that + connection's single-pending-query rule. + """ + + def __init__(self, conn: PGconn, name: str): + self._conn = conn + self.name = name + self._closed = False + + def exec(self, *params: Any) -> Any: + """Bind ``params`` to the statement and execute it, returning + simplified results like ``PGconn.sql()``. This is the libpq equivalent + of psql's ``\\bind_named \\g``.""" + self._conn._check_pending() + return self._conn._exec_prepared_impl(self.name, *params) + + def background_exec(self, *params: Any) -> Future[Any]: + """Execute the statement on the connection's worker thread and return + a Future, with the same semantics and must-consume rule as + ``PGconn.background_sql()``.""" + return self._conn._dispatch_background( + lambda: self._conn._exec_prepared_impl(self.name, *params) + ) + + def close(self) -> None: + """Release the prepared statement (the protocol-level Close message, + like DEALLOCATE). Idempotent, and a no-op if the connection itself is + already gone — the statement died with its session.""" + if self._closed: + return + self._closed = True + if self._conn._handle is None: + return + self._conn._check_pending() + res = self._conn._lib.PQclosePrepared(self._conn._handle, self.name.encode()) + self._conn._result_or_raise( + self._conn._stack.enter_context(PGresult(self._conn._lib, res)) + ) + + def __exit__(self, *exc: object) -> None: + # Mirror PGconn.__exit__: when an exception is already propagating, + # don't let close()'s own errors (e.g. an unconsumed background future + # on the connection) mask it. + if exc[0] is not None: + self._closed = True + return + self.close() + + +def connstr(opts: dict[str, Any]) -> str: + """ + Flattens the provided options into a libpq connection string. Values + are converted to str and quoted/escaped as necessary. + """ + settings: list[str] = [] + + for k, v in opts.items(): + v = str(v) + if not v: + v = "''" + else: + v = v.replace("\\", "\\\\") + v = v.replace("'", "\\'") + + # libpq ends an unquoted value at the first whitespace of any kind + # (not just a space), so wrap in single quotes whenever the value + # contains any whitespace. + if any(c.isspace() for c in v): + v = f"'{v}'" + + settings.append(f"{k}={v}") + + return " ".join(settings) + + +def connect( + libpq_handle: ctypes.CDLL, + stack: contextlib.ExitStack, + **opts: Any, +) -> PGconn: + """ + Connects to a server using the given libpq connection options (host, port, + dbname, connect_timeout, ...) and returns a PGconn registered in ``stack`` + for cleanup. Raises LibpqError if the connection fails. + """ + + conn_p = libpq_handle.PQconnectdb(connstr(opts).encode()) + + if libpq_handle.PQstatus(conn_p) != ConnectionStatus.CONNECTION_OK: + error_msg = libpq_handle.PQerrorMessage(conn_p).decode() + libpq_handle.PQfinish(conn_p) + raise LibpqError(error_msg) + + return stack.enter_context(PGconn(libpq_handle, conn_p, stack=stack)) diff --git a/src/test/pytest/libpq/errors.py b/src/test/pytest/libpq/errors.py new file mode 100644 index 00000000000..8fa206ad5d9 --- /dev/null +++ b/src/test/pytest/libpq/errors.py @@ -0,0 +1,96 @@ +# Copyright (c) 2025, PostgreSQL Global Development Group + +""" +Exception and warning classes for libpq. + +Errors the server reports on a failed result are raised as ``LibpqError``; +NOTICE/WARNING/... messages it reports on a successful result are surfaced as +the ``PostgresMessage`` warning categories. Both sides share +``PostgresDiagnostics`` so a caught error or warning exposes the same +``.detail``/``.hint``/``.sqlstate`` fields. +""" + +from __future__ import annotations + + +class PostgresDiagnostics(Exception): + """Holds the PostgreSQL diagnostic fields (SQLSTATE, detail, hint, ...) the + server attaches to a result. + + The server sends the same set of fields on an error result as on a + NOTICE/WARNING result, so this is mixed into both ``LibpqError`` and + ``PostgresMessage``: a caught notice exposes ``.detail``, + ``.constraint_name``, etc. exactly like a caught error does. + + It roots at ``Exception`` — a shared base of both ``RuntimeError`` and + ``UserWarning`` — so that ``super().__init__(message)`` cooperatively + reaches the real base and stores the message as usual. + """ + + sqlstate: str | None + severity: str | None + primary: str | None + detail: str | None + hint: str | None + schema_name: str | None + table_name: str | None + column_name: str | None + datatype_name: str | None + constraint_name: str | None + position: int | None + context: str | None + + def __init__( + self, + message: str, + *, + sqlstate: str | None = None, + severity: str | None = None, + primary: str | None = None, + detail: str | None = None, + hint: str | None = None, + schema_name: str | None = None, + table_name: str | None = None, + column_name: str | None = None, + datatype_name: str | None = None, + constraint_name: str | None = None, + position: int | None = None, + context: str | None = None, + ): + super().__init__(message) + self.sqlstate = sqlstate + self.severity = severity + self.primary = primary + self.detail = detail + self.hint = hint + self.schema_name = schema_name + self.table_name = table_name + self.column_name = column_name + self.datatype_name = datatype_name + self.constraint_name = constraint_name + self.position = position + self.context = context + + +class LibpqError(PostgresDiagnostics, RuntimeError): + """Exception for libpq errors with PostgreSQL diagnostic fields.""" + + +class PostgresMessage(PostgresDiagnostics, UserWarning): + """Base category for server messages surfaced over libpq as Python warnings. + + Messages the server sends outside of an error result (what psql prints to + stderr) are reported as Python warnings, so tests can assert on them with + ``pytest.warns(..., match=...)``. WARNING and NOTICE map to the subclasses + below; any other level (INFO, LOG, DEBUG, ...) gets this base category. + Like ``LibpqError`` it carries the full diagnostic fields (see + ``PostgresDiagnostics``). + """ + + +class PostgresNotice(PostgresMessage): + """A NOTICE message reported by the server over libpq.""" + + +class PostgresWarning(PostgresMessage): + """A WARNING message reported by the server over libpq.""" diff --git a/src/test/pytest/pypg/__init__.py b/src/test/pytest/pypg/__init__.py new file mode 100644 index 00000000000..8dc1b1dbcb1 --- /dev/null +++ b/src/test/pytest/pypg/__init__.py @@ -0,0 +1,30 @@ +# Copyright (c) 2025, PostgreSQL Global Development Group + +from __future__ import annotations + +from ._env import ( + check_pg_config, + clean_libpq_environment, + pg_test_timeout_default, + require_injection_points, + require_test_extras, + skip_unless_injection_points, + skip_unless_test_extras, +) +from .server import PostgresServer +from .wait import wait_until + +# Clear inherited libpq connection environment variables as soon as the test +# framework is imported, before any server is started or connection is made. +clean_libpq_environment() + +__all__ = [ + "PostgresServer", + "check_pg_config", + "pg_test_timeout_default", + "require_injection_points", + "require_test_extras", + "skip_unless_injection_points", + "skip_unless_test_extras", + "wait_until", +] diff --git a/src/test/pytest/pypg/_env.py b/src/test/pytest/pypg/_env.py new file mode 100644 index 00000000000..91fd4d44adf --- /dev/null +++ b/src/test/pytest/pypg/_env.py @@ -0,0 +1,182 @@ +# Copyright (c) 2025, PostgreSQL Global Development Group + +from __future__ import annotations + +import functools +import logging +import os + +import pytest + +from .paths import INCLUDEDIR_SERVER, SHAREDIR + +logger = logging.getLogger(__name__) + + +# libpq reads many PG* environment variables as connection defaults, and a +# stray value overrides parameters the framework sets explicitly (e.g. GitHub's +# Windows runners preset PGUSER=postgres for their bundled PostgreSQL, breaking +# every connection). Clear them up front. +_LIBPQ_ENV_VARS = ( + "PGAPPNAME", + "PGCLIENTENCODING", + "PGCONNECT_TIMEOUT", + "PGDATA", + "PGDATABASE", + "PGGSSENCMODE", + "PGHOST", + "PGHOSTADDR", + "PGOPTIONS", + "PGPASSFILE", + "PGPASSWORD", + "PGPORT", + "PGREQUIREPEER", + "PGREQUIRESSL", + "PGSERVICE", + "PGSERVICEFILE", + "PGSSLCERT", + "PGSSLCRL", + "PGSSLCRLDIR", + "PGSSLKEY", + "PGSSLMODE", + "PGSSLROOTCERT", + "PGTARGETSESSIONATTRS", + "PGUSER", +) + + +def clean_libpq_environment() -> None: + """Remove inherited libpq connection environment variables (see above).""" + for var in _LIBPQ_ENV_VARS: + os.environ.pop(var, None) + + +def _test_extra_skip_reason(*keys: str) -> str: + return "requires {} to be set in PG_TEST_EXTRA".format(", ".join(keys)) + + +def _has_test_extra(key: str) -> bool: + """ + Returns True if the PG_TEST_EXTRA environment variable contains the given + key. + """ + extra = os.getenv("PG_TEST_EXTRA", "") + return key in extra.split() + + +def require_test_extras(*keys: str) -> pytest.MarkDecorator: + """ + A convenience annotation which will skip tests if all of the required keys + are not present in PG_TEST_EXTRA. + + To skip a particular test function or class: + + @pypg.require_test_extras("ldap") + def test_some_ldap_feature(): + ... + + To skip an entire module: + + pytestmark = pypg.require_test_extra("ssl", "kerberos") + """ + return pytest.mark.skipif( + not all(_has_test_extra(k) for k in keys), + reason=_test_extra_skip_reason(*keys), + ) + + +def skip_unless_test_extras(*keys: str) -> None: + """ + Skip the current test/fixture if any of the required keys are not present + in PG_TEST_EXTRA. Use this inside fixtures where decorators can't be used. + + @pytest.fixture + def my_fixture(): + skip_unless_test_extras("ldap") + ... + """ + if not all(_has_test_extra(k) for k in keys): + pytest.skip(_test_extra_skip_reason(*keys)) + + +_INJECTION_POINTS_SKIP_REASON = "injection points not supported by this build" + + +@functools.cache +def _injection_points_supported() -> bool: + """Return whether the server build supports injection points. + + The ``injection_points`` test extension is only installed when the server + was built with injection point support, so its control file is present + exactly when the feature is available. Checking the filesystem rather than + ``pg_available_extensions`` needs no running node, so this can be used as + a collection-time decorator; the control file is preferred over the shared + library because its name is platform independent. + """ + return (SHAREDIR / "extension" / "injection_points.control").exists() + + +def require_injection_points() -> pytest.MarkDecorator: + """Skip the decorated test/class/module unless the build supports + injection points. + + @pypg.require_injection_points() + def test_some_injection_point(): + ... + + or, for an entire module:: + + pytestmark = pypg.require_injection_points() + """ + return pytest.mark.skipif( + not _injection_points_supported(), + reason=_INJECTION_POINTS_SKIP_REASON, + ) + + +def skip_unless_injection_points() -> None: + """Skip the current test/fixture unless the build supports injection + points. Use this inside fixtures where decorators can't be used; prefer + the ``require_injection_points()`` decorator otherwise. + """ + if not _injection_points_supported(): + pytest.skip(_INJECTION_POINTS_SKIP_REASON) + + +@functools.cache +def _pg_config_h_lines() -> tuple[str, ...]: + """Return the lines of the server build's ``pg_config.h``, stripped. + + Read once and cached, since the build under test does not change during a + session. + """ + path = INCLUDEDIR_SERVER / "pg_config.h" + return tuple(line.strip() for line in path.read_text().splitlines()) + + +def check_pg_config(line: str) -> bool: + """Return whether the server build's ``pg_config.h`` contains a line that + starts with ``line``. + + Use it to gate tests on build-time feature macros, e.g.:: + + if not check_pg_config("#define USE_ICU 1"): + pytest.skip("ICU not supported by this build") + """ + return any(candidate.startswith(line) for candidate in _pg_config_h_lines()) + + +def pg_test_timeout_default() -> int: + """ + Returns the value of the PG_TEST_TIMEOUT_DEFAULT environment variable, in + seconds, or 180 if one was not provided. + """ + default = os.getenv("PG_TEST_TIMEOUT_DEFAULT", "") + if not default: + return 180 + + try: + return int(default) + except ValueError as v: + logger.warning("PG_TEST_TIMEOUT_DEFAULT could not be parsed: " + str(v)) + return 180 diff --git a/src/test/pytest/pypg/bins.py b/src/test/pytest/pypg/bins.py new file mode 100644 index 00000000000..6095d462cef --- /dev/null +++ b/src/test/pytest/pypg/bins.py @@ -0,0 +1,36 @@ +# Copyright (c) 2025, PostgreSQL Global Development Group + +"""Callable handles for every installed PostgreSQL program. + +Importing any name from this module yields a :class:`~pypg.proc.PgBin` for that +program, so ``from pypg.bins import psql`` (or pg_verifybackup, pg_controldata, +...) works for any installed program without a hardcoded list:: + + from pypg.bins import psql, pg_verifybackup + psql("-c", "select 1") + pg_verifybackup.check_standard_options() +""" + +from __future__ import annotations + +import functools + +from .proc import PgBin + + +@functools.cache +def _bin(name: str) -> PgBin: + return PgBin(name) + + +def __getattr__(name: str) -> PgBin: + # PEP 562 module-level __getattr__: any attribute access becomes a cached + # PgBin for that program name. Guard dunders/privates so importlib, + # copy/pickle, and "from pypg.bins import _x" probes raise normally rather + # than fabricating a PgBin("_x"). + # + # NOTE: we don't use the functools.cache decorator directly on this + # function, because that confuses Pyright typechecking. + if name.startswith("_"): + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + return _bin(name) diff --git a/src/test/pytest/pypg/fixtures.py b/src/test/pytest/pypg/fixtures.py new file mode 100644 index 00000000000..a94f2f62196 --- /dev/null +++ b/src/test/pytest/pypg/fixtures.py @@ -0,0 +1,348 @@ +# Copyright (c) 2025, PostgreSQL Global Development Group + +from __future__ import annotations + +import contextlib +import ctypes +import os +import pathlib +import shutil +import tempfile +from collections.abc import Callable, Iterator +from typing import Any + +import pytest +from libpq import PGconn, load_libpq_handle +from libpq import connect as libpq_connect + +from ._env import pg_test_timeout_default +from .paths import BINDIR, LIBDIR +from .server import PostgresServer + +# Stash key for tracking servers for log reporting. +_servers_key = pytest.StashKey[list[PostgresServer]]() + + +def _record_server_for_log_reporting( + request: pytest.FixtureRequest, server: PostgresServer +) -> None: + """Record a server for log reporting on test failure.""" + if _servers_key not in request.node.stash: + request.node.stash[_servers_key] = [] + request.node.stash[_servers_key].append(server) + + +@pytest.fixture(scope="session") +def libpq_handle() -> ctypes.CDLL: + """ + Loads a ctypes handle for libpq. Some common function prototypes are + initialized for general use. + + Session-scoped: the loaded library is immutable, process-global state, so + there is nothing per-module to isolate. + """ + return load_libpq_handle(BINDIR, LIBDIR) + + +@pytest.fixture +def connect(libpq_handle: ctypes.CDLL) -> Iterator[Callable[..., PGconn]]: + """ + Returns a function to connect to PostgreSQL via libpq. + + The returned function accepts connection options as keyword arguments + (host, port, dbname, etc.) and returns a PGconn object. Connections + are automatically cleaned up at the end of the test. + + Example: + conn = connect(host='localhost', port=5432, dbname='postgres') + result = conn.sql("SELECT 1") + """ + with contextlib.ExitStack() as stack: + + def _connect(**opts: object) -> PGconn: + opts.setdefault("connect_timeout", pg_test_timeout_default()) + return libpq_connect(libpq_handle, stack, **opts) + + yield _connect + + +@pytest.fixture(scope="session") +def tmp_check(tmp_path_factory: pytest.TempPathFactory) -> pathlib.Path: + """ + Returns the tmp_check directory that should be used for the tests. If + TESTDATADIR is provided, that will be used; otherwise a new temporary + directory is created in the pytest temp root. + + Session-scoped: this directory is shared by all test files in a pytest + invocation (e.g. ``make check`` sets a single TESTDATADIR for the whole + suite), so everything created in it must have a name that's unique across + the whole suite. Server basedirs get this via _reserve_basedir below. + """ + d = os.getenv("TESTDATADIR") + if d: + d = pathlib.Path(d) + else: + d = tmp_path_factory.mktemp("tmp_check") + + return d + + +@pytest.fixture(scope="module") +def _reserve_basedir( + request: pytest.FixtureRequest, tmp_check: pathlib.Path +) -> Iterator[Callable[[str], pathlib.Path]]: + """ + Returns a function that reserves a basedir for a named server, under which + the server keeps everything it owns (see Server.__init__). The test file + name is included in the directory name, so that test files reusing a + server name never collide inside the suite-wide shared tmp_check directory. + + On teardown the handed-out basedirs are removed again if every test in + this file passed -- a full run otherwise leaves many gigabytes of cluster + data behind, slowing down CI. On any failure, or when PG_TEST_NOCLEAN is + set, the data is kept for debugging. Only this file's + directories are removed -- not all of tmp_check -- so data kept for an + earlier failed file survives later passing ones. This runs after the + dependent server fixtures' teardown has stopped their nodes. + """ + basedirs: list[pathlib.Path] = [] + + def _reserve(name: str) -> pathlib.Path: + d = tmp_check / f"{request.path.stem}_{name}" + basedirs.append(d) + return d + + # Only the Session tracks failures, so detect failures in *this module* + # as the change in its counter over the fixture's lifetime. + failed_before = request.session.testsfailed + yield _reserve + + if ( + request.session.testsfailed == failed_before + and "PG_TEST_NOCLEAN" not in os.environ + ): + for d in basedirs: + shutil.rmtree(d, ignore_errors=True) + + +@pytest.fixture(scope="module") +def basedir(_reserve_basedir: Callable[[str], pathlib.Path]) -> pathlib.Path: + """ + Returns the basedir to use for the pg fixture's server. + """ + + return _reserve_basedir("default") + + +@pytest.fixture(scope="module") +def sockdir() -> Iterator[pathlib.Path]: + """ + Returns the directory name to use as the server's unix_socket_directories + setting. Local client connections use this as the PGHOST. + + Uses tempfile.TemporaryDirectory directly instead of pytest's + tmp_path_factory, because macOS limits Unix socket paths to 104 bytes + and pytest's nested temp directories can exceed that. On Linux the limit is + 108 bytes, but the pytest temp root is usually /tmp while on macOS the root + is much longer. + """ + with tempfile.TemporaryDirectory(prefix="pytest_postgres_sock") as d: + yield pathlib.Path(d) + + +@pytest.fixture(scope="module") +def pg_server_module( + request: pytest.FixtureRequest, + basedir: pathlib.Path, + sockdir: pathlib.Path, + libpq_handle: ctypes.CDLL, +) -> Iterator[PostgresServer]: + """ + Starts a running Postgres server for the test module, listening on + localhost. The HBA initially allows only local UNIX connections from the + same user. + + Module-scoped rather than session-scoped on purpose: meson runs every test + file in its own process, so a session never spans more than one module — + but a local ``pytest pyt/`` run does, and would behave differently if this + were session-scoped. + + Per-test isolation is a separate concern, handled by the ``pg`` fixture, + which opens a per-test cleanup subcontext via ``start_new_test()``. + + Returns a PostgresServer instance with methods for server management, + configuration, and creating test databases/users. + """ + server = PostgresServer("default", basedir, sockdir, libpq_handle) + try: + server.start() + except Exception: + # If startup fails the tests never run, so they never get the chance to + # register the server for log reporting; do it here so the startup logs + # still make it into the failure report. + _record_server_for_log_reporting(request, server) + raise + + yield server + + # Cleanup any test resources, then stop the server. + server.cleanup() + server.stop() + + +@pytest.fixture +def pg( + request: pytest.FixtureRequest, pg_server_module: PostgresServer +) -> Iterator[PostgresServer]: + """ + Per-test server context. Use this fixture to make changes to the server + which will be rolled back at the end of the test (e.g., creating test + users/databases). + + Also captures the PostgreSQL log position at test start so that any new + log entries can be included in the test report on failure. + """ + with pg_server_module.start_new_test() as s: + _record_server_for_log_reporting(request, s) + yield s + + +@pytest.fixture +def create_pg( + request: pytest.FixtureRequest, + sockdir: pathlib.Path, + libpq_handle: ctypes.CDLL, + _reserve_basedir: Callable[[str], pathlib.Path], +) -> Iterator[Callable[..., PostgresServer]]: + """ + Factory fixture to create additional PostgreSQL servers (per-test scope). + + Returns a function that creates new PostgreSQL server instances. + Servers are automatically cleaned up at the end of the test. + + Example: + def test_multiple_servers(create_pg): + node1 = create_pg() + node2 = create_pg() + node3 = create_pg() + """ + servers: list[PostgresServer] = [] + + def _create( + name: str | None = None, start: bool = True, **kwargs: Any + ) -> PostgresServer: + if name is None: + count = len(servers) + 1 + name = f"pg{count}" + + basedir = _reserve_basedir(name) + server = PostgresServer(name, basedir, sockdir, libpq_handle, **kwargs) + servers.append(server) + _record_server_for_log_reporting(request, server) + # Pass start=False when the test must touch the data directory before + # startup (e.g. drop an extra signal file) or expects startup to fail; + # call server.start() yourself afterwards. + if start: + server.start() + return server + + yield _create + + for server in servers: + server.cleanup() + server.stop() + + +@pytest.fixture(scope="module") +def _module_scoped_servers() -> list[PostgresServer]: + """Module-scoped list to track servers created by create_pg_module.""" + return [] + + +@pytest.fixture(scope="module") +def create_pg_module( + request: pytest.FixtureRequest, + sockdir: pathlib.Path, + libpq_handle: ctypes.CDLL, + _reserve_basedir: Callable[[str], pathlib.Path], + _module_scoped_servers: list[PostgresServer], +) -> Iterator[Callable[..., PostgresServer]]: + """ + Factory fixture to create additional PostgreSQL servers (module scope). + + Like create_pg, but servers persist for the entire test module. + Use this when multiple tests in a module can share the same servers. + + A new per-test subcontext is opened on all servers at the start of each + test via the _start_module_server_tests autouse fixture. + + Example: + @pytest.fixture(scope="module") + def shared_nodes(create_pg_module): + return [create_pg_module() for _ in range(3)] + """ + + def _create( + name: str | None = None, start: bool = True, **kwargs: Any + ) -> PostgresServer: + if name is None: + count = len(_module_scoped_servers) + 1 + name = f"pg{count}" + basedir = _reserve_basedir(name) + server = PostgresServer(name, basedir, sockdir, libpq_handle, **kwargs) + _module_scoped_servers.append(server) + _record_server_for_log_reporting(request, server) + if start: + server.start() + return server + + yield _create + + for server in _module_scoped_servers: + server.cleanup() + server.stop() + + +@pytest.fixture(autouse=True) +def _start_module_server_tests( + _module_scoped_servers: list[PostgresServer], +) -> Iterator[None]: + """Opens a per-test subcontext on all module-scoped servers for this test. + + It's hard to reliably detect whether a test uses a module-scoped server or + not. So this simply assumes all tests in the module use the module-scoped + servers. There's little harm in registering servers for tests that don't + use them. + """ + with contextlib.ExitStack() as stack: + for server in _module_scoped_servers: + stack.enter_context(server.start_new_test()) + yield + + +@pytest.hookimpl(wrapper=True, trylast=True) +def pytest_runtest_makereport(item: pytest.Item, call: pytest.CallInfo[None]): + """ + Adds PostgreSQL server logs to the test report sections. + """ + report = yield + + session_servers = item.session.stash.get(_servers_key, []) + + module_node = item.getparent(pytest.Module) + module_servers = module_node.stash.get(_servers_key, []) if module_node else [] + + servers = session_servers + module_servers + item.stash.get(_servers_key, []) + + include_name = len(servers) > 1 + + for server in servers: + content = server.log_content() + if content.strip(): + section_title = f"Postgres log {report.when}" + if include_name: + section_title += f" ({server.name})" + report.sections.append((section_title, content)) + server.reset_log_position() + + return report diff --git a/src/test/pytest/pypg/paths.py b/src/test/pytest/pypg/paths.py new file mode 100644 index 00000000000..44ff72093a8 --- /dev/null +++ b/src/test/pytest/pypg/paths.py @@ -0,0 +1,48 @@ +# Copyright (c) 2025, PostgreSQL Global Development Group + +""" +Install-location discovery for the PostgreSQL build under test. + +The paths are constants for the whole session, so they are plain module +globals, filled from a single ``pg_config`` run at import time. pg_config is +found via the ``PG_CONFIG`` environment variable, falling back to ``PATH``. +The cost is that importing this module fails if pg_config can't be run -- but +no test can do anything useful without an install, so failing collection +loudly is fine. +""" + +from __future__ import annotations + +import os +import pathlib + +from .util import capture + + +def _config_values() -> dict[str, str]: + """All pg_config settings, from a single argument-less pg_config run, which + prints every setting as a ``NAME = value`` line.""" + pg_config = os.environ.get("PG_CONFIG", "pg_config") + values = {} + for line in capture(pg_config, silent=True).splitlines(): + name, sep, value = line.partition(" = ") + if sep: + values[name] = value + return values + + +_values = _config_values() + +BINDIR = pathlib.Path(_values["BINDIR"]) +"""PostgreSQL bin directory (pg_config's ``BINDIR``).""" + +LIBDIR = pathlib.Path(_values["LIBDIR"]) +"""PostgreSQL lib directory (pg_config's ``LIBDIR``).""" + +SHAREDIR = pathlib.Path(_values["SHAREDIR"]) +"""PostgreSQL share directory (pg_config's ``SHAREDIR``).""" + +INCLUDEDIR_SERVER = pathlib.Path(_values["INCLUDEDIR-SERVER"]) +"""PostgreSQL server include directory (pg_config's ``INCLUDEDIR-SERVER``).""" + +del _values diff --git a/src/test/pytest/pypg/portlock.py b/src/test/pytest/pypg/portlock.py new file mode 100644 index 00000000000..876fc41f8b4 --- /dev/null +++ b/src/test/pytest/pypg/portlock.py @@ -0,0 +1,299 @@ +# Copyright (c) 2026, PostgreSQL Global Development Group + +"""Port allocation, shared with the Perl TAP tests. + +This is a port of ``get_free_port()`` and friends from +``src/test/perl/PostgreSQL/Test/Cluster.pm``, deliberately kept mechanically +identical: the same port range, the same bind probe, and the same +``$portdir/$port.rsv`` lock files in the same directory, resolved the same way. +A pytest run and a prove run regularly happen at the same time, this way the +cooperate together when selecting ports. +""" + +from __future__ import annotations + +import atexit +import errno +import os +import pathlib +import random +import socket +import sys + +# Two things Perl gets from its runtime and Python does not: a portable file +# lock and a way to ask whether a pid is alive. Neither module exists on the +# other platform, so these imports cannot be unconditional. +if sys.platform == "win32": + import ctypes + import msvcrt + from ctypes import wintypes +else: + import fcntl + +# Chosen to sit above the range servers typically use on Unix and below the +# range those systems use for ephemeral client ports (Cluster.pm has the same +# two constants and the same reasoning). +PORT_LOWER_BOUND = 10200 +PORT_UPPER_BOUND = 32767 + +_reservation_files: list[pathlib.Path] = [] +# Every port this process has spoken for, so a later search skips it even when +# the server using it is stopped or was never started. Ports that came from +# get_free_port() also have a reservation file; one that a caller picked itself +# does not, which is what this set is really for (see mark_assigned()). +_assigned_ports: set[int] = set() +# Tracking of the last port assigned, to accelerate the search. +_last_port_assigned = random.randint(PORT_LOWER_BOUND, PORT_UPPER_BOUND) + +if sys.platform == "win32": + + class _OVERLAPPED(ctypes.Structure): + """Only ever used zeroed, to lock from offset 0 like Perl does.""" + + _fields_ = [ + ("Internal", wintypes.LPVOID), + ("InternalHigh", wintypes.LPVOID), + ("Offset", wintypes.DWORD), + ("OffsetHigh", wintypes.DWORD), + ("hEvent", wintypes.HANDLE), + ] + + _LOCKFILE_EXCLUSIVE_LOCK = 0x2 + # The length perl's win32_flock() locks (its LK_LEN), and so the length we + # have to lock to collide with a prove run holding the same reservation + # file. + _LK_LEN = 0xFFFF0000 + + # Spelling out the prototypes matters for OpenProcess(): its return value + # is a HANDLE, which ctypes would otherwise truncate to a C int. + _kernel32 = ctypes.WinDLL("kernel32", use_last_error=True) + _kernel32.OpenProcess.restype = wintypes.HANDLE + _kernel32.OpenProcess.argtypes = (wintypes.DWORD, wintypes.BOOL, wintypes.DWORD) + _kernel32.GetExitCodeProcess.argtypes = (wintypes.HANDLE, wintypes.LPDWORD) + _kernel32.CloseHandle.argtypes = (wintypes.HANDLE,) + _kernel32.LockFileEx.argtypes = ( + wintypes.HANDLE, + wintypes.DWORD, + wintypes.DWORD, + wintypes.DWORD, + wintypes.DWORD, + ctypes.POINTER(_OVERLAPPED), + ) + _kernel32.UnlockFileEx.argtypes = ( + wintypes.HANDLE, + wintypes.DWORD, + wintypes.DWORD, + wintypes.DWORD, + ctypes.POINTER(_OVERLAPPED), + ) + + +def _flock_exclusive(fh) -> None: + """Take an exclusive lock on ``fh``, waiting for as long as it takes.""" + if sys.platform != "win32": + # It has to be flock(2), the same primitive Perl's flock uses, and not + # fcntl.lockf(): that is a POSIX record lock, which on Linux does not + # exclude flock(2) holders, so the TAP tests and these would not lock + # against each other at all. + fcntl.flock(fh.fileno(), fcntl.LOCK_EX) + return + + # Windows has no flock(). Call LockFileEx() instead, with the arguments + # perl's own flock() emulation uses [1], so that a pytest run and a prove + # run contend for the same byte range and block on each other rather than + # each taking a lock the other cannot see. + # + # [1] win32_flock() in perl's win32/win32.c: + # https://github.com/Perl/perl5/blob/v5.44.0/win32/win32.c#L3084 + overlapped = _OVERLAPPED() + if not _kernel32.LockFileEx( + msvcrt.get_osfhandle(fh.fileno()), + _LOCKFILE_EXCLUSIVE_LOCK, + 0, + _LK_LEN, + 0, + ctypes.byref(overlapped), + ): + raise ctypes.WinError(ctypes.get_last_error()) + + +def _flock_unlock(fh) -> None: + """Release the lock taken by _flock_exclusive().""" + if sys.platform != "win32": + fcntl.flock(fh.fileno(), fcntl.LOCK_UN) + return + + # The range has to match the one that was locked exactly, or the lock stays + # held until the handle is closed. + overlapped = _OVERLAPPED() + if not _kernel32.UnlockFileEx( + msvcrt.get_osfhandle(fh.fileno()), + 0, + _LK_LEN, + 0, + ctypes.byref(overlapped), + ): + raise ctypes.WinError(ctypes.get_last_error()) + + +def _pid_is_running(pid: int) -> bool: + """Whether ``pid`` is a live process, the question Cluster.pm asks with + ``kill 0``.""" + if sys.platform != "win32": + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + except PermissionError: + # The process exists but is not ours. Cluster.pm treats this as a + # free port too ("process exists and is owned by us" is what makes + # it refuse), so behave the same rather than diverging. + return False + return True + + # os.kill() cannot be used to probe a pid here: on Windows Python documents + # that any signal other than CTRL_C_EVENT/CTRL_BREAK_EVENT unconditionally + # kills the process via TerminateProcess. Probing a reservation must not + # kill whoever holds it. + PROCESS_QUERY_LIMITED_INFORMATION = 0x1000 + STILL_ACTIVE = 259 + + handle = _kernel32.OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, False, pid) + if not handle: + return False + try: + # OpenProcess() alone is not enough: the process object outlives the + # process itself for as long as anybody holds a handle to it, so an + # exited test runner would keep looking alive and leak its port. + exit_code = wintypes.DWORD() + if not _kernel32.GetExitCodeProcess(handle, ctypes.byref(exit_code)): + return False + return exit_code.value == STILL_ACTIVE + finally: + _kernel32.CloseHandle(handle) + + +def _portdir() -> pathlib.Path: + """The lock directory, resolved exactly as Cluster.pm resolves it. + + ``PG_TEST_PORT_DIR`` wins (that is what a buildfarm client sets), then a + ``portlock`` directory at the top of the build tree, and failing that one + under the test's data directory. + """ + portdir = os.environ.get("PG_TEST_PORT_DIR") + if not portdir: + # PostgreSQL::Test::Utils::tmp_check is TESTDATADIR, or "tmp_check". + build_dir = os.environ.get("top_builddir") or os.environ.get( + "TESTDATADIR", "tmp_check" + ) + portdir = f"{build_dir}/portlock" + return pathlib.Path(portdir.replace("\\", "/")) + + +def can_bind(addr: str, port: int) -> bool: + """Whether ``addr:port`` can be bound and listened on right now.""" + family = socket.AF_INET6 if ":" in addr else socket.AF_INET + sock = socket.socket(family, socket.SOCK_STREAM) + try: + # As in the postmaster (and in Cluster.pm's can_bind), don't use + # SO_REUSEADDR on Windows, where it would let us bind a port somebody + # else already has and so report a taken port as free. + if sys.platform != "win32": + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + try: + sock.bind((addr, port)) + sock.listen(socket.SOMAXCONN) + except OSError as e: + # EADDRNOTAVAIL means the address itself is unusable here, which is + # not the port's fault; anything else means the port is taken. + return e.errno == errno.EADDRNOTAVAIL + finally: + sock.close() + return True + + +def _reserve_port(port: int) -> bool: + """Claim ``port`` in the lock directory. False if somebody else holds it. + + The file holds the owning pid, so a reservation left behind by a process + that has since died is reclaimed rather than leaking the port forever. + """ + filename = _portdir() / f"{port}.rsv" + # Open read-write so the lock is not lost by reopening, as Cluster.pm notes. + fd = os.open(filename, os.O_RDWR | os.O_CREAT, 0o644) + with os.fdopen(fd, "r+") as portfile: + _flock_exclusive(portfile) + try: + try: + pid = int((portfile.readline() or "0").strip() or "0") + except ValueError: + pid = 0 + if pid > 0 and _pid_is_running(pid): + return False + portfile.seek(0) + # Fixed width, so a shorter pid cannot leave trailing junk behind. + portfile.write(f"{os.getpid():10d}\n") + portfile.flush() + finally: + _flock_unlock(portfile) + _reservation_files.append(filename) + return True + + +def mark_assigned(port: int) -> None: + """Record that ``port`` is in use by this process, so get_free_port() will + not hand it out again. + + Only needed for a port that did not come from get_free_port(), i.e. one a + caller picked itself: that port has no reservation file, so this set is the + only thing keeping a later search off it. + """ + _assigned_ports.add(port) + + +def get_free_port(addrs: list[str] | None = None) -> int: + """Find a high TCP port nothing is bound to, and reserve it. + + ``addrs`` are the addresses the caller intends to listen on; the port has to + be free on all of them, plus on 127.0.0.1 so the result is usable for the + widest range of purposes. + + The reservation lasts until this process exits, whether the server using it + is running, stopped, or never started at all. + """ + global _last_port_assigned + + probe = ["127.0.0.1", *(addrs or [])] + if sys.platform == "win32": + # Testing 0.0.0.0 covers MSYS, which sets SO_EXCLUSIVEADDRUSE, but is + # not enough for native Windows, hence the individual addresses too. + # Cluster.pm probes exactly these, and only there, for the same reason. + probe += ["127.0.0.2", "127.0.0.3", "0.0.0.0"] + # Preserve order but drop duplicates. + probe = list(dict.fromkeys(probe)) + + _portdir().mkdir(parents=True, exist_ok=True) + + port = _last_port_assigned + while True: + port += 1 + if port > PORT_UPPER_BOUND: + port = PORT_LOWER_BOUND + if port in _assigned_ports: + continue + if not all(can_bind(addr, port) for addr in probe): + continue + if _reserve_port(port): + _last_port_assigned = port + mark_assigned(port) + return port + + +@atexit.register +def _release_reservations() -> None: + """Drop this process's reservations, as Cluster.pm's END block does.""" + for filename in _reservation_files: + try: + filename.unlink() + except OSError: + pass diff --git a/src/test/pytest/pypg/proc.py b/src/test/pytest/pypg/proc.py new file mode 100644 index 00000000000..36e851f949d --- /dev/null +++ b/src/test/pytest/pypg/proc.py @@ -0,0 +1,176 @@ +# Copyright (c) 2025, PostgreSQL Global Development Group + +"""Per-binary callable helpers for invoking installed PostgreSQL programs. + +Each installed program is a :class:`PgBin` instance. Calling it is like +:func:`pypg.util.run` (the program +streams to the console and a nonzero exit raises unless ``check=False``); +:meth:`PgBin.capture` is like :func:`pypg.util.capture` (returns the program's +stdout as text). Both accept ``server=`` to point the program at a running +server. Instances come from :mod:`pypg.bins`, not constructed directly:: + + from pypg.bins import psql, pg_controldata + psql("-c", "select 1", server=pg) # run, raise on failure + state = pg_controldata.capture(pg.datadir) # capture stdout +""" + +from __future__ import annotations + +import os +import pathlib +import re +import shutil +import subprocess +from collections.abc import Sequence +from typing import TYPE_CHECKING, Any + +from . import paths, util + +if TYPE_CHECKING: + from .server import PostgresServer + + +class PgBin: + """A single installed PostgreSQL program, resolved against the test bindir.""" + + def __init__(self, name: str): + self.name = name + # shutil.which rather than Path.exists: it also checks the executable + # bit, and on Windows resolves the implied .exe suffix. + resolved = shutil.which(paths.BINDIR / name) + if resolved is None: + raise FileNotFoundError( + f"program {name!r} is not installed in {paths.BINDIR}" + ) + self.path = pathlib.Path(resolved) + + def __repr__(self) -> str: + return f"PgBin({self.name!r})" + + def _apply_env( + self, + server: PostgresServer | None, + addenv: dict[str, str] | None, + kwargs: dict[str, Any], + ) -> None: + """Layer a server's PG* connection variables and/or ``addenv`` onto the + environment the program will run with.""" + if server is None and addenv is None: + return + env = kwargs.pop("env", None) + env = dict(env if env is not None else os.environ) + if server is not None: + env.update(server.connection_env()) + if addenv is not None: + env.update(addenv) + kwargs["env"] = env + + def __call__( + self, + *args: object, + server: PostgresServer | None = None, + addenv: dict[str, str] | None = None, + **kwargs: Any, + ) -> subprocess.CompletedProcess[Any]: + """Run the program. Like :func:`pypg.util.run`: output is not captured + (it streams) and a nonzero exit raises unless ``check=False``. Pass + ``server=`` to run against a :class:`PostgresServer`.""" + self._apply_env(server, addenv, kwargs) + return util.run(self.path, *args, **kwargs) + + def capture( + self, + *args: object, + server: PostgresServer | None = None, + addenv: dict[str, str] | None = None, + **kwargs: Any, + ) -> str: + """Run the program and return its stdout as text (trailing newline + stripped), like :func:`pypg.util.capture`. Raises on a nonzero exit + unless ``check=False``.""" + self._apply_env(server, addenv, kwargs) + return util.capture(self.path, *args, **kwargs) + + def _capture_both( + self, + *args: object, + server: PostgresServer | None = None, + addenv: dict[str, str] | None = None, + **kwargs: Any, + ) -> subprocess.CompletedProcess[str]: + """Run capturing both stdout and stderr as text (``check=False``), + returning the CompletedProcess. Backs the ``check_*`` helpers.""" + self._apply_env(server, addenv, kwargs) + return util.run( + self.path, + *args, + check=False, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + encoding="utf-8", + **kwargs, + ) + + def check_standard_options(self) -> None: + """Assert the conventions every client program must satisfy. + + ``--help`` and ``--version`` exit 0 writing only to stdout (and --help + keeps its lines within the length limit), and an unknown option exits + nonzero with a stderr message. This bundles what nearly every ``src/bin`` + suite checks. + """ + # The --help convention enforces a maximum line length. This value isn't + # set in stone; it reflects the current project convention (~80). + max_help_line = 95 + + r = self._capture_both("--help") + assert r.returncode == 0, r.stderr + assert r.stdout != "", "--help wrote nothing to stdout" + assert r.stderr == "", f"--help wrote to stderr: {r.stderr}" + too_long = [ln for ln in r.stdout.splitlines() if len(ln) > max_help_line] + assert not too_long, f"--help lines exceed {max_help_line} chars: {too_long}" + + r = self._capture_both("--version") + assert r.returncode == 0, r.stderr + assert r.stdout != "", "--version wrote nothing to stdout" + assert r.stderr == "", f"--version wrote to stderr: {r.stderr}" + + r = self._capture_both("--not-a-valid-option") + assert r.returncode != 0, "expected nonzero exit for an invalid option" + assert r.stderr != "", "expected an error message on stderr" + + def check_all( + self, + *args: object, + exit_code: int = 0, + stdout: str | Sequence[str] = (), + stderr: str | Sequence[str] = (), + server: PostgresServer | None = None, + **kwargs: Any, + ) -> subprocess.CompletedProcess[str]: + """Run the program and assert its exit code and output. + + ``stdout``/``stderr`` is a regex -- or an iterable of regexes -- that + must each be found (``re.search`` with DOTALL and MULTILINE, so ``.`` + spans newlines and ``^``/``$`` match individual lines) in the + respective stream. Returns the completed process so callers can make + further assertions. + """ + if isinstance(stdout, str): + stdout = (stdout,) + if isinstance(stderr, str): + stderr = (stderr,) + r = self._capture_both(*args, server=server, **kwargs) + assert r.returncode == exit_code, ( + f"expected exit {exit_code}, got {r.returncode}\n" + f"stdout: {r.stdout}\nstderr: {r.stderr}" + ) + for pattern in stdout: + assert re.search(pattern, r.stdout, re.DOTALL | re.MULTILINE), ( + f"stdout did not match {pattern!r}\nstdout: {r.stdout}" + ) + for pattern in stderr: + assert re.search(pattern, r.stderr, re.DOTALL | re.MULTILINE), ( + f"stderr did not match {pattern!r}\nstderr: {r.stderr}" + ) + return r diff --git a/src/test/pytest/pypg/server.py b/src/test/pytest/pypg/server.py new file mode 100644 index 00000000000..4d2f2018131 --- /dev/null +++ b/src/test/pytest/pypg/server.py @@ -0,0 +1,1287 @@ +# Copyright (c) 2025, PostgreSQL Global Development Group + +from __future__ import annotations + +import contextlib +import ctypes +import os +import pathlib +import platform +import re +import shutil +import socket +import subprocess +import threading +from collections.abc import Generator, Mapping +from concurrent.futures import Future +from typing import Any + +from libpq import ( + PGconn, + PreparedStatement, +) +from libpq import ( + connect as libpq_connect, +) +from libpq import ( + connstr as libpq_connstr, +) + +from . import bins, portlock +from ._env import pg_test_timeout_default +from .util import shell_path +from .wait import wait_until + +# Database connections land in when neither the caller nor the server's +# default_connection_options names one. +_FALLBACK_DB = "postgres" + + +def _escape_conf_value(value: object) -> str: + """Quote and escape ``value`` as a postgresql.conf single-quoted string — + the inverse of the server's DeescapeQuotedString(), so arbitrary GUC values + round-trip without hand-escaping. A ``bool`` becomes ``on``/``off`` so + boolean GUCs can be set with ``True``/``False``. + """ + if isinstance(value, bool): + value = "on" if value else "off" + value = str(value) + value = value.replace("\\", "\\\\") + value = value.replace("'", "''") + value = value.replace("\n", "\\n") + value = value.replace("\r", "\\r") + value = value.replace("\t", "\\t") + value = value.replace("\b", "\\b") + value = value.replace("\f", "\\f") + return f"'{value}'" + + +class FileBackup(contextlib.AbstractContextManager): + """Context manager that snapshots a file's contents on entry and restores + them on exit. Used to roll back per-test edits to the server's config + files; the snapshot is held in memory rather than a sidecar file, so + nothing is left behind in the data directory. + """ + + def __init__(self, path: pathlib.Path): + self._path = path + self._contents: str | None = None + + def __enter__(self) -> FileBackup: + self._contents = self._path.read_text() + return self + + def __exit__(self, *exc: object) -> None: + assert self._contents is not None # set by __enter__ + self._path.write_text(self._contents) + + +class PostgresServer: + """ + Represents a running PostgreSQL server instance with management utilities. + Provides methods for configuration, user/database creation, and server control. + """ + + def __init__( + self, + name: str, + basedir: pathlib.Path, + sockdir: pathlib.Path, + libpq_handle: ctypes.CDLL, + *, + hostaddr: str | None = None, + port: int | None = None, + initdb_opts: list[str] | None = None, + from_backup: pathlib.Path | None = None, + streaming_primary: PostgresServer | None = None, + allows_streaming: bool | str = False, + archiving: bool = False, + restoring: PostgresServer | None = None, + restoring_standby: bool = True, + conf: dict[str, Any] | None = None, + ): + """ + Initialize a PostgreSQL server instance. Call start() to actually + start the server. + + Args: + name: The name of this server instance (for logging purposes) + basedir: Directory holding everything belonging to this server: + the data directory (``pgdata``), base backups taken from it + (``backup``), and archived WAL (``archives``). Must be unique + per server *and* per test file, since under autoconf's make + check the whole suite shares one tmp_check directory. + sockdir: Path to directory for Unix sockets + libpq_handle: ctypes handle to libpq + hostaddr: If provided, use this specific address (e.g., "127.0.0.2") + port: If provided, use this port instead of finding a free one, + is currently only allowed if hostaddr is also provided + initdb_opts: Extra arguments to pass to initdb (e.g. + ["--locale=C", "--encoding=LATIN1"]). When provided the fast + INITDB_TEMPLATE copy is bypassed and a real initdb is run, since + the template was created with the default locale/encoding. + from_backup: Path to a base backup (as produced by + ``PostgresServer.backup()``) to copy into the data directory + instead of running initdb. Use this to build a standby or a + point-in-time-recovery node. + streaming_primary: When building from a backup, the upstream server + to stream WAL from. Sets ``primary_conninfo`` and creates a + ``standby.signal`` file so the node starts as a streaming + standby. Its ``application_name`` is this node's name, so the + primary can ``wait_for_catchup()`` on it by name. + allows_streaming: Configure this server as a replication primary + (``wal_log_hints`` plus generous ``max_wal_senders`` / + ``max_replication_slots``). Pass ``"logical"`` to set + ``wal_level = logical`` instead of ``replica``; any other + truthy value configures plain physical streaming. + archiving: Enable WAL archiving (``archive_mode = on`` plus an + ``archive_command`` that copies segments into this server's + ``archive_dir``). + restoring: When building from a backup, the upstream server whose + ``archive_dir`` to restore WAL from (sets ``restore_command``). + By default a ``standby.signal`` is dropped (the node keeps + replaying archived WAL as a standby); pass + ``restoring_standby=False`` for a ``recovery.signal`` instead + (archive recovery that promotes when WAL runs out). + restoring_standby: See ``restoring``. + conf: Extra GUC settings (a ``{name: value}`` dict) to append to + postgresql.conf before the first start. Use for settings that + must be present at startup (e.g. a recovery target), since + ``create_pg`` starts the server immediately. + """ + + if hostaddr is None and port is not None: + raise NotImplementedError("port was provided without hostaddr") + + self.name = name + # Everything belonging to this server lives under its basedir, so + # cleanup is a single rmtree and anything a test places next to the + # datadir (tablespaces, COPY files, ...) is scoped to this server + # automatically. + self.basedir = basedir + self.datadir = basedir / "pgdata" + self._backup_root = basedir / "backup" + self.archive_dir = basedir / "archives" + self.sockdir = sockdir + self.libpq_handle = libpq_handle + # The log deliberately lives outside the data directory: pg_basebackup + # copies unknown files in pgdata, so a log in there would leak the + # primary's log lines into backups and confuse log searches on nodes + # built from them. + self.log = basedir / "postgresql.log" + self._log_start_pos = 0 + basedir.mkdir(parents=True, exist_ok=True) + + # ExitStack for cleanup callbacks + self._cleanup_stack = contextlib.ExitStack() + + # Cached connection reused by sql() (see _get_default_conn). Closed + # (and reopened lazily on next use) via close_default_conn: on + # stop()/restart(), at per-test teardown, or explicitly by a test + # after it made the backend die. + self._default_conn: PGconn | None = None + + # Connection options connect() and everything built on it use when the + # caller doesn't pass them (see the default_connection_options + # property). An unset dbname falls back to _FALLBACK_DB. + self._default_conn_opts: dict[str, Any] = {} + + # Per-test config save/restore state (see _snapshot_conf_if_needed). + # Tracking is only armed inside start_new_test(); config written during + # __init__ and by module-scoped setup is deliberately left untracked so + # that it persists for the lifetime of the server. + self._track_conf_changes = False + self._conf_snapshotted = False + # None until the test applies an edit; then "reload" or "restart" (the + # strongest apply the test did), which is how the edit is reverted. + self._conf_restore_mode: str | None = None + + # Determine whether to use Unix sockets + use_unix_sockets = platform.system() != "Windows" and hostaddr is None + + # The backup carries the primary's config; the conf appended below + # (port, sockets, ...) overrides it since later entries win. + if from_backup is not None: + shutil.copytree(from_backup, self.datadir) + os.chmod(self.datadir, 0o700) + # Use INITDB_TEMPLATE if available (much faster than running initdb), + # unless caller-supplied initdb options require a real initdb. + elif (initdb_template := os.environ.get("INITDB_TEMPLATE")) and ( + not initdb_opts and os.path.isdir(initdb_template) + ): + shutil.copytree(initdb_template, self.datadir) + else: + if platform.system() == "Windows": + auth_method = "trust" + else: + auth_method = "peer" + bins.initdb( + "--no-sync", + "--auth", + auth_method, + "--pgdata", + self.datadir, + *(initdb_opts or []), + ) + + # Figure out which addresses to listen on, then take a port that is + # free on all of them. Ports come from pypg.portlock, which reserves + # them the same way (and in the same directory) as the Perl TAP tests, + # so the two harnesses never pick the same one. + if hostaddr is not None: + # Explicit address provided + addrs: list[str] = [hostaddr] + elif socket.has_dualstack_ipv6(): + hostaddr = "::1" + addrs = [hostaddr, "127.0.0.1"] + else: + hostaddr = "127.0.0.1" + addrs = [hostaddr] + + if port is None: + port = portlock.get_free_port(addrs) + else: + portlock.mark_assigned(port) + + # Store the computed values + self.hostaddr = hostaddr + self.port = port + # Including the host to use for connections - either the socket + # directory or TCP address + if use_unix_sockets: + self.host = str(sockdir) + else: + self.host = hostaddr + + self.append_conf( + # An empty value disables Unix sockets when using TCP, avoiding + # lock conflicts. + unix_socket_directories=sockdir.as_posix() if use_unix_sockets else "", + listen_addresses=",".join(addrs), + port=port, + log_connections="all", + fsync=False, + datestyle="ISO", + timezone="UTC", + # With the 5s default, tests that repeatedly restart a primary + # under synchronous replication stall ~5s per reconnect before the + # standby reattaches, dominating the run time. + wal_retrieve_retry_interval="500ms", + ) + + # Replication-primary settings. Most already default to streaming- + # capable values; wal_log_hints (off by default) is the one that + # matters for pg_rewind-style tests. + if allows_streaming: + self.append_conf( + wal_level="logical" if allows_streaming == "logical" else "replica", + max_wal_senders=10, + max_replication_slots=10, + wal_log_hints=True, + hot_standby=True, + max_wal_size="128MB", + ) + if platform.system() == "Windows": + copy_cmd = "copy" + else: + copy_cmd = "cp" + + # Configure streaming replication from the primary: set + # primary_conninfo and drop a standby.signal so the node comes up as a + # streaming standby. + if streaming_primary is not None: + conninfo = streaming_primary.connstr(application_name=self.name) + self.append_conf(primary_conninfo=conninfo) + (self.datadir / "standby.signal").touch() + + # Enable WAL archiving. archive_mode is a postmaster GUC, so this must + # be configured before the first start. + if archiving: + os.makedirs(self.archive_dir, exist_ok=True) + archive_target = shell_path(self.archive_dir / "%f") + self.append_conf( + archive_mode=True, + archive_command=f'{copy_cmd} "%p" "{archive_target}"', + wal_level="replica", + ) + + # Restore WAL from an upstream server's archive. + if restoring is not None: + archive_source = shell_path(restoring.archive_dir / "%f") + self.append_conf( + restore_command=f'{copy_cmd} "{archive_source}" "%p"', + ) + signal = "standby.signal" if restoring_standby else "recovery.signal" + (self.datadir / signal).touch() + + # Caller-supplied startup config (e.g. a recovery target). + if conf: + self.append_conf(**conf) + + def start(self) -> None: + """Start the server using pg_ctl.""" + self.pg_ctl("start") + self.pid = self._read_postmaster_pid() + + def _read_postmaster_pid(self) -> int: + """Read the postmaster PID from the server's postmaster.pid file.""" + with open(self.datadir / "postmaster.pid") as f: + return int(f.readline().strip()) + + def is_running(self) -> bool: + """Whether the postmaster looks to be running, based on the presence of + its postmaster.pid file (pg_ctl removes it on a clean stop).""" + return (self.datadir / "postmaster.pid").exists() + + def reload(self) -> None: + """Reload postgresql.conf and pg_hba.conf via ``pg_ctl reload`` (SIGHUP). + + Only settings that can change at SIGHUP take effect; postmaster-level + settings (shared_buffers, archive_mode, ...) need restart(). + + When this applies a config edit made during a test (see + start_new_test), it records that the edit must be reverted with a reload + at the end of the test. A restart, if the test also did one, takes + precedence. + """ + self.pg_ctl("reload") + if ( + self._track_conf_changes + and self._conf_snapshotted + and self._conf_restore_mode is None + ): + self._conf_restore_mode = "reload" + + def restart(self, mode: str = "fast") -> None: + """Restart the server via ``pg_ctl restart`` and refresh the postmaster + PID. + + When this applies a config edit made during a test (see + start_new_test), it records that reverting the edit will also need a + restart, since a reload cannot undo a postmaster-level GUC. (A restart + before any edit just re-reads the baseline, so it does not count.) + """ + self.close_default_conn() + self.pg_ctl("restart", "--mode", mode) + self.pid = self._read_postmaster_pid() + if self._track_conf_changes and self._conf_snapshotted: + self._conf_restore_mode = "restart" + + def promote(self) -> None: + """Promote a standby/recovery node to a primary, waiting for the + promotion to finish (pg_ctl promote -w).""" + self.pg_ctl("promote", "-w") + + def enable_streaming(self, primary: PostgresServer) -> None: + """Reconfigure this (stopped) node to stream from ``primary`` as a + standby: set ``primary_conninfo`` and drop a ``standby.signal``. Use it + to re-attach a former primary as a standby of a newly-promoted node (a + role swap); call ``start()`` afterwards. The standby's + ``application_name`` is this node's name so the new primary can + ``wait_for_catchup()`` on it by name. + """ + conninfo = primary.connstr(application_name=self.name) + self.append_conf(primary_conninfo=conninfo) + (self.datadir / "standby.signal").touch() + + def current_log_position(self) -> int: + """Get the current end position of the log file.""" + if self.log.exists(): + return self.log.stat().st_size + return 0 + + def reset_log_position(self) -> None: + """Mark current log position as start for log_content().""" + self._log_start_pos = self.current_log_position() + + @contextlib.contextmanager + def start_new_test(self) -> Generator[PostgresServer, None, None]: + """ + Prepare server for a new test. + + Resets log position and enters a cleanup subcontext. Within that + subcontext config edits are tracked: the first config edit snapshots + the server's config files, and they are restored when the test finishes. + If the test applied its edit (reload/restart), the restore re-applies + the same way; if it never did, the restore is just the file rollback, + since the running server never picked the edit up. See + _snapshot_conf_if_needed. + + The cached sql() connection is dropped first. sql() closes it when the + cleanup context it was opened in ends, so a connection opened outside a + test -- by a module-scoped fixture's setup, most likely -- belongs to + the module and would otherwise be shared by every test in it, carrying + session state (a GUC, an open transaction, a temporary table) from one + test into the next. Closing here means each test's first sql() opens a + connection registered against that test. + """ + self.reset_log_position() + self.close_default_conn() + with self.subcontext(): + self._conf_snapshotted = False + self._conf_restore_mode = None + self._track_conf_changes = True + try: + yield self + finally: + # Disarm before the subcontext unwinds, so the restore callback + # it runs (via restart()) does not re-escalate the mode. + self._track_conf_changes = False + + def psql(self, *args: object) -> None: + """Run psql with the given arguments.""" + bins.psql("-w", *args, server=self) + + def sql(self, query: str, *params: Any, simplify_result: bool = True) -> Any: + """Execute a SQL query via libpq. Returns simplified results, or a list + of row tuples with ``simplify_result=False`` (see ``PGconn.sql``). + + Runs on a cached "default" connection reused across calls (see + ``_get_default_conn``). Two consequences of the reuse: + + - Session state (SET, temp objects, prepared statements) persists + across sql() calls on the same node. Use ``sql_oneshot()`` or an + explicit ``connect()`` for a fresh session or non-default connection + options. + - If the backend dies underneath the cached connection (a crash, + ``pg_terminate_backend``), there is deliberately no automatic + reconnect: the error keeps propagating until the connection is + explicitly invalidated by ``stop()``/``restart()``, per-test + teardown, or ``close_default_conn()``. + """ + return self._get_default_conn().sql( + query, *params, simplify_result=simplify_result + ) + + def sql_oneshot(self, query: str, *params: Any, **connection_opts: Any) -> Any: + """Execute a SQL query on a fresh, single-use connection. + + Any keyword arguments are passed through to ``connect()`` as connection + options, e.g. ``sql_oneshot(q, dbname="mydb")``. Use this over ``sql()`` + when the query must run in its own session: connecting as another + user/database, or when the disconnect itself matters (temporary slots + or objects dropped at session end, stats flushed on exit, picking up a + reloaded setting immediately, ...). + """ + with self.connect(**connection_opts) as conn: + return conn.sql(query, *params) + + def sql_batch(self, *queries: str) -> list[Any]: + """Run several statements like consecutive ``sql()`` calls and return a + list with every statement's simplified result (see ``PGconn.sql_batch``). + + Like ``sql()`` this runs on the cached default connection, so session + state the batch establishes (SET, temp objects, an unfinished BEGIN) + persists into later ``sql()``/``sql_batch()`` calls on the node. A + batch that scopes session state to its final statement (``SET ROLE; + UPDATE ...``) should use ``sql_batch_oneshot()`` instead, so the state + dies with the connection. + """ + return self._get_default_conn().sql_batch(*queries) + + def sql_batch_oneshot(self, *queries: str, **connection_opts: Any) -> list[Any]: + """Run several statements on a fresh, single-use connection and return + a list with every statement's simplified result. + + Any keyword arguments are passed through to ``connect()`` as connection + options, e.g. ``sql_batch_oneshot(q1, q2, dbname="mydb")``. Use this + over ``sql_batch()`` when the batch builds up session state that must + not leak into later calls on the node, or when the disconnect itself + matters (see ``sql_oneshot``). + """ + with self.connect(**connection_opts) as conn: + return conn.sql_batch(*queries) + + def prepare(self, query: str, *, name: str | None = None) -> PreparedStatement: + """Parse ``query`` into a named prepared statement on the cached + default connection and return the ``PreparedStatement`` (see + ``PGconn.prepare``). + + Like ``sql()`` this uses the connection shared by all default-conn + methods, so ``stmt.exec()`` sees session state from earlier ``sql()`` + calls, an unconsumed ``stmt.background_exec()`` future blocks + ``node.sql()``, and the statement survives until the connection is + invalidated (stop/restart, per-test teardown, ``close_default_conn``). + """ + return self._get_default_conn().prepare(query, name=name) + + def background_sql( + self, query: str, *params: Any, simplify_result: bool = True + ) -> Future[Any]: + """Dispatch a query that is expected to *block* and return its Future + (see ``PGconn.background_sql``). + + Runs on the cached default connection, which refuses further queries + while the future is unconsumed — any ``node.sql()`` before collecting + ``.result()`` will raise. Use ``background_sql_oneshot()`` or an + explicit ``connect()`` when the default connection must stay usable + while the query is blocked. + """ + return self._get_default_conn().background_sql( + query, *params, simplify_result=simplify_result + ) + + def background_sql_oneshot( + self, + query: str, + *params: Any, + simplify_result: bool = True, + **connection_opts: Any, + ) -> Future[Any]: + """Dispatch a blocking query on its own fresh, single-use connection + and return its Future (see ``PGconn.background_sql``). + + Keyword arguments are passed through to ``connect()`` as connection + options. Use this over ``background_sql()`` when the test needs + ``node.sql()`` to keep working while the query is blocked. The + connection cannot be closed before this returns — the query is still + running on it — so the worker thread finishes it as soon as the query + completes. The future must still be consumed (call ``.result()``). + """ + conn = self.connect(**connection_opts) + return conn.background_sql( + query, *params, close_when_done=True, simplify_result=simplify_result + ) + + def _get_default_conn(self) -> PGconn: + """Return the cached connection used by sql(), opening it lazily on + first use. + + connect() registers the connection in whatever cleanup stack is + current (a per-test subcontext when called inside a test), so + close_default_conn is registered there too: when that context tears + down, the cache is forgotten along with it and the next sql() + reconnects.""" + if self._default_conn is None: + self._default_conn = self.connect() + self._cleanup_stack.callback(self.close_default_conn) + return self._default_conn + + @property + def default_connection_options(self) -> dict[str, Any]: + """The connection options that ``connect()`` — and so ``sql()``, + ``poll_query_until()`` and everything else built on it — applies when + the caller passes none. + + Assigning to it closes the cached ``sql()`` connection, so the next + query reconnects with the new options:: + + pg.sql("CREATE DATABASE mydb") + pg.default_connection_options = {"dbname": "mydb"} + pg.sql("CREATE EXTENSION thing") # runs in mydb + + pg.default_connection_options = {"user": "regress_user"} + pg.sql("SELECT current_user") # runs as regress_user + + Use it when a test does most of its work under options other than the + defaults — in another database, as another role, or with ``options`` + (libpq's PGOPTIONS) — instead of repeating them at every call, or + opening a connection by hand just to carry them. An explicit keyword to + ``connect()``/``sql_oneshot()`` still wins, so a test can reach back to + the default at any point. + + Assignment replaces the whole mapping rather than merging into it; the + getter returns a copy, so a test that wants to add one option can read, + update and assign back. ``dbname`` falls back to ``postgres`` whenever + the mapping does not set it. + + Deliberately *not* applied to connection_env(): that describes the + server to a client program run with those PG* variables, where silently + substituting a different database or role would be a surprise. + """ + return dict(self._default_conn_opts) + + @default_connection_options.setter + def default_connection_options(self, opts: Mapping[str, Any]) -> None: + # The cached connection carries the old options, so it has to go for the + # change to mean anything for sql(). + self.close_default_conn() + self._default_conn_opts = dict(opts) + + def close_default_conn(self) -> None: + """Close and forget the cached connection used by sql(), if any; the + next sql() call opens a fresh one. + + Called by stop()/restart() (the cached session would otherwise outlive + the backend it was talking to). sql() never reconnects on its own, so + call this from a test to restore sql() connectivity after the backend + died some other way — a crash, pg_terminate_backend, ... — once that + error has been asserted. + """ + if self._default_conn is not None: + self._default_conn.close() + self._default_conn = None + + def append_conf(self, **gucs: object) -> None: + """Append GUC settings to postgresql.conf. + + Each keyword is written as ``name = 'value'`` with the value escaped + for a postgresql.conf single-quoted string, so callers never have to + quote or escape values themselves (a ``bool`` becomes ``on``/``off``). + For GUCs whose names are not valid Python identifiers — the dotted + names of extension GUCs — unpack a dict:: + + node.append_conf(primary_conninfo=conninfo, work_mem="4MB") + node.append_conf(**{"basebackup_to_shell.command": cmd}) + """ + self._snapshot_conf_if_needed() + with open(self.datadir / "postgresql.conf", "a") as f: + f.writelines( + f"{name} = {_escape_conf_value(value)}\n" + for name, value in gucs.items() + ) + + def adjust_conf(self, **gucs: object) -> None: + """Set each given GUC in postgresql.conf, replacing any existing (or + commented-out) line for it; a value of ``None`` removes the setting. + + Values are escaped exactly like append_conf, but unlike append_conf + this leaves a single clean line per setting rather than relying on + later-line-wins, and does not reload the server:: + + node.adjust_conf(work_mem="8MB", fsync=False) + node.adjust_conf(autovacuum=None) # remove the setting + node.adjust_conf(**{"auto_explain.log_min_duration": 0}) + """ + self._snapshot_conf_if_needed() + path = self.datadir / "postgresql.conf" + lines = path.read_text().splitlines() + for setting, value in gucs.items(): + pat = re.compile(rf"^\s*#?\s*{re.escape(setting)}\s*=") + lines = [ln for ln in lines if not pat.match(ln)] + if value is not None: + lines.append(f"{setting} = {_escape_conf_value(value)}") + path.write_text("\n".join(lines) + "\n") + + def _config_files(self) -> list[pathlib.Path]: + """The config files that are rolled back between tests.""" + return [ + self.datadir / name + for name in ( + "postgresql.conf", + "postgresql.auto.conf", + "pg_hba.conf", + "pg_ident.conf", + ) + ] + + def _snapshot_conf_if_needed(self) -> None: + """Back up every config file before the first config edit of a test, so + the edits can be rolled back when the test finishes. + + A no-op unless tracking is armed (start_new_test); config written + during __init__ or module-scoped setup stays untracked on purpose, so + it persists for the lifetime of the server. All config files are backed + up together on the first edit — they are tiny and it keeps a single + restore point. Only the config-editing helpers arm this, so edits made + purely through SQL (e.g. ``ALTER SYSTEM``) are not covered. + """ + if not self._track_conf_changes or self._conf_snapshotted: + return + self._conf_snapshotted = True + # Pushed before the FileBackups so it unwinds last: the files are + # restored first, and only then does the server reload/restart. + self._cleanup_stack.callback(self._reapply_conf) + for path in self._config_files(): + self._cleanup_stack.enter_context(FileBackup(path)) + + def _reapply_conf(self) -> None: + """Bring the server to a running state with the config files just + restored by the FileBackups in effect. Runs as the last step of a test + that edited config; see _snapshot_conf_if_needed. + """ + if not self.is_running(): + # The test left the server stopped; start it so the next test + # finds a running server. A fresh start reads the restored config + # in full, so the reload-vs-restart distinction does not apply. + self.start() + elif self._conf_restore_mode == "restart": + self.restart() + elif self._conf_restore_mode == "reload": + self.reload() + # else: the test never applied its edit, so the running server never + # picked it up; restoring the files is enough. + + def reset_hba(self, database: str, role: str, method: str) -> None: + """Replace pg_hba.conf with a single local rule and reload the server. + + The rule is written across a continuation line, so every use also + exercises pg_hba.conf continuation-line parsing. + """ + self._snapshot_conf_if_needed() + hba = self.datadir / "pg_hba.conf" + hba.write_text(f"local {database} {role}\\\n {method}\n") + self.reload() + + def reset_ident(self, map_name: str, system_user: str, pg_user: str) -> None: + """Replace pg_ident.conf with a single user-name-map entry and + reload.""" + self._snapshot_conf_if_needed() + ident = self.datadir / "pg_ident.conf" + ident.write_text(f"{map_name} {system_user} {pg_user}\n") + self.reload() + + def poll_query_until( + self, + query: str, + *params: Any, + expected: Any = True, + dbname: str | None = None, + timeout: float | None = None, + ) -> Any: + """Run ``query`` repeatedly until it returns ``expected``. + + Any positional ``params`` after ``query`` are bound to its + ``$1, $2, ...`` placeholders, like ``sql()``. The comparison is against + the simplified Python result of ``sql()`` (so ``expected`` is ``True`` + for a boolean ``t`` probe, an ``int`` for a count, a tuple for a + multi-column row, and so on) rather than psql text. Returns the matching + result, or raises ``TimeoutError`` once the timeout (defaulting to + PG_TEST_TIMEOUT_DEFAULT) is exhausted. + """ + # Close the polling connection on return rather than leaking it until + # teardown; a lingering connection to ``dbname`` would otherwise block + # e.g. CREATE DATABASE WITH TEMPLATE on that database. + with self.connect(dbname=dbname) as conn: + for _ in wait_until( + f"query never returned {expected!r}: {query}", timeout=timeout + ): + result = conn.sql(query, *params) + if result == expected: + return result + + def pg_ctl(self, *args: object) -> None: + """Run pg_ctl with the given arguments.""" + # Many tests bounce the server through a raw pg_ctl call instead of + # stop()/restart(); the cached sql() connection would be left pointing + # at a dead backend, so invalidate it here too. + if "stop" in args or "restart" in args: + self.close_default_conn() + bins.pg_ctl("--pgdata", self.datadir, "--log", self.log, *args, server=self) + + def connection_env(self) -> dict[str, str]: + """Return the PG* environment variables that point a client program at + this server. + + Use this to run an installed client program (createdb, vacuumdb, ...) + against this server while capturing its output, e.g. via + ``vacuumdb(..., server=pg)`` from :mod:`pypg.bins`. + """ + return { + "PGHOST": str(self.host), + "PGPORT": str(self.port), + "PGDATABASE": "postgres", + "PGDATA": str(self.datadir), + } + + def connstr(self, *, dbname: str | None = None, **opts: object) -> str: + """Return a libpq connection string pointing at this server. + + As in ``connect()``, a ``dbname`` of None means this server's + ``default_connection_options``. Extra keyword options (e.g. + ``application_name``) are appended. Used for ``primary_conninfo`` on + standbys and by replication clients. + """ + return libpq_connstr(self._conn_opts(dbname=dbname, **opts)) + + def backup( + self, + backup_name: str = "my_backup", + backup_options: list[str] | None = None, + ) -> pathlib.Path: + """Take a base backup of this (running) server with pg_basebackup. + + The backup is written under a per-server backups directory and the path + is returned, suitable for passing as ``from_backup`` when creating a + standby. + """ + backup_path = self._backup_root / backup_name + backup_path.parent.mkdir(parents=True, exist_ok=True) + bins.pg_basebackup( + "--no-sync", + "--pgdata", + backup_path, + "--host", + self.host, + "--port", + str(self.port), + "--checkpoint", + "fast", + *(backup_options or []), + ) + return backup_path + + def pg_recvlogical_upto( + self, + slot_name: str, + endpos: str, + *, + dbname: str | None = None, + timeout: float | None = None, + options: dict[str, str] | None = None, + ) -> str: + """Stream a logical slot's changes up to ``endpos`` with pg_recvlogical. + + Runs ``pg_recvlogical --start`` (which confirms the changes it reads, + advancing the slot) and returns its stdout as text with the trailing + newline stripped. ``options`` is a dict of plugin output options, each + passed as ``--option name=value``. + """ + args = [ + "--slot", + slot_name, + "--dbname", + self.connstr(dbname=dbname), + "--endpos", + endpos, + "--file", + "-", + "--no-loop", + "--start", + ] + for k, v in (options or {}).items(): + args.append("--option") + args.append(f"{k}={v}") + return bins.pg_recvlogical.capture(*args, timeout=timeout) + + def advance_wal(self, num: int) -> None: + """Advance WAL by ``num`` segments. + + Emits an empty logical message and forces a segment switch ``num`` + times. ``pg_switch_wal()`` flushes WAL, so ``pg_logical_emit_message()`` + is safe in non-transactional mode. + """ + with self.connect() as conn: + for _ in range(num): + conn.sql("SELECT pg_logical_emit_message(false, '', 'foo')") + conn.sql("SELECT pg_switch_wal()") + + def _get_insert_lsn(self) -> int: + """Return the current insert LSN of this server, in bytes.""" + return int(self.sql("SELECT pg_current_wal_insert_lsn() - '0/0'")) + + def emit_wal(self, size: int) -> int: + """Emit a transactional logical message of ``size`` bytes and return the + resulting end LSN, in bytes.""" + return int( + self.sql( + "SELECT pg_logical_emit_message(true, '', repeat('a', $1)) - '0/0'", + size, + ) + ) + + def write_wal( + self, tli: int, lsn: int, segment_size: int, data: bytes + ) -> pathlib.Path: + """Write ``data`` (bytes) into the WAL segment file at byte ``lsn`` on + timeline ``tli``, returning the segment path. Used to corrupt WAL on a + stopped server.""" + segment = lsn // segment_size + offset = lsn % segment_size + path = pathlib.Path(self.datadir) / "pg_wal" / f"{tli:08X}{0:08X}{segment:08X}" + with open(path, "r+b") as f: + f.seek(offset) + f.write(data) + return path + + def advance_wal_out_of_record_splitting_zone(self, wal_block_size: int) -> int: + """Advance WAL to a safe distance from the end of a page (enough to fit + a couple of small records), returning the end LSN in bytes.""" + page_threshold = wal_block_size // 4 + end_lsn = self._get_insert_lsn() + page_offset = end_lsn % wal_block_size + while page_offset >= wal_block_size - page_threshold: + self.emit_wal(page_threshold) + end_lsn = self._get_insert_lsn() + page_offset = end_lsn % wal_block_size + return end_lsn + + def advance_wal_to_record_splitting_zone(self, wal_block_size: int) -> int: + """Advance WAL so close to the end of a page that an XLogRecordHeader + would not fit on it, returning the end LSN in bytes.""" + record_header_size = 24 + end_lsn = self._get_insert_lsn() + page_offset = end_lsn % wal_block_size + + # Get fairly close to the end of a page in big steps. + while page_offset <= wal_block_size - 512: + self.emit_wal(wal_block_size - page_offset - 256) + end_lsn = self._get_insert_lsn() + page_offset = end_lsn % wal_block_size + + # Calibrate the message size so we can get closer 8 bytes at a time. + message_size = wal_block_size - 80 + while page_offset <= wal_block_size - record_header_size: + self.emit_wal(message_size) + end_lsn = self._get_insert_lsn() + old_offset = page_offset + page_offset = end_lsn % wal_block_size + # Adjust the message size until it causes 8-byte changes in offset, + # enough to be able to split a record header. + delta = page_offset - old_offset + if delta > 8: + message_size -= 8 + elif delta <= 0: + message_size += 8 + return end_lsn + + def backup_fs_cold(self, backup_name: str = "cold_backup") -> pathlib.Path: + """Take a filesystem-level cold backup of this (stopped) server. + + Copies the whole data directory, including WAL, into a per-server + backups directory and returns the path, suitable for ``from_backup``. + The server must be stopped, as no attempt is made to handle concurrent + writes; a node restored from such a backup enters crash recovery before + switching to archive recovery. + """ + backup_path = self._backup_root / backup_name + backup_path.parent.mkdir(parents=True, exist_ok=True) + shutil.copytree( + self.datadir, + backup_path, + ignore=shutil.ignore_patterns("postmaster.pid", "postmaster.opts"), + ) + return backup_path + + def lsn(self, mode: str = "write") -> str: + """Return a current WAL LSN of this server as a string. + + ``mode`` selects the function: ``insert``/``flush``/``write`` on a + primary, ``receive``/``replay`` on a standby. + """ + funcs = { + "insert": "pg_current_wal_insert_lsn()", + "flush": "pg_current_wal_flush_lsn()", + "write": "pg_current_wal_lsn()", + "receive": "pg_last_wal_receive_lsn()", + "replay": "pg_last_wal_replay_lsn()", + } + return self.sql(f"SELECT {funcs[mode]}") + + def wait_for_catchup( + self, + standby_name: str | PostgresServer, + mode: str = "replay", + target_lsn: str | None = None, + ) -> None: + """Wait until a streaming standby has caught up to ``target_lsn``. + + Polls pg_stat_replication on this (upstream) server until the standby's + ``_lsn`` has reached ``target_lsn`` (the upstream's current write + LSN by default) while in the ``streaming`` state. ``standby_name`` is + matched against the standby's ``application_name``, which the streaming + helpers set to the node name. + """ + if isinstance(standby_name, PostgresServer): + standby_name = standby_name.name + if target_lsn is None: + # On a standby (e.g. a standby acting as a publisher) the write LSN + # isn't available; use the replay LSN. + if self.sql("SELECT pg_is_in_recovery()"): + target_lsn = self.lsn("replay") + else: + target_lsn = self.lsn("write") + # mode names a pg_stat_replication LSN column (sent/write/flush/replay), + # so it is interpolated as an identifier; the values are bound as params. + query = ( + f"SELECT $1 <= {mode}_lsn AND state = 'streaming' " + "FROM pg_catalog.pg_stat_replication " + "WHERE application_name = $2" + ) + self.poll_query_until(query, target_lsn, standby_name) + + def wait_for_slot_catchup( + self, + slot_name: str, + mode: str = "restart", + target_lsn: str | None = None, + ) -> None: + """Wait until a replication slot's ``_lsn`` reaches ``target_lsn``. + + Polls pg_replication_slots on this server. ``mode`` is ``restart`` or + ``confirmed_flush``. + """ + assert target_lsn is not None, "target lsn must be specified" + assert mode in ("restart", "confirmed_flush") + # mode names a pg_replication_slots LSN column, so it is interpolated as + # an identifier; the values are bound as params. + self.poll_query_until( + f"SELECT $1 <= {mode}_lsn " + "FROM pg_catalog.pg_replication_slots WHERE slot_name = $2", + target_lsn, + slot_name, + ) + + def wait_for_subscription_sync( + self, + publisher: PostgresServer | None = None, + subname: str | None = None, + dbname: str | None = None, + ) -> None: + """Wait for a subscription's initial table sync to finish, then for the + subscriber to catch up to the publisher. + + Called on the subscriber: polls pg_subscription_rel until every table is + synced (``r``/``s``). If ``publisher`` and ``subname`` are given, also + waits for the publisher's walsender (named after the subscription) to + catch up. + """ + self.poll_query_until( + "SELECT count(1) = 0 FROM pg_subscription_rel " + "WHERE srsubstate NOT IN ('r', 's')", + dbname=dbname, + ) + if publisher is not None and subname is not None: + publisher.wait_for_catchup(subname) + + @contextlib.contextmanager + def repeat_query( + self, query: str, interval: float = 0.1, dbname: str | None = None + ) -> Generator[None, None, None]: + """Context manager that runs ``query`` repeatedly in the background on + its own connection until the block exits, like psql's ``\\watch``. + + Used to keep generating activity (e.g. transactions producing running + xact records) while the test does other work. Errors from a connection + torn down by a deliberate stop/restart are swallowed. + """ + stop = threading.Event() + conn = self.connect(dbname=dbname) + + def loop(): + while not stop.is_set(): + try: + conn.sql(query) + except Exception: + return + stop.wait(interval) + + worker = threading.Thread(target=loop, daemon=True) + worker.start() + try: + yield + finally: + stop.set() + worker.join(timeout=10) + try: + conn.close() + except Exception: + pass + + def log_standby_snapshot(self, standby: PostgresServer, slot_name: str) -> None: + """Emit the ``xl_running_xacts`` record a standby's logical slot + creation is waiting for. + + Called on the primary: waits until the standby slot's ``restart_lsn`` is + determined, then runs ``pg_log_standby_snapshot()``. + """ + standby.poll_query_until( + "SELECT restart_lsn IS NOT NULL FROM pg_catalog.pg_replication_slots " + "WHERE slot_name = $1", + slot_name, + ) + self.sql("SELECT pg_log_standby_snapshot()") + + def create_logical_slot_on_standby( + self, primary: PostgresServer, slot_name: str, dbname: str | None = None + ) -> None: + """Create a logical replication slot on this standby. + + Logical slot creation on a standby blocks until an ``xl_running_xacts`` + record arrives, so it is driven from a background ``pg_recvlogical + --create-slot`` while the primary is asked to log a standby snapshot. + """ + recv = subprocess.Popen( + [ + str(bins.pg_recvlogical.path), + "--dbname", + self.connstr(dbname=dbname), + "--plugin", + "test_decoding", + "--slot", + slot_name, + "--create-slot", + ], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + # Arrange for the xl_running_xacts record pg_recvlogical waits for. + primary.log_standby_snapshot(self, slot_name) + recv.wait() + assert ( + self.sql( + "SELECT slot_type FROM pg_catalog.pg_replication_slots " + "WHERE slot_name = $1", + slot_name, + ) + == "logical" + ), f"{slot_name} on standby created" + + def wait_for_event(self, backend_type: str, wait_event: str) -> None: + """Wait until some backend is parked on a given wait event. + + Polls pg_stat_activity until a backend of ``backend_type`` reports + ``wait_event``. Use it after dispatching a blocking query with + ``PGconn.background_sql()`` to confirm it has reached the expected + wait point. + """ + self.poll_query_until( + "SELECT count(*) > 0 FROM pg_stat_activity " + "WHERE backend_type = $1 AND wait_event = $2", + backend_type, + wait_event, + ) + + def wait_for_injection_point(self, name: str) -> None: + """Wait until some backend is parked at the named injection point. + + Polls pg_stat_activity for a backend whose wait event is the injection + point (``wait_event_type = 'InjectionPoint'``). Use after dispatching a + query with ``PGconn.background_sql()`` that is expected to block on + a point attached in ``'wait'`` mode. Unlike ``wait_for_event()`` it + does not constrain the backend type, so it also catches background + workers (autovacuum, checkpointer, ...) parked at the point. + """ + self.poll_query_until( + "SELECT count(*) > 0 FROM pg_stat_activity " + "WHERE wait_event_type = 'InjectionPoint' AND wait_event = $1", + name, + ) + + @contextlib.contextmanager + def subcontext(self) -> Generator[PostgresServer, None, None]: + """ + Create a new cleanup context for per-test isolation. + + Temporarily replaces the cleanup stack so that any cleanup callbacks + registered within this context will be cleaned up when the context exits. + """ + old_stack = self._cleanup_stack + self._cleanup_stack = contextlib.ExitStack() + try: + self._cleanup_stack.__enter__() + yield self + finally: + self._cleanup_stack.__exit__(None, None, None) + self._cleanup_stack = old_stack + + def stop(self, mode: str = "fast") -> None: + """ + Stop the PostgreSQL server instance. + + Ignores failures if the server is already stopped. + """ + self.close_default_conn() + try: + self.pg_ctl("stop", "--mode", mode) + except subprocess.CalledProcessError: + # Server may have already been stopped + pass + + def log_content(self) -> str: + """Return log content from the current context's start position.""" + return self.log_since(self._log_start_pos) + + def log_since(self, offset: int) -> str: + """Return log content written since the given byte offset. + + Pair with current_log_position() to capture exactly the log a single + operation produces:: + + offset = pg.current_log_position() + conn.sql("...") + assert "..." in pg.log_since(offset) + """ + if not self.log.exists(): + return "" + # Read as bytes and decode leniently: offsets are byte positions, and + # the server log is not guaranteed to be UTF-8 (e.g. log_connections + # records a LATIN1 database name verbatim). A strict decode would crash + # the failure-report hook on such logs. + with open(self.log, "rb") as f: + f.seek(offset) + return f.read().decode("utf-8", errors="replace") + + def wait_for_log( + self, pattern: str, offset: int = 0, timeout: float | None = None + ) -> int | None: + """Wait until the log written since ``offset`` matches ``pattern``. + + Returns the log's end offset once the regex matches, so chained waits + can continue from there. Raises ``TimeoutError`` otherwise. + """ + for _ in wait_until(f"log never matched {pattern!r}", timeout=timeout): + if re.search(pattern, self.log_since(offset)): + return self.current_log_position() + + @contextlib.contextmanager + def log_contains( + self, pattern: str, times: int | None = None + ) -> Generator[None, None, None]: + """ + Context manager that checks if the log matches pattern during the block. + + Args: + pattern: The regex pattern to search for. + times: If None, any number of matches is accepted. + If a number, exactly that many matches are required. + """ + start_pos = self.current_log_position() + yield + # See log_since(): decode leniently, the log may contain non-UTF-8 + # bytes. + with open(self.log, "rb") as f: + f.seek(start_pos) + content = f.read().decode("utf-8", errors="replace") + if times is None: + assert re.search(pattern, content), f"Pattern {pattern!r} not found in log" + else: + match_count = len(re.findall(pattern, content)) + assert match_count == times, ( + f"Expected {times} matches of {pattern!r}, found {match_count}" + ) + + def cleanup(self) -> None: + """Run all registered cleanup callbacks.""" + self.close_default_conn() + self._cleanup_stack.close() + + def _conn_opts(self, **opts: Any) -> dict[str, Any]: + """The connection options that point at this server, with anything the + caller passed taking precedence. Shared by connect() and connstr(), so + a connection and a connection string describing it cannot drift apart. + + An option passed as None means "whatever this server defaults to", so + the helpers taking an optional dbname (or user, ...) can pass theirs + straight through instead of resolving the fallback themselves. + """ + for name, value in self._default_conn_opts.items(): + if opts.get(name) is None: + opts[name] = value + if opts.get("dbname") is None: + opts["dbname"] = _FALLBACK_DB + opts.setdefault("host", self.host) + opts.setdefault("port", self.port) + return opts + + def connect(self, **opts: Any) -> PGconn: + """ + Creates a connection to this PostgreSQL server instance. + + Args: + **opts: Additional connection options (can override defaults) + + Returns: + PGconn: Connected database connection + + Example: + conn = pg.connect() + conn = pg.connect(dbname='mydb') + """ + opts = self._conn_opts(**opts) + # Only meaningful for a connection this process makes, so it is not + # part of _conn_opts(): connstr() hands its string to a standby or a + # client program, which have their own ideas about connect timeouts. + opts.setdefault("connect_timeout", pg_test_timeout_default()) + + return libpq_connect(self.libpq_handle, self._cleanup_stack, **opts) diff --git a/src/test/pytest/pypg/util.py b/src/test/pytest/pypg/util.py new file mode 100644 index 00000000000..ce4331a8647 --- /dev/null +++ b/src/test/pytest/pypg/util.py @@ -0,0 +1,126 @@ +# Copyright (c) 2025, PostgreSQL Global Development Group + +from __future__ import annotations + +import os +import platform +import shlex +import stat +import subprocess +import sys +from typing import Any + + +def shell_path(path: str | os.PathLike[str]) -> str: + """Render ``path`` for embedding in a shell command that the *server* runs + (archive_command, restore_command, ...). This needs backslashes on + Windows, even in MinGW environments. + + A plain ``str.replace`` rather than ``pathlib.PureWindowsPath`` because + the MinGW (MSYS2) Python swaps pathlib's separator (and ``os.sep``) to + "/", so pathlib/os.path still emit forward slashes there even though + ``platform.system() == "Windows"`` [1][2]. A literal replace is + separator-agnostic either way. + + [1] https://bugs.python.org/issue44778 + [2] https://sourceforge.net/p/mingw-w64/bugs/912/ + """ + if platform.system() == "Windows": + return str(path).replace("/", "\\") + return str(path) + + +def eprint(*args: object, **kwargs: Any) -> None: + """eprint prints to stderr""" + print(*args, file=sys.stderr, **kwargs) + + +def check_mode_recursive( + root: str | os.PathLike[str], dir_mode: int, file_mode: int +) -> list[str]: + """Check permissions of a directory tree (usually a data directory), + returning a list of paths whose mode differs from the expected one -- + empty if everything matches, so tests can assert on the result and get + the offending paths in the failure message. Files that vanish mid-walk + are ignored: a running server can remove files (e.g. in pg_stat) while we + are walking. + """ + violations = [] + + def check(path: str, expected: int) -> None: + try: + mode = stat.S_IMODE(os.stat(path).st_mode) + except FileNotFoundError: + return + if mode != expected: + violations.append(f"{path}: mode {oct(mode)} != {oct(expected)}") + + check(os.fspath(root), dir_mode) + for dirpath, dirnames, filenames in os.walk(root): + for d in dirnames: + check(os.path.join(dirpath, d), dir_mode) + for f in filenames: + check(os.path.join(dirpath, f), file_mode) + return violations + + +def run( + *command: object, + check: bool = True, + shell: bool | None = None, + silent: bool = False, + **kwargs: Any, +) -> subprocess.CompletedProcess[Any]: + """run runs the given command and prints it to stderr""" + + __tracebackhide__ = True # Don't show in pytest stack traces + + if shell is None: + shell = len(command) == 1 and isinstance(command[0], str) + + # A shell command is a single string; everything else is a list of + # stringified argv elements. Build it into a fresh local rather than + # rebinding the *command parameter (whose static type is a tuple). + cmd: str | list[str] + if shell: + # The shell auto-detection above only sets shell when the single + # argument is a str; an explicit shell=True is the caller's promise of + # the same, so command[0] is the shell command line. + assert isinstance(command[0], str) + cmd = command[0] + else: + cmd = [str(c) for c in command] + + if not silent: + if shell: + eprint(f"+ {cmd}") + else: + eprint(f"+ {shlex.join(cmd)}") + + if silent: + kwargs.setdefault("stdout", subprocess.DEVNULL) + + result = subprocess.run(cmd, check=False, shell=shell, **kwargs) + + # Manually throw CalledProcessError to avoid subprocess.run's huge body + # poluting stack traces. + if check and result.returncode: + raise subprocess.CalledProcessError( + result.returncode, cmd, result.stdout, result.stderr + ) + + return result + + +def capture( + command: object, + *args: object, + stdout: int = subprocess.PIPE, + encoding: str = "utf-8", + **kwargs: Any, +) -> str: + __tracebackhide__ = True # Don't pollute pytest stack traces + + return run( + command, *args, stdout=stdout, encoding=encoding, **kwargs + ).stdout.removesuffix("\n") diff --git a/src/test/pytest/pypg/wait.py b/src/test/pytest/pypg/wait.py new file mode 100644 index 00000000000..449a6f7c9bd --- /dev/null +++ b/src/test/pytest/pypg/wait.py @@ -0,0 +1,48 @@ +# Copyright (c) 2025, PostgreSQL Global Development Group + +from __future__ import annotations + +import time +from collections.abc import Iterator + +from ._env import pg_test_timeout_default + + +def wait_until( + error_message: str = "Did not complete", + timeout: float | None = None, + interval: float | None = None, +) -> Iterator[None]: + """ + Loop until the timeout is reached. If the timeout is reached, raise an + exception with the given error message. + + Use it to poll for a condition, breaking out once it holds:: + + for _ in wait_until("standby did not catch up"): + if standby.sql("SELECT ...") == expected: + break + + The timeout defaults to PG_TEST_TIMEOUT_DEFAULT. + + By default the sleep between attempts starts at 1ms and doubles up to + 100ms. Pass ``interval`` to poll at a fixed rate instead (e.g. when each + attempt is itself expensive). + """ + if timeout is None: + timeout = pg_test_timeout_default() + + start = time.time() + end = start + timeout + last_printed_progress = start + sleep_for = interval if interval is not None else 0.001 + while time.time() < end: + if timeout > 5 and time.time() - last_printed_progress > 5: + last_printed_progress = time.time() + print(f"{error_message} in {time.time() - start} seconds - will retry") + yield + time.sleep(sleep_for) + if interval is None: + sleep_for = min(sleep_for * 2, 0.1) + + raise TimeoutError(error_message + " in time") -- 2.54.0