"""
Base classes and shared logic for pycopg.
Contains abstract base classes and mixins used by Database and AsyncDatabase.
"""
from __future__ import annotations
from abc import ABC
from collections.abc import Mapping, Sequence
from pathlib import Path
from typing import Any
from pycopg.config import Config
from pycopg.utils import validate_identifiers, validate_timestamp
# Public type alias for query parameters — single source of truth exported via pycopg.__all__
Params = Sequence[Any] | Mapping[str, Any] | None
# Public surface consumed across module boundaries — DatabaseBase/QueryMixin by
# database.py+async_database.py, Params by __init__.py+maint.py+database.py+
# async_database.py, the build_* helpers by admin.py/backup.py,
# _decode_server_version by pool.py+database.py+async_database.py.
__all__ = [
"Params",
"DatabaseBase",
"QueryMixin",
"SessionMixin",
"build_pg_dump_cmd",
"build_pg_restore_cmd",
"build_role_options",
"_decode_server_version",
]
def _decode_server_version(v: int) -> str:
"""Convert a psycopg server_version integer to a dotted version string.
psycopg exposes the server version as a single integer using PostgreSQL's
encoding scheme: ``major * 10000 + minor * 100 + patch`` (pre-10) or
``major * 10000 + patch`` (PostgreSQL ≥ 10, where ``minor`` is always 0).
Parameters
----------
v : int
Integer server version from ``connection.info.server_version``
(e.g. ``160001`` for PostgreSQL 16.1, ``90605`` for 9.6.5).
Returns
-------
str
Dotted version string. PostgreSQL ≥ 10 uses two parts
(``"16.1"``); pre-10 uses three parts (``"9.6.5"``).
Examples
--------
>>> _decode_server_version(160001)
'16.1'
>>> _decode_server_version(90605)
'9.6.5'
"""
major = v // 10000
if major >= 10:
return f"{major}.{v % 10000}"
minor = (v % 10000) // 100
return f"{major}.{minor}.{v % 100}"
def _validate_cli_pattern(value: str, context: str) -> None:
"""Guard CLI pattern values passed to pg_dump/pg_restore -t/-T/-n flags.
pg_dump accepts table/schema PATTERNS (wildcards, 'schema.table'), so
strict SQL identifier validation is too restrictive. The minimum guard
is to reject values that start with '-' (flag injection) or contain
control characters and newlines (which could subvert shell logging).
Parameters
----------
value : str
Pattern value to validate.
context : str
Human-readable context for the error message (e.g. 'tables').
Raises
------
ValueError
If the value looks like a flag or contains unsafe characters.
"""
if not value:
raise ValueError(f"Empty value in {context} list")
if value.startswith("-"):
raise ValueError(
f"Invalid {context} value {value!r}: values starting with '-' are "
"not permitted (possible flag injection)"
)
# Reject control characters (NUL, newlines, tabs, etc.) as they are never
# valid in pg_dump pattern arguments.
for ch in value:
if ord(ch) < 0x20:
raise ValueError(
f"Invalid {context} value {value!r}: contains control character "
f"U+{ord(ch):04X} which is not permitted"
)
[docs]
class DatabaseBase(ABC):
"""Abstract base class for Database and AsyncDatabase.
Provides shared configuration, factory methods, and SQL query constants.
Subclasses must implement the actual execution methods.
"""
[docs]
def __init__(self, config: Config):
"""Initialize database with configuration.
Parameters
----------
config : Config
Database configuration.
"""
self.config = config
[docs]
@classmethod
def from_env(cls, dotenv_path: str | Path | None = None) -> DatabaseBase:
"""Create database from environment variables.
Parameters
----------
dotenv_path : str or Path, optional
Path to .env file.
Returns
-------
DatabaseBase
Database instance.
"""
return cls(Config.from_env(dotenv_path))
[docs]
@classmethod
def from_url(cls, url: str) -> DatabaseBase:
"""Create database from connection URL.
Parameters
----------
url : str
PostgreSQL connection URL.
Returns
-------
DatabaseBase
Database instance.
"""
return cls(Config.from_url(url))
def __repr__(self) -> str:
"""Return string representation of the database instance."""
class_name = self.__class__.__name__
return f"{class_name}({self.config.database!r} @ {self.config.host}:{self.config.port})"
[docs]
class QueryMixin:
"""Mixin providing common query building utilities.
Used by both sync and async database classes for consistent
SQL generation and validation.
"""
@staticmethod
def _build_insert_sql(
table: str,
columns: list[str],
schema: str = "public",
on_conflict: str | None = None,
) -> tuple[str, str]:
"""Build INSERT SQL template.
Parameters
----------
table : str
Table name.
columns : list of str
Column names.
schema : str, optional
Schema name, by default "public".
on_conflict : str, optional
ON CONFLICT clause.
Returns
-------
tuple of (str, str)
Tuple of (sql_template, columns_str).
"""
validate_identifiers(table, schema, *columns)
cols_str = ", ".join(columns)
placeholders = ", ".join(["%s"] * len(columns))
conflict_clause = f" ON CONFLICT {on_conflict}" if on_conflict else ""
sql = f"INSERT INTO {schema}.{table} ({cols_str}) VALUES ({placeholders}){conflict_clause}"
return sql, cols_str
@staticmethod
def _build_batch_insert_sql(
table: str,
columns: list[str],
rows: list[dict[str, Any]],
schema: str = "public",
on_conflict: str | None = None,
) -> tuple[str, list[Any]]:
"""Build batch INSERT SQL with multiple VALUES.
Parameters
----------
table : str
Table name.
columns : list of str
Column names.
rows : list of dict
List of row dicts.
schema : str, optional
Schema name, by default "public".
on_conflict : str, optional
ON CONFLICT clause.
Returns
-------
tuple of (str, list)
Tuple of (sql, params).
"""
validate_identifiers(table, schema, *columns)
cols_str = ", ".join(columns)
conflict_clause = f" ON CONFLICT {on_conflict}" if on_conflict else ""
placeholders = []
params: list[Any] = []
for row in rows:
row_placeholders = ", ".join(["%s"] * len(columns))
placeholders.append(f"({row_placeholders})")
params.extend(row.get(col) for col in columns)
values_str = ", ".join(placeholders)
sql = f"INSERT INTO {schema}.{table} ({cols_str}) VALUES {values_str}{conflict_clause}"
return sql, params
@staticmethod
def _build_select_sql(
table: str,
columns: list[str] | None = None,
schema: str = "public",
where: str | None = None,
order_by: str | None = None,
limit: int | None = None,
offset: int | None = None,
) -> str:
"""Build SELECT SQL.
Parameters
----------
table : str
Table name.
columns : list of str, optional
Column names, by default None (selects all columns).
schema : str, optional
Schema name, by default "public".
where : str, optional
WHERE clause (without WHERE keyword).
order_by : str, optional
ORDER BY clause.
limit : int, optional
LIMIT value.
offset : int, optional
OFFSET value.
Returns
-------
str
SQL string.
"""
validate_identifiers(table, schema)
if columns:
validate_identifiers(*columns)
cols_str = ", ".join(columns) if columns else "*"
sql = f"SELECT {cols_str} FROM {schema}.{table}"
if where:
sql += f" WHERE {where}"
if order_by:
sql += f" ORDER BY {order_by}"
if limit is not None:
sql += f" LIMIT {int(limit)}"
if offset is not None:
sql += f" OFFSET {int(offset)}"
return sql
@staticmethod
def _build_where_dict(where: dict[str, Any]) -> tuple[str, list[Any]]:
"""Build WHERE fragment from a dict of equality conditions.
Parameters
----------
where : dict
Mapping of column names to values. Keys become validated SQL
identifiers; values become positional ``%s`` parameters.
Must be non-empty — callers are responsible for rejecting an
empty dict before invoking this builder.
Returns
-------
tuple of (str, list)
Tuple of (fragment, params) where ``fragment`` is an AND-ed
``col = %s`` string suitable for ``WHERE {fragment}`` and
``params`` is the list of values in dict-insertion order.
Raises
------
InvalidIdentifierError
If any column key fails identifier validation.
"""
validate_identifiers(*where.keys())
fragment = " AND ".join(f"{col} = %s" for col in where)
params = list(where.values())
return fragment, params
@staticmethod
def _resolve_predicate(
where: dict[str, Any] | None,
where_sql: str | None,
where_params: Sequence[Any] | None,
) -> tuple[str, list[Any]]:
"""Resolve a predicate fragment from exactly one source (D-01..D-04).
Two predicate sources are supported: the equality ``where`` dict
(delegated to :meth:`_build_where_dict`) and a raw, caller-authored
``where_sql`` fragment. The raw fragment is a **documented, intentional
escape hatch** (C-02) for conditions the dict form cannot express
(e.g. ``"age > %s AND status = ANY(%s)"``) — it is passed through
**verbatim** and is **never** run through ``validate_identifiers``.
Only its bound ``where_params`` are treated as untrusted values
(positional ``%s`` parameters). Callers own the SQL text they pass as
``where_sql``; this method only enforces that the two sources are
never mixed or desynced.
Parameters
----------
where : dict, optional
Equality conditions as column-name to value mapping. Mutually
exclusive with ``where_sql``.
where_sql : str, optional
Raw SQL WHERE fragment (without the ``WHERE`` keyword), passed
through verbatim. Mutually exclusive with ``where``.
where_params : Sequence, optional
Positional ``%s`` parameters bound against ``where_sql``. Requires
``where_sql`` to be present; a placeholder-free ``where_sql`` with
``where_params=None`` is legal.
Returns
-------
tuple of (str, list)
Tuple of ``(fragment, params)``. An empty string fragment means no
predicate source was supplied (legal for read helpers with an
optional predicate; destructive helpers must reject an empty
fragment themselves).
Raises
------
ValueError
If ``where_params`` is supplied without ``where_sql`` (D-04); if
both ``where`` and ``where_sql`` are supplied (D-02); or if
``where_sql`` is present but empty/whitespace-only (D-03).
InvalidIdentifierError
If ``where`` is supplied and any of its keys fail identifier
validation (delegated to :meth:`_build_where_dict`).
"""
if where_params is not None and where_sql is None:
raise ValueError("where_params requires where_sql.")
if where and where_sql is not None:
raise ValueError("Pass either 'where' (dict) or 'where_sql', not both.")
if where_sql is not None:
if not where_sql.strip():
raise ValueError("where_sql must be a non-empty SQL fragment.")
return where_sql, list(where_params or [])
if where:
return QueryMixin._build_where_dict(where)
return "", []
@staticmethod
def _build_keyset_where(
order_cols: list[str], after: Sequence[Any], descending: bool
) -> tuple[str, list[Any]]:
"""Build a row-value keyset (seek) WHERE fragment (D-06/D-07).
Emits a PostgreSQL row-value comparison — ``(c1, c2) > (%s, %s)`` —
the same lexicographic tuple-compare shape shipped by the ETL
composite-watermark builder
(:func:`pycopg.etl._build_incremental_extract_sql`). Row-value
comparison is index-friendly: a btree index on ``(c1, c2)`` supports
it as a range scan. **Caveat (IN-01):** ``order_cols`` should be
``NOT NULL`` (or the caller must guarantee no ``NULL`` values) — the
comparison follows SQL ``NULL`` semantics, so a row whose ordered
column is ``NULL`` makes the comparison evaluate to ``NULL`` and the
row is silently excluded from the seek predicate.
Parameters
----------
order_cols : list of str
Ordered column names the page is sorted by. Validated via
``validate_identifiers`` (C-02) before interpolation — this is
the line-1 guard for this builder.
after : Sequence
Cursor values, positionally keyed to ``order_cols``. Bound as
``%s`` parameters, never interpolated (T-50-05).
descending : bool
When ``True``, the comparison operator flips from ``>`` to
``<`` (D-07) to match a ``DESC`` ordered page. The caller is
responsible for keeping the ``ORDER BY`` direction in sync with
this flip (Pitfall 3).
Returns
-------
tuple of (str, list)
Tuple of ``(fragment, params)`` where ``fragment`` is
``"(c1, c2) > (%s, %s)"`` (or ``<`` when descending) suitable for
``WHERE {fragment}``, and ``params`` is ``list(after)``.
Raises
------
InvalidIdentifierError
If any column in ``order_cols`` fails identifier validation.
"""
validate_identifiers(*order_cols)
op = "<" if descending else ">"
lhs = "(" + ", ".join(order_cols) + ")"
rhs = "(" + ", ".join(["%s"] * len(order_cols)) + ")"
return f"{lhs} {op} {rhs}", list(after)
[docs]
class SessionMixin:
"""Mixin providing session/connection reuse capabilities.
Allows keeping a connection open for multiple operations,
reducing connection overhead for batch operations.
"""
_session_connection: Any = None
_in_session: bool = False
def _get_session_connection(self) -> Any:
"""Get the current session connection if in session mode."""
if self._in_session and self._session_connection is not None:
return self._session_connection
return None
def _is_in_session(self) -> bool:
"""Check if currently in session mode."""
return self._in_session
[docs]
def build_pg_dump_cmd(
host: str,
port: int,
user: str,
database: str,
output_file: str | Path,
format: str = "custom",
schema_only: bool = False,
data_only: bool = False,
tables: list[str] | None = None,
exclude_tables: list[str] | None = None,
schemas: list[str] | None = None,
compress: int = 6,
jobs: int = 1,
) -> list[str]:
"""Build a pg_dump command argv list.
Constructs the argument list for invoking pg_dump. Pure function: no
I/O, no environment access, no secrets. The caller runs the command
and manages credentials via environment variables.
Parameters
----------
host : str
Database host.
port : int
Database port.
user : str
Database user.
database : str
Database name.
output_file : str or Path
Output file path.
format : str, optional
Dump format — 'plain', 'custom', 'directory', or 'tar', by default "custom".
schema_only : bool, optional
Dump only schema, no data, by default False.
data_only : bool, optional
Dump only data, no schema, by default False.
tables : list, optional
Only dump these tables.
exclude_tables : list, optional
Exclude these tables.
schemas : list, optional
Only dump these schemas.
compress : int, optional
Compression level (0-9, for custom format), by default 6.
jobs : int, optional
Parallel jobs (for directory format), by default 1.
Returns
-------
list
List of strings (argv) suitable for passing to a process runner.
"""
output_file = Path(output_file)
cmd = ["pg_dump"]
# Connection params
cmd.extend(["-h", host])
cmd.extend(["-p", str(port)])
cmd.extend(["-U", user])
cmd.extend(["-d", database])
# Format
format_map = {"plain": "p", "custom": "c", "directory": "d", "tar": "t"}
cmd.extend(["-F", format_map[format]])
# Options
if schema_only:
cmd.append("--schema-only")
if data_only:
cmd.append("--data-only")
if compress and format == "custom":
cmd.extend(["-Z", str(compress)])
if jobs > 1 and format == "directory":
cmd.extend(["-j", str(jobs)])
# Tables — guard against flag injection (CR-04)
if tables:
for table in tables:
_validate_cli_pattern(table, "tables")
cmd.extend(["-t", table])
if exclude_tables:
for table in exclude_tables:
_validate_cli_pattern(table, "exclude_tables")
cmd.extend(["-T", table])
if schemas:
for schema in schemas:
_validate_cli_pattern(schema, "schemas")
cmd.extend(["-n", schema])
# Output
cmd.extend(["-f", str(output_file)])
return cmd
[docs]
def build_pg_restore_cmd(
host: str,
port: int,
user: str,
database: str,
input_file: str | Path,
clean: bool = False,
if_exists: bool = True,
create: bool = False,
data_only: bool = False,
schema_only: bool = False,
tables: list[str] | None = None,
schemas: list[str] | None = None,
jobs: int = 1,
no_owner: bool = False,
no_privileges: bool = False,
) -> list[str]:
"""Build a pg_restore command argv list.
Constructs the argument list for invoking pg_restore. Pure function: no
I/O, no environment access, no secrets. The caller runs the command
and manages credentials via environment variables.
The .sql / non-existent file early-return branch (which delegates to
_psql_restore) is the caller's responsibility; this builder assumes a
binary-format input file.
Parameters
----------
host : str
Database host.
port : int
Database port.
user : str
Database user.
database : str
Database name.
input_file : str or Path
Backup file path.
clean : bool, optional
Drop objects before recreating, by default False.
if_exists : bool, optional
Use IF EXISTS with clean (prevents errors), by default True.
create : bool, optional
Create database before restoring, by default False.
data_only : bool, optional
Restore only data, by default False.
schema_only : bool, optional
Restore only schema, by default False.
tables : list, optional
Only restore these tables.
schemas : list, optional
Only restore these schemas.
jobs : int, optional
Parallel jobs, by default 1.
no_owner : bool, optional
Don't restore ownership, by default False.
no_privileges : bool, optional
Don't restore privileges, by default False.
Returns
-------
list
List of strings (argv) suitable for passing to a process runner.
"""
cmd = ["pg_restore"]
# Connection params
cmd.extend(["-h", host])
cmd.extend(["-p", str(port)])
cmd.extend(["-U", user])
cmd.extend(["-d", database])
# Options
if clean:
cmd.append("--clean")
if if_exists:
cmd.append("--if-exists")
if create:
cmd.append("--create")
if data_only:
cmd.append("--data-only")
if schema_only:
cmd.append("--schema-only")
if jobs > 1:
cmd.extend(["-j", str(jobs)])
if no_owner:
cmd.append("--no-owner")
if no_privileges:
cmd.append("--no-privileges")
# Tables/Schemas — guard against flag injection (CR-04)
if tables:
for table in tables:
_validate_cli_pattern(table, "tables")
cmd.extend(["-t", table])
if schemas:
for schema in schemas:
_validate_cli_pattern(schema, "schemas")
cmd.extend(["-n", schema])
cmd.append(str(Path(input_file)))
return cmd
[docs]
def build_role_options(
login: bool = True,
superuser: bool = False,
createdb: bool = False,
createrole: bool = False,
inherit: bool = True,
replication: bool = False,
connection_limit: int = -1,
password: Any = None,
valid_until: str | None = None,
) -> list[str]:
"""Build the options list for a CREATE ROLE statement.
Constructs the SQL option tokens for CREATE ROLE … WITH <options>.
The password VALUE is never stored in or returned from this builder
(D-04); when password is truthy the literal placeholder "PASSWORD %s"
is appended so the caller can bind the actual secret via a parameterized
execute call.
Parameters
----------
login : bool, optional
Can log in (True = user, False = group role), by default True.
superuser : bool, optional
Is superuser, by default False.
createdb : bool, optional
Can create databases, by default False.
createrole : bool, optional
Can create other roles, by default False.
inherit : bool, optional
Inherits privileges from member roles, by default True.
replication : bool, optional
Can initiate streaming replication, by default False.
connection_limit : int, optional
Max concurrent connections (-1 = unlimited), by default -1.
password : optional
Truthiness flag only — when truthy, appends "PASSWORD %s"
placeholder. The actual secret must be bound by the caller.
valid_until : str, optional
Password expiration (e.g. '2025-12-31'). Validated
against the timestamp whitelist before use.
Returns
-------
list
List of SQL option token strings.
"""
options = []
if login:
options.append("LOGIN")
else:
options.append("NOLOGIN")
if superuser:
options.append("SUPERUSER")
if createdb:
options.append("CREATEDB")
if createrole:
options.append("CREATEROLE")
if not inherit:
options.append("NOINHERIT")
if replication:
options.append("REPLICATION")
if connection_limit != -1:
# Runtime int guard — consistent with number_partitions in timescale.py (WR-01)
if isinstance(connection_limit, bool) or not isinstance(connection_limit, int):
raise TypeError(
f"connection_limit must be an int, "
f"got {type(connection_limit).__name__!r}"
)
options.append(f"CONNECTION LIMIT {connection_limit}")
if password:
# Append placeholder only; caller binds the actual secret value
options.append("PASSWORD %s")
if valid_until:
validate_timestamp(valid_until)
options.append(f"VALID UNTIL '{valid_until}'")
return options