diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 3558ab54d..260870066 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -114,3 +114,92 @@ jobs: run: bash tests/smoke_postgres_container_logging.sh - name: Run legacy SQL cleanup integration smoke run: bash tests/smoke_legacy_sql_cleanup.sh + - name: Set up Python 3.14 for candidate parity + uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 + with: + python-version: "3.14" + - name: Set up uv for candidate parity + uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 + with: + prune-cache: true + - name: Install locked project dependencies + run: uv sync --locked + - name: Download exact pg8000 candidate dependency closure + run: >- + python -m pip download --no-deps --only-binary=:all: + --dest /tmp/pg8000-candidate + pg8000==1.31.5 + python-dateutil==2.9.0.post0 + scramp==1.4.17 + asn1crypto==1.5.1 + six==1.17.0 + - name: Verify pg8000 candidate dependency digests + shell: bash + run: | + cat <<'EOF' | sha256sum --check --strict + 0af2c1926b153307639868d2ee5cef6cd3a7d07448e12736989b10e1d491e201 /tmp/pg8000-candidate/pg8000-1.31.5-py3-none-any.whl + a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427 /tmp/pg8000-candidate/python_dateutil-2.9.0.post0-py2.py3-none-any.whl + a4e3fd2e8169461a28a13777a166d3da94274454f0714a7d3023fee124474ac8 /tmp/pg8000-candidate/scramp-1.4.17-py3-none-any.whl + db4e40728b728508912cbb3d44f19ce188f218e9eba635821bb4b68564f8fd67 /tmp/pg8000-candidate/asn1crypto-1.5.1-py2.py3-none-any.whl + 4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274 /tmp/pg8000-candidate/six-1.17.0-py2.py3-none-any.whl + EOF + - name: Verify pg8000 candidate dependency licenses + run: python tools/verify_candidate_wheel_licenses.py /tmp/pg8000-candidate + - name: Install exact candidate closure into the CI environment + run: >- + uv pip install --python .venv/bin/python --no-deps + /tmp/pg8000-candidate/pg8000-1.31.5-py3-none-any.whl + /tmp/pg8000-candidate/python_dateutil-2.9.0.post0-py2.py3-none-any.whl + /tmp/pg8000-candidate/scramp-1.4.17-py3-none-any.whl + /tmp/pg8000-candidate/asn1crypto-1.5.1-py2.py3-none-any.whl + /tmp/pg8000-candidate/six-1.17.0-py2.py3-none-any.whl + - name: Verify candidate environment dependency consistency + run: uv pip check --python .venv/bin/python + - name: Start candidate PostgreSQL runtime + shell: bash + run: | + candidate_password="$(python -c 'import secrets; print(secrets.token_urlsafe(32))')" + echo "::add-mask::$candidate_password" + password_file="$(mktemp)" + printf '%s' "$candidate_password" > "$password_file" + chmod 600 "$password_file" + container="pg-llm-batch-pg8000-${GITHUB_RUN_ID}-${GITHUB_RUN_ATTEMPT}" + echo "PG8000_CANDIDATE_PASSWORD_FILE=$password_file" >> "$GITHUB_ENV" + echo "PG8000_CANDIDATE_CONTAINER=$container" >> "$GITHUB_ENV" + docker run --detach --name "$container" \ + --publish 127.0.0.1:5432:5432 \ + --mount "type=bind,source=$password_file,target=/run/secrets/postgres_password,readonly" \ + --env POSTGRES_USER=pgllm \ + --env POSTGRES_DB=pgllm \ + --env POSTGRES_PASSWORD_FILE=/run/secrets/postgres_password \ + pg-llm-batch-postgres:ci + - name: Wait for candidate PostgreSQL health contract + shell: bash + run: | + for attempt in $(seq 1 90); do + if docker exec "$PG8000_CANDIDATE_CONTAINER" \ + pg_isready -U pgllm -d pgllm >/dev/null 2>&1; then + ready="$(docker exec "$PG8000_CANDIDATE_CONTAINER" \ + psql -U pgllm -d pgllm -tAc \ + "SELECT bool_and(is_ready) FROM pg_llm_batch_health_check() WHERE component IN ('database','pg_tiktoken','com_config')" \ + 2>/dev/null || true)" + if [ "$ready" = "t" ]; then + exit 0 + fi + fi + sleep 2 + done + docker logs "$PG8000_CANDIDATE_CONTAINER" + exit 1 + - name: Run real pg8000 candidate PostgreSQL smoke + run: uv run --no-sync python tests/smoke_pg8000_candidate_postgres.py + - name: Tear down candidate PostgreSQL runtime + if: ${{ always() }} + shell: bash + run: | + if [ -n "${PG8000_CANDIDATE_CONTAINER:-}" ]; then + docker rm --force "$PG8000_CANDIDATE_CONTAINER" >/dev/null 2>&1 || true + fi + if [ -n "${PG8000_CANDIDATE_PASSWORD_FILE:-}" ]; then + rm -f "$PG8000_CANDIDATE_PASSWORD_FILE" + fi diff --git a/docs/product-technical-gap-baseline.md b/docs/product-technical-gap-baseline.md new file mode 100644 index 000000000..d31753f22 --- /dev/null +++ b/docs/product-technical-gap-baseline.md @@ -0,0 +1,56 @@ +# Product and technical gap baseline + +This document records shipped truth separately from active-PR evidence. Exact PR heads, checks, reviews, rulesets, and release identities must always be read live before merge or release decisions; this file is not a substitute for GitHub evidence. + +## Product boundary + +pg-llm-batch owns durable PostgreSQL-backed asynchronous LLM batch preparation, token/size accounting, provider-neutral `BatchInferencePort` lifecycle state, tenant/RLS enforcement, result ingestion, audit/provenance, and recovery. Model/provider discovery and routing remain contextual-orchestrator authority. Foreign product truth is consumed through released versioned contracts and anti-corruption layers; source copying, mutable-branch production dependencies, and cross-service application-table SQL are out of bounds. + +## Protected-main truth + +The protected integration branch is `main`. At the latest refresh it was `bdff1273d3885dedc5187632e1c8838b470c9b6d`. The package remains version `0.1.0`, and its production dependency graph still includes `psycopg[binary]>=3.1`. Therefore issue #322, replacement of the LGPL-family Psycopg runtime dependency, remains an open commercial-policy defect. No public release should claim that the current `pip install .` runtime graph is commercially clean while that defect remains. + +The repository has no immutable GitHub release at the latest refresh. A release is not ready merely because a branch is green: one exact protected head must pass the repository's applicable CI, security, coverage/docstring, package, SBOM/provenance, reproducibility, migration/rollback/recovery, operability, and review gates before version/tag/publication evidence is promoted. + +## Active delivery lanes + +PR #233 remains the dependency-root delivery lane and must be judged from its live head and live base, not predecessor evidence. + +PR #323 is the active Draft migration lane for issue #322. It establishes a driver-neutral PostgreSQL anti-corruption port, retains Psycopg only as the current baseline adapter, and evaluates pg8000 1.31.5 as candidate evidence without promoting it into the production manifest. The lane already exercises parameter binding, tuple-row normalization, row-count semantics, transaction and cleanup precedence, forced-RLS tenant scope, JSONB/UUID/timestamp behavior, exact candidate dependency hashes, and real PostgreSQL candidate smoke tests. + +The current candidate supply-chain work also verifies license metadata for the exact five-wheel pg8000 candidate closure before installation. The verifier reads bounded wheel `METADATA` without importing candidate code, rejects GPL/LGPL/AGPL-family declarations, requires positive reviewed permissive-license evidence for every exact package/version, and rejects an unexpected wheel set. This strengthens candidate admission but does not itself approve a production driver replacement. + +## Highest-priority gaps + +| Gap | Current state | Required next evidence | +| --- | --- | --- | +| Commercial PostgreSQL runtime dependency | P0 / active | Complete issue #322: preserve shipped DB semantics while removing every disallowed GPL/LGPL/AGPL-family runtime package from the committed dependency graph. | +| Candidate driver contract parity | Active Draft | Close conninfo URI/keyword/service-selector compatibility, driver-level JSONB/error classification, concurrency/recovery, timeout/health, schema/restore, and package-installed behavior with realistic PostgreSQL evidence. | +| Candidate supply-chain admission | Active / strengthened | Exact wheel hashes and license metadata are now gated; complete vulnerability/SBOM/provenance and final runtime-graph evidence before promotion. | +| Exact-head CI execution | External control-plane dependency plus local continuation | Current required jobs must acquire a runner and check out the exact current head. Queued/pre-checkout evidence is non-passing; central `.github#712` owns the organization runner-admission diagnosis. | +| Immutable product release | Not yet published | After the production dependency replacement and all gates pass on one integrated protected head, perform version/CHANGELOG/tag/package/SBOM/provenance/reproducibility/rollback publication and verify artifact identity. | +| Context Graph / EA projection | Candidate-only until released authority exists | `context-graph-contracts` and `enterprise-architecture-core` currently expose no immutable GitHub release. Do not pin mutable producer heads. Continue pg-owned release-readiness seams and adopt only a verified released contract. | + +## Commercial acceptance for issue #322 + +Completion requires all of the following on the final production graph, not only candidate fixtures: + +- parameterized SQL and injection-safe bindings remain intact; +- commit, rollback, context-manager, cleanup-error precedence, cancellation/recovery, and connection lifecycle remain deterministic; +- tenant authority and transaction-local `set_config` behavior remain correct under forced RLS and restricted roles; +- JSON/JSONB, UUID, timestamp, row, row-count, and relevant PostgreSQL error semantics remain compatible; +- DSN parsing/rendering preserves the repository's supported URI, keyword, and service-selector contract without credential leakage into argv or logs; +- concurrency, idempotency, checkpoint, schema application, logical restore, health, and finite-connect behavior pass realistic PostgreSQL tests; +- the committed runtime graph and built artifacts contain no disallowed GPL/LGPL/AGPL-family package; +- package, license, vulnerability, SBOM, provenance, and reproducibility evidence bind the same immutable artifacts; +- the final unchanged head passes exact-source required checks and then-live review/ruleset requirements without self-approval or gate weakening. + +## Context Fabric boundary + +`context-graph-contracts` remains a contract-only Shared Kernel for canonical object/authority references, truth origin/status, bitemporal semantics, provenance, Context Assertion, and CloudEvents/schema/conformance/admission contracts. `enterprise-architecture-core` remains the EA Decision Plane. While their dedicated Context Fabric writer is active, pg-llm-batch treats both repositories as read-only source dependencies and advances their existing owner paths with exact consumer RED/GREEN criteria instead of creating competing writers. + +Prompt, response, batch-result, and user data remain pg/product-domain data and are not copied into EA authoritative architecture tables. Deployable service/API/worker/database/runtime/provider/version and lifecycle/risk/ownership/remediation changes may be projected only through a verified released Context Graph contract with provenance. + +## Evidence discipline + +Queued, pending, skipped-required, cancelled, absent, predecessor-head, model-only, and status-only evidence is non-passing. A current blocker is the next work item at its actual owner: pg-owned causes require a realistic RED, the smallest causal repair, focused/full GREEN, and exact-head refetch; foreign-owned causes require advancement of the existing owner path followed by independent pg work. A report, comment, handoff, or documentation-only change is never completion while executable code/test/release work remains. diff --git a/pg_llm_batch/checkpoint_store.py b/pg_llm_batch/checkpoint_store.py index 4eda50195..398ec9cbb 100644 --- a/pg_llm_batch/checkpoint_store.py +++ b/pg_llm_batch/checkpoint_store.py @@ -10,14 +10,14 @@ from .db import ( DEFAULT_TENANT_SCOPE, - _require_psycopg, _set_transaction_tenant_scope, - psycopg, validate_endpoint_alias, validate_remote_resource_id, validate_tenant_scope, ) from .exceptions import ConfigError, PgLlmBatchError, ValidationError +from .postgres_driver_port import PostgresDriverPort +from .postgres_driver_runtime import retained_postgres_driver from .result_streaming import BatchResultCheckpoint MIGRATION_PATH = ( @@ -85,6 +85,21 @@ def _validated_postgres_dsn(value: Any) -> str: return value +def _connect_postgres( + postgres_dsn: str, + postgres_driver: PostgresDriverPort | None, +) -> Any: + """Connect through the selected PostgreSQL driver boundary. + + Explicitly injected migration drivers remain available for parity and + degraded-mode tests. Without an injection, the centralized runtime selector + supplies the retained implementation so checkpoint persistence no longer + imports or constructs Psycopg directly. + """ + selected_driver = postgres_driver or retained_postgres_driver() + return selected_driver.connect(postgres_dsn) + + def _validated_checkpoint(value: Any, field: str) -> BatchResultCheckpoint: """Require one immutable checkpoint whose counters fit PostgreSQL storage.""" if not isinstance(value, BatchResultCheckpoint): @@ -171,13 +186,14 @@ def _checkpoint_values(checkpoint: BatchResultCheckpoint) -> tuple[Any, ...]: def apply_result_checkpoint_schema( postgres_dsn: str, migration_path: Optional[str] = None, + *, + postgres_driver: PostgresDriverPort | None = None, ) -> None: """Apply the idempotent durable result-checkpoint migration.""" dsn = _validated_postgres_dsn(postgres_dsn) - _require_psycopg() path = Path(migration_path) if migration_path else MIGRATION_PATH sql = path.read_text(encoding="utf-8") - with psycopg.connect(dsn) as conn: + with _connect_postgres(dsn, postgres_driver) as conn: with conn.cursor() as cur: cur.execute(sql) conn.commit() @@ -191,9 +207,11 @@ def __init__( postgres_dsn: str, *, tenant_scope: str = DEFAULT_TENANT_SCOPE, + postgres_driver: PostgresDriverPort | None = None, ) -> None: - """Bind one explicit database and trusted local tenant scope to the store.""" + """Bind one explicit database, tenant scope, and optional driver port.""" self.postgres_dsn = _validated_postgres_dsn(postgres_dsn) + self._postgres_driver = postgres_driver try: self.tenant_scope = validate_tenant_scope(tenant_scope) except ValidationError as exc: @@ -210,8 +228,7 @@ def load( endpoint_alias: str, ) -> Optional[BatchResultCheckpoint]: """Load the current checkpoint in one package-owned transaction.""" - _require_psycopg() - with psycopg.connect(self.postgres_dsn) as conn: + with _connect_postgres(self.postgres_dsn, self._postgres_driver) as conn: with conn.cursor() as cur: return self.load_in_transaction( cur, @@ -257,8 +274,7 @@ def save( expected_previous: Optional[BatchResultCheckpoint] = None, ) -> BatchResultCheckpoint: """Create or advance a checkpoint in one package-owned transaction.""" - _require_psycopg() - with psycopg.connect(self.postgres_dsn) as conn: + with _connect_postgres(self.postgres_dsn, self._postgres_driver) as conn: with conn.cursor() as cur: saved = self.save_in_transaction( cur, diff --git a/pg_llm_batch/cli.py b/pg_llm_batch/cli.py index 8d09cded3..0333235f0 100644 --- a/pg_llm_batch/cli.py +++ b/pg_llm_batch/cli.py @@ -33,17 +33,16 @@ import sys import warnings from contextlib import ExitStack +from functools import partial from typing import List, Optional -from psycopg import ProgrammingError -from psycopg.conninfo import conninfo_to_dict - -from . import db +from . import db, postgres_driver_runtime from .batch_api_client import BatchAPIClient, config_credentials_provider from .bootstrap import resolve_dsn, resolve_secret_key from .config import PostgresConfigStore, SecretStore from .exceptions import ConfigError, PgLlmBatchError from .health import check_health, serve_healthz +from .postgres_driver_port import PostgresDriverPort from .token_counter import TokenCounter MAX_SECRET_INPUT_CHARACTERS = 65_536 @@ -73,28 +72,53 @@ def error(self, message: str) -> None: super().error(redacted_message) -def _validate_cli_dsn(value: str) -> str: - """Accept valid libpq selectors while refusing credential-bearing argv data.""" +def _default_postgres_driver() -> PostgresDriverPort: + """Delegate concrete-driver construction to the canonical runtime selector. + + CLI parsing owns credential-safe argument validation, not the concrete + PostgreSQL client choice. A single selector keeps migration cutover atomic + across CLI, service, persistence, and Compose surfaces. + """ + return postgres_driver_runtime.retained_postgres_driver() + + +def _validate_cli_dsn( + value: str, + *, + postgres_driver: PostgresDriverPort | None = None, +) -> str: + """Accept valid PostgreSQL selectors without concrete-driver coupling.""" + driver = ( + postgres_driver + if postgres_driver is not None + else _default_postgres_driver() + ) try: - parameters = conninfo_to_dict(value) - except ProgrammingError: - raise argparse.ArgumentTypeError( - "Postgres DSN must be valid libpq connection information" - ) from None + parameters = driver.parse_conninfo(value) + except Exception as exc: + if driver.is_invalid_conninfo(exc): + raise argparse.ArgumentTypeError( + "Postgres DSN must be valid connection information" + ) from None + raise if CLI_DSN_SENSITIVE_PARAMETERS.intersection(parameters): raise argparse.ArgumentTypeError( "Credential-bearing Postgres DSNs are not accepted in --dsn; " - "use libpq secret mechanisms outside process argv" + "use PostgreSQL secret mechanisms outside process argv" ) return value -def _add_common(parser: argparse.ArgumentParser) -> None: +def _add_common( + parser: argparse.ArgumentParser, + *, + postgres_driver: PostgresDriverPort | None = None, +) -> None: """Add the shared credential-free ``--dsn`` selector to a subcommand parser.""" parser.add_argument( "--dsn", default=None, - type=_validate_cli_dsn, + type=partial(_validate_cli_dsn, postgres_driver=postgres_driver), help=( "Credential-free Postgres selector " "(else PG_LLM_BATCH_DSN bootstrap env var)" @@ -186,40 +210,42 @@ def _read_token_input() -> str: raise ConfigError("Token input must be valid UTF-8") from None -def build_parser() -> argparse.ArgumentParser: - """Build the command-line parser and all supported subcommands.""" +def build_parser( + *, + postgres_driver: PostgresDriverPort | None = None, +) -> argparse.ArgumentParser: + """Build the command-line parser with an injectable PostgreSQL DSN parser.""" parser = _RedactingArgumentParser( prog="pg_llm_batch", description="Standalone Postgres LLM batch engine", ) sub = parser.add_subparsers(dest="command", required=True) - p_init = sub.add_parser("init-db", help="Apply batch schema (idempotent)") - _add_common(p_init) + _add_common(p_init, postgres_driver=postgres_driver) p_cfg = sub.add_parser("config", help="Manage KV config and secrets") cfg_sub = p_cfg.add_subparsers(dest="config_command", required=True) p_set = cfg_sub.add_parser("set", help="Set a config value") - _add_common(p_set) + _add_common(p_set, postgres_driver=postgres_driver) p_set.add_argument("category") p_set.add_argument("key") p_set.add_argument("value") p_get = cfg_sub.add_parser("get", help="Get a config value") - _add_common(p_get) + _add_common(p_get, postgres_driver=postgres_driver) p_get.add_argument("category") p_get.add_argument("key") p_secret = cfg_sub.add_parser( "set-secret", help="Store a secret from a no-echo prompt or standard input", ) - _add_common(p_secret) + _add_common(p_secret, postgres_driver=postgres_driver) p_secret.add_argument("secret_key") p_count = sub.add_parser( "count-tokens", help="Count bounded UTF-8 stdin content without exposing it in argv", ) - _add_common(p_count) + _add_common(p_count, postgres_driver=postgres_driver) p_count.add_argument("--model", required=True) p_count.add_argument( "--stdin", @@ -229,38 +255,38 @@ def build_parser() -> argparse.ArgumentParser: ) p_submit = sub.add_parser("submit", help="Upload payload + create batch job") - _add_common(p_submit) + _add_common(p_submit, postgres_driver=postgres_driver) p_submit.add_argument("--endpoint", required=True, help="Endpoint alias") p_submit.add_argument("--file-path", required=True, help="memory://") p_submit.add_argument("--batch-endpoint", default="/v1/chat/completions") p_poll = sub.add_parser("poll", help="Poll a batch job status once") - _add_common(p_poll) + _add_common(p_poll, postgres_driver=postgres_driver) p_poll.add_argument("--endpoint", required=True) p_poll.add_argument("--batch-id", required=True) p_wait = sub.add_parser("wait", help="Wait for a terminal batch status") - _add_common(p_wait) + _add_common(p_wait, postgres_driver=postgres_driver) p_wait.add_argument("--endpoint", required=True) p_wait.add_argument("--batch-id", required=True) p_wait.add_argument("--poll-interval", type=float, default=5.0) p_wait.add_argument("--timeout", type=float, default=3600.0) p_retrieve = sub.add_parser("retrieve", help="Download batch results") - _add_common(p_retrieve) + _add_common(p_retrieve, postgres_driver=postgres_driver) p_retrieve.add_argument("--endpoint", required=True) p_retrieve.add_argument("--batch-id", required=True) p_cancel = sub.add_parser("cancel", help="Cancel a provider batch job") - _add_common(p_cancel) + _add_common(p_cancel, postgres_driver=postgres_driver) p_cancel.add_argument("--endpoint", required=True) p_cancel.add_argument("--batch-id", required=True) p_health = sub.add_parser("health", help="Print readiness report") - _add_common(p_health) + _add_common(p_health, postgres_driver=postgres_driver) p_serve = sub.add_parser("serve-healthz", help="Serve GET /healthz") - _add_common(p_serve) + _add_common(p_serve, postgres_driver=postgres_driver) p_serve.add_argument("--host", default="127.0.0.1") p_serve.add_argument("--port", type=int, default=8080) @@ -430,4 +456,4 @@ async def _go() -> int: if __name__ == "__main__": # pragma: no cover - sys.exit(main()) + sys.exit(main()) \ No newline at end of file diff --git a/pg_llm_batch/compose_bootstrap.py b/pg_llm_batch/compose_bootstrap.py index b3987f406..16615c60c 100644 --- a/pg_llm_batch/compose_bootstrap.py +++ b/pg_llm_batch/compose_bootstrap.py @@ -5,8 +5,8 @@ The Compose profile keeps the database password out of committed configuration, process arguments, and the credential-free bootstrap DSN. This module reads the single explicitly mounted secret, combines it with the bootstrap target only in -process memory using psycopg's conninfo quoting, and hands the result directly to -the existing health server. +process memory through the selected PostgreSQL driver boundary, and hands the +result directly to the existing health server. """ from __future__ import annotations @@ -15,16 +15,27 @@ from pathlib import Path from typing import Sequence -from psycopg.conninfo import make_conninfo - +from . import postgres_driver_runtime from .bootstrap import resolve_dsn from .exceptions import ConfigError from .health import serve_healthz +from .postgres_driver_port import PostgresDriverPort _DEFAULT_PASSWORD_FILE = Path("/run/secrets/postgres_password") _MAX_PASSWORD_BYTES = 65_536 +def _default_postgres_driver() -> PostgresDriverPort: + """Delegate retained-driver construction to the canonical runtime selector. + + Compose owns secret-file handling and private DSN assembly, not concrete + database-client selection. Routing the default through one runtime owner + keeps a future commercially admitted replacement atomic across package + surfaces instead of leaving a hidden Psycopg construction path here. + """ + return postgres_driver_runtime.retained_postgres_driver() + + def _load_database_password(password_file: Path) -> str: """Read one bounded UTF-8 password from an explicitly mounted secret file.""" try: @@ -50,20 +61,58 @@ def _load_database_password(password_file: Path) -> str: return password -def _build_private_dsn(base_dsn: str, password: str) -> str: - """Add the password to a validated DSN using psycopg's conninfo quoting.""" +def _build_private_dsn( + base_dsn: str, + password: str, + *, + postgres_driver: PostgresDriverPort | None = None, +) -> str: + """Add the mounted password through the selected reviewed conninfo renderer. + + The selected driver parses the credential-free selector and renders a fresh + parameter snapshot containing the mounted password. The retained concrete + implementation is selected only by ``postgres_driver_runtime`` while the + commercial migration is incomplete, so this module has no independent + concrete-driver conninfo authority. Parser or renderer diagnostics are + normalized so secret material never escapes this bootstrap boundary. + """ + driver = ( + postgres_driver + if postgres_driver is not None + else _default_postgres_driver() + ) try: - return make_conninfo(base_dsn, password=password) + parameters = dict(driver.parse_conninfo(base_dsn)) + parameters["password"] = password + return driver.make_conninfo(parameters) + except ConfigError: + raise except Exception: raise ConfigError("The PostgreSQL bootstrap target is invalid.") from None -def run_compose_health(password_file: Path = _DEFAULT_PASSWORD_FILE) -> None: - """Serve health checks using the credential-free DSN plus mounted secret.""" +def run_compose_health( + password_file: Path = _DEFAULT_PASSWORD_FILE, + *, + postgres_driver: PostgresDriverPort | None = None, +) -> None: + """Serve readiness with one driver owning private DSN assembly and database I/O.""" base_dsn = resolve_dsn(None) password = _load_database_password(password_file) - private_dsn = _build_private_dsn(base_dsn, password) - serve_healthz(private_dsn, host="0.0.0.0", port=8080) + private_dsn = _build_private_dsn( + base_dsn, + password, + postgres_driver=postgres_driver, + ) + if postgres_driver is None: + serve_healthz(private_dsn, host="0.0.0.0", port=8080) + return + serve_healthz( + private_dsn, + host="0.0.0.0", + port=8080, + postgres_driver=postgres_driver, + ) def _password_file_from_args(argv: Sequence[str] | None) -> Path: diff --git a/pg_llm_batch/config.py b/pg_llm_batch/config.py index 1a33a4d03..e81837e64 100644 --- a/pg_llm_batch/config.py +++ b/pg_llm_batch/config.py @@ -21,11 +21,11 @@ from typing import Any, Dict, Iterable, Optional, Tuple, Type from .exceptions import ConfigError - -try: # pragma: no cover - optional dependency - import psycopg # type: ignore -except ImportError: # pragma: no cover - psycopg = None # type: ignore +from .postgres_driver_port import PostgresDriverPort +from .postgres_driver_runtime import ( + PostgresDriverUnavailableError, + retained_postgres_driver, +) try: # pragma: no cover - optional dependency from cryptography.fernet import Fernet # type: ignore @@ -155,23 +155,56 @@ def _split_full_key(full_key: str) -> Tuple[str, str]: return "global", full_key +def _connect_store_database( + dsn: str, + postgres_driver: PostgresDriverPort | None, + *, + missing_dependency_message: str, +) -> Any: + """Open one config-store connection through the shared driver selector. + + Explicit driver injection remains authoritative for candidate or host-owned + adapters. Otherwise the centralized runtime boundary selects the retained + concrete implementation. Missing retained-driver support is translated to + the store's existing bounded ``ConfigError`` contract. + """ + if postgres_driver is None: + try: + postgres_driver = retained_postgres_driver() + except PostgresDriverUnavailableError as exc: + raise ConfigError(missing_dependency_message) from exc + return postgres_driver.connect(dsn) + + +def _set_store_autocommit(connection: Any) -> None: + """Enable explicit store autocommit through the driver-neutral connection.""" + connection.set_autocommit(True) + + class PostgresConfigStore: """PostgreSQL-backed KV configuration store (``com_config`` table).""" TABLE_NAME = "com_config" - def __init__(self, dsn: str) -> None: - """Connect to PostgreSQL and initialize the configuration cache.""" - if psycopg is None: - raise ConfigError("psycopg is required for PostgresConfigStore") + def __init__( + self, + dsn: str, + *, + postgres_driver: PostgresDriverPort | None = None, + ) -> None: + """Connect through the selected driver and initialize the config cache.""" if not dsn: raise ConfigError( "A Postgres DSN must be provided explicitly (no os.getenv for config)" ) self.dsn = dsn - self._conn = psycopg.connect(self.dsn) + self._conn = _connect_store_database( + self.dsn, + postgres_driver, + missing_dependency_message="psycopg is required for PostgresConfigStore", + ) try: - self._conn.autocommit = True + _set_store_autocommit(self._conn) self.cache: Dict[str, Dict[str, Any]] = {} self._ensure_table() self._ensure_defaults() @@ -301,10 +334,9 @@ def __init__( fernet_key: Optional[str] = None, *, require_encryption: bool = False, + postgres_driver: PostgresDriverPort | None = None, ) -> None: - """Connect using optional Fernet encryption or fail when it is required.""" - if psycopg is None: - raise ConfigError("psycopg is required for SecretStore") + """Connect through the selected driver with the requested secret policy.""" if not dsn: raise ConfigError("A Postgres DSN must be provided explicitly") if require_encryption and not fernet_key: @@ -316,9 +348,13 @@ def __init__( "Fernet encryption requires the optional cryptography dependency" ) self.dsn = dsn - self._conn = psycopg.connect(self.dsn) + self._conn = _connect_store_database( + self.dsn, + postgres_driver, + missing_dependency_message="psycopg is required for SecretStore", + ) try: - self._conn.autocommit = True + _set_store_autocommit(self._conn) self._fernet = None if fernet_key and Fernet is not None: self._fernet = Fernet(fernet_key.encode("utf-8")) @@ -413,4 +449,4 @@ def close(self) -> None: def get_config_store(dsn: str) -> PostgresConfigStore: """Construct a config store. DSN must be passed explicitly (no getenv).""" - return PostgresConfigStore(dsn) \ No newline at end of file + return PostgresConfigStore(dsn) diff --git a/pg_llm_batch/db.py b/pg_llm_batch/db.py index 564677604..a483ba8fe 100644 --- a/pg_llm_batch/db.py +++ b/pg_llm_batch/db.py @@ -19,11 +19,8 @@ from typing import Any, Dict, Optional from .exceptions import ValidationError - -try: # pragma: no cover - optional dependency - import psycopg # type: ignore -except ImportError: # pragma: no cover - psycopg = None # type: ignore +from .postgres_driver_port import PostgresDriverPort +from .postgres_driver_runtime import retained_postgres_driver logger = logging.getLogger(__name__) @@ -90,26 +87,42 @@ def __init__(self) -> None: super().__init__("Stored virtual payload failed integrity validation") -def _require_psycopg() -> None: - """Raise a clear error when the optional psycopg dependency is unavailable.""" - if psycopg is None: # pragma: no cover - raise RuntimeError("psycopg is required for database access") +def _connect_database( + dsn: str, + postgres_driver: PostgresDriverPort | None, +) -> Any: + """Open one connection through the selected PostgreSQL driver boundary. + + Explicitly injected migration drivers remain authoritative for candidate and + degraded-mode tests. When no driver is injected, one centralized runtime + selector supplies the retained implementation so bounded contexts no longer + import or construct Psycopg directly. + """ + selected_driver = postgres_driver or retained_postgres_driver() + return selected_driver.connect(dsn) -def apply_schema(dsn: str) -> None: - """Apply the package-owned idempotent schema to one PostgreSQL database.""" - _require_psycopg() +def apply_schema( + dsn: str, + *, + postgres_driver: PostgresDriverPort | None = None, +) -> None: + """Apply the package-owned schema through the selected PostgreSQL driver.""" sql = SCHEMA_PATH.read_text(encoding="utf-8") - with psycopg.connect(dsn) as conn: + with _connect_database(dsn, postgres_driver) as conn: with conn.cursor() as cur: cur.execute(sql) conn.commit() -def load_virtual_payload(dsn: str, file_id: str) -> Optional[str]: - """Load one canonical package-owned JSONL payload or fail closed.""" - _require_psycopg() - with psycopg.connect(dsn) as conn: +def load_virtual_payload( + dsn: str, + file_id: str, + *, + postgres_driver: PostgresDriverPort | None = None, +) -> Optional[str]: + """Load canonical package JSONL through the selected PostgreSQL driver.""" + with _connect_database(dsn, postgres_driver) as conn: with conn.cursor() as cur: cur.execute( "SELECT content FROM llm_batch_file_payloads WHERE file_id = %s", @@ -395,11 +408,17 @@ def normalize_provider_metadata(value: Any) -> Dict[str, Any]: return _provider_metadata(value)[0] -def reserve_remote_batch_observation_order(dsn: str) -> int: - """Reserve and return one positive database-owned lifecycle order.""" - _require_psycopg() - with psycopg.connect(dsn) as conn: +def reserve_remote_batch_observation_order( + dsn: str, + *, + tenant_scope: str = DEFAULT_TENANT_SCOPE, + postgres_driver: PostgresDriverPort | None = None, +) -> int: + """Reserve one positive lifecycle order after binding validated tenant scope.""" + normalized_tenant_scope = validate_tenant_scope(tenant_scope) + with _connect_database(dsn, postgres_driver) as conn: with conn.cursor() as cur: + _set_transaction_tenant_scope(cur, normalized_tenant_scope) cur.execute("SELECT nextval('llm_remote_batch_observation_sequence')") row = cur.fetchone() if ( @@ -422,6 +441,20 @@ def _set_transaction_tenant_scope(cursor: Any, tenant_scope: str) -> None: ) +def _cursor_row_count( + cursor: Any, + _postgres_driver: PostgresDriverPort | None, +) -> int | None: + """Read an exact affected-row count through a driver-neutral cursor surface.""" + row_count = getattr(cursor, "row_count", None) + value = row_count() if callable(row_count) else getattr(cursor, "rowcount", None) + if value is None or value == -1: + return None + if type(value) is not int or value < 0: + return None + return value + + def _normalize_remote_batch_snapshot( tenant_scope: str, endpoint_alias: str, @@ -525,8 +558,9 @@ def _persist_remote_batch_state( observation_order: int, *, observed_at: Optional[datetime] = None, + postgres_driver: PostgresDriverPort | None = None, ) -> Dict[str, Any]: - """Persist one validated tenant-qualified lifecycle projection.""" + """Persist one validated tenant lifecycle projection through the driver port.""" snapshot, metadata_json = _normalize_remote_batch_snapshot( tenant_scope, endpoint_alias, @@ -651,12 +685,12 @@ def _persist_remote_batch_state( terminal_at, observed, ) - _require_psycopg() - with psycopg.connect(dsn) as conn: + with _connect_database(dsn, postgres_driver) as conn: with conn.cursor() as cur: _set_transaction_tenant_scope(cur, snapshot["tenant_scope"]) cur.execute(sql, params) - if getattr(cur, "rowcount", None) == 0: + affected_rows = _cursor_row_count(cur, postgres_driver) + if affected_rows in (None, 0): cur.execute( """ SELECT tenant_scope, @@ -707,8 +741,9 @@ def persist_remote_batch_state( observation_order: int, *, observed_at: Optional[datetime] = None, + postgres_driver: PostgresDriverPort | None = None, ) -> Dict[str, Any]: - """Persist one standalone projection without changing its return shape.""" + """Persist one standalone projection through the selected PostgreSQL driver.""" snapshot = _persist_remote_batch_state( dsn, DEFAULT_TENANT_SCOPE, @@ -716,6 +751,7 @@ def persist_remote_batch_state( provider_batch, observation_order, observed_at=observed_at, + postgres_driver=postgres_driver, ) snapshot.pop("tenant_scope", None) snapshot.pop("total_requests_known", None) @@ -730,8 +766,9 @@ def persist_tenant_remote_batch_state( observation_order: int, *, observed_at: Optional[datetime] = None, + postgres_driver: PostgresDriverPort | None = None, ) -> Dict[str, Any]: - """Persist one lifecycle projection for an explicit trusted tenant scope.""" + """Persist one trusted-tenant lifecycle projection through the driver port.""" snapshot = _persist_remote_batch_state( dsn, tenant_scope, @@ -739,6 +776,7 @@ def persist_tenant_remote_batch_state( provider_batch, observation_order, observed_at=observed_at, + postgres_driver=postgres_driver, ) snapshot.pop("total_requests_known", None) return snapshot @@ -749,8 +787,10 @@ def get_tenant_remote_batch_state( tenant_scope: str, endpoint_alias: str, remote_batch_id: str, + *, + postgres_driver: PostgresDriverPort | None = None, ) -> Optional[Dict[str, Any]]: - """Return one lifecycle projection visible to a validated tenant scope.""" + """Return one tenant-visible lifecycle projection through the driver port.""" normalized_tenant_scope = validate_tenant_scope(tenant_scope) normalized_alias = validate_endpoint_alias(endpoint_alias) normalized_remote_batch_id = validate_remote_resource_id( @@ -780,8 +820,7 @@ def get_tenant_remote_batch_state( AND endpoint_alias = %s AND remote_batch_id = %s """ - _require_psycopg() - with psycopg.connect(dsn) as conn: + with _connect_database(dsn, postgres_driver) as conn: with conn.cursor() as cur: _set_transaction_tenant_scope(cur, normalized_tenant_scope) cur.execute( @@ -804,31 +843,40 @@ def get_remote_batch_state( dsn: str, endpoint_alias: str, remote_batch_id: str, + *, + postgres_driver: PostgresDriverPort | None = None, ) -> Optional[Dict[str, Any]]: - """Return one lifecycle projection from the standalone tenant scope.""" + """Return one standalone lifecycle projection through the driver port.""" return get_tenant_remote_batch_state( dsn, DEFAULT_TENANT_SCOPE, endpoint_alias, remote_batch_id, + postgres_driver=postgres_driver, ) -def get_model_metadata(dsn: Optional[str], model_id: str) -> Optional[Dict[str, Any]]: - """Fetch model mode and tokenizer metadata for a model identifier. +def get_model_metadata( + dsn: Optional[str], + model_id: str, + *, + postgres_driver: PostgresDriverPort | None = None, +) -> Optional[Dict[str, Any]]: + """Fetch model metadata through the selected PostgreSQL driver boundary. Args: dsn: Optional PostgreSQL connection string. model_id: Provider model identifier to resolve. + postgres_driver: Optional migration driver retained only for database I/O. Returns: A dictionary containing normalized ``mode`` and ``tokenizer_model`` when found, otherwise ``None``. """ - if not dsn or psycopg is None or not model_id: + if not dsn or not model_id: return None try: - with psycopg.connect(dsn) as conn: + with _connect_database(dsn, postgres_driver) as conn: with conn.cursor() as cur: cur.execute( """ diff --git a/pg_llm_batch/health.py b/pg_llm_batch/health.py index 6abed36a7..ab22e043d 100644 --- a/pg_llm_batch/health.py +++ b/pg_llm_batch/health.py @@ -13,10 +13,11 @@ import logging from typing import Any, Dict, List -try: # pragma: no cover - optional dependency - import psycopg # type: ignore -except ImportError: # pragma: no cover - psycopg = None # type: ignore +from .postgres_driver_port import PostgresDriverPort +from .postgres_driver_runtime import ( + PostgresDriverUnavailableError, + retained_postgres_driver, +) logger = logging.getLogger(__name__) @@ -24,18 +25,46 @@ REQUIRED_COMPONENTS = {"database", "pg_tiktoken", "com_config"} -def check_health(dsn: str) -> Dict[str, Any]: - """Return a readiness report ``{ready: bool, components: [...]}``.""" - if psycopg is None: - return { - "ready": False, - "components": [ - {"component": "psycopg", "is_ready": False, "detail": "not installed"} - ], - } +def _connect_health_database( + dsn: str, + postgres_driver: PostgresDriverPort | None, +) -> Any: + """Open a bounded readiness connection through the shared driver selector. + + An explicitly injected driver remains authoritative for candidate and + degraded-mode checks. Otherwise the centralized runtime boundary selects the + retained concrete client. Both paths receive the same five-second connection + budget so migration cannot silently weaken readiness liveness semantics. + """ + if postgres_driver is None: + try: + postgres_driver = retained_postgres_driver() + except PostgresDriverUnavailableError: + return None + return postgres_driver.connect(dsn, connect_timeout_seconds=5) + + +def check_health( + dsn: str, + *, + postgres_driver: PostgresDriverPort | None = None, +) -> Dict[str, Any]: + """Return a readiness report using the injected or retained database driver.""" components: List[Dict[str, Any]] = [] try: - with psycopg.connect(dsn, connect_timeout=5) as conn: + connection = _connect_health_database(dsn, postgres_driver) + if connection is None: + return { + "ready": False, + "components": [ + { + "component": "psycopg", + "is_ready": False, + "detail": "not installed", + } + ], + } + with connection as conn: with conn.cursor() as cur: cur.execute( "SELECT component, is_ready, detail FROM pg_llm_batch_health_check()" @@ -48,11 +77,16 @@ def check_health(dsn: str) -> Dict[str, Any]: "detail": detail, } ) - except Exception as exc: + except Exception: + logger.debug("Database readiness check failed") return { "ready": False, "components": [ - {"component": "database", "is_ready": False, "detail": str(exc)} + { + "component": "database", + "is_ready": False, + "detail": "database readiness check failed", + } ], } @@ -120,8 +154,14 @@ def public_health_report(report: Dict[str, Any]) -> Dict[str, Any]: } -def serve_healthz(dsn: str, host: str = "0.0.0.0", port: int = 8080) -> None: - """Serve a minimal ``/healthz`` endpoint (blocking).""" +def serve_healthz( + dsn: str, + host: str = "0.0.0.0", + port: int = 8080, + *, + postgres_driver: PostgresDriverPort | None = None, +) -> None: + """Serve ``/healthz`` using the selected PostgreSQL driver boundary.""" from http.server import BaseHTTPRequestHandler, HTTPServer class _Handler(BaseHTTPRequestHandler): @@ -133,7 +173,9 @@ def do_GET(self) -> None: # noqa: N802 (stdlib naming) self.send_response(404) self.end_headers() return - report = public_health_report(check_health(dsn)) + report = public_health_report( + check_health(dsn, postgres_driver=postgres_driver) + ) body = json.dumps(report).encode("utf-8") self.send_response(200 if report["ready"] else 503) self.send_header("Content-Type", "application/json") diff --git a/pg_llm_batch/orchestrator.py b/pg_llm_batch/orchestrator.py index 11f2aabda..318e32bf5 100644 --- a/pg_llm_batch/orchestrator.py +++ b/pg_llm_batch/orchestrator.py @@ -20,15 +20,10 @@ from . import db from .config import PostgresConfigStore from .exceptions import ValidationError +from .postgres_driver_port import PostgresDriverPort +from .postgres_driver_runtime import retained_postgres_driver from .token_counter import BatchAccumulator, TokenCounter -try: # pragma: no cover - optional dependency - import psycopg # type: ignore - from psycopg.types.json import Jsonb # type: ignore -except ImportError: # pragma: no cover - psycopg = None # type: ignore - Jsonb = None # type: ignore - def _validate_effective_token_limit(value: Optional[int]) -> Optional[int]: """Validate an optional stricter runtime token limit.""" @@ -60,11 +55,40 @@ class BatchPayload: class PostgresBatchOrchestrator: """Assemble and persist JSONL batch payloads from queued requests.""" - def __init__(self, dsn: str) -> None: - """Initialize the orchestrator with an explicit PostgreSQL DSN.""" - if not dsn or psycopg is None: - raise RuntimeError("A Postgres DSN and psycopg are required") + def __init__( + self, + dsn: str, + *, + postgres_driver: PostgresDriverPort | None = None, + ) -> None: + """Initialize with an explicit DSN through the centralized driver boundary. + + Candidate and host adapters may be injected explicitly. Ordinary runtime + construction uses the retained selector so this bounded context does not + keep a second concrete Psycopg import, JSONB adapter, or connection path. + """ + if not dsn: + raise RuntimeError("A Postgres DSN is required") self.dsn = dsn + self._postgres_driver = ( + postgres_driver if postgres_driver is not None else retained_postgres_driver() + ) + + def _connect_database(self) -> Any: + """Open one orchestrator connection through the selected driver boundary.""" + return self._postgres_driver.connect(self.dsn) + + def _set_autocommit(self, connection: Any, enabled: bool) -> None: + """Set transaction mode through the driver-neutral connection contract.""" + connection.set_autocommit(enabled) + + def _adapt_jsonb(self, value: object) -> object: + """Adapt JSONB through the selected driver while preserving exact payload data.""" + return self._postgres_driver.jsonb(value) + + def _cursor_row_count(self, cursor: Any) -> int | None: + """Read affected-row evidence through the selected cursor contract.""" + return cursor.row_count() def _resolve_batch_uuid(self, batch_key: str) -> Optional[str]: """Resolve an exact string batch UUID or input-file-path selector.""" @@ -79,7 +103,7 @@ def _resolve_batch_uuid(self, batch_key: str) -> Optional[str]: return batch_key except ValueError: pass - with psycopg.connect(self.dsn) as conn: + with self._connect_database() as conn: with conn.cursor() as cur: cur.execute( "SELECT batch_uuid FROM llm_batches " @@ -115,7 +139,7 @@ def prepare_batches( ), ) - with psycopg.connect(self.dsn) as conn: + with self._connect_database() as conn: with conn.cursor() as cur: cur.execute( """ @@ -130,9 +154,16 @@ def prepare_batches( ) rows: List[Tuple] = cur.fetchall() - config = PostgresConfigStore(self.dsn) + config = PostgresConfigStore( + self.dsn, + postgres_driver=self._postgres_driver, + ) try: - counter = TokenCounter(self.dsn, config=config) + counter = TokenCounter( + self.dsn, + config=config, + postgres_driver=self._postgres_driver, + ) try: if validated_token_limit is not None: counter.effective_limit = min( @@ -156,7 +187,11 @@ def _assemble_payloads( payloads: List[Dict[str, Any]] = [] for (request_uuid, system_prompt, user_prompt, model_name) in rows: - metadata = db.get_model_metadata(self.dsn, model_name) + metadata = db.get_model_metadata( + self.dsn, + model_name, + postgres_driver=self._postgres_driver, + ) mode = str((metadata or {}).get("mode") or "").lower() system_for_tokens = system_prompt if mode != "embedding" else "" @@ -278,8 +313,8 @@ def _persist_payloads( immediate_limit = counter.azure_max_files_per_job lock_key = self._batch_lock_key(batch_uuid) - with psycopg.connect(self.dsn) as conn: - conn.autocommit = False + with self._connect_database() as conn: + self._set_autocommit(conn, False) with conn.cursor() as cur: cur.execute("SELECT pg_advisory_xact_lock(%s)", (lock_key,)) cur.execute( @@ -331,11 +366,7 @@ def _persist_payloads( request_ids = [str(item) for item in meta.get("request_ids", [])] content = "\n".join(lines) + ("\n" if lines else "") payload_doc = {"text": content, "line_count": len(lines)} - adapted = ( - Jsonb(payload_doc) - if Jsonb is not None - else json.dumps(payload_doc) - ) + adapted = self._adapt_jsonb(payload_doc) cur.execute( """ INSERT INTO llm_batch_file_payloads (file_id, content) @@ -399,7 +430,7 @@ def _persist_payloads( """, (file_uuid, batch_uuid, request_ids), ) - if cur.rowcount != len(request_ids): + if self._cursor_row_count(cur) != len(request_ids): raise ValidationError( field="request_ids", value=request_ids, diff --git a/pg_llm_batch/pg8000_candidate_driver_port.py b/pg_llm_batch/pg8000_candidate_driver_port.py new file mode 100644 index 000000000..22a5262c6 --- /dev/null +++ b/pg_llm_batch/pg8000_candidate_driver_port.py @@ -0,0 +1,424 @@ +"""Candidate-only pg8000 implementation of the PostgreSQL driver port. + +The commercial migration needs a connection factory, not only cursor wrappers. +pg8000 1.31.5 accepts explicit DB-API connection keyword arguments but does not +provide libpq conninfo or service-file parsing. This adapter therefore owns a +bounded anti-corruption parser for the single-host PostgreSQL URI and keyword +conninfo forms already needed by pg-llm-batch. A caller may inject a service +resolver so service-file I/O and precedence remain outside the concrete driver; +only the resolved parameter subset already admitted by this candidate reaches +pg8000. Query options, multi-host/socket forms, and other libpq-only semantics +remain fail closed until they have separate compatibility evidence. + +The module does not import pg8000. An exact candidate DB-API module must be +injected after artifact, license, integrity, and environment admission, keeping +pg8000 out of the committed production dependency graph while issue #322 remains +open. +""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from types import ModuleType +from typing import cast +from urllib.parse import quote, unquote, urlsplit + +from .pg8000_driver_candidate_adapter import ( + Pg8000CandidateAdapterError, + validate_pg8000_dbapi_module, +) +from .pg8000_driver_candidate_errors import ( + is_pg8000_candidate_undefined_function, +) +from .pg8000_driver_candidate_jsonb import adapt_pg8000_jsonb +from .pg8000_thread_affine_candidate_adapter import ( + Pg8000ThreadAffineCandidateConnectionAdapter, +) +from .postgres_driver_port import PostgresConnectionPort, PostgresDriverPort + + +_HEX_DIGITS = frozenset("0123456789abcdefABCDEF") +_ALLOWED_PARAMETER_KEYS = frozenset({"user", "password", "host", "port", "dbname"}) +_KEYWORD_PARAMETER_KEYS = _ALLOWED_PARAMETER_KEYS | frozenset({"service"}) +_DEFAULT_PORT = 5432 +_MIN_PORT = 1 +_MAX_PORT = 65_535 +_AMBIGUOUS_HOST_TOKENS = frozenset("/?,#@[]\\%") +_KEYWORD_SEPARATOR = " " + +ServiceResolver = Callable[[str], Mapping[str, str]] + + +class Pg8000CandidateInvalidConninfoError(Pg8000CandidateAdapterError): + """Identify only failures at the candidate connection-selector boundary. + + The error message never reflects the supplied selector or credentials. A + failure means the candidate cannot represent that PostgreSQL selector under + the currently proved URI/keyword contract; it does not mean the database + rejected a connection attempt. + """ + + +def _invalid_selector(*, unsupported: bool = False) -> Pg8000CandidateInvalidConninfoError: + """Create one non-content-bearing selector error for malformed or missing data.""" + if unsupported: + return Pg8000CandidateInvalidConninfoError( + "PostgreSQL connection selector is unsupported" + ) + return Pg8000CandidateInvalidConninfoError( + "PostgreSQL connection selector is invalid" + ) + + +def _contains_control(value: str) -> bool: + """Return whether text contains ASCII control or DEL framing characters.""" + return any(ord(character) < 0x20 or ord(character) == 0x7F for character in value) + + +def _validate_percent_encoding(value: str) -> None: + """Reject incomplete or non-hex percent escapes before URI decoding.""" + index = 0 + while index < len(value): + if value[index] != "%": + index += 1 + continue + if ( + index + 2 >= len(value) + or value[index + 1] not in _HEX_DIGITS + or value[index + 2] not in _HEX_DIGITS + ): + raise _invalid_selector() + index += 3 + + +def _decode_component(value: str, *, allow_empty: bool = False) -> str: + """Decode one UTF-8 URI component after strict percent and framing checks.""" + _validate_percent_encoding(value) + try: + decoded = unquote(value, encoding="utf-8", errors="strict") + except (UnicodeDecodeError, UnicodeError): + raise _invalid_selector() from None + if (not decoded and not allow_empty) or _contains_control(decoded) or "\x00" in decoded: + raise _invalid_selector() + return decoded + + +def _parse_port(value: object) -> int: + """Return one exact TCP port while rejecting bools and out-of-range values.""" + if type(value) is int: + port = value + elif type(value) is str and value.isascii() and value.isdigit(): + port = int(value) + else: + raise _invalid_selector() + if port < _MIN_PORT or port > _MAX_PORT: + raise _invalid_selector() + return port + + +def _validate_host(host: str) -> str: + """Keep one TCP host while rejecting forms that imply unsupported selectors. + + Commas would turn PostgreSQL URI authority into a multi-host selector; + percent escapes can encode Unix-socket paths or IPv6 zone identifiers; and + whitespace/backslashes or URI delimiters make rendering ambiguous. Those + contracts require separate compatibility evidence and therefore fail closed + instead of being passed to pg8000 as a misleading single hostname. + """ + if ( + not host + or _contains_control(host) + or any(character.isspace() for character in host) + or any(token in host for token in _AMBIGUOUS_HOST_TOKENS) + ): + raise _invalid_selector() + return host + + +def _parse_postgresql_uri(dsn: str) -> dict[str, str]: + """Parse the candidate's reviewed single-host PostgreSQL URI subset. + + Service selectors, query parameters, fragments, multi-host forms, Unix-socket + selectors, and libpq-specific options remain outside this bounded slice. They + fail closed rather than being approximated by pg8000 connection arguments. + """ + if type(dsn) is not str or not dsn or _contains_control(dsn): + raise _invalid_selector() + if not (dsn.startswith("postgresql://") or dsn.startswith("postgres://")): + raise _invalid_selector(unsupported=True) + + try: + parsed = urlsplit(dsn) + except ValueError: + raise _invalid_selector() from None + + if parsed.scheme not in {"postgresql", "postgres"}: + raise _invalid_selector(unsupported=True) + if parsed.query or parsed.fragment: + raise _invalid_selector(unsupported=True) + if not parsed.netloc or parsed.username is None or parsed.hostname is None: + raise _invalid_selector() + if not parsed.path.startswith("/") or len(parsed.path) <= 1: + raise _invalid_selector() + raw_database = parsed.path[1:] + if "/" in raw_database: + raise _invalid_selector(unsupported=True) + + user = _decode_component(parsed.username) + password = ( + _decode_component(parsed.password, allow_empty=True) + if parsed.password is not None + else None + ) + host = _validate_host(parsed.hostname) + database = _decode_component(raw_database) + try: + port = parsed.port if parsed.port is not None else _DEFAULT_PORT + except ValueError: + raise _invalid_selector() from None + port = _parse_port(port) + + result = { + "user": user, + "host": host, + "port": str(port), + "dbname": database, + } + if password is not None: + result["password"] = password + return result + + +def _read_keyword_value(dsn: str, start: int) -> tuple[str, int]: + """Read one bounded libpq-style keyword value and return the next offset. + + The candidate deliberately accepts only ASCII-space separators. Single-quoted + values and backslash escaping follow the portable conninfo forms needed by + current deployments, while tabs/newlines and unterminated escapes fail closed + instead of acquiring implicit libpq parser semantics. + """ + if start >= len(dsn): + return "", start + + quoted = dsn[start] == "'" + index = start + 1 if quoted else start + characters: list[str] = [] + while index < len(dsn): + character = dsn[index] + if quoted and character == "'": + index += 1 + if index < len(dsn) and dsn[index] != _KEYWORD_SEPARATOR: + raise _invalid_selector() + return "".join(characters), index + if not quoted and character == _KEYWORD_SEPARATOR: + return "".join(characters), index + if character == "\\": + index += 1 + if index >= len(dsn): + raise _invalid_selector() + characters.append(dsn[index]) + index += 1 + continue + if character == "'" and not quoted: + raise _invalid_selector() + characters.append(character) + index += 1 + + if quoted: + raise _invalid_selector() + return "".join(characters), index + + +def _parse_keyword_fields(dsn: str) -> dict[str, str]: + """Parse the reviewed keyword grammar without yet resolving service authority. + + ``service`` is grammar-recognized so a separately injected resolver can own + service-file I/O and precedence. Other unknown parameters remain unsupported; + duplicate keys are rejected instead of relying on libpq's last-value-wins + behavior because duplicated authority is ambiguous at the migration boundary. + """ + if type(dsn) is not str or not dsn or _contains_control(dsn): + raise _invalid_selector() + if any(character.isspace() and character != _KEYWORD_SEPARATOR for character in dsn): + raise _invalid_selector() + + index = 0 + params: dict[str, str] = {} + while index < len(dsn): + while index < len(dsn) and dsn[index] == _KEYWORD_SEPARATOR: + index += 1 + if index >= len(dsn): + break + + key_start = index + while index < len(dsn) and dsn[index] not in {_KEYWORD_SEPARATOR, "="}: + index += 1 + key = dsn[key_start:index] + if not key: + raise _invalid_selector() + while index < len(dsn) and dsn[index] == _KEYWORD_SEPARATOR: + index += 1 + if index >= len(dsn) or dsn[index] != "=": + raise _invalid_selector() + index += 1 + while index < len(dsn) and dsn[index] == _KEYWORD_SEPARATOR: + index += 1 + + if key not in _KEYWORD_PARAMETER_KEYS: + raise _invalid_selector(unsupported=True) + if key in params: + raise _invalid_selector() + + value, index = _read_keyword_value(dsn, index) + params[key] = value + + return params + + +def _copy_parameter_mapping(params: Mapping[str, str]) -> dict[str, str]: + """Copy exact built-in values while rejecting non-candidate parameters.""" + if not isinstance(params, Mapping): + raise _invalid_selector() + copied: dict[str, str] = {} + for key, value in params.items(): + if type(key) is not str or key not in _ALLOWED_PARAMETER_KEYS: + raise _invalid_selector(unsupported=True) + if type(value) is not str or _contains_control(value) or "\x00" in value: + raise _invalid_selector() + copied[key] = value + return copied + + +def _validate_parameter_mapping(params: Mapping[str, str]) -> dict[str, str]: + """Normalize the complete candidate connection parameter set.""" + copied = _copy_parameter_mapping(params) + if "user" not in copied or "host" not in copied or "dbname" not in copied: + raise _invalid_selector() + if not copied["user"] or not copied["dbname"]: + raise _invalid_selector() + copied["host"] = _validate_host(copied["host"]) + copied["port"] = str(_parse_port(copied.get("port", str(_DEFAULT_PORT)))) + return copied + + +def _render_host(host: str) -> str: + """Render a validated single host, adding URI brackets only for IPv6 form.""" + if ":" in host: + return f"[{host}]" + return host + + +class Pg8000CandidateDriverAdapter(PostgresDriverPort): + """Prove pg8000 on bounded single-host PostgreSQL connection selectors. + + The injected module must already be the exact candidate artifact. This class + supplies no artifact discovery, dependency installation, or fallback. It + converts only reviewed URI/keyword fields to pg8000 DB-API keyword arguments. + Service-file lookup is deliberately an injected anti-corruption boundary: the + driver never reads process environment variables or filesystem service files + itself, and direct conninfo values override resolver-provided values before + the merged parameter set is validated. + """ + + def __init__( + self, + dbapi_module: ModuleType, + *, + service_resolver: ServiceResolver | None = None, + ) -> None: + """Bind one admitted DB-API module and optional explicit service resolver.""" + validate_pg8000_dbapi_module(dbapi_module) + connect = vars(dbapi_module).get("connect") + if not callable(connect): + raise Pg8000CandidateAdapterError( + "PostgreSQL driver connection factory is incompatible" + ) + self._dbapi_module = dbapi_module + self._connect = connect + self._service_resolver = service_resolver + + def connect( + self, + dsn: str, + *, + connect_timeout_seconds: int | None = None, + ) -> PostgresConnectionPort: + """Open one candidate connection from a proved selector and finite timeout. + + The original DSN is never forwarded to pg8000. Parsed values are supplied + as explicit DB-API keywords so unsupported libpq selector semantics cannot + be silently inherited or misrepresented. + """ + if connect_timeout_seconds is not None and ( + type(connect_timeout_seconds) is not int or connect_timeout_seconds <= 0 + ): + raise Pg8000CandidateInvalidConninfoError( + "PostgreSQL driver timeout is invalid" + ) + params = self.parse_conninfo(dsn) + kwargs: dict[str, object] = { + "user": params["user"], + "host": params["host"], + "port": _parse_port(params["port"]), + "database": params["dbname"], + } + if "password" in params: + kwargs["password"] = params["password"] + if connect_timeout_seconds is not None: + kwargs["timeout"] = connect_timeout_seconds + raw_connection = self._connect(**kwargs) + return Pg8000ThreadAffineCandidateConnectionAdapter(raw_connection) + + def parse_conninfo(self, dsn: str) -> Mapping[str, str]: + """Parse admitted URI/keyword selectors and resolve service authority. + + A service selector is accepted only when a resolver was explicitly + injected. The resolver returns ordinary connection parameters; this + adapter then applies any direct conninfo overrides and validates the final + merged subset. This mirrors PostgreSQL's service-then-direct precedence + without silently acquiring service-file or environment authority. + """ + if type(dsn) is not str or not dsn: + raise _invalid_selector() + if dsn.startswith("postgresql://") or dsn.startswith("postgres://"): + return _parse_postgresql_uri(dsn) + + parsed = _parse_keyword_fields(dsn) + service_name = parsed.pop("service", None) + if service_name is None: + return _validate_parameter_mapping(parsed) + if not service_name: + raise _invalid_selector() + if self._service_resolver is None: + raise _invalid_selector(unsupported=True) + + resolved = _copy_parameter_mapping(self._service_resolver(service_name)) + resolved.update(parsed) + return _validate_parameter_mapping(resolved) + + def make_conninfo(self, params: Mapping[str, str]) -> str: + """Render the proved parameter subset as a safely percent-encoded URI.""" + copied = _validate_parameter_mapping(params) + user = quote(copied["user"], safe="") + password = copied.get("password") + credentials = user + if password is not None: + credentials += f":{quote(password, safe='')}" + host = _render_host(copied["host"]) + database = quote(copied["dbname"], safe="") + return f"postgresql://{credentials}@{host}:{copied['port']}/{database}" + + def jsonb(self, value: object) -> object: + """Use the separately admitted candidate JSONB serialization boundary.""" + return adapt_pg8000_jsonb(value) + + def is_invalid_conninfo(self, error: BaseException) -> bool: + """Recognize only errors emitted by this candidate selector boundary.""" + return isinstance(error, Pg8000CandidateInvalidConninfoError) + + def is_undefined_function(self, error: BaseException) -> bool: + """Classify SQLSTATE 42883 through the exact injected DB-API authority.""" + return is_pg8000_candidate_undefined_function( + error, + dbapi_module=cast(object, self._dbapi_module), + ) diff --git a/pg_llm_batch/pg8000_candidate_service_file.py b/pg_llm_batch/pg8000_candidate_service_file.py new file mode 100644 index 000000000..428440d53 --- /dev/null +++ b/pg_llm_batch/pg8000_candidate_service_file.py @@ -0,0 +1,213 @@ +"""Bounded caller-selected ``pg_service.conf`` resolver for the pg8000 candidate. + +PostgreSQL service files are an INI-like indirection from a service name to +connection parameters. pg8000 does not implement libpq's service-file lookup, +and the database driver must not silently acquire process-environment or +filesystem-discovery authority while pg-llm-batch evaluates a replacement for +Psycopg. This candidate resolver therefore reads exactly one caller-selected +file, applies a finite byte budget, and returns only the exact target stanza. + +The parser intentionally does not implement libpq LDAP lookup or ambient +``PGSERVICEFILE``/user/system search precedence. Those capabilities require +separate security and compatibility evidence. The returned mapping is validated +again by :class:`Pg8000CandidateDriverAdapter`, so unsupported PostgreSQL +connection parameters remain fail closed at the driver boundary. +""" + +from __future__ import annotations + +import os +import stat +from pathlib import Path + +from .pg8000_candidate_driver_port import Pg8000CandidateInvalidConninfoError + + +_MAX_SERVICE_FILE_BYTES = 64 * 1024 + + +def _invalid_service_file(*, unsupported: bool = False) -> Pg8000CandidateInvalidConninfoError: + """Return one non-content-bearing error for service-file resolution failures.""" + if unsupported: + return Pg8000CandidateInvalidConninfoError( + "PostgreSQL connection selector is unsupported" + ) + return Pg8000CandidateInvalidConninfoError( + "PostgreSQL connection selector is invalid" + ) + + +def _has_disallowed_control(value: str) -> bool: + """Reject framing controls while allowing ordinary horizontal whitespace.""" + return any( + (ord(character) < 0x20 and character != "\t") or ord(character) == 0x7F + for character in value + ) + + +def _validate_service_name(service_name: object) -> str: + """Validate one exact service identity without normalizing caller authority.""" + if ( + type(service_name) is not str + or not service_name + or service_name != service_name.strip() + or _has_disallowed_control(service_name) + or "[" in service_name + or "]" in service_name + ): + raise _invalid_service_file() + return service_name + + +def _service_file_snapshot(observed: os.stat_result) -> tuple[int, int, int, int, int, int]: + """Capture metadata that must remain stable while service bytes are retained.""" + return ( + observed.st_dev, + observed.st_ino, + observed.st_mode, + observed.st_size, + observed.st_mtime_ns, + observed.st_ctime_ns, + ) + + +def _read_bounded_utf8(path: Path) -> str: + """Read one explicit regular service file under a finite UTF-8 byte budget. + + The caller-selected path is opened nonblocking where the platform supports + it, then the retained descriptor is required to name one stable regular + file before and after the bounded read. This prevents a FIFO/device path or + in-place mutation from becoming connection-selector authority while bytes + are being inspected. + """ + flags = ( + os.O_RDONLY + | getattr(os, "O_BINARY", 0) + | getattr(os, "O_CLOEXEC", 0) + | getattr(os, "O_NONBLOCK", 0) + ) + try: + descriptor = os.open(path, flags) + except (OSError, ValueError): + raise _invalid_service_file() from None + + primary_error: BaseException | None = None + try: + before = os.fstat(descriptor) + if not stat.S_ISREG(before.st_mode): + raise _invalid_service_file() + before_snapshot = _service_file_snapshot(before) + + chunks: list[bytes] = [] + remaining = _MAX_SERVICE_FILE_BYTES + 1 + while remaining > 0: + chunk = os.read(descriptor, remaining) + if not chunk: + break + chunks.append(chunk) + remaining -= len(chunk) + payload = b"".join(chunks) + + after = os.fstat(descriptor) + if ( + _service_file_snapshot(after) != before_snapshot + or len(payload) != after.st_size + ): + raise _invalid_service_file() + except Pg8000CandidateInvalidConninfoError as exc: + primary_error = exc + raise + except (OSError, ValueError) as exc: + primary_error = exc + raise _invalid_service_file() from None + finally: + try: + os.close(descriptor) + except OSError: + if primary_error is None: + raise _invalid_service_file() from None + + if len(payload) > _MAX_SERVICE_FILE_BYTES: + raise _invalid_service_file() + try: + text = payload.decode("utf-8", errors="strict") + except UnicodeDecodeError: + raise _invalid_service_file() from None + if "\x00" in text: + raise _invalid_service_file() + return text + + +class Pg8000CandidateServiceFileResolver: + """Resolve one service stanza from an explicit local service-file capability. + + ``service_file`` is selected by the caller and retained as a concrete path; + this object never discovers user/system files and never reads environment + variables. Duplicate section/key authority and malformed target lines fail + closed. Non-target stanza contents are not promoted into the selected + connection parameters. + """ + + def __init__(self, service_file: Path) -> None: + """Retain exactly one caller-selected service file after validating its path type.""" + if not isinstance(service_file, Path): + raise _invalid_service_file() + self._service_file = service_file + + def __call__(self, service_name: str) -> dict[str, str]: + """Return the exact target stanza or fail without reflecting file content.""" + target = _validate_service_name(service_name) + text = _read_bounded_utf8(self._service_file) + sections: set[str] = set() + target_found = False + target_active = False + parameters: dict[str, str] = {} + + for raw_line in text.splitlines(): + if _has_disallowed_control(raw_line): + raise _invalid_service_file() + stripped = raw_line.strip() + if not stripped or stripped.startswith("#"): + continue + + if stripped.startswith("["): + if ( + not stripped.endswith("]") + or stripped.count("[") != 1 + or stripped.count("]") != 1 + ): + raise _invalid_service_file() + section_name = stripped[1:-1].strip() + if ( + not section_name + or _has_disallowed_control(section_name) + or section_name in sections + ): + raise _invalid_service_file() + sections.add(section_name) + target_active = section_name == target + if target_active: + target_found = True + continue + + if not target_active: + continue + if stripped.lower().startswith("ldap://"): + raise _invalid_service_file(unsupported=True) + if "=" not in stripped: + raise _invalid_service_file() + key, value = stripped.split("=", 1) + key = key.strip() + value = value.strip() + if ( + not key + or _has_disallowed_control(key) + or _has_disallowed_control(value) + or key in parameters + ): + raise _invalid_service_file() + parameters[key] = value + + if not target_found: + raise _invalid_service_file() + return parameters diff --git a/pg_llm_batch/pg8000_driver_candidate_adapter.py b/pg_llm_batch/pg8000_driver_candidate_adapter.py new file mode 100644 index 000000000..78d184418 --- /dev/null +++ b/pg_llm_batch/pg8000_driver_candidate_adapter.py @@ -0,0 +1,376 @@ +"""Candidate-only pg8000 DB-API adapters for commercial migration evidence. + +This module intentionally stops short of a production ``PostgresDriverPort``. +pg8000 1.31.5 documents the DB-API cursor, transaction, autocommit, parameter +binding, module-only connection thread sharing, and ``-1`` unknown-row-count +behavior needed by part of the current port, but pg-llm-batch has not yet proved +its full conninfo/service-selector, JSONB adaptation, PostgreSQL +error-classification, Python 3.14, RLS, transport failure recovery, package, +SBOM, and provenance contract on one exact artifact. Keeping this adapter +candidate-only lets those portable semantics be exercised without making an +unreleased or unverified runtime dependency canonical. +""" + +from __future__ import annotations + +from types import ModuleType +from typing import Any, cast + +from .postgres_driver_port import PostgresConnectionPort, PostgresCursorPort + + +class Pg8000CandidateAdapterError(RuntimeError): + """Report a candidate-boundary mismatch without exposing database content. + + A mismatch means the candidate cannot yet be promoted through the shared + PostgreSQL port. The error is deliberately separate from pg8000's database + exceptions so callers cannot mistake missing adapter evidence for a server + or application failure. + """ + + +def validate_pg8000_dbapi_module(dbapi_module: object) -> None: + """Fail closed unless imported pg8000 DB-API metadata matches package use. + + pg8000 exposes ``paramstyle`` as mutable module state. pg-llm-batch's current + SQL uses DB-API ``format`` placeholders, so a future production candidate + factory must run this guard immediately after importing the exact admitted + pg8000 artifact and before creating adapters or executing SQL. The documented + ``threadsafety == 1`` value is also part of this boundary: code may share the + module across threads but must not infer that one connection is shareable. + Metadata is read from an exact ``ModuleType`` dictionary rather than through + arbitrary shaped objects whose attribute access could execute caller-controlled + code. + + Raises: + Pg8000CandidateAdapterError: If DB-API level, parameter style, or thread + sharing semantics differ from the exact reviewed candidate contract. + """ + if type(dbapi_module) is not ModuleType: + raise Pg8000CandidateAdapterError("PostgreSQL driver module identity is invalid") + + module = cast(ModuleType, dbapi_module) + metadata = vars(module) + api_level = metadata.get("apilevel") + parameter_style = metadata.get("paramstyle") + thread_safety = metadata.get("threadsafety") + + if type(api_level) is not str or api_level != "2.0": + raise Pg8000CandidateAdapterError("PostgreSQL driver API level is incompatible") + if type(parameter_style) is not str or parameter_style != "format": + raise Pg8000CandidateAdapterError( + "PostgreSQL driver parameter style is incompatible" + ) + if type(thread_safety) is not int or thread_safety != 1: + raise Pg8000CandidateAdapterError( + "PostgreSQL driver thread safety is incompatible" + ) + + +class Pg8000CandidateCursorAdapter(PostgresCursorPort): + """Exercise pg8000 DB-API cursor semantics behind the canonical cursor port. + + The raw cursor remains dependency-injected because this candidate slice must + not add pg8000 to the production dependency graph before the exact artifact + passes license, security, Python, PostgreSQL, and recovery admission. Query + text and bound parameters are forwarded unchanged; materialized list rows are + normalized to the tuple representation already used by pg-llm-batch. + """ + + def __init__(self, cursor: Any) -> None: + """Retain one already-admitted raw cursor without acquiring connection authority.""" + self._cursor = cursor + + @staticmethod + def _normalize_result_row(row: object) -> tuple[object, ...]: + """Normalize one DB-API positional row while rejecting ambiguous shapes. + + pg8000 documents list-like result rows. The package canonicalizes exact + ``list`` and ``tuple`` containers only; mapping or custom containers are + rejected so a driver-specific row factory cannot silently change domain + indexing or equality semantics. + """ + if type(row) is tuple: + return row + if type(row) is list: + return tuple(row) + raise Pg8000CandidateAdapterError("PostgreSQL driver result row is invalid") + + def execute( + self, + query: str, + params: object | None = None, + ) -> Pg8000CandidateCursorAdapter: + """Forward package-authored SQL and bound parameters without interpolation. + + pg8000's DB-API interface defaults to ``format`` parameter style, which + matches the existing ``%s`` package SQL. This candidate method forwards + both objects unchanged so later real-driver tests can detect any semantic + mismatch rather than hiding it in an adapter rewrite. + """ + self._cursor.execute(query, params) + return self + + def executemany( + self, + query: str, + params_seq: object, + ) -> Pg8000CandidateCursorAdapter: + """Forward one statement and parameter sequence without implicit commits. + + Transaction ownership remains with the retained connection. The adapter + therefore does not commit between items or transform the supplied + sequence into independently executed application operations. + """ + self._cursor.executemany(query, params_seq) + return self + + def fetchone(self) -> tuple[object, ...] | None: + """Return one canonical tuple row or ``None`` at end of results. + + Only row-container normalization belongs here. Field-count, type, tenant, + and domain validation remain responsibilities of the consuming bounded + context after the database adapter returns. + """ + row = self._cursor.fetchone() + if row is None: + return None + return self._normalize_result_row(row) + + def fetchmany(self, size: int) -> list[tuple[object, ...]]: + """Return a bounded result page and reject invalid caller or driver budgets. + + ``bool`` is rejected even though it subclasses ``int`` because an + accidental truth value must not become a one-row resource budget. The + adapter also verifies that the concrete DB-API candidate honors that + budget; over-delivery is a candidate-contract failure rather than extra + data the application may silently materialize. + """ + if type(size) is not int or size <= 0: + raise Pg8000CandidateAdapterError("PostgreSQL driver fetch size is invalid") + rows = self._cursor.fetchmany(size) + try: + returned_count = len(rows) + except (TypeError, ValueError, OverflowError): + raise Pg8000CandidateAdapterError( + "PostgreSQL driver fetch result is invalid" + ) from None + if returned_count > size: + raise Pg8000CandidateAdapterError( + "PostgreSQL driver fetch result exceeds requested size" + ) + return [self._normalize_result_row(row) for row in rows] + + def fetchall(self) -> list[tuple[object, ...]]: + """Normalize all rows from an already bounded package-authored query. + + This preserves current package compatibility but does not authorize new + unbounded queries; untrusted-result paths must continue to enforce their + own finite SQL and fetch budgets. + """ + return [self._normalize_result_row(row) for row in self._cursor.fetchall()] + + def row_count(self) -> int | None: + """Normalize pg8000's documented ``-1`` unknown row count to ``None``. + + Exact non-negative counts remain usable as mutation evidence. Any other + negative sentinel or non-integer value fails closed because the package + must not interpret an undocumented driver state as exact write success. + """ + value = self._cursor.rowcount + if type(value) is not int: + raise Pg8000CandidateAdapterError("PostgreSQL driver row count is invalid") + if value == -1: + return None + if value < 0: + raise Pg8000CandidateAdapterError("PostgreSQL driver row count is invalid") + return value + + def __enter__(self) -> Pg8000CandidateCursorAdapter: + """Enter the package cursor context without requiring a driver extension. + + Python DB-API 2.0 standardizes ``Cursor.close()`` but not a cursor context + manager. The anti-corruption adapter therefore owns context entry instead + of making an undocumented pg8000 ``__enter__`` method part of the product + contract. + """ + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + traceback: object | None, + ) -> bool | None: + """Close the DB-API cursor without replacing an active application error. + + Cursor exit owns resource cleanup only; transaction commit or rollback + remains a connection-level responsibility. If cleanup fails while an + application exception is already in flight, the application exception + remains primary. A close-only failure still propagates. Returning + ``False`` preserves ordinary context-manager exception propagation while + avoiding a driver-specific cursor context-manager dependency. + """ + try: + self._cursor.close() + except BaseException: + if exc is not None: + raise exc from None + raise + return False + + +class Pg8000CandidateConnectionAdapter(PostgresConnectionPort): + """Exercise portable pg8000 DB-API connection semantics on one raw connection. + + This adapter proves only the connection/cursor portion of the migration port. + It never opens a connection itself and therefore cannot bypass the still-open + DSN/conninfo/service-selector admission problem. All operations stay on the + injected raw connection so transaction-local RLS state cannot migrate to an + implicit second session. The candidate's DB-API thread level does not permit + callers to infer that this retained connection is safe to share across threads; + that package-level concurrency boundary remains a separate admission gate. + """ + + def __init__(self, connection: Any) -> None: + """Retain one admitted raw connection and initialize terminal-state tracking.""" + self._connection = connection + self._closed = False + + def cursor(self) -> Pg8000CandidateCursorAdapter: + """Create a candidate cursor on this exact retained database connection. + + No second connection or hidden pool is introduced. Later PostgreSQL + acceptance must prove the real driver preserves the same session for + tenant-local ``set_config`` and lifecycle SQL. + """ + return Pg8000CandidateCursorAdapter(self._connection.cursor()) + + def execute( + self, + query: str, + params: object | None = None, + ) -> Pg8000CandidateCursorAdapter: + """Execute on this connection and release the owned cursor on failure. + + DB-API does not require ``Connection.execute``. The adapter therefore + creates the cursor itself and owns it until a successful execution hands + the cursor back to the caller. If execution fails before that handoff, + cleanup is attempted immediately so a database error cannot strand an + unreachable cursor. A secondary close failure never replaces the primary + execution failure. + """ + cursor = self.cursor() + try: + cursor.execute(query, params) + except BaseException as execution_error: + try: + cursor.__exit__( + type(execution_error), + execution_error, + execution_error.__traceback__, + ) + except BaseException: + raise execution_error from None + raise + return cursor + + def commit(self) -> None: + """Commit the current local transaction through the raw DB-API connection. + + The adapter adds no retry or distributed-delivery semantics; higher + bounded contexts retain responsibility for replay and idempotency. + """ + self._connection.commit() + + def rollback(self) -> None: + """Roll back the current local transaction and propagate driver failures. + + Rollback errors remain visible because hiding them would make recovery + evidence claim a clean transaction boundary that PostgreSQL did not prove. + """ + self._connection.rollback() + + def set_autocommit(self, enabled: bool) -> None: + """Set pg8000's documented autocommit property from an exact boolean only. + + Rejecting integer truthiness prevents configuration mistakes from being + normalized into transaction-policy changes at the infrastructure edge. + """ + if type(enabled) is not bool: + raise Pg8000CandidateAdapterError("PostgreSQL driver autocommit is invalid") + self._connection.autocommit = enabled + + def is_closed(self) -> bool: + """Report whether this adapter has successfully closed its raw connection. + + DB-API 2.0 requires ``close()`` but not a portable public liveness flag. + The anti-corruption layer therefore tracks only the state it owns instead + of reading a pg8000 implementation detail. This is intentionally not a + network health probe; unexpected transport failure remains an operation + error that recovery tests must prove is discarded and reconnected. + """ + return self._closed + + def close(self) -> None: + """Release the raw connection while preserving pg8000's terminal-close state. + + pg8000 1.31.5 documents that its underlying socket is closed even if the + PostgreSQL protocol-level close reports an error. The adapter therefore + records this capability as locally closed in ``finally`` while allowing + the original close failure to propagate. A failed close must never leave + a definitively released connection eligible for reuse. + """ + try: + self._connection.close() + finally: + self._closed = True + + def __enter__(self) -> Pg8000CandidateConnectionAdapter: + """Enter the package transaction context without a driver-only extension. + + The anti-corruption layer owns the existing pg-llm-batch connection + context contract: successful exit commits, exceptional exit rolls back, + and both paths close the physical connection. Entry itself must not open a + second session or mutate transaction policy. + """ + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + traceback: object | None, + ) -> bool | None: + """Commit or roll back, close, and preserve the highest-priority failure. + + pg8000's public DB-API contract documents ``commit``, ``rollback``, and + ``close`` but does not require a connection context-manager extension. + Owning the package policy here removes that undocumented dependency while + retaining the transaction semantics required by candidate admission. A + commit or rollback failure remains primary over both an application error + and later cleanup failure. If rollback succeeds, the application error + remains primary over a later close failure. A close-only failure still + propagates on an otherwise successful exit. + """ + transaction_error: BaseException | None = None + try: + if exc_type is None: + self.commit() + else: + self.rollback() + except BaseException as error: + transaction_error = error + + try: + self.close() + except BaseException: + if transaction_error is not None: + raise transaction_error from None + if exc is not None: + raise exc from None + raise + + if transaction_error is not None: + raise transaction_error + return False diff --git a/pg_llm_batch/pg8000_driver_candidate_errors.py b/pg_llm_batch/pg8000_driver_candidate_errors.py new file mode 100644 index 000000000..bf88c0883 --- /dev/null +++ b/pg_llm_batch/pg8000_driver_candidate_errors.py @@ -0,0 +1,79 @@ +"""Candidate-only pg8000 PostgreSQL error classification. + +The production package still uses Psycopg while the commercial driver migration +is incomplete. This module proves only narrow pg8000 error semantics needed by +``PostgresDriverPort`` without importing pg8000 into the committed runtime +dependency graph. Callers must inject the exact admitted DB-API module from the +candidate environment; message text is never used as authority. +""" + +from __future__ import annotations + +from types import ModuleType + + +_UNDEFINED_FUNCTION_SQLSTATE = "42883" + + +class Pg8000CandidateErrorEvidenceError(RuntimeError): + """Reject malformed candidate exception authority before classification. + + Candidate metadata participates in a commercial dependency decision. An + invalid module or exception-class authority therefore fails closed instead + of being interpreted as a PostgreSQL server error or a successful parity + result. + """ + + +def _programming_error_type(dbapi_module: object) -> type[BaseException]: + """Return the exact DB-API ProgrammingError class from an admitted module. + + ``ModuleType`` identity is required so shaped objects cannot execute custom + attribute access while supplying security-relevant error metadata. The + exported class must be an actual ``BaseException`` subtype before any + candidate exception is inspected. + """ + if type(dbapi_module) is not ModuleType: + raise Pg8000CandidateErrorEvidenceError( + "PostgreSQL candidate DB-API module authority is invalid" + ) + programming_error = vars(dbapi_module).get("ProgrammingError") + if ( + type(programming_error) is not type + or not issubclass(programming_error, BaseException) + ): + raise Pg8000CandidateErrorEvidenceError( + "PostgreSQL candidate ProgrammingError authority is invalid" + ) + return programming_error + + +def is_pg8000_candidate_undefined_function( + error: BaseException, + *, + dbapi_module: object, +) -> bool: + """Recognize only PostgreSQL SQLSTATE 42883 from the exact candidate class. + + pg8000 server errors carry a PostgreSQL response mapping as the sole + ``ProgrammingError`` argument. Classification requires the exact injected + DB-API exception type, an exact built-in ``dict`` payload, and an exact + string SQLSTATE. Severity and message text are intentionally ignored, so + translated or attacker-controlled diagnostics cannot manufacture the + undefined-function fallback signal used by token-counting code. + + Raises: + Pg8000CandidateErrorEvidenceError: If the injected DB-API module does not + expose a trustworthy ``ProgrammingError`` class authority. + """ + programming_error = _programming_error_type(dbapi_module) + if type(error) is not programming_error: + return False + arguments = error.args + if type(arguments) is not tuple or len(arguments) != 1: + return False + payload = arguments[0] + if type(payload) is not dict: + return False + sqlstate = payload.get("C") + return type(sqlstate) is str and sqlstate == _UNDEFINED_FUNCTION_SQLSTATE diff --git a/pg_llm_batch/pg8000_driver_candidate_jsonb.py b/pg_llm_batch/pg8000_driver_candidate_jsonb.py new file mode 100644 index 000000000..f8128d722 --- /dev/null +++ b/pg_llm_batch/pg8000_driver_candidate_jsonb.py @@ -0,0 +1,58 @@ +"""Candidate-only JSONB adaptation for the pg8000 migration lane. + +pg8000 1.31.5's DB-API contract sends JSON as serialized text and returns JSON +values deserialized. The production runtime still uses Psycopg; this module only +proves the JSONB parameter adaptation needed before a permissively licensed +candidate can implement ``PostgresDriverPort.jsonb``. It does not promote +pg8000 into the runtime dependency graph or bypass the remaining conninfo, RLS, +recovery, package, SBOM, and provenance gates. +""" + +from __future__ import annotations + +import json + + +class Pg8000CandidateJsonbError(RuntimeError): + """Report an invalid candidate JSONB value without reflecting payload content. + + The error intentionally contains no serialized value because batch payloads + may contain purpose-bound user or provider content. Callers can classify the + candidate contract failure without turning diagnostics into a content leak. + """ + + +def adapt_pg8000_jsonb(value: object) -> str: + """Serialize one validated value for pg8000 DB-API JSONB parameter binding. + + PostgreSQL JSON/JSONB does not admit IEEE non-finite numeric literals, and an + isolated Unicode surrogate cannot be encoded as the UTF-8 client text sent to + PostgreSQL. The candidate therefore fails closed before database I/O for + either case and for objects that Python's JSON encoder cannot represent. + Non-ASCII text is retained as Unicode rather than escaped so the adapter can + exercise the same client-encoding boundary used by real multilingual batch + payloads. + + Args: + value: A caller-validated JSON-compatible Python value. + + Returns: + Compact UTF-8-encodable JSON text suitable for a pg8000 DB-API parameter. + + Raises: + Pg8000CandidateJsonbError: If the value is not finite JSON or cannot be + represented as UTF-8 JSON text. + """ + try: + serialized = json.dumps( + value, + ensure_ascii=False, + allow_nan=False, + separators=(",", ":"), + ) + serialized.encode("utf-8") + except (TypeError, ValueError, UnicodeEncodeError, RecursionError): + raise Pg8000CandidateJsonbError( + "PostgreSQL driver JSONB value is invalid" + ) from None + return serialized diff --git a/pg_llm_batch/pg8000_thread_affine_candidate_adapter.py b/pg_llm_batch/pg8000_thread_affine_candidate_adapter.py new file mode 100644 index 000000000..92861e6e2 --- /dev/null +++ b/pg_llm_batch/pg8000_thread_affine_candidate_adapter.py @@ -0,0 +1,153 @@ +"""Thread-affine pg8000 candidate adapters for concurrency admission. + +pg8000 1.31.5 declares DB-API ``threadsafety == 1``: threads may share the +module, but not connections. The portable candidate adapters intentionally do +not invent a stronger guarantee. This layer binds each candidate connection and +cursor to the thread that created it and fails before raw driver access when a +capability crosses that boundary. + +The layer remains candidate-only. It is exercised by the exact-artifact +PostgreSQL smoke test and must not be treated as production dependency approval +until the remaining conninfo, error, recovery, package, SBOM, and provenance +gates pass on one immutable artifact. +""" + +from __future__ import annotations + +from threading import get_ident +from typing import Any + +from .pg8000_driver_candidate_adapter import ( + Pg8000CandidateAdapterError, + Pg8000CandidateConnectionAdapter, + Pg8000CandidateCursorAdapter, +) + + +_THREAD_AFFINITY_ERROR = "PostgreSQL driver connection must not be shared across threads" +_CURSOR_THREAD_AFFINITY_ERROR = "PostgreSQL driver cursor must not be shared across threads" + + +class Pg8000ThreadAffineCandidateCursorAdapter(Pg8000CandidateCursorAdapter): + """Bind one candidate cursor capability to its creating thread. + + The base adapter owns DB-API row, parameter, fetch-budget, and cleanup + normalization. This subclass adds only the concurrency invariant required by + pg8000's declared thread-safety level and performs the check before every raw + cursor access. + """ + + def __init__(self, cursor: Any) -> None: + """Retain one candidate cursor and bind its capability to this thread.""" + super().__init__(cursor) + self._owner_thread_id = get_ident() + + def _require_owner_thread(self) -> None: + """Reject cross-thread cursor use before touching driver-owned state.""" + if get_ident() != self._owner_thread_id: + raise Pg8000CandidateAdapterError(_CURSOR_THREAD_AFFINITY_ERROR) + + def execute( + self, + query: str, + params: object | None = None, + ) -> Pg8000ThreadAffineCandidateCursorAdapter: + """Execute only on the thread that owns the raw candidate cursor.""" + self._require_owner_thread() + super().execute(query, params) + return self + + def executemany( + self, + query: str, + params_seq: object, + ) -> Pg8000ThreadAffineCandidateCursorAdapter: + """Execute a parameter sequence only on the cursor owner thread.""" + self._require_owner_thread() + super().executemany(query, params_seq) + return self + + def fetchone(self) -> tuple[object, ...] | None: + """Fetch one row only from the cursor owner thread.""" + self._require_owner_thread() + return super().fetchone() + + def fetchmany(self, size: int) -> list[tuple[object, ...]]: + """Fetch one bounded page only from the cursor owner thread.""" + self._require_owner_thread() + return super().fetchmany(size) + + def fetchall(self) -> list[tuple[object, ...]]: + """Fetch an already bounded result only from the cursor owner thread.""" + self._require_owner_thread() + return super().fetchall() + + def row_count(self) -> int | None: + """Read affected-row evidence only from the cursor owner thread.""" + self._require_owner_thread() + return super().row_count() + + def __enter__(self) -> Pg8000ThreadAffineCandidateCursorAdapter: + """Enter the cursor context only on the owner thread.""" + self._require_owner_thread() + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + traceback: object | None, + ) -> bool | None: + """Release the raw cursor only on the thread that owns it.""" + self._require_owner_thread() + return super().__exit__(exc_type, exc, traceback) + + +class Pg8000ThreadAffineCandidateConnectionAdapter(Pg8000CandidateConnectionAdapter): + """Bind one candidate connection and its cursors to the creating thread. + + pg8000's DB-API metadata explicitly permits module sharing but not connection + sharing. Serializing a shared connection with a lock would still exceed that + contract, so this adapter rejects cross-thread connection access rather than + treating mutual exclusion as proof of portability. + """ + + def __init__(self, connection: Any) -> None: + """Retain one candidate connection and bind its session to this thread.""" + super().__init__(connection) + self._owner_thread_id = get_ident() + + def _require_owner_thread(self) -> None: + """Reject cross-thread connection use before touching raw driver state.""" + if get_ident() != self._owner_thread_id: + raise Pg8000CandidateAdapterError(_THREAD_AFFINITY_ERROR) + + def cursor(self) -> Pg8000ThreadAffineCandidateCursorAdapter: + """Create a thread-affine cursor on the exact owned connection.""" + self._require_owner_thread() + return Pg8000ThreadAffineCandidateCursorAdapter(self._connection.cursor()) + + def commit(self) -> None: + """Commit only on the thread that owns the candidate connection.""" + self._require_owner_thread() + super().commit() + + def rollback(self) -> None: + """Roll back only on the thread that owns the candidate connection.""" + self._require_owner_thread() + super().rollback() + + def set_autocommit(self, enabled: bool) -> None: + """Change transaction mode only on the candidate connection owner thread.""" + self._require_owner_thread() + super().set_autocommit(enabled) + + def close(self) -> None: + """Close the raw candidate connection only from its owner thread.""" + self._require_owner_thread() + super().close() + + def __enter__(self) -> Pg8000ThreadAffineCandidateConnectionAdapter: + """Enter the candidate transaction context only on the owner thread.""" + self._require_owner_thread() + return self diff --git a/pg_llm_batch/postgres_driver_candidate.py b/pg_llm_batch/postgres_driver_candidate.py new file mode 100644 index 000000000..dabea5884 --- /dev/null +++ b/pg_llm_batch/postgres_driver_candidate.py @@ -0,0 +1,362 @@ +"""Fail-closed commercial acceptance for PostgreSQL driver candidates. + +The repository must replace its current LGPL-family Psycopg runtime dependency +without turning an unverified alternative into production authority. This module +revalidates a bounded candidate snapshot and decides only whether a candidate +has enough permissive-license, immutable license/artifact identity, Python-version, +vulnerability, and capability evidence to enter parity validation. Production +approval remains a later gate that requires a concrete adapter plus +PostgreSQL/RLS/recovery/package evidence. +""" + +from __future__ import annotations + +from dataclasses import dataclass +import re +import unicodedata + + +REQUIRED_POSTGRES_DRIVER_CAPABILITIES = frozenset( + { + "autocommit_state", + "connection_closed_state", + "connection_context", + "connection_context_commit_rollback", + "connection_thread_affinity", + "conninfo_keyword_parse_render", + "conninfo_service_selector", + "conninfo_uri_parse_render", + "cursor_context", + "finite_connect_timeout", + "invalid_conninfo_classification", + "jsonb", + "parameterized_sql", + "result_row_semantics", + "row_count", + "sql_parameter_style_adaptation", + "transaction_commit_rollback", + "undefined_function_classification", + "uuid_timestamp_adaptation", + } +) +"""Capabilities a replacement driver must evidence before parity validation.""" + +REQUIRED_POSTGRES_DRIVER_PYTHON_VERSIONS = frozenset( + {"3.10", "3.11", "3.12", "3.13", "3.14"} +) +"""Repository-supported Python minors a replacement driver must evidence explicitly.""" + +POSTGRES_DRIVER_CANDIDATE_EVIDENCE_SCHEMA_VERSION = "2" +"""Version of the candidate-evidence receipt interpreted by this evaluator.""" + +_APPROVED_PERMISSIVE_LICENSES = frozenset( + { + "Apache-2.0", + "BSD-2-Clause", + "BSD-3-Clause", + "ISC", + "MIT", + "PostgreSQL", + } +) +_MAX_IDENTITY_EVIDENCE_BYTES = 256 +_MAX_PYTHON_VERSION_EVIDENCE_ITEMS = 32 +_MAX_VULNERABILITY_EVIDENCE_ITEMS = 256 +_MINOR_PYTHON_VERSION = re.compile(r"^[1-9][0-9]*\.[0-9]+$") +_PYPA_DISTRIBUTION_NAME = re.compile( + r"[A-Za-z0-9](?:[A-Za-z0-9._-]*[A-Za-z0-9])?\Z" +) +_SOURCE_COMMIT_SHA = re.compile(r"^[0-9a-f]{40}$") +_ARTIFACT_SHA256 = re.compile(r"^[0-9a-f]{64}$") +_VULNERABILITY_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$") + + +class PostgresDriverCandidateEvidenceError(ValueError): + """Reject malformed or mutable evidence before commercial evaluation. + + Candidate metadata can influence a supply-chain migration decision, so the + evaluator accepts only immutable primitive evidence with exact digest and + version shapes. It never repairs or guesses malformed package metadata. + """ + + +def _validate_identity_text(label: str, value: object) -> None: + """Require one finite, exact package-identity token without normalization. + + Package name, version, and SPDX evidence participate in an acquisition + decision and can arrive from untrusted package metadata. Rejecting whitespace, + controls, malformed Unicode, and Unicode format characters keeps one evidence + value from becoming multiple visual or line-oriented identities, while the + UTF-8 byte ceiling bounds malformed metadata without inventing a + package-manager-specific grammar. + """ + if type(value) is not str or not value: + raise PostgresDriverCandidateEvidenceError( + f"PostgreSQL driver {label} evidence is invalid" + ) + try: + encoded_value = value.encode("utf-8") + except UnicodeEncodeError: + raise PostgresDriverCandidateEvidenceError( + f"PostgreSQL driver {label} evidence is invalid" + ) from None + if ( + len(encoded_value) > _MAX_IDENTITY_EVIDENCE_BYTES + or any( + character.isspace() + or ord(character) < 32 + or ord(character) == 127 + or unicodedata.category(character) == "Cf" + for character in value + ) + ): + raise PostgresDriverCandidateEvidenceError( + f"PostgreSQL driver {label} evidence is invalid" + ) + + +def _validate_vulnerability_ids(values: object) -> tuple[str, ...]: + """Validate bounded advisory identifiers without normalizing scan evidence. + + The tuple may be empty only when the bound vulnerability report found no + known advisories. Identifiers remain opaque CVE/GHSA/vendor tokens; the + evaluator constrains their representation and cardinality rather than + inventing a namespace or letting untrusted scan metadata amplify evaluation + work and decision receipts without bound. + """ + if ( + type(values) is not tuple + or len(values) > _MAX_VULNERABILITY_EVIDENCE_ITEMS + ): + raise PostgresDriverCandidateEvidenceError( + "PostgreSQL driver vulnerability evidence is invalid" + ) + if any( + type(value) is not str or _VULNERABILITY_ID.fullmatch(value) is None + for value in values + ): + raise PostgresDriverCandidateEvidenceError( + "PostgreSQL driver vulnerability evidence is invalid" + ) + if len(set(values)) != len(values): + raise PostgresDriverCandidateEvidenceError( + "PostgreSQL driver vulnerability evidence is invalid" + ) + return values + + +@dataclass(frozen=True, slots=True) +class PostgresDriverCandidateEvidence: + """Describe one validated PostgreSQL-driver package candidate. + + ``source_commit_sha`` identifies the reviewed source revision, + ``license_report_sha256`` binds the exact license evidence used for + ``license_spdx``, ``artifact_sha256`` identifies the exact distributable, + ``vulnerability_report_sha256`` binds the exact vulnerability evidence used + for the decision, and ``capability_report_sha256`` binds the exact parity + capability report from which ``capabilities`` is derived. + ``known_vulnerability_ids`` records unresolved advisories from the bound + vulnerability report. ``python_versions`` and ``capabilities`` must contain + explicit evidence rather than inferred support from a nearby release or + similar database driver. ``evidence_schema_version`` prevents a future + receipt shape from being silently interpreted under today's semantics. + Evaluation revalidates a fresh snapshot because Python's frozen dataclasses + do not make ``object.__setattr__`` an authority boundary. + """ + + package_name: str + package_version: str + license_spdx: str + license_report_sha256: str + python_versions: tuple[str, ...] + source_commit_sha: str + artifact_sha256: str + vulnerability_report_sha256: str + capability_report_sha256: str + known_vulnerability_ids: tuple[str, ...] + capabilities: frozenset[str] + evidence_schema_version: str = POSTGRES_DRIVER_CANDIDATE_EVIDENCE_SCHEMA_VERSION + + def __post_init__(self) -> None: + """Validate candidate evidence without normalizing ambiguous inputs.""" + for label, value in ( + ("package name", self.package_name), + ("package version", self.package_version), + ("license", self.license_spdx), + ): + _validate_identity_text(label, value) + if _PYPA_DISTRIBUTION_NAME.fullmatch(self.package_name) is None: + raise PostgresDriverCandidateEvidenceError( + "PostgreSQL driver package name evidence is invalid" + ) + if ( + type(self.evidence_schema_version) is not str + or self.evidence_schema_version + != POSTGRES_DRIVER_CANDIDATE_EVIDENCE_SCHEMA_VERSION + ): + raise PostgresDriverCandidateEvidenceError( + "PostgreSQL driver candidate evidence schema version is unsupported" + ) + if ( + type(self.license_report_sha256) is not str + or _ARTIFACT_SHA256.fullmatch(self.license_report_sha256) is None + ): + raise PostgresDriverCandidateEvidenceError( + "PostgreSQL driver license report evidence is invalid" + ) + if ( + type(self.python_versions) is not tuple + or not self.python_versions + or len(self.python_versions) > _MAX_PYTHON_VERSION_EVIDENCE_ITEMS + ): + raise PostgresDriverCandidateEvidenceError( + "PostgreSQL driver Python version evidence is invalid" + ) + if any( + type(version) is not str or _MINOR_PYTHON_VERSION.fullmatch(version) is None + for version in self.python_versions + ): + raise PostgresDriverCandidateEvidenceError( + "PostgreSQL driver Python version evidence is invalid" + ) + if len(set(self.python_versions)) != len(self.python_versions): + raise PostgresDriverCandidateEvidenceError( + "PostgreSQL driver Python version evidence is invalid" + ) + if ( + type(self.source_commit_sha) is not str + or _SOURCE_COMMIT_SHA.fullmatch(self.source_commit_sha) is None + ): + raise PostgresDriverCandidateEvidenceError( + "PostgreSQL driver source commit evidence is invalid" + ) + if ( + type(self.artifact_sha256) is not str + or _ARTIFACT_SHA256.fullmatch(self.artifact_sha256) is None + ): + raise PostgresDriverCandidateEvidenceError( + "PostgreSQL driver artifact digest evidence is invalid" + ) + if ( + type(self.vulnerability_report_sha256) is not str + or _ARTIFACT_SHA256.fullmatch(self.vulnerability_report_sha256) is None + ): + raise PostgresDriverCandidateEvidenceError( + "PostgreSQL driver vulnerability report evidence is invalid" + ) + if ( + type(self.capability_report_sha256) is not str + or _ARTIFACT_SHA256.fullmatch(self.capability_report_sha256) is None + ): + raise PostgresDriverCandidateEvidenceError( + "PostgreSQL driver capability report evidence is invalid" + ) + _validate_vulnerability_ids(self.known_vulnerability_ids) + if type(self.capabilities) is not frozenset or not self.capabilities: + raise PostgresDriverCandidateEvidenceError( + "PostgreSQL driver capability evidence is invalid" + ) + if any(type(capability) is not str for capability in self.capabilities): + raise PostgresDriverCandidateEvidenceError( + "PostgreSQL driver capability evidence is invalid" + ) + if len(self.capabilities) > len(REQUIRED_POSTGRES_DRIVER_CAPABILITIES): + raise PostgresDriverCandidateEvidenceError( + "PostgreSQL driver capability evidence contains an unknown capability" + ) + unknown_capabilities = self.capabilities - REQUIRED_POSTGRES_DRIVER_CAPABILITIES + if unknown_capabilities: + raise PostgresDriverCandidateEvidenceError( + "PostgreSQL driver capability evidence contains an unknown capability" + ) + + +@dataclass(frozen=True, slots=True) +class PostgresDriverCandidateDecision: + """Record whether validated evidence permits candidate parity validation. + + ``production_approved`` is deliberately always false in this stage. A package + that clears this evaluator still needs a concrete ``PostgresDriverPort`` + adapter and realistic PostgreSQL/RLS/concurrency/recovery/package gates. + """ + + eligible_for_parity_validation: bool + production_approved: bool + reasons: tuple[str, ...] + + +def _validated_candidate_snapshot( + evidence: PostgresDriverCandidateEvidence, +) -> PostgresDriverCandidateEvidence: + """Capture and revalidate exact package evidence before policy evaluation. + + Candidate evidence crosses a supply-chain decision boundary. Requiring the + exact package type before member access prevents candidate-shaped objects from + executing caller-controlled accessors, while reconstruction reapplies every + primitive/container invariant after any post-construction mutation. Deleted + slots are normalized to the package's fixed evidence error. + """ + if type(evidence) is not PostgresDriverCandidateEvidence: + raise PostgresDriverCandidateEvidenceError( + "PostgreSQL driver candidate evidence is invalid" + ) + try: + return PostgresDriverCandidateEvidence( + package_name=evidence.package_name, + package_version=evidence.package_version, + license_spdx=evidence.license_spdx, + license_report_sha256=evidence.license_report_sha256, + python_versions=evidence.python_versions, + source_commit_sha=evidence.source_commit_sha, + artifact_sha256=evidence.artifact_sha256, + vulnerability_report_sha256=evidence.vulnerability_report_sha256, + capability_report_sha256=evidence.capability_report_sha256, + known_vulnerability_ids=evidence.known_vulnerability_ids, + capabilities=evidence.capabilities, + evidence_schema_version=evidence.evidence_schema_version, + ) + except AttributeError: + raise PostgresDriverCandidateEvidenceError( + "PostgreSQL driver candidate evidence is invalid" + ) from None + + +def evaluate_postgres_driver_candidate( + evidence: PostgresDriverCandidateEvidence, +) -> PostgresDriverCandidateDecision: + """Evaluate one candidate without promoting it to a production dependency. + + The decision first revalidates one exact package-owned snapshot, then fails + closed when the bound vulnerability report contains a known advisory, the + SPDX identifier is not in the repository's explicitly reviewed permissive + set, any repository-required Python runtime is not evidenced, or any runtime + capability required by the migration port is absent. DSN evidence remains + split across URI, keyword, and service selectors so a driver cannot claim + generic conninfo support while silently dropping a shipped selector family. + Reasons are deterministic so CI and acquisition diligence can compare exact + evidence. + """ + snapshot = _validated_candidate_snapshot(evidence) + reasons = [ + f"known_vulnerability:{vulnerability_id}" + for vulnerability_id in sorted(snapshot.known_vulnerability_ids) + ] + if snapshot.license_spdx not in _APPROVED_PERMISSIVE_LICENSES: + reasons.append("license_not_approved") + missing_python_versions = REQUIRED_POSTGRES_DRIVER_PYTHON_VERSIONS - set( + snapshot.python_versions + ) + reasons.extend( + f"missing_python_version:{version}" + for version in sorted(missing_python_versions) + ) + missing_capabilities = REQUIRED_POSTGRES_DRIVER_CAPABILITIES - snapshot.capabilities + reasons.extend( + f"missing_capability:{capability}" + for capability in sorted(missing_capabilities) + ) + return PostgresDriverCandidateDecision( + eligible_for_parity_validation=not reasons, + production_approved=False, + reasons=tuple(reasons), + ) diff --git a/pg_llm_batch/postgres_driver_port.py b/pg_llm_batch/postgres_driver_port.py new file mode 100644 index 000000000..61a432fdc --- /dev/null +++ b/pg_llm_batch/postgres_driver_port.py @@ -0,0 +1,271 @@ +"""Provider-neutral PostgreSQL driver contracts for runtime decoupling. + +The package currently has direct Psycopg coupling at several infrastructure +boundaries. These abstract ports describe the database capabilities those +callers actually need without choosing a concrete PostgreSQL driver. Concrete +adapters remain infrastructure concerns and must preserve parameterized SQL, +transaction semantics, connection-string handling, JSONB adaptation, and +PostgreSQL error classification. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import Mapping + + +class PostgresCursorPort(ABC): + """Describe the synchronous cursor surface used by pg-llm-batch. + + Implementations must preserve parameter binding rather than interpolating + SQL text themselves. Materialized result rows use positional tuples as the + canonical package representation because existing bounded contexts use tuple + indexing and equality. A concrete driver returning another row container must + normalize it inside its adapter before exposing it through this port. + """ + + @abstractmethod + def execute( + self, + query: str, + params: object | None = None, + ) -> PostgresCursorPort: + """Execute one parameterized PostgreSQL operation and retain the cursor. + + ``query`` is package-authored SQL and ``params`` carries bound values. + Implementations must not downgrade this call into string formatting or + another transport that changes PostgreSQL parameter semantics. + """ + + @abstractmethod + def executemany(self, query: str, params_seq: object) -> PostgresCursorPort: + """Execute one package-authored operation for a parameter sequence. + + Concrete adapters are responsible for preserving the driver's normal + transactional behavior and must not silently commit between entries. + """ + + @abstractmethod + def fetchone(self) -> tuple[object, ...] | None: + """Return one canonical tuple row, or ``None`` when no row remains. + + The adapter owns only the row-container normalization. The consuming + bounded context remains responsible for validating field count, primitive + types, and semantic meaning before database evidence becomes trusted. + """ + + @abstractmethod + def fetchmany(self, size: int) -> list[tuple[object, ...]]: + """Return at most ``size`` canonical tuple rows. + + Bounded contexts use this operation when an explicit row budget is part + of the product contract; adapters must preserve that finite request and + fail closed rather than silently dropping malformed materialized rows. + """ + + @abstractmethod + def fetchall(self) -> list[tuple[object, ...]]: + """Return canonical tuple rows for an already bounded result set. + + This method exists for compatibility with current package code. New + untrusted-result paths should prefer a bounded query and ``fetchmany``. + """ + + @abstractmethod + def row_count(self) -> int | None: + """Return an exact affected-row count or ``None`` when it is unknown. + + Existing persistence paths use the count to detect missing updates and + partial batch membership writes. Concrete adapters must normalize native + unknown sentinels at this boundary so consumers never mistake a + driver-specific integer such as ``-1`` for exact success evidence. + """ + + @abstractmethod + def __enter__(self) -> PostgresCursorPort: + """Enter the cursor context without changing transaction ownership. + + Connection-level transaction semantics remain owned by the surrounding + connection port rather than being hidden in cursor entry. + """ + + @abstractmethod + def __exit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + traceback: object | None, + ) -> bool | None: + """Leave the cursor context and release adapter-owned cursor resources. + + Returning a truthy value may suppress an exception, so concrete + adapters should preserve their underlying driver's normal behavior. + """ + + +class PostgresConnectionPort(ABC): + """Describe the synchronous PostgreSQL connection capability the package uses. + + The port deliberately keeps transaction mode and local closed-state + inspection explicit because current token-counting, configuration, and batch + assembly paths depend on those semantics. A replacement driver must not + weaken RLS, replay, or recovery behavior by hiding them behind an opaque + wrapper. + """ + + @abstractmethod + def cursor(self) -> PostgresCursorPort: + """Create a cursor bound to this connection's current transaction. + + The returned cursor must implement ``PostgresCursorPort`` semantics and + must not open a second implicit connection. + """ + + @abstractmethod + def execute( + self, + query: str, + params: object | None = None, + ) -> PostgresCursorPort: + """Execute one parameterized statement using this exact connection. + + This convenience operation must retain the same transaction and session + state, including transaction-local tenant ``set_config`` values. + """ + + @abstractmethod + def commit(self) -> None: + """Commit the current local PostgreSQL transaction. + + A successful return means only that the concrete driver reported local + transaction commit; it does not imply distributed exactly-once delivery. + """ + + @abstractmethod + def rollback(self) -> None: + """Roll back the current local PostgreSQL transaction. + + Adapters must preserve PostgreSQL rollback behavior and must not convert + rollback failures into a successful application outcome. + """ + + @abstractmethod + def set_autocommit(self, enabled: bool) -> None: + """Select explicit connection autocommit behavior without attribute leakage. + + Existing runtime paths intentionally use both autocommit and explicit + transactions. Concrete adapters translate this operation into their + native API while preserving PostgreSQL transaction and session scope. + """ + + @abstractmethod + def is_closed(self) -> bool: + """Report whether this adapter knows the connection was locally closed. + + Cached owners use this signal to avoid intentionally reusing a connection + whose capability has already been released. This method is not a network + liveness probe: an unexpected transport or server failure must surface + from the attempted database operation and enter the bounded recovery path + rather than being guessed from driver-private state. + """ + + @abstractmethod + def close(self) -> None: + """Release the concrete database connection and its session authority. + + Implementations must not keep an implicit reusable connection alive + after callers intentionally close this capability. + """ + + @abstractmethod + def __enter__(self) -> PostgresConnectionPort: + """Enter the package-owned transaction context on this exact connection. + + Entry must not open a second session, implicitly commit, or change the + caller's transaction mode. The canonical context policy is defined at + this port so a replacement driver need not expose a proprietary context + manager extension merely to preserve pg-llm-batch behavior. + """ + + @abstractmethod + def __exit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + traceback: object | None, + ) -> bool | None: + """Commit normal exit, roll back exceptional exit, close, and propagate. + + Adapters must preserve these package transaction semantics even when the + concrete DB-API does not supply a connection context manager. Commit or + rollback failures remain visible, cleanup still runs, and application + exceptions must not be suppressed. + """ + + +class PostgresDriverPort(ABC): + """Define the PostgreSQL-driver anti-corruption layer required by the package. + + This port owns no model discovery, provider routing, LLM credentials, or + batch-provider selection. It exists solely to let pg-llm-batch replace a + concrete PostgreSQL client while keeping its database and tenant contracts + stable and testable. + """ + + @abstractmethod + def connect( + self, + dsn: str, + *, + connect_timeout_seconds: int | None = None, + ) -> PostgresConnectionPort: + """Open one synchronous PostgreSQL connection for a validated DSN. + + The concrete adapter must honor the requested finite positive timeout in + whole seconds when supplied and return a connection whose transaction and + session behavior conforms to ``PostgresConnectionPort``. + """ + + @abstractmethod + def parse_conninfo(self, dsn: str) -> Mapping[str, str]: + """Parse a PostgreSQL connection selector without exposing credentials. + + The returned mapping is used only for deterministic policy decisions; + callers remain responsible for rejecting credential-bearing selectors + where their boundary requires a credential-free DSN. + """ + + @abstractmethod + def make_conninfo(self, params: Mapping[str, str]) -> str: + """Render validated PostgreSQL connection parameters safely. + + Concrete adapters must use their reviewed conninfo quoting rules rather + than ad-hoc concatenation when values may contain PostgreSQL syntax. + """ + + @abstractmethod + def jsonb(self, value: object) -> object: + """Adapt one validated Python value for a PostgreSQL JSONB parameter. + + The adapter may return a driver-specific wrapper, but that wrapper must + remain confined behind this infrastructure boundary and out of domain + models and public package contracts. + """ + + @abstractmethod + def is_invalid_conninfo(self, error: BaseException) -> bool: + """Classify only a PostgreSQL connection-selector grammar failure. + + CLI and bootstrap boundaries need to normalize malformed DSN syntax + without importing one concrete driver's exception type. Implementations + must not broaden this category to connection, authentication, or runtime + database failures. + """ + + @abstractmethod + def is_undefined_function(self, error: BaseException) -> bool: + """Classify the PostgreSQL undefined-function error without leaking it. + + Token-counting fallback logic needs this narrow database-error category; + adapters must not broaden it to unrelated provider or application errors. + """ diff --git a/pg_llm_batch/postgres_driver_runtime.py b/pg_llm_batch/postgres_driver_runtime.py new file mode 100644 index 000000000..586a7dcfc --- /dev/null +++ b/pg_llm_batch/postgres_driver_runtime.py @@ -0,0 +1,40 @@ +"""Runtime selection for the retained PostgreSQL driver implementation. + +Concrete database-client authority belongs at one infrastructure boundary while +pg-llm-batch migrates away from Psycopg. Bounded contexts consume only +:class:`PostgresDriverPort`; this module lazily constructs the retained adapter +until a commercially admitted replacement is ready. Keeping the import lazy +also preserves explicit driver injection for candidate and degraded-mode tests. +""" + +from __future__ import annotations + +from .postgres_driver_port import PostgresDriverPort + + +class PostgresDriverUnavailableError(RuntimeError): + """Report that the retained PostgreSQL client cannot be constructed. + + The fixed diagnostic deliberately omits import paths, environment details, + DSNs, and credentials. Import failures unrelated to Psycopg are re-raised so + packaging defects are not misclassified as an optional-client absence. + """ + + +def retained_postgres_driver() -> PostgresDriverPort: + """Return the currently retained concrete driver behind the neutral port. + + Psycopg remains a temporary migration baseline only. The import lives here + so callers do not acquire a second concrete-driver dependency and the future + production replacement can be switched at one reviewed runtime boundary. + """ + try: + from .psycopg_driver_adapter import PsycopgDriverAdapter + except ModuleNotFoundError as exc: + missing_name = exc.name or "" + if missing_name != "psycopg" and not missing_name.startswith("psycopg."): + raise + raise PostgresDriverUnavailableError( + "Retained PostgreSQL driver is unavailable" + ) from None + return PsycopgDriverAdapter() diff --git a/pg_llm_batch/psycopg_driver_adapter.py b/pg_llm_batch/psycopg_driver_adapter.py new file mode 100644 index 000000000..a1aa3bf98 --- /dev/null +++ b/pg_llm_batch/psycopg_driver_adapter.py @@ -0,0 +1,260 @@ +"""Psycopg adapter for the provider-neutral PostgreSQL driver port. + +This module is a migration baseline, not the commercial replacement itself. It +encapsulates the PostgreSQL-client behavior that existing pg-llm-batch code +currently receives from Psycopg so a future permissively licensed adapter can be +verified against the same transaction, parameter-binding, conninfo, JSONB, and +error-classification contract before the LGPL-family runtime dependency is +removed. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +import psycopg +from psycopg import ProgrammingError +from psycopg.conninfo import conninfo_to_dict, make_conninfo +from psycopg.errors import UndefinedFunction +from psycopg.types.json import Jsonb + +from .postgres_driver_port import ( + PostgresConnectionPort, + PostgresCursorPort, + PostgresDriverPort, +) + + +class PsycopgDriverAdapterError(RuntimeError): + """Report a fixed adapter-contract failure without reflecting database data. + + The adapter uses this error only when a driver-facing primitive violates the + migration port itself, such as a non-boolean autocommit state, unsupported row + container, or a row-count value with the wrong Python type. PostgreSQL + execution errors continue to propagate through Psycopg so existing bounded + contexts can classify them. + """ + + +class PsycopgInvalidConninfoError(PsycopgDriverAdapterError): + """Identify conninfo grammar failures created at the adapter parsing boundary. + + Psycopg's public ``ProgrammingError`` class also represents server-side SQL + errors such as undefined tables and malformed statements. Wrapping only + failures raised by conninfo parsing/rendering prevents those unrelated + database errors from being misclassified as an invalid DSN. + """ + + +class PsycopgCursorAdapter(PostgresCursorPort): + """Wrap one PostgreSQL cursor while preserving canonical tuple row semantics. + + Package-authored query text and bound parameters are handed to the retained + cursor unchanged. Result rows are normalized from exact tuple/list containers + to tuples because current pg-llm-batch bounded contexts use positional tuple + identity and equality. This keeps a future DB-API driver that returns list rows + from silently changing application behavior. + """ + + def __init__(self, cursor: Any) -> None: + """Retain one Psycopg cursor behind the driver-neutral cursor contract.""" + self._cursor = cursor + + @staticmethod + def _normalize_result_row(row: object) -> tuple[object, ...]: + """Normalize one materialized DB-API row to the positional tuple contract.""" + if type(row) is tuple: + return row + if type(row) is list: + return tuple(row) + raise PsycopgDriverAdapterError("PostgreSQL driver result row is invalid") + + def execute( + self, + query: str, + params: object | None = None, + ) -> PsycopgCursorAdapter: + """Execute one query with Psycopg parameter binding and retain this wrapper.""" + self._cursor.execute(query, params) + return self + + def executemany( + self, + query: str, + params_seq: object, + ) -> PsycopgCursorAdapter: + """Execute one query for a parameter sequence without implicit commits.""" + self._cursor.executemany(query, params_seq) + return self + + def fetchone(self) -> tuple[object, ...] | None: + """Return one canonical tuple row, or ``None`` only for end-of-results.""" + row = self._cursor.fetchone() + if row is None: + return None + return self._normalize_result_row(row) + + def fetchmany(self, size: int) -> list[tuple[object, ...]]: + """Return at most the requested finite page and reject malformed results.""" + if type(size) is not int or size <= 0: + raise PsycopgDriverAdapterError("PostgreSQL driver fetch size is invalid") + rows = self._cursor.fetchmany(size) + try: + returned_count = len(rows) + except (TypeError, ValueError, OverflowError): + raise PsycopgDriverAdapterError( + "PostgreSQL driver fetch result is invalid" + ) from None + if returned_count > size: + raise PsycopgDriverAdapterError( + "PostgreSQL driver fetch result exceeds requested size" + ) + return [self._normalize_result_row(row) for row in rows] + + def fetchall(self) -> list[tuple[object, ...]]: + """Return bounded query results while rejecting malformed rows.""" + return [self._normalize_result_row(row) for row in self._cursor.fetchall()] + + def row_count(self) -> int | None: + """Return an exact non-negative count or normalize Psycopg's unknown sentinel.""" + value = self._cursor.rowcount + if type(value) is not int: + raise PsycopgDriverAdapterError("PostgreSQL driver row count is invalid") + if value == -1: + return None + if value < 0: + raise PsycopgDriverAdapterError("PostgreSQL driver row count is invalid") + return value + + def __enter__(self) -> PsycopgCursorAdapter: + """Enter the retained cursor context while preserving wrapper identity.""" + self._cursor.__enter__() + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + traceback: object | None, + ) -> bool | None: + """Delegate cursor cleanup and exception propagation to Psycopg.""" + return self._cursor.__exit__(exc_type, exc, traceback) + + +class PsycopgConnectionAdapter(PostgresConnectionPort): + """Wrap one Psycopg connection while preserving its session and transaction. + + The same retained raw connection backs cursor creation, direct execution, + commit, rollback, autocommit selection, and closed-state inspection. The + adapter therefore cannot silently move tenant-local ``set_config`` state to + another connection. + """ + + def __init__(self, connection: Any) -> None: + """Retain one Psycopg connection as the exact session capability.""" + self._connection = connection + + def cursor(self) -> PsycopgCursorAdapter: + """Create a wrapped cursor on this exact retained PostgreSQL connection.""" + return PsycopgCursorAdapter(self._connection.cursor()) + + def execute( + self, + query: str, + params: object | None = None, + ) -> PsycopgCursorAdapter: + """Execute through this connection without opening an implicit second one.""" + return PsycopgCursorAdapter(self._connection.execute(query, params)) + + def commit(self) -> None: + """Commit the current local PostgreSQL transaction through Psycopg.""" + self._connection.commit() + + def rollback(self) -> None: + """Roll back the current local PostgreSQL transaction through Psycopg.""" + self._connection.rollback() + + def set_autocommit(self, enabled: bool) -> None: + """Set Psycopg autocommit only from an exact boolean policy decision.""" + if type(enabled) is not bool: + raise PsycopgDriverAdapterError("PostgreSQL driver autocommit is invalid") + self._connection.autocommit = enabled + + def is_closed(self) -> bool: + """Return Psycopg's public closed-state signal without truthiness coercion.""" + value = self._connection.closed + if type(value) is not bool: + raise PsycopgDriverAdapterError("PostgreSQL driver closed state is invalid") + return value + + def close(self) -> None: + """Close the retained PostgreSQL connection and release its session state.""" + self._connection.close() + + def __enter__(self) -> PsycopgConnectionAdapter: + """Enter Psycopg's connection context while preserving wrapper identity.""" + self._connection.__enter__() + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + traceback: object | None, + ) -> bool | None: + """Delegate transaction-context exit and error propagation to Psycopg.""" + return self._connection.__exit__(exc_type, exc, traceback) + + +class PsycopgDriverAdapter(PostgresDriverPort): + """Expose existing Psycopg behavior through the migration anti-corruption port. + + This class establishes a parity baseline only. It does not make Psycopg an + approved commercial dependency and it does not acquire model/provider routing + authority from contextual-orchestrator. + """ + + def connect( + self, + dsn: str, + *, + connect_timeout_seconds: int | None = None, + ) -> PsycopgConnectionAdapter: + """Connect with an optional exact positive libpq timeout in whole seconds.""" + kwargs: dict[str, int] = {} + if connect_timeout_seconds is not None: + if type(connect_timeout_seconds) is not int or connect_timeout_seconds <= 0: + raise PsycopgDriverAdapterError("PostgreSQL driver timeout is invalid") + kwargs["connect_timeout"] = connect_timeout_seconds + return PsycopgConnectionAdapter(psycopg.connect(dsn, **kwargs)) + + def parse_conninfo(self, dsn: str) -> Mapping[str, str]: + """Parse conninfo and narrow Psycopg's broad ProgrammingError category.""" + try: + return conninfo_to_dict(dsn) + except ProgrammingError: + raise PsycopgInvalidConninfoError( + "PostgreSQL connection selector is invalid" + ) from None + + def make_conninfo(self, params: Mapping[str, str]) -> str: + """Render conninfo and narrow Psycopg's broad ProgrammingError category.""" + try: + return make_conninfo(**dict(params)) + except ProgrammingError: + raise PsycopgInvalidConninfoError( + "PostgreSQL connection selector is invalid" + ) from None + + def jsonb(self, value: object) -> Jsonb: + """Wrap a validated Python value in Psycopg's JSONB parameter adapter.""" + return Jsonb(value) + + def is_invalid_conninfo(self, error: BaseException) -> bool: + """Recognize only errors wrapped at the conninfo grammar boundary.""" + return isinstance(error, PsycopgInvalidConninfoError) + + def is_undefined_function(self, error: BaseException) -> bool: + """Recognize only Psycopg's PostgreSQL undefined-function error category.""" + return isinstance(error, UndefinedFunction) diff --git a/pg_llm_batch/token_counter.py b/pg_llm_batch/token_counter.py index e31b3d65d..e083a41bc 100644 --- a/pg_llm_batch/token_counter.py +++ b/pg_llm_batch/token_counter.py @@ -16,21 +16,17 @@ import logging from dataclasses import dataclass from io import StringIO +from threading import RLock from typing import Any, Dict, List, Optional, Tuple from .db import get_model_metadata from .exceptions import TokenLimitExceededError, ValidationError from .models import BatchRequest +from .postgres_driver_port import PostgresDriverPort +from .postgres_driver_runtime import retained_postgres_driver logger = logging.getLogger(__name__) -try: # pragma: no cover - optional dependency - import psycopg # type: ignore - from psycopg.errors import UndefinedFunction # type: ignore -except ImportError: # pragma: no cover - psycopg = None # type: ignore - UndefinedFunction = Exception # type: ignore - @dataclass(frozen=True) class _EncoderInfo: @@ -57,8 +53,15 @@ def __init__( *, config: Optional[Any] = None, buffer_percentage: Optional[int] = None, + postgres_driver: PostgresDriverPort | None = None, ) -> None: - """Initialize PostgreSQL token counting and configured batch limits.""" + """Initialize token counting through the centralized PostgreSQL driver boundary. + + Explicitly injected drivers remain authoritative for candidate and test + paths. Ordinary runtime construction acquires the retained implementation + from :mod:`postgres_driver_runtime`, so this bounded context no longer + owns a second concrete Psycopg import or connection fallback. + """ if not postgres_dsn: raise ValidationError( field="postgres_dsn", @@ -67,7 +70,11 @@ def __init__( ) self.postgres_dsn = postgres_dsn self.config = config - self._pg_conn: Optional["psycopg.Connection"] = None + self._postgres_driver = ( + postgres_driver if postgres_driver is not None else retained_postgres_driver() + ) + self._pg_conn: Optional[Any] = None + self._pg_connection_lock = RLock() self._pg_available: bool = False self._encoder_cache: Dict[str, _EncoderInfo] = {} @@ -119,8 +126,7 @@ def __init__( ), ) - if psycopg is not None: - self._pg_available = self._ensure_pg_tiktoken() + self._pg_available = self._ensure_pg_tiktoken() @staticmethod def _require_positive_limit(field: str, value: Any) -> int: @@ -159,17 +165,28 @@ def get_encoder(self, model: str) -> _EncoderInfo: return info def count_tokens(self, text: str, model: str) -> int: - """Count tokens through pg_tiktoken or fail when it is unavailable.""" + """Count tokens while serializing use of the retained PostgreSQL session. + + A replacement DB-API driver may permit module sharing without permitting + concurrent use of one connection. The counter intentionally retains one + autocommit session for repeated pg_tiktoken calls, so the lock protects + that exact session through execution, error classification, and cleanup + rather than assuming stronger driver thread semantics. + """ if not text: return 0 - if self._pg_available: - try: - return self._count_tokens_postgres(text, model) - except UndefinedFunction: - self._pg_available = False - logger.warning("pg_tiktoken extension/functions unavailable") - except Exception: # pragma: no cover - runtime DB variance - logger.debug("PostgreSQL token counting failed") + with self._pg_connection_lock: + if self._pg_available: + try: + return self._count_tokens_postgres(text, model) + except Exception as error: # pragma: no cover - runtime DB variance + if self._is_undefined_function(error): + self._pg_available = False + self.close() + logger.warning("pg_tiktoken extension/functions unavailable") + else: + self.close() + logger.debug("PostgreSQL token counting failed") raise RuntimeError( "Token counting requires pg_tiktoken. Enable the extension and pass a " "valid DSN." @@ -262,15 +279,16 @@ def split_oversized_batch( return batches def close(self) -> None: - """Close and clear the cached PostgreSQL token-counting connection.""" - conn = self._pg_conn - self._pg_conn = None - if conn is None: - return - try: - conn.close() - except Exception: - pass + """Close and clear the cached PostgreSQL token-counting connection safely.""" + with self._pg_connection_lock: + conn = self._pg_conn + self._pg_conn = None + if conn is None: + return + try: + conn.close() + except Exception: + pass def _resolve_config_value(self, category: str, key: str, default: Any) -> Any: """Read a config value from the KV store, returning the default on any failure.""" @@ -284,63 +302,72 @@ def _resolve_config_value(self, category: str, key: str, default: Any) -> Any: def _ensure_pg_tiktoken(self) -> bool: """Verify the pre-provisioned pg_tiktoken extension and functions read-only.""" - if psycopg is None: - return False - try: - conn = self._get_pg_conn() - with conn.cursor() as cur: - cur.execute( - """ - SELECT EXISTS ( - SELECT 1 - FROM pg_extension - WHERE extname = %s - ), - to_regprocedure('tiktoken_count(text,text)') IS NOT NULL, - to_regprocedure('tiktoken_encode(text,text)') IS NOT NULL - """, - ("pg_tiktoken",), - ) - row = cur.fetchone() - return bool(row and row == (True, True, True)) - except Exception: - self.close() - return False - - def _get_pg_conn(self) -> "psycopg.Connection": - """Return a cached autocommit PostgreSQL connection, reconnecting if closed.""" - assert psycopg is not None - if self._pg_conn is None or self._pg_conn.closed: - self._pg_conn = psycopg.connect(self.postgres_dsn) - self._pg_conn.autocommit = True - return self._pg_conn + with self._pg_connection_lock: + try: + conn = self._get_pg_conn() + with conn.cursor() as cur: + cur.execute( + """ + SELECT EXISTS ( + SELECT 1 + FROM pg_extension + WHERE extname = %s + ), + to_regprocedure('tiktoken_count(text,text)') IS NOT NULL, + to_regprocedure('tiktoken_encode(text,text)') IS NOT NULL + """, + ("pg_tiktoken",), + ) + row = cur.fetchone() + return bool(row and row == (True, True, True)) + except Exception: + self.close() + return False + + def _get_pg_conn(self) -> Any: + """Return a cached autocommit connection under the session reuse lock.""" + with self._pg_connection_lock: + if self._pg_conn is not None and not self._pg_conn.is_closed(): + return self._pg_conn + self._pg_conn = self._postgres_driver.connect(self.postgres_dsn) + self._pg_conn.set_autocommit(True) + return self._pg_conn + + def _is_undefined_function(self, error: BaseException) -> bool: + """Classify undefined-function failures through the selected driver boundary.""" + return self._postgres_driver.is_undefined_function(error) def _count_tokens_postgres(self, text: str, model: str) -> int: - """Count tokens for text via pg_tiktoken, falling back to tiktoken_encode.""" - if psycopg is None: - raise RuntimeError("PostgreSQL integration is unavailable") - conn = self._get_pg_conn() - tiktoken_name = self.get_encoder(model).tokenizer_name - with conn.cursor() as cur: - try: - cur.execute("SELECT tiktoken_count(%s, %s)", (tiktoken_name, text)) - row = cur.fetchone() - if row and row[0] is not None: - return int(row[0]) - except UndefinedFunction: - cur.execute( - "SELECT COUNT(*) FROM tiktoken_encode(%s, %s)", - (tiktoken_name, text), - ) - row = cur.fetchone() - if row and row[0] is not None: - return int(row[0]) - raise - return 0 + """Count tokens while retaining one non-concurrent PostgreSQL session.""" + with self._pg_connection_lock: + conn = self._get_pg_conn() + tiktoken_name = self.get_encoder(model).tokenizer_name + with conn.cursor() as cur: + try: + cur.execute("SELECT tiktoken_count(%s, %s)", (tiktoken_name, text)) + row = cur.fetchone() + if row and row[0] is not None: + return int(row[0]) + except Exception as error: + if not self._is_undefined_function(error): + raise + cur.execute( + "SELECT COUNT(*) FROM tiktoken_encode(%s, %s)", + (tiktoken_name, text), + ) + row = cur.fetchone() + if row and row[0] is not None: + return int(row[0]) + raise + return 0 def _get_tokenizer_from_db(self, model: str) -> Optional[str]: """Return the tokenizer model recorded in model metadata, or None if unset.""" - metadata = get_model_metadata(self.postgres_dsn, model) + metadata = get_model_metadata( + self.postgres_dsn, + model, + postgres_driver=self._postgres_driver, + ) if metadata and metadata.get("tokenizer_model"): return str(metadata["tokenizer_model"]) return None @@ -390,72 +417,43 @@ def reset(self) -> None: self.entries: List[Tuple[str, str, int]] = [] self.total_tokens = 0 self.record_count = 0 - self.byte_size = 0 - - def compute_tokens( - self, system_prompt: str, user_prompt: str - ) -> Tuple[int, int, int]: - """Return total, system, and user token counts for one prompt pair.""" - system_tokens = self.token_counter.count_tokens(system_prompt or "", self.model) - user_tokens = self.token_counter.count_tokens(user_prompt or "", self.model) - return system_tokens + user_tokens, system_tokens, user_tokens - - @staticmethod - def compute_byte_size(json_line: str) -> int: - """Return the UTF-8 byte size including the JSONL newline.""" - return len(json_line.encode("utf-8")) + 1 + self.total_bytes = 0 + self._payload = StringIO() - def would_exceed(self, tokens: int, byte_size: int) -> bool: - """Report whether adding a line would exceed any active limit.""" - if self.record_count == 0: + def can_add(self, jsonl_line: str, tokens: int) -> bool: + """Return whether an entry would fit every configured resource ceiling.""" + if type(jsonl_line) is not str: + return False + if type(tokens) is not int or tokens < 0: return False - if self.total_tokens + tokens > self.token_limit: - return True - if self.byte_size + byte_size > self.max_bytes: - return True + line_bytes = len((jsonl_line + "\n").encode("utf-8")) if self.record_count + 1 > self.max_records: - return True - return False + return False + if self.total_bytes + line_bytes > self.max_bytes: + return False + return self.total_tokens + tokens <= self.token_limit - def add_entry( - self, request_id: str, json_line: str, tokens: int, byte_size: int - ) -> None: - """Append one valid record and update aggregate counters.""" - if tokens > self.token_limit: - raise TokenLimitExceededError( - current_tokens=tokens, - limit_tokens=self.token_limit, - batch_id=request_id, - ) - if byte_size > self.max_bytes: - raise ValidationError( - field="byte_size", - value=byte_size, - reason=f"single JSONL record exceeds max_bytes={self.max_bytes}", - ) - self.entries.append((request_id, json_line, tokens)) + def add(self, request_id: str, jsonl_line: str, tokens: int) -> bool: + """Append one validated JSONL line when all resource ceilings permit it.""" + if not self.can_add(jsonl_line, tokens): + return False + line_bytes = len((jsonl_line + "\n").encode("utf-8")) + self.entries.append((request_id, jsonl_line, tokens)) + self._payload.write(jsonl_line) + self._payload.write("\n") self.total_tokens += tokens self.record_count += 1 - self.byte_size += byte_size - - def drain(self) -> Dict[str, Any]: - """Return accumulated metadata and reset the accumulator.""" - if not self.entries: - return {} - metadata = { - "record_count": self.record_count, - "total_tokens": self.total_tokens, - "request_ids": [rid for rid, _, _ in self.entries], - "lines": [line for _, line, _ in self.entries], - "byte_size": self.byte_size, - } - self.reset() - return metadata - - def to_jsonl(self) -> str: - """Return accumulated lines as newline-terminated JSONL text (in-memory).""" - buffer = StringIO() - for _, line, _ in self.entries: - buffer.write(line) - buffer.write("\n") - return buffer.getvalue() \ No newline at end of file + self.total_bytes += line_bytes + return True + + def content(self) -> str: + """Return the canonical newline-terminated JSONL payload.""" + return self._payload.getvalue() + + def is_empty(self) -> bool: + """Return whether no request has been accumulated.""" + return not self.entries + + def __len__(self) -> int: + """Return the number of accumulated requests.""" + return self.record_count diff --git a/tests/fake_postgres_driver_port.py b/tests/fake_postgres_driver_port.py new file mode 100644 index 000000000..980efed83 --- /dev/null +++ b/tests/fake_postgres_driver_port.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Driver-port adapter for legacy in-memory PostgreSQL unit-test fakes. + +Production bounded contexts no longer import Psycopg directly while the +commercial driver migration is in progress. These wrappers let the existing +in-memory SQL fake exercise the same ``PostgresDriverPort`` contract without +reintroducing concrete-client authority into product code. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +from pg_llm_batch.postgres_driver_port import PostgresConnectionPort + + +class _FakeCursorPort: + """Adapt one legacy fake cursor to the driver-neutral cursor contract.""" + + def __init__(self, cursor: Any) -> None: + self._cursor = cursor + + def execute(self, query: str, params: object | None = None) -> _FakeCursorPort: + self._cursor.execute(query, params) + return self + + def executemany(self, query: str, params_seq: object) -> _FakeCursorPort: + self._cursor.executemany(query, params_seq) + return self + + def fetchone(self) -> tuple[object, ...] | None: + row = self._cursor.fetchone() + return None if row is None else tuple(row) + + def fetchmany(self, size: int) -> list[tuple[object, ...]]: + fetchmany = getattr(self._cursor, "fetchmany", None) + if not callable(fetchmany): + return self.fetchall()[:size] + return [tuple(row) for row in fetchmany(size)] + + def fetchall(self) -> list[tuple[object, ...]]: + return [tuple(row) for row in self._cursor.fetchall()] + + def row_count(self) -> int | None: + value = getattr(self._cursor, "rowcount", None) + if value is None or value == -1: + return None + return int(value) + + def __enter__(self) -> _FakeCursorPort: + self._cursor.__enter__() + return self + + def __exit__(self, *exc: object) -> object: + return self._cursor.__exit__(*exc) + + +class _FakeConnectionPort: + """Adapt one legacy fake connection while preserving exact session identity.""" + + def __init__(self, connection: Any) -> None: + self._connection = connection + + def cursor(self) -> _FakeCursorPort: + return _FakeCursorPort(self._connection.cursor()) + + def execute(self, query: str, params: object | None = None) -> _FakeCursorPort: + cursor = self.cursor() + return cursor.execute(query, params) + + def commit(self) -> None: + self._connection.commit() + + def rollback(self) -> None: + rollback = getattr(self._connection, "rollback", None) + if callable(rollback): + rollback() + + def set_autocommit(self, enabled: bool) -> None: + self._connection.autocommit = enabled + + def is_closed(self) -> bool: + return bool(self._connection.closed) + + def close(self) -> None: + self._connection.close() + + def __enter__(self) -> _FakeConnectionPort: + self._connection.__enter__() + return self + + def __exit__(self, *exc: object) -> object: + return self._connection.__exit__(*exc) + + +class FakePsycopgDriverPort: + """Expose ``tests.conftest.FakePsycopg`` through ``PostgresDriverPort``. + + The fake intentionally implements only deterministic unit-test semantics. + Real PostgreSQL compatibility, RLS, recovery, and candidate admission remain + covered by their dedicated integration lanes rather than being inferred from + this in-memory adapter. + """ + + def __init__(self, psycopg_fake: Any) -> None: + self._psycopg_fake = psycopg_fake + + def connect( + self, + dsn: str, + *, + connect_timeout_seconds: int | None = None, + ) -> PostgresConnectionPort: + kwargs: dict[str, object] = {} + if connect_timeout_seconds is not None: + kwargs["connect_timeout"] = connect_timeout_seconds + return _FakeConnectionPort(self._psycopg_fake.connect(dsn, **kwargs)) + + def parse_conninfo(self, dsn: str) -> Mapping[str, str]: + return {"dsn": dsn} + + def make_conninfo(self, params: Mapping[str, str]) -> str: + return str(params.get("dsn", "")) + + def jsonb(self, value: object) -> object: + return value + + def is_invalid_conninfo(self, error: BaseException) -> bool: + return False + + def is_undefined_function(self, error: BaseException) -> bool: + return isinstance(error, self._psycopg_fake.errors.UndefinedFunction) diff --git a/tests/smoke_pg8000_candidate_postgres.py b/tests/smoke_pg8000_candidate_postgres.py new file mode 100644 index 000000000..e70da637a --- /dev/null +++ b/tests/smoke_pg8000_candidate_postgres.py @@ -0,0 +1,375 @@ +"""Exercise the exact pg8000 candidate against a real PostgreSQL boundary. + +This script is intentionally outside pytest discovery. CI installs one immutable +pg8000 candidate artifact and runs this smoke against the repository PostgreSQL +image without adding the candidate to the production dependency graph. The +checks cover the candidate URI, keyword, and explicit service connection +selectors, portable connection/cursor ACL, thread-affine connection use, +transaction, parameter, JSONB, UUID/timestamp, affected-row, narrow PostgreSQL +error classification, restore-catalog inspection, and transaction-local tenant +semantics that must be proven before candidate promotion. +""" + +from __future__ import annotations + +from datetime import datetime, timezone +from importlib import metadata +import os +from pathlib import Path +import uuid + +from pg8000 import dbapi + +from pg_llm_batch.pg8000_candidate_driver_port import Pg8000CandidateDriverAdapter +from pg_llm_batch.pg8000_candidate_service_file import Pg8000CandidateServiceFileResolver +from pg_llm_batch.pg8000_driver_candidate_jsonb import adapt_pg8000_jsonb +from pg_llm_batch.postgres_restore_acceptance import inspect_postgres_restore_catalog + +_EXPECTED_VERSION = "1.31.5" +_EXPECTED_DATABASE = "pgllm" +_EXPECTED_USER = "pgllm" +_CREDENTIAL_FREE_DSN = "postgresql://pgllm@127.0.0.1:5432/pgllm" + + +def _candidate_driver() -> Pg8000CandidateDriverAdapter: + """Bind the exact admitted pg8000 DB-API module to the candidate driver port.""" + return Pg8000CandidateDriverAdapter(dbapi) + + +def _candidate_password() -> str: + """Read the ephemeral CI credential without placing it in process arguments.""" + password_file = os.environ.get("PG8000_CANDIDATE_PASSWORD_FILE") + if not password_file: + raise RuntimeError("PG8000_CANDIDATE_PASSWORD_FILE is required") + try: + password = Path(password_file).read_text(encoding="utf-8") + except (OSError, UnicodeError): + raise RuntimeError("PG8000 candidate password file could not be read") from None + if not password: + raise RuntimeError("PG8000 candidate password file is empty") + return password + + +def _connection() -> object: + """Open one finite candidate connection from a private in-memory URI selector.""" + driver = _candidate_driver() + parameters = dict(driver.parse_conninfo(_CREDENTIAL_FREE_DSN)) + parameters["password"] = _candidate_password() + private_dsn = driver.make_conninfo(parameters) + return driver.connect(private_dsn, connect_timeout_seconds=5) + + +def _assert_keyword_and_service_selector_connections() -> None: + """Prove exact-artifact connection parity beyond the URI-only happy path. + + Unit tests establish grammar and precedence, but issue #322 requires the + replacement to preserve the selectors used by deployed PostgreSQL clients. + This probe therefore opens real PostgreSQL sessions through both the bounded + keyword grammar and the explicit caller-selected service-file resolver. The + temporary service file remains credential-free; the ephemeral password is + injected by the trusted in-process resolver so it is never written to disk. + """ + password = _candidate_password() + keyword_driver = _candidate_driver() + keyword_connection = keyword_driver.connect( + "host=127.0.0.1 port=5432 dbname=pgllm user=pgllm " + f"password={password}", + connect_timeout_seconds=5, + ) + try: + with keyword_connection.cursor() as cursor: + cursor.execute("SELECT current_database(), current_user") + if cursor.fetchone() != (_EXPECTED_DATABASE, _EXPECTED_USER): + raise AssertionError("candidate keyword selector connection changed") + finally: + keyword_connection.close() + + service_file = Path(os.environ["PG8000_CANDIDATE_PASSWORD_FILE"]).with_name( + "pg8000_candidate_service.conf" + ) + service_file.write_text( + "[candidate]\n" + "host=127.0.0.1\n" + "port=5432\n" + "dbname=pgllm\n" + "user=pgllm\n", + encoding="utf-8", + ) + try: + file_resolver = Pg8000CandidateServiceFileResolver(service_file) + + def resolve_service(service_name: str) -> dict[str, str]: + parameters = file_resolver(service_name) + parameters["password"] = password + return parameters + + service_driver = Pg8000CandidateDriverAdapter( + dbapi, + service_resolver=resolve_service, + ) + service_connection = service_driver.connect( + "service=candidate", + connect_timeout_seconds=5, + ) + try: + with service_connection.cursor() as cursor: + cursor.execute("SELECT current_database(), current_user") + if cursor.fetchone() != (_EXPECTED_DATABASE, _EXPECTED_USER): + raise AssertionError("candidate service selector connection changed") + finally: + service_connection.close() + finally: + service_file.unlink(missing_ok=True) + + +def _cleanup() -> None: + """Remove candidate-only database objects even after a prior interrupted smoke.""" + connection = _connection() + try: + connection.set_autocommit(True) + with connection.cursor() as cursor: + cursor.execute("DROP TABLE IF EXISTS pg8000_candidate_contract") + cursor.execute("DROP ROLE IF EXISTS pg8000_candidate_reader") + finally: + connection.close() + + +def _assert_restore_catalog_inspection() -> None: + """Prove the candidate can inspect the packaged restore catalog exactly. + + The production restore acceptance query binds finite Python lists through + ``ANY(%s)`` and consumes catalog booleans and tuple rows. Running that exact + query through pg8000 closes a driver-parity gap that unit adapters cannot + prove, without turning the candidate into the production runtime. + """ + connection = _connection() + try: + connection.set_autocommit(True) + evidence = inspect_postgres_restore_catalog(connection) + if evidence.required_table_count != 11: + raise AssertionError("candidate restore catalog table evidence changed") + if evidence.required_index_count != 2: + raise AssertionError("candidate restore catalog index evidence changed") + if evidence.lifecycle_rls_forced is not True: + raise AssertionError("candidate restore catalog RLS evidence changed") + finally: + connection.close() + + +def _prepare_rls_fixture() -> tuple[uuid.UUID, datetime]: + """Create an ephemeral RLS fixture and return exact typed evidence values.""" + evidence_uuid = uuid.uuid4() + evidence_time = datetime.now(timezone.utc).replace(microsecond=0) + connection = _connection() + try: + connection.set_autocommit(True) + with connection.cursor() as cursor: + cursor.execute("CREATE ROLE pg8000_candidate_reader NOLOGIN") + cursor.execute( + """ + CREATE TABLE pg8000_candidate_contract ( + tenant_scope TEXT NOT NULL, + evidence_uuid UUID NOT NULL, + evidence_time TIMESTAMPTZ NOT NULL, + evidence_json JSONB NOT NULL + ) + """ + ) + cursor.execute( + "ALTER TABLE pg8000_candidate_contract ENABLE ROW LEVEL SECURITY" + ) + cursor.execute( + "ALTER TABLE pg8000_candidate_contract FORCE ROW LEVEL SECURITY" + ) + cursor.execute( + """ + CREATE POLICY candidate_tenant_scope ON pg8000_candidate_contract + USING ( + tenant_scope = current_setting( + 'pg_llm_batch.tenant_scope', true + ) + ) + """ + ) + cursor.execute( + "GRANT SELECT ON pg8000_candidate_contract TO pg8000_candidate_reader" + ) + cursor.execute( + """ + INSERT INTO pg8000_candidate_contract + (tenant_scope, evidence_uuid, evidence_time, evidence_json) + VALUES (%s, %s, %s, %s), (%s, %s, %s, %s) + """, + ( + "tenant-a", + evidence_uuid, + evidence_time, + adapt_pg8000_jsonb({"candidate": "pg8000", "visible": True}), + "tenant-b", + uuid.uuid4(), + evidence_time, + adapt_pg8000_jsonb({"candidate": "pg8000", "visible": False}), + ), + ) + if cursor.row_count() != 2: + raise AssertionError("pg8000 candidate row-count evidence is not exact") + finally: + connection.close() + return evidence_uuid, evidence_time + + +def _assert_transaction_commit() -> None: + """Prove a successful package connection context commits a real write.""" + connection = _connection() + with connection as transaction: + with transaction.cursor() as cursor: + cursor.execute( + """ + UPDATE pg8000_candidate_contract + SET evidence_json = %s + WHERE tenant_scope = %s + """, + (adapt_pg8000_jsonb({"committed": True}), "tenant-b"), + ) + if cursor.row_count() != 1: + raise AssertionError("candidate commit probe did not update one row") + + verification = _connection() + try: + with verification.cursor() as cursor: + cursor.execute( + """ + SELECT evidence_json + FROM pg8000_candidate_contract + WHERE tenant_scope = %s + """, + ("tenant-b",), + ) + if cursor.fetchone() != ({"committed": True},): + raise AssertionError("candidate connection context did not commit") + finally: + verification.close() + + +def _assert_transaction_rollback() -> None: + """Prove an exceptional package connection context rolls a real write back.""" + connection = _connection() + try: + with connection as transaction: + with transaction.cursor() as cursor: + cursor.execute( + """ + UPDATE pg8000_candidate_contract + SET evidence_json = %s + WHERE tenant_scope = %s + """, + (adapt_pg8000_jsonb({"rolled_back": True}), "tenant-a"), + ) + if cursor.row_count() != 1: + raise AssertionError("candidate rollback probe did not update one row") + raise RuntimeError("candidate rollback probe") + except RuntimeError as exc: + if str(exc) != "candidate rollback probe": + raise + + verification = _connection() + try: + with verification.cursor() as cursor: + cursor.execute( + """ + SELECT evidence_json + FROM pg8000_candidate_contract + WHERE tenant_scope = %s + """, + ("tenant-a",), + ) + if cursor.fetchone() != ({"candidate": "pg8000", "visible": True},): + raise AssertionError("candidate connection context did not roll back") + finally: + verification.close() + + +def _assert_undefined_function_classification() -> None: + """Prove SQLSTATE-based undefined-function classification on real PostgreSQL.""" + driver = _candidate_driver() + connection = _connection() + try: + with connection.cursor() as cursor: + try: + cursor.execute("SELECT pg_llm_batch_candidate_missing_function()") + except BaseException as error: + if not driver.is_undefined_function(error): + raise AssertionError( + "candidate undefined-function classification changed" + ) from error + else: + raise AssertionError("candidate undefined-function probe unexpectedly exists") + connection.rollback() + finally: + connection.close() + + +def _assert_typed_rls_read( + expected_uuid: uuid.UUID, + expected_time: datetime, +) -> None: + """Prove transaction-local tenant scope and typed result semantics together.""" + connection = _connection() + with connection as transaction: + with transaction.cursor() as cursor: + cursor.execute("SET ROLE pg8000_candidate_reader") + cursor.execute( + "SELECT set_config('pg_llm_batch.tenant_scope', %s, true)", + ("tenant-a",), + ) + cursor.execute( + """ + SELECT tenant_scope, evidence_uuid, evidence_time, evidence_json + FROM pg8000_candidate_contract + ORDER BY tenant_scope + """ + ) + rows = cursor.fetchmany(1) + if len(rows) != 1: + raise AssertionError("RLS candidate read exceeded one visible tenant row") + tenant_scope, evidence_uuid, evidence_time, evidence_json = rows[0] + if tenant_scope != "tenant-a": + raise AssertionError("transaction-local tenant scope was not preserved") + if evidence_uuid != expected_uuid: + raise AssertionError("UUID adaptation changed candidate evidence") + if evidence_time != expected_time: + raise AssertionError("timestamp adaptation changed candidate evidence") + if evidence_json != {"candidate": "pg8000", "visible": True}: + raise AssertionError("JSONB adaptation changed candidate evidence") + if cursor.fetchmany(1): + raise AssertionError("RLS exposed another tenant through the candidate") + + +def main() -> None: + """Run exact-artifact and real-PostgreSQL candidate acceptance probes.""" + if metadata.version("pg8000") != _EXPECTED_VERSION: + raise AssertionError("unexpected pg8000 candidate version") + _candidate_driver() + _assert_keyword_and_service_selector_connections() + + connection = _connection() + with connection as transaction: + with transaction.cursor() as cursor: + cursor.execute("SELECT current_database(), current_user, %s::text", ("bound",)) + if cursor.fetchone() != (_EXPECTED_DATABASE, _EXPECTED_USER, "bound"): + raise AssertionError("candidate parameter/result semantics changed") + + _assert_restore_catalog_inspection() + _assert_undefined_function_classification() + _cleanup() + try: + evidence_uuid, evidence_time = _prepare_rls_fixture() + _assert_transaction_commit() + _assert_transaction_rollback() + _assert_typed_rls_read(evidence_uuid, evidence_time) + finally: + _cleanup() + + +if __name__ == "__main__": + main() diff --git a/tests/test_batch_assembly.py b/tests/test_batch_assembly.py index fc814d8a9..b631ed8d7 100644 --- a/tests/test_batch_assembly.py +++ b/tests/test_batch_assembly.py @@ -14,16 +14,23 @@ from pg_llm_batch.orchestrator import BatchPayload, PostgresBatchOrchestrator from pg_llm_batch.token_counter import TokenCounter from tests.conftest import FakePsycopg +from tests.fake_postgres_driver_port import FakePsycopgDriverPort @pytest.fixture() def fake_pg(monkeypatch): + """Route assembly tests through one shared driver-port fake.""" fake = FakePsycopg() - monkeypatch.setattr(tc_mod, "psycopg", fake) - monkeypatch.setattr(tc_mod, "UndefinedFunction", fake.errors.UndefinedFunction) - monkeypatch.setattr(db_mod, "psycopg", fake) - monkeypatch.setattr(orch_mod, "psycopg", fake) - monkeypatch.setattr(db_mod, "get_model_metadata", lambda dsn, model: None) + driver = FakePsycopgDriverPort(fake) + fake.driver = driver + monkeypatch.setattr(tc_mod, "retained_postgres_driver", lambda: driver) + monkeypatch.setattr(db_mod, "retained_postgres_driver", lambda: driver) + monkeypatch.setattr(orch_mod, "retained_postgres_driver", lambda: driver) + monkeypatch.setattr( + db_mod, + "get_model_metadata", + lambda dsn, model, *, postgres_driver=None: None, + ) return fake @@ -78,12 +85,9 @@ def test_assemble_payloads_splits_on_token_limit(fake_pg): assert all(p["record_count"] == 1 for p in payloads) -def test_orchestrator_requires_dsn_and_driver(monkeypatch, fake_pg): - with pytest.raises(RuntimeError, match="DSN and psycopg"): +def test_orchestrator_requires_dsn(fake_pg): + with pytest.raises(RuntimeError, match="Postgres DSN"): PostgresBatchOrchestrator("") - monkeypatch.setattr(orch_mod, "psycopg", None) - with pytest.raises(RuntimeError, match="DSN and psycopg"): - PostgresBatchOrchestrator("postgresql://x") def test_resolve_batch_uuid_direct_and_lookup(monkeypatch, fake_pg): @@ -175,8 +179,9 @@ def close(self): config = Config() class Counter: - def __init__(self, dsn, config): + def __init__(self, dsn, config, *, postgres_driver=None): assert (dsn, config) == ("postgresql://x", globals_config) + assert postgres_driver is fake_pg.driver self.effective_limit = 100 def close(self): @@ -185,7 +190,11 @@ def close(self): globals_config = config orch = PostgresBatchOrchestrator("postgresql://x") monkeypatch.setattr(fake_pg, "connect", lambda _dsn: Connection()) - monkeypatch.setattr(orch_mod, "PostgresConfigStore", lambda _dsn: globals_config) + monkeypatch.setattr( + orch_mod, + "PostgresConfigStore", + lambda _dsn, *, postgres_driver=None: globals_config, + ) monkeypatch.setattr(orch_mod, "TokenCounter", Counter) monkeypatch.setattr(orch, "_resolve_batch_uuid", lambda _key: "resolved") monkeypatch.setattr( @@ -219,7 +228,9 @@ def test_assemble_payloads_handles_empty_model_switch_and_null_user(fake_pg, mon monkeypatch.setattr( db_mod, "get_model_metadata", - lambda _dsn, model: {"mode": "embedding"} if model == "embed" else None, + lambda _dsn, model, *, postgres_driver=None: ( + {"mode": "embedding"} if model == "embed" else None + ), ) rows = [ ("r1", "ignored system", "vector", "embed"), @@ -315,7 +326,7 @@ def commit(self): def test_persist_payloads_separates_ready_and_overflow(monkeypatch, fake_pg): connection, executions, many = _persistence_connection() monkeypatch.setattr(fake_pg, "connect", lambda _dsn: connection) - monkeypatch.setattr(orch_mod, "Jsonb", lambda value: ("jsonb", value)) + monkeypatch.setattr(fake_pg.driver, "jsonb", lambda value: ("jsonb", value)) counter = TokenCounter("postgresql://x") counter.azure_max_files_per_job = 1 batch_uuid = "11111111-1111-1111-1111-111111111111" diff --git a/tests/test_batch_persistence_failures.py b/tests/test_batch_persistence_failures.py index bf91aa178..791d62bcf 100644 --- a/tests/test_batch_persistence_failures.py +++ b/tests/test_batch_persistence_failures.py @@ -3,8 +3,6 @@ from __future__ import annotations -from types import SimpleNamespace - import pytest from pg_llm_batch import orchestrator as orch_mod @@ -81,13 +79,77 @@ def commit(self): self.commits += 1 +class _CursorPort: + """Expose failure-mode row-count evidence through the driver cursor contract.""" + + def __init__(self, cursor: Cursor) -> None: + self._cursor = cursor + + def __enter__(self): + self._cursor.__enter__() + return self + + def __exit__(self, *exc): + return self._cursor.__exit__(*exc) + + def execute(self, sql, params=None): + self._cursor.execute(sql, params) + return self + + def executemany(self, sql, params): + self._cursor.executemany(sql, params) + return self + + def fetchone(self): + return self._cursor.fetchone() + + def fetchall(self): + return self._cursor.fetchall() + + def row_count(self): + return self._cursor.rowcount + + +class _ConnectionPort: + """Preserve one failure-mode transaction while satisfying the driver port.""" + + def __init__(self, connection: Connection) -> None: + self._connection = connection + + def __enter__(self): + self._connection.__enter__() + return self + + def __exit__(self, *exc): + return self._connection.__exit__(*exc) + + def cursor(self): + return _CursorPort(self._connection.cursor()) + + def set_autocommit(self, enabled): + self._connection.autocommit = enabled + + def commit(self): + self._connection.commit() + + +class _Driver: + """Return one deterministic failure connection through the driver boundary.""" + + def __init__(self, connection: Connection) -> None: + self.connection = connection + + def connect(self, _dsn, **_kwargs): + return _ConnectionPort(self.connection) + + def jsonb(self, value): + return value + + def _orchestrator(monkeypatch, mode: str): connection = Connection(mode) - monkeypatch.setattr( - orch_mod, - "psycopg", - SimpleNamespace(connect=lambda _dsn: connection), - ) + driver = _Driver(connection) + monkeypatch.setattr(orch_mod, "retained_postgres_driver", lambda: driver) return PostgresBatchOrchestrator("postgresql://x"), connection diff --git a/tests/test_candidate_wheel_license_verifier.py b/tests/test_candidate_wheel_license_verifier.py new file mode 100644 index 000000000..8088a3c40 --- /dev/null +++ b/tests/test_candidate_wheel_license_verifier.py @@ -0,0 +1,182 @@ +"""Regression tests for immutable candidate-wheel license admission. + +The pg8000 migration is intended to remove an LGPL-family production dependency. +Hash-pinning a candidate wheel closure proves artifact identity but does not prove +that every installed transitive dependency satisfies the repository's inbound +license policy. These tests keep that commercial-policy decision executable. +""" + +from __future__ import annotations + +import importlib.util +from pathlib import Path +import zipfile + +import pytest + + +_REPOSITORY_ROOT = Path(__file__).parents[1] +_TOOL_PATH = _REPOSITORY_ROOT / "tools" / "verify_candidate_wheel_licenses.py" + + +def _load_verifier(): + """Load the repository-owned verifier without turning tools into a package.""" + spec = importlib.util.spec_from_file_location("candidate_license_verifier", _TOOL_PATH) + if spec is None or spec.loader is None: + raise AssertionError("candidate license verifier could not be loaded") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _write_wheel(directory: Path, filename: str, *, name: str, version: str, license_lines: tuple[str, ...]) -> None: + """Write the smallest synthetic wheel metadata fixture needed by the verifier.""" + metadata = [ + "Metadata-Version: 2.4", + f"Name: {name}", + f"Version: {version}", + *license_lines, + "", + ] + dist_info = filename.split("-", 1)[0].replace("-", "_") + with zipfile.ZipFile(directory / filename, "w") as archive: + archive.writestr(f"{dist_info}-{version}.dist-info/METADATA", "\n".join(metadata)) + + +def _write_valid_closure(directory: Path) -> None: + """Create license metadata matching the exact pg8000 candidate closure.""" + _write_wheel( + directory, + "pg8000-1.31.5-py3-none-any.whl", + name="pg8000", + version="1.31.5", + license_lines=("License-Expression: BSD-3-Clause",), + ) + _write_wheel( + directory, + "python_dateutil-2.9.0.post0-py2.py3-none-any.whl", + name="python-dateutil", + version="2.9.0.post0", + license_lines=( + "License: Dual License", + "Classifier: License :: OSI Approved :: Apache Software License", + "Classifier: License :: OSI Approved :: BSD License", + ), + ) + _write_wheel( + directory, + "scramp-1.4.17-py3-none-any.whl", + name="scramp", + version="1.4.17", + license_lines=("License-Expression: MIT-0",), + ) + _write_wheel( + directory, + "asn1crypto-1.5.1-py2.py3-none-any.whl", + name="asn1crypto", + version="1.5.1", + license_lines=("License: MIT",), + ) + _write_wheel( + directory, + "six-1.17.0-py2.py3-none-any.whl", + name="six", + version="1.17.0", + license_lines=("License: MIT",), + ) + + +def test_exact_candidate_closure_requires_permissive_license_evidence(tmp_path: Path) -> None: + verifier = _load_verifier() + _write_valid_closure(tmp_path) + + verifier.verify_candidate_wheel_licenses(tmp_path) + + +def test_candidate_license_cli_accepts_platform_path_subclass(tmp_path: Path) -> None: + verifier = _load_verifier() + _write_valid_closure(tmp_path) + + assert verifier.main([str(tmp_path)]) == 0 + + +def test_candidate_closure_rejects_gpl_family_metadata_even_with_permissive_marker(tmp_path: Path) -> None: + verifier = _load_verifier() + _write_valid_closure(tmp_path) + wheel_path = tmp_path / "scramp-1.4.17-py3-none-any.whl" + wheel_path.unlink() + _write_wheel( + tmp_path, + wheel_path.name, + name="scramp", + version="1.4.17", + license_lines=( + "License-Expression: MIT-0", + "Classifier: License :: OSI Approved :: GNU Lesser General Public License v3 (LGPLv3)", + ), + ) + + with pytest.raises(verifier.CandidateWheelLicenseError, match="disallowed license"): + verifier.verify_candidate_wheel_licenses(tmp_path) + + +def test_candidate_closure_rejects_permissive_token_inside_unapproved_word(tmp_path: Path) -> None: + verifier = _load_verifier() + _write_valid_closure(tmp_path) + wheel_path = tmp_path / "six-1.17.0-py2.py3-none-any.whl" + wheel_path.unlink() + _write_wheel( + tmp_path, + wheel_path.name, + name="six", + version="1.17.0", + license_lines=("License: Limited proprietary terms",), + ) + + with pytest.raises(verifier.CandidateWheelLicenseError, match="approved license evidence"): + verifier.verify_candidate_wheel_licenses(tmp_path) + + +def test_candidate_closure_rejects_missing_positive_license_evidence(tmp_path: Path) -> None: + verifier = _load_verifier() + _write_valid_closure(tmp_path) + wheel_path = tmp_path / "six-1.17.0-py2.py3-none-any.whl" + wheel_path.unlink() + _write_wheel( + tmp_path, + wheel_path.name, + name="six", + version="1.17.0", + license_lines=("License: UNKNOWN",), + ) + + with pytest.raises(verifier.CandidateWheelLicenseError, match="approved license evidence"): + verifier.verify_candidate_wheel_licenses(tmp_path) + + +def test_candidate_closure_rejects_unexpected_wheel_set(tmp_path: Path) -> None: + verifier = _load_verifier() + _write_valid_closure(tmp_path) + _write_wheel( + tmp_path, + "unexpected-1.0-py3-none-any.whl", + name="unexpected", + version="1.0", + license_lines=("License: MIT",), + ) + + with pytest.raises(verifier.CandidateWheelLicenseError, match="wheel set is invalid"): + verifier.verify_candidate_wheel_licenses(tmp_path) + + +def test_candidate_license_gate_runs_before_candidate_install() -> None: + """CI must verify license metadata before any candidate wheel is installed.""" + workflow = (_REPOSITORY_ROOT / ".github" / "workflows" / "ci.yml").read_text( + encoding="utf-8" + ) + verification_step = "- name: Verify pg8000 candidate dependency licenses" + install_step = "- name: Install exact candidate closure into the CI environment" + + assert verification_step in workflow + assert "python tools/verify_candidate_wheel_licenses.py /tmp/pg8000-candidate" in workflow + assert workflow.index(verification_step) < workflow.index(install_step) diff --git a/tests/test_checkpoint_store_driver_port.py b/tests/test_checkpoint_store_driver_port.py new file mode 100644 index 000000000..844b51e2c --- /dev/null +++ b/tests/test_checkpoint_store_driver_port.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Regression tests for checkpoint persistence through the PostgreSQL driver port.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import pytest + +import pg_llm_batch.checkpoint_store as checkpoint_store +from pg_llm_batch.checkpoint_store import ( + PostgresBatchResultCheckpointStore, + apply_result_checkpoint_schema, +) + + +class _PortCursor: + """Provide the minimal cursor behavior needed by the checkpoint boundary.""" + + def __init__(self, calls: list[tuple[str, object | None]]) -> None: + self.calls = calls + self.result: object | None = None + + def execute(self, query: str, params: object | None = None) -> "_PortCursor": + """Record parameterized SQL without interpolating caller values.""" + self.calls.append((query, params)) + self.result = (params[0],) if query.startswith("SELECT set_config") and params else None + return self + + def fetchone(self) -> object | None: + """Return the bounded result from the previous fake operation.""" + return self.result + + def __enter__(self) -> "_PortCursor": + """Retain cursor identity across the context-manager boundary.""" + return self + + def __exit__(self, *_exc: object) -> None: + """Release no external resources in the deterministic fake.""" + return None + + +class _PortConnection: + """Expose one retained fake connection for driver-port regression tests.""" + + def __init__(self, calls: list[tuple[str, object | None]]) -> None: + self.calls = calls + self.commit_count = 0 + + def cursor(self) -> _PortCursor: + """Create a cursor bound to this exact fake connection.""" + return _PortCursor(self.calls) + + def commit(self) -> None: + """Record explicit package-owned transaction commit authority.""" + self.commit_count += 1 + + def __enter__(self) -> "_PortConnection": + """Retain connection identity across the context-manager boundary.""" + return self + + def __exit__(self, *_exc: object) -> None: + """Release no external resources in the deterministic fake.""" + return None + + +class _DriverPortFake: + """Connect checkpoint operations without exposing a concrete database client.""" + + def __init__(self) -> None: + self.dsns: list[str] = [] + self.calls: list[tuple[str, object | None]] = [] + self.connections: list[_PortConnection] = [] + + def connect(self, dsn: str, **_kwargs: Any) -> _PortConnection: + """Return one connection and preserve the exact validated DSN.""" + self.dsns.append(dsn) + connection = _PortConnection(self.calls) + self.connections.append(connection) + return connection + + +def _deny_default_driver_path(monkeypatch: pytest.MonkeyPatch) -> None: + """Fail if an explicitly injected driver silently reacquires the runtime default.""" + + def fail_default_driver() -> None: + raise AssertionError("default PostgreSQL runtime driver was reached") + + monkeypatch.setattr( + checkpoint_store, + "retained_postgres_driver", + fail_default_driver, + ) + + +def test_checkpoint_store_load_uses_injected_driver_port_without_default_driver( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A migrated store must reach tenant SQL only through its injected driver.""" + _deny_default_driver_path(monkeypatch) + driver = _DriverPortFake() + store = PostgresBatchResultCheckpointStore( + "postgresql://unit", + tenant_scope="tenant-a", + postgres_driver=driver, # type: ignore[arg-type] + ) + + assert store.load("worker-a", "batch-1", "default") is None + assert driver.dsns == ["postgresql://unit"] + assert driver.calls[0][1] == ("tenant-a",) + assert driver.calls[1][1] == ("tenant-a", "worker-a", "default", "batch-1") + + +def test_checkpoint_schema_application_uses_injected_driver_port_without_default_driver( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Schema migration must run through the injected driver without fallback.""" + _deny_default_driver_path(monkeypatch) + migration = tmp_path / "checkpoint.sql" + migration.write_text("CREATE TABLE checkpoint_probe (probe_id BIGINT);", encoding="utf-8") + driver = _DriverPortFake() + + apply_result_checkpoint_schema( + "postgresql://unit", + str(migration), + postgres_driver=driver, # type: ignore[arg-type] + ) + + assert driver.dsns == ["postgresql://unit"] + assert driver.calls == [("CREATE TABLE checkpoint_probe (probe_id BIGINT);", None)] + assert driver.connections[0].commit_count == 1 diff --git a/tests/test_cli_postgres_driver_port.py b/tests/test_cli_postgres_driver_port.py new file mode 100644 index 000000000..c50343f9f --- /dev/null +++ b/tests/test_cli_postgres_driver_port.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CLI regressions for the PostgreSQL driver-migration boundary.""" + +from __future__ import annotations + +import ast +from pathlib import Path + +import pytest + +from pg_llm_batch import cli +from pg_llm_batch import postgres_driver_runtime + + +class _CandidateConninfoError(Exception): + """Candidate-driver parse failure used to exercise CLI normalization.""" + + +class _CandidateDriver: + """Minimal candidate parser double for the CLI migration seam.""" + + def __init__(self) -> None: + self.parsed_values: list[str] = [] + + def __bool__(self) -> bool: + """Remain falsy so dependency selection cannot rely on object truthiness.""" + return False + + def parse_conninfo(self, value: str) -> dict[str, str]: + """Parse bounded fixtures without delegating to Psycopg.""" + self.parsed_values.append(value) + if value == "invalid": + raise _CandidateConninfoError + if value == "credential-bearing": + return {"host": "db.internal", "password": "redacted-fixture"} + return {"host": "db.internal", "dbname": "batch"} + + def is_invalid_conninfo(self, error: BaseException) -> bool: + """Classify only this double's explicit conninfo failure.""" + return type(error) is _CandidateConninfoError + + +def test_cli_parser_uses_injected_postgres_driver_for_dsn() -> None: + """Candidate validation must not route DSN parsing back through Psycopg.""" + driver = _CandidateDriver() + + parser = cli.build_parser(postgres_driver=driver) + args = parser.parse_args(["health", "--dsn", "candidate-selector"]) + + assert args.dsn == "candidate-selector" + assert driver.parsed_values == ["candidate-selector"] + + +def test_cli_parser_normalizes_candidate_conninfo_failure() -> None: + """Driver-specific parse errors remain bounded argparse diagnostics.""" + driver = _CandidateDriver() + parser = cli.build_parser(postgres_driver=driver) + + with pytest.raises(SystemExit) as exc_info: + parser.parse_args(["health", "--dsn", "invalid"]) + + assert exc_info.value.code == 2 + + +def test_cli_parser_rejects_candidate_reported_credential_fields() -> None: + """Driver migration must preserve the credential-free argv contract.""" + driver = _CandidateDriver() + parser = cli.build_parser(postgres_driver=driver) + + with pytest.raises(SystemExit) as exc_info: + parser.parse_args(["health", "--dsn", "credential-bearing"]) + + assert exc_info.value.code == 2 + + +def test_cli_default_driver_delegates_to_runtime_owner(monkeypatch) -> None: + """The CLI must not retain a second concrete-driver construction authority.""" + driver = _CandidateDriver() + monkeypatch.setattr( + postgres_driver_runtime, + "retained_postgres_driver", + lambda: driver, + ) + + assert cli._default_postgres_driver() is driver + + +def test_cli_module_has_no_eager_psycopg_import() -> None: + """Importing the CLI must not itself require the retained legacy driver.""" + source = Path(cli.__file__).read_text(encoding="utf-8") + tree = ast.parse(source) + + direct_psycopg_imports = [ + node + for node in ast.walk(tree) + if ( + isinstance(node, ast.Import) + and any( + alias.name == "psycopg" or alias.name.startswith("psycopg.") + for alias in node.names + ) + ) + or ( + isinstance(node, ast.ImportFrom) + and node.module is not None + and (node.module == "psycopg" or node.module.startswith("psycopg.")) + ) + ] + + assert direct_psycopg_imports == [] diff --git a/tests/test_compose_bootstrap_driver_port.py b/tests/test_compose_bootstrap_driver_port.py new file mode 100644 index 000000000..cebde1b58 --- /dev/null +++ b/tests/test_compose_bootstrap_driver_port.py @@ -0,0 +1,139 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Regression tests for Compose bootstrap through the PostgreSQL driver port.""" + +from __future__ import annotations + +from collections.abc import Mapping +from pathlib import Path + +from pg_llm_batch import compose_bootstrap +from pg_llm_batch import postgres_driver_runtime + + +class _BootstrapDriver: + """Capture conninfo and health-service use without importing a concrete driver.""" + + def __init__(self) -> None: + self.parsed: list[str] = [] + self.rendered: list[dict[str, str]] = [] + + def parse_conninfo(self, dsn: str) -> Mapping[str, str]: + """Return the credential-free target fields expected by the bootstrap.""" + self.parsed.append(dsn) + return { + "user": "pgllm", + "host": "postgres", + "port": "5432", + "dbname": "pgllm", + } + + def make_conninfo(self, params: Mapping[str, str]) -> str: + """Record the exact private parameter map and return an opaque DSN.""" + snapshot = dict(params) + self.rendered.append(snapshot) + return "driver-private-dsn" + + +def test_build_private_dsn_uses_injected_driver_without_legacy_renderer() -> None: + """Mounted-secret assembly must be usable after the Psycopg renderer is removed.""" + driver = _BootstrapDriver() + + private_dsn = compose_bootstrap._build_private_dsn( + "postgresql://pgllm@postgres:5432/pgllm", + "private-password", + postgres_driver=driver, # type: ignore[arg-type] + ) + + assert private_dsn == "driver-private-dsn" + assert driver.parsed == ["postgresql://pgllm@postgres:5432/pgllm"] + assert driver.rendered == [ + { + "user": "pgllm", + "host": "postgres", + "port": "5432", + "dbname": "pgllm", + "password": "private-password", + } + ] + + +def test_build_private_dsn_uses_default_driver_boundary_without_direct_renderer( + monkeypatch, +) -> None: + """Default secret assembly must select the driver port rather than Psycopg helpers.""" + driver = _BootstrapDriver() + monkeypatch.setattr(compose_bootstrap, "_default_postgres_driver", lambda: driver) + + private_dsn = compose_bootstrap._build_private_dsn( + "postgresql://pgllm@postgres:5432/pgllm", + "private-password", + ) + + assert private_dsn == "driver-private-dsn" + assert driver.parsed == ["postgresql://pgllm@postgres:5432/pgllm"] + assert driver.rendered == [ + { + "user": "pgllm", + "host": "postgres", + "port": "5432", + "dbname": "pgllm", + "password": "private-password", + } + ] + + +def test_compose_default_driver_delegates_to_runtime_owner(monkeypatch) -> None: + """Compose must not retain a second concrete-driver construction authority.""" + driver = _BootstrapDriver() + monkeypatch.setattr( + postgres_driver_runtime, + "retained_postgres_driver", + lambda: driver, + ) + + assert compose_bootstrap._default_postgres_driver() is driver + + +def test_run_compose_health_forwards_one_driver_to_dsn_and_health_boundaries( + tmp_path: Path, + monkeypatch, +) -> None: + """One selected driver must own both secret-safe DSN assembly and readiness I/O.""" + password_file = tmp_path / "database-password" + password_file.write_text("private-password", encoding="utf-8") + driver = _BootstrapDriver() + observed: dict[str, object] = {} + + monkeypatch.setattr( + compose_bootstrap, + "resolve_dsn", + lambda explicit=None: "postgresql://pgllm@postgres:5432/pgllm", + ) + + def capture_health( + dsn: str, + host: str, + port: int, + *, + postgres_driver=None, + ) -> None: + observed.update( + dsn=dsn, + host=host, + port=port, + postgres_driver=postgres_driver, + ) + + monkeypatch.setattr(compose_bootstrap, "serve_healthz", capture_health) + + compose_bootstrap.run_compose_health( + password_file, + postgres_driver=driver, # type: ignore[arg-type] + ) + + assert observed == { + "dsn": "driver-private-dsn", + "host": "0.0.0.0", + "port": 8080, + "postgres_driver": driver, + } diff --git a/tests/test_config.py b/tests/test_config.py index 859c55415..d1fc15a5d 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -3,6 +3,8 @@ from __future__ import annotations +from typing import Any + import pytest from pg_llm_batch import config as cfg_mod @@ -11,10 +13,45 @@ from tests.conftest import FakePsycopg +class _FakeDriverConnection: + """Adapt the legacy unit-test connection to the driver-neutral store surface.""" + + def __init__(self, connection: Any) -> None: + self._connection = connection + + @property + def closed(self) -> bool: + """Expose the underlying deterministic closed state for cleanup assertions.""" + return self._connection.closed + + def cursor(self): + """Return the exact cursor backed by the shared fake KV state.""" + return self._connection.cursor() + + def set_autocommit(self, enabled: bool) -> None: + """Apply the store's explicit autocommit policy to the legacy fake.""" + self._connection.autocommit = enabled + + def close(self) -> None: + """Close the underlying fake connection.""" + self._connection.close() + + +class _FakeDriver: + """Provide the shared FakePsycopg state through the runtime driver contract.""" + + def __init__(self, client: FakePsycopg) -> None: + self.client = client + + def connect(self, *args: Any, **kwargs: Any) -> _FakeDriverConnection: + """Return one driver-neutral wrapper around the existing deterministic fake.""" + return _FakeDriverConnection(self.client.connect(*args, **kwargs)) + + @pytest.fixture() def fake_pg(monkeypatch): fake = FakePsycopg() - monkeypatch.setattr(cfg_mod, "psycopg", fake) + monkeypatch.setattr(cfg_mod, "retained_postgres_driver", lambda: _FakeDriver(fake)) return fake @@ -25,16 +62,13 @@ def test_config_requires_dsn(fake_pg): def test_config_defaults_seeded_and_typed(fake_pg): store = PostgresConfigStore("postgresql://x") - # int coercion from stored string assert store.get("token_limits", "per_batch") == 5_000_000_000 - # bool coercion assert store.get("optimization", "smart_batching") is True def test_config_set_get_roundtrip(fake_pg): store = PostgresConfigStore("postgresql://x") store.set("gateway", "base_url", "https://gw.example/v1") - # bypass cache to prove it persisted to the backing table store.cache.clear() assert store.get("gateway", "base_url") == "https://gw.example/v1" @@ -42,7 +76,6 @@ def test_config_set_get_roundtrip(fake_pg): def test_secret_store_base64_without_key(fake_pg, caplog): store = SecretStore("postgresql://x", fernet_key=None) store.set_secret("gateway_api_key.default", "sk-secret-123") - # stored obfuscated, not plaintext stored_value = fake_pg.store.secrets["gateway_api_key.default"][0] assert stored_value != "sk-secret-123" assert store.get_secret("gateway_api_key.default") == "sk-secret-123" @@ -57,11 +90,9 @@ def test_secret_store_fernet_encrypts_at_rest(fake_pg): value, is_encrypted = fake_pg.store.secrets["gateway_api_key.default"] assert is_encrypted is True assert value != "sk-abc" - # a wrong key cannot decrypt other = SecretStore("postgresql://x", fernet_key=Fernet.generate_key().decode()) with pytest.raises(Exception): other.get_secret("gateway_api_key.default") - # the right key can assert store.get_secret("gateway_api_key.default") == "sk-abc" @@ -121,13 +152,22 @@ def test_config_missing_lookup_show_and_factory(fake_pg): def test_store_constructor_requires_dependency_and_dsn(monkeypatch, fake_pg): - monkeypatch.setattr(cfg_mod, "psycopg", None) + def unavailable_driver(): + raise cfg_mod.PostgresDriverUnavailableError( + "Retained PostgreSQL driver is unavailable" + ) + + monkeypatch.setattr(cfg_mod, "retained_postgres_driver", unavailable_driver) with pytest.raises(ConfigError, match="psycopg is required"): PostgresConfigStore("postgresql://x") with pytest.raises(ConfigError, match="psycopg is required"): SecretStore("postgresql://x") - monkeypatch.setattr(cfg_mod, "psycopg", fake_pg) + monkeypatch.setattr( + cfg_mod, + "retained_postgres_driver", + lambda: _FakeDriver(fake_pg), + ) with pytest.raises(ConfigError, match="DSN"): PostgresConfigStore("") with pytest.raises(ConfigError, match="DSN"): diff --git a/tests/test_config_driver_port.py b/tests/test_config_driver_port.py new file mode 100644 index 000000000..9aedb7eaa --- /dev/null +++ b/tests/test_config_driver_port.py @@ -0,0 +1,112 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Regression tests for configuration persistence through the PostgreSQL port.""" + +from __future__ import annotations + +from typing import Any + +import pytest + +from pg_llm_batch.config import PostgresConfigStore, SecretStore +from tests.conftest import FakeCursor, FakeKVStore + + +class _ConfigConnection: + """Expose KV persistence through the driver-neutral connection surface.""" + + def __init__(self, store: FakeKVStore) -> None: + self.store = store + self.autocommit_values: list[bool] = [] + self.closed = False + + def cursor(self) -> FakeCursor: + """Return a cursor backed by the shared deterministic KV store.""" + return FakeCursor(self.store) + + def set_autocommit(self, enabled: bool) -> None: + """Record the store's explicit autocommit policy decision.""" + self.autocommit_values.append(enabled) + + def close(self) -> None: + """Record deterministic release of this connection capability.""" + self.closed = True + + +class _ConfigDriver: + """Return distinct connections backed by one deterministic database state.""" + + def __init__(self) -> None: + self.store = FakeKVStore() + self.dsns: list[str] = [] + self.connections: list[_ConfigConnection] = [] + + def connect(self, dsn: str, **_kwargs: Any) -> _ConfigConnection: + """Capture one validated DSN and create a connection on shared state.""" + self.dsns.append(dsn) + connection = _ConfigConnection(self.store) + self.connections.append(connection) + return connection + + +def test_config_store_uses_injected_driver_without_concrete_import() -> None: + """Configuration CRUD must not require the retained concrete client when injected.""" + driver = _ConfigDriver() + + store = PostgresConfigStore( + "postgresql://example", + postgres_driver=driver, # type: ignore[arg-type] + ) + try: + store.set("batch_size", "default", 321) + assert store.get("batch_size", "default") == 321 + assert driver.dsns == ["postgresql://example"] + assert driver.connections[0].autocommit_values == [True] + finally: + store.close() + + assert driver.connections[0].closed is True + + +def test_secret_store_uses_injected_driver_without_concrete_import() -> None: + """Secret persistence must retain the same DB seam without concrete imports.""" + driver = _ConfigDriver() + + store = SecretStore( + "postgresql://example", + postgres_driver=driver, # type: ignore[arg-type] + ) + try: + store.set_secret("provider.key", "secret-value") + assert store.get_secret("provider.key") == "secret-value" + assert driver.dsns == ["postgresql://example"] + assert driver.connections[0].autocommit_values == [True] + finally: + store.close() + + assert driver.connections[0].closed is True + + +def test_config_store_closes_connection_when_autocommit_setup_fails() -> None: + """A replacement-driver setup failure must not leak the opened DB connection.""" + class _BrokenConnection(_ConfigConnection): + def set_autocommit(self, enabled: bool) -> None: + """Fail after connection creation to exercise constructor cleanup.""" + super().set_autocommit(enabled) + raise RuntimeError("autocommit setup failed") + + class _BrokenDriver(_ConfigDriver): + def connect(self, dsn: str, **_kwargs: Any) -> _BrokenConnection: + """Return the observable broken connection for cleanup verification.""" + self.dsns.append(dsn) + connection = _BrokenConnection(self.store) + self.connections.append(connection) + return connection + + driver = _BrokenDriver() + with pytest.raises(RuntimeError, match="autocommit setup failed"): + PostgresConfigStore( + "postgresql://example", + postgres_driver=driver, # type: ignore[arg-type] + ) + + assert driver.connections[0].closed is True diff --git a/tests/test_connection_lifecycle.py b/tests/test_connection_lifecycle.py index 7f51474e5..8a4fabd88 100644 --- a/tests/test_connection_lifecycle.py +++ b/tests/test_connection_lifecycle.py @@ -3,6 +3,7 @@ from __future__ import annotations +from threading import RLock from types import SimpleNamespace from typing import Any @@ -54,7 +55,7 @@ class _InitializationConnection: """Expose whether a partially initialized store releases its connection.""" def __init__(self) -> None: - """Start open with autocommit disabled like a new psycopg connection.""" + """Start open with autocommit disabled like a new database connection.""" self.autocommit = False self.closed = False @@ -63,14 +64,43 @@ def close(self) -> None: self.closed = True +class _InitializationConnectionPort: + """Expose initialization cleanup through the driver-neutral connection API.""" + + def __init__(self, connection: _InitializationConnection) -> None: + self._connection = connection + + def set_autocommit(self, enabled: bool) -> None: + self._connection.autocommit = enabled + + def close(self) -> None: + self._connection.close() + + +class _InitializationDriver: + """Return one observable initialization connection through the runtime port.""" + + def __init__(self, connection: _InitializationConnection) -> None: + self._connection = connection + + def connect(self, _dsn: str, **_kwargs: Any) -> _InitializationConnectionPort: + return _InitializationConnectionPort(self._connection) + + class _OwnedConfig: """Record whether the orchestrator closes its owned config store.""" instances: list["_OwnedConfig"] = [] - def __init__(self, dsn: str) -> None: - """Record the explicit DSN and register this owned store.""" + def __init__( + self, + dsn: str, + *, + postgres_driver: object | None = None, + ) -> None: + """Record the explicit DSN, shared driver, and owned store lifecycle.""" self.dsn = dsn + self.postgres_driver = postgres_driver self.closed = False self.instances.append(self) @@ -84,10 +114,17 @@ class _OwnedCounter: instances: list["_OwnedCounter"] = [] - def __init__(self, dsn: str, *, config: _OwnedConfig) -> None: + def __init__( + self, + dsn: str, + *, + config: _OwnedConfig, + postgres_driver: object | None = None, + ) -> None: """Record constructor ownership and expose the runtime token limit.""" self.dsn = dsn self.config = config + self.postgres_driver = postgres_driver self.effective_limit = 100 self.closed = False self.instances.append(self) @@ -102,7 +139,7 @@ def _prepare_owned_orchestrator(monkeypatch: Any) -> PostgresBatchOrchestrator: _OwnedConfig.instances = [] _OwnedCounter.instances = [] driver = SimpleNamespace(connect=lambda _dsn: _QueryConnection()) - monkeypatch.setattr(orchestrator_module, "psycopg", driver) + monkeypatch.setattr(orchestrator_module, "retained_postgres_driver", lambda: driver) monkeypatch.setattr(orchestrator_module, "PostgresConfigStore", _OwnedConfig) monkeypatch.setattr(orchestrator_module, "TokenCounter", _OwnedCounter) orchestrator = PostgresBatchOrchestrator("postgresql://example") @@ -154,8 +191,14 @@ def test_prepare_batches_closes_config_when_counter_construction_fails( """Config ownership must be released even when token setup cannot finish.""" orchestrator = _prepare_owned_orchestrator(monkeypatch) - def fail_counter(_dsn: str, *, config: _OwnedConfig) -> _OwnedCounter: + def fail_counter( + _dsn: str, + *, + config: _OwnedConfig, + postgres_driver: object | None = None, + ) -> _OwnedCounter: assert config is _OwnedConfig.instances[0] + assert postgres_driver is orchestrator._postgres_driver raise RuntimeError("counter construction failed") monkeypatch.setattr(orchestrator_module, "TokenCounter", fail_counter) @@ -172,11 +215,8 @@ def test_config_store_constructor_closes_connection_after_setup_failure( ) -> None: """A failed config-store setup must release the connection it already acquired.""" connection = _InitializationConnection() - monkeypatch.setattr( - config_module, - "psycopg", - SimpleNamespace(connect=lambda _dsn: connection), - ) + driver = _InitializationDriver(connection) + monkeypatch.setattr(config_module, "retained_postgres_driver", lambda: driver) def fail_table_setup(_store: PostgresConfigStore) -> None: raise RuntimeError("config setup failed") @@ -194,11 +234,8 @@ def test_secret_store_constructor_closes_connection_after_setup_failure( ) -> None: """A failed secret-store setup must release the connection it already acquired.""" connection = _InitializationConnection() - monkeypatch.setattr( - config_module, - "psycopg", - SimpleNamespace(connect=lambda _dsn: connection), - ) + driver = _InitializationDriver(connection) + monkeypatch.setattr(config_module, "retained_postgres_driver", lambda: driver) def fail_table_setup(_store: SecretStore) -> None: raise RuntimeError("secret setup failed") @@ -215,6 +252,7 @@ def test_token_counter_close_releases_cached_connection() -> None: """Closing a counter must release and clear its cached PostgreSQL connection.""" closed: list[str] = [] counter = object.__new__(TokenCounter) + counter._pg_connection_lock = RLock() counter._pg_conn = SimpleNamespace(close=lambda: closed.append("closed")) counter.close() @@ -227,6 +265,7 @@ def test_token_counter_close_releases_cached_connection() -> None: def test_token_counter_close_clears_connection_after_driver_failure() -> None: """Driver cleanup failure must not retain the unusable cached connection.""" counter = object.__new__(TokenCounter) + counter._pg_connection_lock = RLock() def fail_close() -> None: raise RuntimeError("driver close failed") diff --git a/tests/test_db.py b/tests/test_db.py index 57130e3ad..bf6988191 100644 --- a/tests/test_db.py +++ b/tests/test_db.py @@ -42,7 +42,9 @@ def commit(self): self.driver.commits += 1 -class _Psycopg: +class _Driver: + """Minimal database driver fake for default and injected port paths.""" + def __init__(self, row=None, error=None): self.row = row self.error = error @@ -57,9 +59,23 @@ def connect(self, dsn): return _Connection(self) +def _use_default_driver(monkeypatch: pytest.MonkeyPatch, driver: _Driver) -> None: + """Route the package default through the canonical runtime selector seam.""" + monkeypatch.setattr(db, "retained_postgres_driver", lambda: driver) + + +def _deny_default_driver(monkeypatch: pytest.MonkeyPatch) -> None: + """Fail if an explicitly injected database operation reacquires the default.""" + + def fail_default_driver(): + raise AssertionError("default PostgreSQL runtime driver was reached") + + monkeypatch.setattr(db, "retained_postgres_driver", fail_default_driver) + + def test_apply_schema_executes_packaged_file(monkeypatch, tmp_path): - driver = _Psycopg() - monkeypatch.setattr(db, "psycopg", driver) + driver = _Driver() + _use_default_driver(monkeypatch, driver) schema = tmp_path / "schema.sql" schema.write_text("CREATE TABLE snake_case_name (id int);", encoding="utf-8") monkeypatch.setattr(db, "SCHEMA_PATH", schema) @@ -70,10 +86,25 @@ def test_apply_schema_executes_packaged_file(monkeypatch, tmp_path): assert driver.commits == 1 +def test_apply_schema_uses_injected_driver_without_default_driver(monkeypatch, tmp_path): + """Schema bootstrap must honor an injected driver without hidden reacquisition.""" + driver = _Driver() + _deny_default_driver(monkeypatch) + schema = tmp_path / "schema.sql" + schema.write_text("CREATE TABLE snake_case_name (id int);", encoding="utf-8") + monkeypatch.setattr(db, "SCHEMA_PATH", schema) + + db.apply_schema("postgresql://x", postgres_driver=driver) + + assert driver.connections == ["postgresql://x"] + assert driver.executions == [("CREATE TABLE snake_case_name (id int);", None)] + assert driver.commits == 1 + + def test_apply_schema_refuses_caller_selected_sql(monkeypatch, tmp_path): """Caller-controlled local files must not acquire arbitrary SQL authority.""" - driver = _Psycopg() - monkeypatch.setattr(db, "psycopg", driver) + driver = _Driver() + _use_default_driver(monkeypatch, driver) untrusted_schema = tmp_path / "untrusted.sql" untrusted_schema.write_text("DROP TABLE llm_requests;", encoding="utf-8") @@ -93,20 +124,39 @@ def test_apply_schema_refuses_caller_selected_sql(monkeypatch, tmp_path): ], ) def test_load_virtual_payload_preserves_canonical_jsonl(monkeypatch, stored, expected): - driver = _Psycopg((stored,)) - monkeypatch.setattr(db, "psycopg", driver) + driver = _Driver((stored,)) + _use_default_driver(monkeypatch, driver) assert db.load_virtual_payload("postgresql://x", "file-1") == expected assert driver.executions[0][1] == ("file-1",) +def test_load_virtual_payload_uses_injected_driver_without_default_driver(monkeypatch): + """Virtual payload reads must honor the explicit driver boundary.""" + stored = {"text": '{"id":1}\n', "line_count": 1} + driver = _Driver((stored,)) + _deny_default_driver(monkeypatch) + + assert ( + db.load_virtual_payload( + "postgresql://x", + "file-1", + postgres_driver=driver, + ) + == '{"id":1}\n' + ) + assert driver.connections == ["postgresql://x"] + assert driver.executions[0][1] == ("file-1",) + + def test_load_virtual_payload_returns_none_when_missing(monkeypatch): - monkeypatch.setattr(db, "psycopg", _Psycopg(None)) + driver = _Driver(None) + _use_default_driver(monkeypatch, driver) assert db.load_virtual_payload("postgresql://x", "missing") is None def test_model_metadata_normalizes_mode_and_handles_absence(monkeypatch): - driver = _Psycopg((" CHAT ", "o200k_base")) - monkeypatch.setattr(db, "psycopg", driver) + driver = _Driver((" CHAT ", "o200k_base")) + _use_default_driver(monkeypatch, driver) assert db.get_model_metadata("postgresql://x", "gpt-4o") == { "mode": "chat", "tokenizer_model": "o200k_base", @@ -123,14 +173,36 @@ def test_model_metadata_normalizes_mode_and_handles_absence(monkeypatch): assert db.get_model_metadata("postgresql://x", "") is None +def test_model_metadata_uses_injected_driver_without_default_driver(monkeypatch): + """Tokenizer metadata lookup must honor the explicit driver boundary.""" + driver = _Driver((" CHAT ", "o200k_base")) + _deny_default_driver(monkeypatch) + + assert db.get_model_metadata( + "postgresql://x", + "gpt-4o", + postgres_driver=driver, + ) == { + "mode": "chat", + "tokenizer_model": "o200k_base", + } + assert driver.connections == ["postgresql://x"] + + def test_model_metadata_driver_failure_is_nonfatal(monkeypatch, caplog): - monkeypatch.setattr(db, "psycopg", _Psycopg(error=OSError("database down"))) + driver = _Driver(error=OSError("database down")) + _use_default_driver(monkeypatch, driver) with caplog.at_level("DEBUG"): assert db.get_model_metadata("postgresql://x", "gpt-4o") is None assert "database down" in caplog.text -def test_database_access_requires_psycopg(monkeypatch): - monkeypatch.setattr(db, "psycopg", None) - with pytest.raises(RuntimeError, match="psycopg is required"): +def test_database_access_requires_runtime_driver(monkeypatch): + """Default database access must fail when the canonical runtime selector fails.""" + + def fail_default_driver(): + raise RuntimeError("PostgreSQL runtime driver is required") + + monkeypatch.setattr(db, "retained_postgres_driver", fail_default_driver) + with pytest.raises(RuntimeError, match="runtime driver is required"): db.load_virtual_payload("postgresql://x", "file") diff --git a/tests/test_health.py b/tests/test_health.py index e93625805..b245256c6 100644 --- a/tests/test_health.py +++ b/tests/test_health.py @@ -43,24 +43,25 @@ def cursor(self): return _Cursor(self._rows) -class _Psycopg: - """Minimal psycopg facade returning fixed health rows.""" +class _Driver: + """Minimal retained-driver facade returning fixed health rows.""" def __init__(self, rows): self._rows = rows - def connect(self, _dsn, *, connect_timeout): - assert connect_timeout == 5 + def connect(self, _dsn, *, connect_timeout_seconds): + assert connect_timeout_seconds == 5 return _Connection(self._rows) +def _select_driver(monkeypatch, driver) -> None: + """Bind one retained readiness driver without exposing concrete-client imports.""" + monkeypatch.setattr(health, "retained_postgres_driver", lambda: driver) + + def test_missing_required_component_is_reported_not_ready(monkeypatch): """A partial health function result must never pass by vacuous truth.""" - monkeypatch.setattr( - health, - "psycopg", - _Psycopg([("database", True, "connected")]), - ) + _select_driver(monkeypatch, _Driver([("database", True, "connected")])) report = health.check_health("postgresql://example") @@ -76,9 +77,14 @@ def test_missing_required_component_is_reported_not_ready(monkeypatch): } -def test_health_dependency_and_database_failures_include_reason(monkeypatch): - """Dependency and connection failures are explicit rather than hidden.""" - monkeypatch.setattr(health, "psycopg", None) +def test_health_dependency_and_database_failures_are_bounded(monkeypatch): + """Dependency absence is explicit while runtime failures stay content-free.""" + def unavailable_driver(): + raise health.PostgresDriverUnavailableError( + "Retained PostgreSQL driver is unavailable" + ) + + monkeypatch.setattr(health, "retained_postgres_driver", unavailable_driver) report = health.check_health("postgresql://example") assert report == { "ready": False, @@ -87,15 +93,26 @@ def test_health_dependency_and_database_failures_include_reason(monkeypatch): ], } - class BrokenPsycopg: + class BrokenDriver: @staticmethod - def connect(_dsn, *, connect_timeout): - raise OSError(f"connection refused after {connect_timeout}s") + def connect(_dsn, *, connect_timeout_seconds): + raise OSError( + f"private-dsn-sentinel after {connect_timeout_seconds}s" + ) - monkeypatch.setattr(health, "psycopg", BrokenPsycopg()) + _select_driver(monkeypatch, BrokenDriver()) report = health.check_health("postgresql://example") - assert report["ready"] is False - assert "connection refused after 5s" in report["components"][0]["detail"] + assert report == { + "ready": False, + "components": [ + { + "component": "database", + "is_ready": False, + "detail": "database readiness check failed", + } + ], + } + assert "private-dsn-sentinel" not in repr(report) def test_health_requires_every_required_component(monkeypatch): @@ -106,17 +123,19 @@ def test_health_requires_every_required_component(monkeypatch): ("com_config", True, "ready"), ("optional_metrics", False, "disabled"), ] - monkeypatch.setattr(health, "psycopg", _Psycopg(rows)) + _select_driver(monkeypatch, _Driver(rows)) assert health.check_health("postgresql://example")["ready"] is True rows[1] = ("pg_tiktoken", False, "extension unavailable") - monkeypatch.setattr(health, "psycopg", _Psycopg(rows)) + _select_driver(monkeypatch, _Driver(rows)) assert health.check_health("postgresql://example")["ready"] is False def test_serve_healthz_reports_status_body_and_not_found(monkeypatch): - """The HTTP wrapper emits JSON readiness and a strict 404 elsewhere.""" + """The HTTP wrapper emits JSON readiness and forwards the selected DB driver.""" events = [] + observed = {} + driver = object() class FakeHTTPServer: def __init__(self, address, handler_class): @@ -142,13 +161,19 @@ def serve_forever(self): handler = self.handler_class.__new__(self.handler_class) assert handler.log_message("ignored") is None + def fake_check_health(_dsn, *, postgres_driver=None): + observed["driver"] = postgres_driver + return {"ready": False, "components": []} + monkeypatch.setattr("http.server.HTTPServer", FakeHTTPServer) - monkeypatch.setattr( - health, - "check_health", - lambda _dsn: {"ready": False, "components": []}, + monkeypatch.setattr(health, "check_health", fake_check_health) + health.serve_healthz( + "postgresql://example", + host="127.0.0.1", + port=8090, + postgres_driver=driver, # type: ignore[arg-type] ) - health.serve_healthz("postgresql://example", host="127.0.0.1", port=8090) + assert observed["driver"] is driver assert ("address", ("127.0.0.1", 8090)) in events assert ("/other", "status", 404) in events assert ("/healthz/", "status", 503) in events diff --git a/tests/test_health_driver_port.py b/tests/test_health_driver_port.py new file mode 100644 index 000000000..1f4657369 --- /dev/null +++ b/tests/test_health_driver_port.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Regression tests for readiness checks through the PostgreSQL driver port.""" + +from __future__ import annotations + +from typing import Any + +from pg_llm_batch import health + + +class _HealthCursor: + """Expose deterministic readiness rows through the driver-neutral cursor shape.""" + + def __init__(self) -> None: + self.executions: list[tuple[str, object | None]] = [] + + def execute(self, query: str, params: object | None = None) -> "_HealthCursor": + """Record the package-authored health query without modifying it.""" + self.executions.append((query, params)) + return self + + def fetchall(self) -> list[tuple[object, ...]]: + """Return all required healthy component rows as canonical tuples.""" + return [ + ("database", True, "connected"), + ("pg_tiktoken", True, "installed"), + ("com_config", True, "ready"), + ] + + def __enter__(self) -> "_HealthCursor": + """Retain the cursor identity during the readiness transaction.""" + return self + + def __exit__(self, *_exc: object) -> None: + """Release no external resources in the deterministic fake.""" + return None + + +class _HealthConnection: + """Retain one cursor for the driver-neutral readiness connection.""" + + def __init__(self) -> None: + self.cursor_value = _HealthCursor() + + def cursor(self) -> _HealthCursor: + """Return the cursor bound to this exact connection.""" + return self.cursor_value + + def __enter__(self) -> "_HealthConnection": + """Retain connection identity across context-manager entry.""" + return self + + def __exit__(self, *_exc: object) -> None: + """Release no external resources in the deterministic fake.""" + return None + + +class _HealthDriver: + """Capture readiness connection parameters without a concrete client import.""" + + def __init__(self) -> None: + self.calls: list[tuple[str, int | None]] = [] + self.connection = _HealthConnection() + + def connect( + self, + dsn: str, + *, + connect_timeout_seconds: int | None = None, + ) -> _HealthConnection: + """Preserve the exact DSN and bounded five-second readiness timeout.""" + self.calls.append((dsn, connect_timeout_seconds)) + return self.connection + + +def test_check_health_uses_injected_driver_without_concrete_import() -> None: + """Readiness must work entirely through an explicitly injected driver seam.""" + driver = _HealthDriver() + + report = health.check_health( + "postgresql://example", + postgres_driver=driver, # type: ignore[arg-type] + ) + + assert report["ready"] is True + assert driver.calls == [("postgresql://example", 5)] + assert driver.connection.cursor_value.executions == [ + ("SELECT component, is_ready, detail FROM pg_llm_batch_health_check()", None) + ] + + +def test_check_health_bounds_injected_driver_failures() -> None: + """Replacement-driver failures remain bounded without reflecting connection data.""" + secret_sentinel = "postgresql://user:private-password@db.example/batch" + + class _BrokenDriver: + def connect(self, _dsn: str, **_kwargs: Any) -> None: + raise OSError(f"connection refused for {secret_sentinel}") + + report = health.check_health( + "postgresql://example", + postgres_driver=_BrokenDriver(), # type: ignore[arg-type] + ) + + assert report == { + "ready": False, + "components": [ + { + "component": "database", + "is_ready": False, + "detail": "database readiness check failed", + } + ], + } + assert secret_sentinel not in repr(report) diff --git a/tests/test_orchestrator_batch_key_authority.py b/tests/test_orchestrator_batch_key_authority.py index 8258177d9..cc5e95c1e 100644 --- a/tests/test_orchestrator_batch_key_authority.py +++ b/tests/test_orchestrator_batch_key_authority.py @@ -8,7 +8,6 @@ import pytest -from pg_llm_batch import orchestrator as orchestrator_module from pg_llm_batch.exceptions import ValidationError from pg_llm_batch.orchestrator import PostgresBatchOrchestrator @@ -33,7 +32,6 @@ def __str__(self) -> str: ], ) def test_non_string_batch_key_fails_before_database_io( - monkeypatch: pytest.MonkeyPatch, invalid_key: object, ) -> None: """Only exact strings may choose UUID/path authority before PostgreSQL access.""" @@ -43,12 +41,11 @@ def forbidden_connect(_dsn: str) -> Any: database_calls.append("connect") raise AssertionError("database I/O must not occur for a non-string selector") - monkeypatch.setattr( - orchestrator_module, - "psycopg", - SimpleNamespace(connect=forbidden_connect), + driver = SimpleNamespace(connect=forbidden_connect) + orchestrator = PostgresBatchOrchestrator( + "postgresql://example", + postgres_driver=driver, ) - orchestrator = PostgresBatchOrchestrator("postgresql://example") with pytest.raises(ValidationError) as caught: orchestrator._resolve_batch_uuid(invalid_key) # type: ignore[arg-type] @@ -58,20 +55,17 @@ def forbidden_connect(_dsn: str) -> Any: assert database_calls == [] -def test_valid_uuid_string_preserves_exact_selector_without_database_io( - monkeypatch: pytest.MonkeyPatch, -) -> None: +def test_valid_uuid_string_preserves_exact_selector_without_database_io() -> None: """A valid exact UUID string must remain the direct authoritative selector.""" database_calls: list[str] = [] - monkeypatch.setattr( - orchestrator_module, - "psycopg", - SimpleNamespace( - connect=lambda _dsn: database_calls.append("connect") - or (_ for _ in ()).throw(AssertionError("unexpected database lookup")) - ), + driver = SimpleNamespace( + connect=lambda _dsn: database_calls.append("connect") + or (_ for _ in ()).throw(AssertionError("unexpected database lookup")) + ) + orchestrator = PostgresBatchOrchestrator( + "postgresql://example", + postgres_driver=driver, ) - orchestrator = PostgresBatchOrchestrator("postgresql://example") batch_key = "00000000-0000-0000-0000-000000000123" assert orchestrator._resolve_batch_uuid(batch_key) == batch_key @@ -120,19 +114,14 @@ def cursor(self) -> _LookupCursor: return _LookupCursor(self.lookup_values) -def test_exact_non_uuid_string_is_used_as_path_key_without_rewriting( - monkeypatch: pytest.MonkeyPatch, -) -> None: +def test_exact_non_uuid_string_is_used_as_path_key_without_rewriting() -> None: """An exact non-UUID string must retain byte-for-byte path-key identity.""" lookup_values: list[str] = [] - monkeypatch.setattr( - orchestrator_module, - "psycopg", - SimpleNamespace( - connect=lambda _dsn: _LookupConnection(lookup_values), - ), + driver = SimpleNamespace(connect=lambda _dsn: _LookupConnection(lookup_values)) + orchestrator = PostgresBatchOrchestrator( + "postgresql://example", + postgres_driver=driver, ) - orchestrator = PostgresBatchOrchestrator("postgresql://example") batch_key = "memory://batch/Case-Sensitive-Key" assert orchestrator._resolve_batch_uuid(batch_key) == ( diff --git a/tests/test_orchestrator_driver_port.py b/tests/test_orchestrator_driver_port.py new file mode 100644 index 000000000..129dcd431 --- /dev/null +++ b/tests/test_orchestrator_driver_port.py @@ -0,0 +1,302 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Driver-port regressions for PostgreSQL batch assembly and persistence.""" + +from __future__ import annotations + +from typing import Any + +import pytest + +from pg_llm_batch import db +from pg_llm_batch import orchestrator as orchestrator_module +from pg_llm_batch.orchestrator import PostgresBatchOrchestrator + + +class _Cursor: + """Expose deterministic rows for orchestrator driver-port acceptance.""" + + def __init__(self, driver: _Driver) -> None: + self.driver = driver + self.next_one: tuple[object, ...] | None = None + self.next_all: list[tuple[object, ...]] = [] + self._row_count = 0 + + def __enter__(self) -> _Cursor: + return self + + def __exit__(self, *_exc: Any) -> None: + return None + + def execute(self, query: str, params: object | None = None) -> _Cursor: + self.driver.executions.append((query, params)) + self._row_count = 0 + if "input_file_path" in query: + self.next_one = ("22222222-2222-2222-2222-222222222222",) + elif "FROM llm_requests" in query and "ORDER BY created_at" in query: + self.next_all = list(self.driver.request_rows) + elif "SELECT queue_uuid" in query: + self.next_one = ("33333333-3333-3333-3333-333333333333",) + elif "SELECT file_path" in query: + self.next_all = [] + elif "RETURNING file_uuid" in query: + self.next_one = ("44444444-4444-4444-4444-444444444444",) + elif "UPDATE llm_requests" in query: + assert isinstance(params, tuple) + self._row_count = len(params[2]) + return self + + def executemany(self, query: str, params_seq: object) -> _Cursor: + self.driver.many.append((query, list(params_seq))) # type: ignore[arg-type] + return self + + def fetchone(self) -> tuple[object, ...] | None: + return self.next_one + + def fetchall(self) -> list[tuple[object, ...]]: + return list(self.next_all) + + def row_count(self) -> int: + return self._row_count + + +class _Connection: + """Retain one fake transaction and explicit autocommit state.""" + + def __init__(self, driver: _Driver) -> None: + self.driver = driver + self.autocommit_values: list[bool] = [] + self.commits = 0 + self.closed = False + + def __enter__(self) -> _Connection: + return self + + def __exit__(self, *_exc: Any) -> None: + return None + + def cursor(self) -> _Cursor: + return _Cursor(self.driver) + + def set_autocommit(self, enabled: bool) -> None: + self.autocommit_values.append(enabled) + + def commit(self) -> None: + self.commits += 1 + + def rollback(self) -> None: + return None + + def is_closed(self) -> bool: + return self.closed + + def close(self) -> None: + self.closed = True + + +class _Driver: + """Minimal concrete-driver-free port fake for the orchestrator boundary.""" + + def __init__(self) -> None: + self.executions: list[tuple[str, object | None]] = [] + self.many: list[tuple[str, list[object]]] = [] + self.connections: list[_Connection] = [] + self.request_rows: list[tuple[object, ...]] = [] + self.jsonb_values: list[object] = [] + + def connect( + self, + dsn: str, + *, + connect_timeout_seconds: int | None = None, + ) -> _Connection: + assert dsn == "postgresql://x" + assert connect_timeout_seconds is None + connection = _Connection(self) + self.connections.append(connection) + return connection + + def jsonb(self, value: object) -> object: + self.jsonb_values.append(value) + return ("jsonb", value) + + +def _deny_default_driver(monkeypatch: pytest.MonkeyPatch) -> None: + """Fail if explicit orchestrator injection silently reacquires runtime default.""" + + def fail_default_driver(): + raise AssertionError("default PostgreSQL runtime driver was reached") + + monkeypatch.setattr( + orchestrator_module, + "retained_postgres_driver", + fail_default_driver, + ) + + +def test_orchestrator_accepts_injected_driver_without_default_driver( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Batch lookup must remain usable through an explicit replacement driver.""" + driver = _Driver() + _deny_default_driver(monkeypatch) + + orchestrator = PostgresBatchOrchestrator( + "postgresql://x", + postgres_driver=driver, + ) + + assert orchestrator._resolve_batch_uuid("input.jsonl") == ( + "22222222-2222-2222-2222-222222222222" + ) + assert len(driver.connections) == 1 + + +def test_assemble_payloads_passes_driver_to_model_metadata( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Model metadata reads must preserve the explicitly selected driver.""" + driver = _Driver() + _deny_default_driver(monkeypatch) + metadata_calls: list[tuple[str, str, object]] = [] + + def _metadata( + dsn: str, + model: str, + *, + postgres_driver: object = None, + ) -> dict[str, str]: + metadata_calls.append((dsn, model, postgres_driver)) + return {"mode": "embedding"} + + monkeypatch.setattr(db, "get_model_metadata", _metadata) + + class _Counter: + effective_limit = 100 + azure_max_records_per_file = 10 + azure_max_bytes_per_file = 10_000 + + def count_tokens(self, text: str, model: str) -> int: + return 1 if text else 0 + + orchestrator = PostgresBatchOrchestrator( + "postgresql://x", + postgres_driver=driver, + ) + payloads = orchestrator._assemble_payloads( + _Counter(), # type: ignore[arg-type] + [("req-1", "ignored", "content", "embed-model")], + ) + + assert metadata_calls == [("postgresql://x", "embed-model", driver)] + assert payloads[0]["record_count"] == 1 + + +def test_prepare_batches_propagates_driver_to_store_and_token_counter( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Preparation must keep one explicit driver boundary across its DB helpers.""" + driver = _Driver() + _deny_default_driver(monkeypatch) + calls: list[tuple[str, object]] = [] + + class _Config: + def close(self) -> None: + calls.append(("config_close", self)) + + config = _Config() + + def _config_factory( + dsn: str, + *, + postgres_driver: object = None, + ) -> _Config: + assert dsn == "postgresql://x" + calls.append(("config_driver", postgres_driver)) + return config + + class _Counter: + effective_limit = 100 + azure_max_files_per_job = 1 + + def __init__( + self, + dsn: str, + *, + config: object, + postgres_driver: object = None, + ) -> None: + assert dsn == "postgresql://x" + assert config is globals_config + calls.append(("counter_driver", postgres_driver)) + + def close(self) -> None: + calls.append(("counter_close", self)) + + globals_config = config + monkeypatch.setattr(orchestrator_module, "PostgresConfigStore", _config_factory) + monkeypatch.setattr(orchestrator_module, "TokenCounter", _Counter) + + orchestrator = PostgresBatchOrchestrator( + "postgresql://x", + postgres_driver=driver, + ) + monkeypatch.setattr( + orchestrator, + "_resolve_batch_uuid", + lambda _batch_key: "11111111-1111-1111-1111-111111111111", + ) + monkeypatch.setattr(orchestrator, "_assemble_payloads", lambda _counter, _rows: []) + monkeypatch.setattr( + orchestrator, + "_persist_payloads", + lambda _payloads, batch_key, _counter: { + "ready": [batch_key], + "overflow": [], + }, + ) + + result = orchestrator.prepare_batches(batch_uuid="source.jsonl") + + assert result["ready"] == ["11111111-1111-1111-1111-111111111111"] + assert ("config_driver", driver) in calls + assert ("counter_driver", driver) in calls + assert any("FROM llm_requests" in query for query, _params in driver.executions) + assert sum(name.endswith("_close") for name, _value in calls) == 2 + + +def test_persist_payloads_uses_driver_jsonb_transaction_and_row_count( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Persistence must use only the selected driver's JSONB and row-count seams.""" + driver = _Driver() + _deny_default_driver(monkeypatch) + orchestrator = PostgresBatchOrchestrator( + "postgresql://x", + postgres_driver=driver, + ) + + class _Counter: + azure_max_files_per_job = 1 + + result = orchestrator._persist_payloads( + [ + { + "part_index": 0, + "record_count": 1, + "total_tokens": 2, + "request_ids": ["aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"], + "lines": ['{"custom_id":"req-1"}'], + } + ], + "11111111-1111-1111-1111-111111111111", + _Counter(), # type: ignore[arg-type] + ) + + assert len(result["ready"]) == 1 + assert result["overflow"] == [] + assert driver.jsonb_values == [ + {"text": '{"custom_id":"req-1"}\n', "line_count": 1} + ] + assert driver.connections[0].autocommit_values == [False] + assert driver.connections[0].commits == 1 + assert len(driver.many) == 1 diff --git a/tests/test_pg8000_candidate_driver_port.py b/tests/test_pg8000_candidate_driver_port.py new file mode 100644 index 000000000..44942d6ca --- /dev/null +++ b/tests/test_pg8000_candidate_driver_port.py @@ -0,0 +1,345 @@ +"""Candidate driver-port regressions for the permissive PostgreSQL migration. + +These tests pin the bounded URI and keyword-conninfo selector slices that pg8000 +can exercise without libpq. Service-file I/O and libpq-only options remain +explicit fail-closed gaps; passing this suite must not be interpreted as full +issue #322 admission or production dependency approval. +""" + +from __future__ import annotations + +from types import ModuleType + +import pytest + +from pg_llm_batch.pg8000_candidate_driver_port import ( + Pg8000CandidateDriverAdapter, + Pg8000CandidateInvalidConninfoError, +) +from pg_llm_batch.pg8000_thread_affine_candidate_adapter import ( + Pg8000ThreadAffineCandidateConnectionAdapter, +) + + +class _ProgrammingError(Exception): + """Stand in for the exact candidate DB-API ProgrammingError authority.""" + + +def _candidate_module() -> tuple[ModuleType, dict[str, object]]: + """Build one DB-API-shaped module and capture exact connection arguments.""" + module = ModuleType("candidate_dbapi") + module.apilevel = "2.0" + module.paramstyle = "format" + module.threadsafety = 1 + module.ProgrammingError = _ProgrammingError + captured: dict[str, object] = {} + + def connect(**kwargs: object) -> object: + captured.update(kwargs) + return object() + + module.connect = connect + return module, captured + + +def test_candidate_uri_conninfo_round_trip_preserves_encoded_identity() -> None: + """Preserve URI identity while decoding values only at the driver boundary.""" + module, _ = _candidate_module() + driver = Pg8000CandidateDriverAdapter(module) + + params = driver.parse_conninfo( + "postgresql://batch%20user:p%40ss@db.example:6543/batch%2Fqueue" + ) + + assert params == { + "user": "batch user", + "password": "p@ss", + "host": "db.example", + "port": "6543", + "dbname": "batch/queue", + } + assert driver.parse_conninfo(driver.make_conninfo(params)) == params + + +def test_candidate_keyword_conninfo_round_trip_preserves_quoted_identity() -> None: + """Accept the bounded libpq keyword form needed by existing CLI deployments.""" + module, _ = _candidate_module() + driver = Pg8000CandidateDriverAdapter(module) + + params = driver.parse_conninfo( + "host=db.example port=6543 dbname='batch queue' " + "user='batch user' password='p@ss word'" + ) + + assert params == { + "host": "db.example", + "port": "6543", + "dbname": "batch queue", + "user": "batch user", + "password": "p@ss word", + } + assert driver.parse_conninfo(driver.make_conninfo(params)) == params + + +def test_candidate_keyword_conninfo_supports_bare_escaped_and_empty_values() -> None: + """Preserve bounded libpq escaping without delegating parsing back to libpq.""" + module, _ = _candidate_module() + driver = Pg8000CandidateDriverAdapter(module) + + params = driver.parse_conninfo( + r" host = db.example user=batch dbname=batch\ queue password= " + ) + + assert params == { + "host": "db.example", + "user": "batch", + "dbname": "batch queue", + "password": "", + "port": "5432", + } + assert driver.parse_conninfo(driver.make_conninfo(params)) == params + + +def test_candidate_service_selector_uses_injected_resolver_and_direct_overrides() -> None: + """Resolve service authority outside the driver and preserve direct overrides.""" + module, _ = _candidate_module() + resolved_names: list[str] = [] + + def resolve_service(service_name: str) -> dict[str, str]: + resolved_names.append(service_name) + return { + "host": "service.example", + "port": "5432", + "dbname": "service_db", + "user": "service_user", + "password": "service-secret", + } + + driver = Pg8000CandidateDriverAdapter(module, service_resolver=resolve_service) + + params = driver.parse_conninfo( + "service=analytics port=6543 dbname='override db'" + ) + + assert resolved_names == ["analytics"] + assert params == { + "host": "service.example", + "port": "6543", + "dbname": "override db", + "user": "service_user", + "password": "service-secret", + } + + +def test_candidate_service_connect_uses_resolved_parameters_not_raw_selector() -> None: + """Keep service-file transport outside pg8000 while connecting with resolved values.""" + module, captured = _candidate_module() + + def resolve_service(service_name: str) -> dict[str, str]: + assert service_name == "analytics" + return { + "host": "127.0.0.1", + "port": "5544", + "dbname": "queue", + "user": "batch", + } + + driver = Pg8000CandidateDriverAdapter(module, service_resolver=resolve_service) + connection = driver.connect("service=analytics", connect_timeout_seconds=5) + + assert isinstance(connection, Pg8000ThreadAffineCandidateConnectionAdapter) + assert captured == { + "user": "batch", + "host": "127.0.0.1", + "port": 5544, + "database": "queue", + "timeout": 5, + } + assert "service" not in captured + assert "dsn" not in captured + + +def test_candidate_conninfo_rejects_unproved_service_and_libpq_options() -> None: + """Keep selectors outside the proved portable subset fail closed instead of guessing.""" + module, _ = _candidate_module() + driver = Pg8000CandidateDriverAdapter(module) + + for dsn in ( + "service=production", + "host=db.example dbname=batch user=batch sslmode=require", + "postgresql://batch@db.example/batch?sslmode=require", + "postgresql://batch@db.example/batch#fragment", + ): + with pytest.raises( + Pg8000CandidateInvalidConninfoError, + match="PostgreSQL connection selector is unsupported", + ): + driver.parse_conninfo(dsn) + + +def test_candidate_service_resolution_rejects_empty_or_unproved_parameters() -> None: + """Do not let a service resolver expand the candidate beyond admitted fields.""" + module, _ = _candidate_module() + + with pytest.raises(Pg8000CandidateInvalidConninfoError): + Pg8000CandidateDriverAdapter( + module, + service_resolver=lambda _: { + "host": "db.example", + "dbname": "batch", + "user": "batch", + }, + ).parse_conninfo("service=''") + + driver = Pg8000CandidateDriverAdapter( + module, + service_resolver=lambda _: { + "host": "db.example", + "dbname": "batch", + "user": "batch", + "sslmode": "verify-full", + }, + ) + with pytest.raises( + Pg8000CandidateInvalidConninfoError, + match="PostgreSQL connection selector is unsupported", + ): + driver.parse_conninfo("service=production") + + +def test_candidate_keyword_conninfo_rejects_ambiguous_or_malformed_grammar() -> None: + """Reject duplicated authority, malformed quoting, and unproved separators.""" + module, _ = _candidate_module() + driver = Pg8000CandidateDriverAdapter(module) + + for dsn in ( + "host=db.example user=batch dbname=queue host=other.example", + "host=db.example user=batch dbname=queue password='unterminated", + "host=db.example user=batch dbname=queue password=trailing\\", + "host=db.example user=batch dbname=que'ue", + "host=db.example user=batch dbname='queue'x", + "host db.example user=batch dbname=queue", + "=db.example user=batch dbname=queue", + "host=db.example\u00a0user=batch dbname=queue", + ): + with pytest.raises(Pg8000CandidateInvalidConninfoError): + driver.parse_conninfo(dsn) + + with pytest.raises(Pg8000CandidateInvalidConninfoError): + driver.parse_conninfo(123) # type: ignore[arg-type] + with pytest.raises(Pg8000CandidateInvalidConninfoError): + driver.parse_conninfo("") + + +def test_candidate_conninfo_rejects_ambiguous_or_unproved_host_forms() -> None: + """Reject multi-host, socket, zone-id, whitespace, and delimiter host forms.""" + module, _ = _candidate_module() + driver = Pg8000CandidateDriverAdapter(module) + + for dsn in ( + "postgresql://batch@db-a.example,db-b.example/batch", + "postgresql://batch@db%20name.example/batch", + "postgresql://batch@%2Fvar%2Frun%2Fpostgresql/batch", + "postgresql://batch@[fe80::1%25eth0]/batch", + "postgresql://batch@db\\name.example/batch", + ): + with pytest.raises(Pg8000CandidateInvalidConninfoError): + driver.parse_conninfo(dsn) + + +def test_candidate_conninfo_rejects_malformed_percent_port_and_control_data() -> None: + """Reject malformed selectors through one non-content-bearing error contract.""" + module, _ = _candidate_module() + driver = Pg8000CandidateDriverAdapter(module) + + for dsn in ( + "postgresql://batch%ZZ@db.example/batch", + "postgresql://batch@db.example:70000/batch", + "postgresql://batch@db.example/batch\nservice=other", + ): + with pytest.raises(Pg8000CandidateInvalidConninfoError): + driver.parse_conninfo(dsn) + + +def test_candidate_make_conninfo_rejects_unknown_or_non_string_parameters() -> None: + """Prevent driver-specific or truthy values from entering the URI renderer.""" + module, _ = _candidate_module() + driver = Pg8000CandidateDriverAdapter(module) + baseline: dict[str, object] = { + "user": "batch", + "host": "db.example", + "port": "5432", + "dbname": "batch", + } + + with pytest.raises(Pg8000CandidateInvalidConninfoError): + driver.make_conninfo({**baseline, "service": "production"}) # type: ignore[arg-type] + with pytest.raises(Pg8000CandidateInvalidConninfoError): + driver.make_conninfo({**baseline, "port": 5432}) # type: ignore[arg-type] + + +def test_candidate_connect_maps_uri_and_finite_timeout_without_raw_dsn_forwarding() -> None: + """Translate the proved URI subset to pg8000 kwargs and retain thread affinity.""" + module, captured = _candidate_module() + driver = Pg8000CandidateDriverAdapter(module) + + connection = driver.connect( + "postgresql://batch:secret@127.0.0.1:5544/queue", + connect_timeout_seconds=5, + ) + + assert isinstance(connection, Pg8000ThreadAffineCandidateConnectionAdapter) + assert captured == { + "user": "batch", + "password": "secret", + "host": "127.0.0.1", + "port": 5544, + "database": "queue", + "timeout": 5, + } + assert "dsn" not in captured + + +def test_candidate_connect_maps_keyword_conninfo_without_raw_dsn_forwarding() -> None: + """Translate keyword conninfo through the same explicit pg8000 argument boundary.""" + module, captured = _candidate_module() + driver = Pg8000CandidateDriverAdapter(module) + + connection = driver.connect("host=127.0.0.1 user=batch dbname=queue") + + assert isinstance(connection, Pg8000ThreadAffineCandidateConnectionAdapter) + assert captured == { + "user": "batch", + "host": "127.0.0.1", + "port": 5432, + "database": "queue", + } + assert "dsn" not in captured + + +@pytest.mark.parametrize("timeout", [True, False, 0, -1, 1.5, "5"]) +def test_candidate_connect_rejects_non_positive_or_non_integer_timeout(timeout: object) -> None: + """Do not coerce booleans, floats, text, or non-positive values into policy.""" + module, _ = _candidate_module() + driver = Pg8000CandidateDriverAdapter(module) + + with pytest.raises( + Pg8000CandidateInvalidConninfoError, + match="PostgreSQL driver timeout is invalid", + ): + driver.connect( + "postgresql://batch@127.0.0.1/queue", + connect_timeout_seconds=timeout, # type: ignore[arg-type] + ) + + +def test_candidate_invalid_conninfo_classifier_is_narrow() -> None: + """Classify only errors created by the candidate selector boundary.""" + module, _ = _candidate_module() + driver = Pg8000CandidateDriverAdapter(module) + + error = Pg8000CandidateInvalidConninfoError( + "PostgreSQL connection selector is invalid" + ) + assert driver.is_invalid_conninfo(error) is True + assert driver.is_invalid_conninfo(RuntimeError("database unavailable")) is False diff --git a/tests/test_pg8000_candidate_secret_boundary.py b/tests/test_pg8000_candidate_secret_boundary.py new file mode 100644 index 000000000..f7651bb37 --- /dev/null +++ b/tests/test_pg8000_candidate_secret_boundary.py @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Protect the pg8000 candidate smoke credential from duplicate env propagation.""" + +from __future__ import annotations + +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[1] + + +def _read(relative_path: str) -> str: + """Read one repository-owned candidate acceptance artifact.""" + return (ROOT / relative_path).read_text(encoding="utf-8") + + +def test_pg8000_candidate_credential_uses_only_ephemeral_password_file() -> None: + """Keep the candidate password in one masked file boundary, not GITHUB_ENV.""" + workflow = _read(".github/workflows/ci.yml") + smoke = _read("tests/smoke_pg8000_candidate_postgres.py") + + assert "PG8000_CANDIDATE_PASSWORD_FILE=$password_file" in workflow + assert "PG_LLM_BATCH_POSTGRES_PASSWORD=$candidate_password" not in workflow + assert 'os.environ.get("PG_LLM_BATCH_POSTGRES_PASSWORD")' not in smoke + assert 'os.environ.get("PG8000_CANDIDATE_PASSWORD_FILE")' in smoke + assert ".read_text(encoding=\"utf-8\")" in smoke diff --git a/tests/test_pg8000_candidate_service_file.py b/tests/test_pg8000_candidate_service_file.py new file mode 100644 index 000000000..34d76abf5 --- /dev/null +++ b/tests/test_pg8000_candidate_service_file.py @@ -0,0 +1,279 @@ +"""Candidate service-file resolver regressions for the PostgreSQL migration.""" + +from __future__ import annotations + +import os +import stat +from pathlib import Path + +import pytest + +from pg_llm_batch.pg8000_candidate_driver_port import Pg8000CandidateInvalidConninfoError +from pg_llm_batch.pg8000_candidate_service_file import Pg8000CandidateServiceFileResolver + + +def test_candidate_service_file_resolves_exact_section_without_ambient_state( + tmp_path: Path, +) -> None: + """Read only caller-selected service-file authority and preserve exact values.""" + service_file = tmp_path / "pg_service.conf" + service_file.write_text( + "# unrelated service\n" + "[other]\n" + "host=other.example\n" + "\n" + "[analytics]\n" + "host=db.example\n" + "port=6543\n" + "dbname=batch queue\n" + "user=batch user\n" + "password=service-secret\n", + encoding="utf-8", + ) + + resolver = Pg8000CandidateServiceFileResolver(service_file) + + assert resolver("analytics") == { + "host": "db.example", + "port": "6543", + "dbname": "batch queue", + "user": "batch user", + "password": "service-secret", + } + + +def test_candidate_service_file_rejects_missing_duplicate_or_empty_target( + tmp_path: Path, +) -> None: + """Fail closed when a service identity has no single authoritative stanza.""" + service_file = tmp_path / "pg_service.conf" + service_file.write_text( + "[analytics]\nhost=db.example\n[analytics]\nhost=other.example\n", + encoding="utf-8", + ) + resolver = Pg8000CandidateServiceFileResolver(service_file) + + with pytest.raises(Pg8000CandidateInvalidConninfoError): + resolver("analytics") + with pytest.raises(Pg8000CandidateInvalidConninfoError): + resolver("") + with pytest.raises(Pg8000CandidateInvalidConninfoError): + resolver("missing") + + +def test_candidate_service_file_rejects_duplicate_keys_and_malformed_lines( + tmp_path: Path, +) -> None: + """Do not invent last-value-wins or permissive grammar for target authority.""" + duplicate_key = tmp_path / "duplicate.conf" + duplicate_key.write_text( + "[analytics]\nhost=db.example\nhost=other.example\n", + encoding="utf-8", + ) + with pytest.raises(Pg8000CandidateInvalidConninfoError): + Pg8000CandidateServiceFileResolver(duplicate_key)("analytics") + + malformed = tmp_path / "malformed.conf" + malformed.write_text("[analytics]\nhost db.example\n", encoding="utf-8") + with pytest.raises(Pg8000CandidateInvalidConninfoError): + Pg8000CandidateServiceFileResolver(malformed)("analytics") + + +def test_candidate_service_file_rejects_ldap_and_control_data( + tmp_path: Path, +) -> None: + """Keep network lookup and framed data outside the local candidate resolver.""" + ldap_file = tmp_path / "ldap.conf" + ldap_file.write_text( + "[analytics]\nldap://directory.example/dc=example?description?one?(cn=db)\n", + encoding="utf-8", + ) + with pytest.raises( + Pg8000CandidateInvalidConninfoError, + match="PostgreSQL connection selector is unsupported", + ): + Pg8000CandidateServiceFileResolver(ldap_file)("analytics") + + control_file = tmp_path / "control.conf" + control_file.write_bytes(b"[analytics]\nhost=db.example\x00evil\n") + with pytest.raises(Pg8000CandidateInvalidConninfoError): + Pg8000CandidateServiceFileResolver(control_file)("analytics") + + +def test_candidate_service_file_rejects_non_utf8_or_oversized_input( + tmp_path: Path, +) -> None: + """Bound service metadata before decoding or parsing it.""" + invalid_utf8 = tmp_path / "invalid-utf8.conf" + invalid_utf8.write_bytes(b"[analytics]\nhost=\xff\n") + with pytest.raises(Pg8000CandidateInvalidConninfoError): + Pg8000CandidateServiceFileResolver(invalid_utf8)("analytics") + + oversized = tmp_path / "oversized.conf" + oversized.write_bytes(b"#" * 65_537) + with pytest.raises(Pg8000CandidateInvalidConninfoError): + Pg8000CandidateServiceFileResolver(oversized)("analytics") + + +def test_candidate_service_file_rejects_non_file_and_non_string_service( + tmp_path: Path, +) -> None: + """Normalize missing-file and invalid service identities at the candidate boundary.""" + with pytest.raises(Pg8000CandidateInvalidConninfoError): + Pg8000CandidateServiceFileResolver(tmp_path / "missing.conf")("analytics") + + service_file = tmp_path / "pg_service.conf" + service_file.write_text("[analytics]\nhost=db.example\n", encoding="utf-8") + resolver = Pg8000CandidateServiceFileResolver(service_file) + with pytest.raises(Pg8000CandidateInvalidConninfoError): + resolver(7) # type: ignore[arg-type] + + +def test_candidate_service_file_uses_nonblocking_descriptor_open( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Reach the service-file byte budget without a blocking special-file open.""" + service_file = tmp_path / "pg_service.conf" + service_file.write_text("[analytics]\nhost=db.example\n", encoding="utf-8") + observed_flags: list[int] = [] + real_open = os.open + + def recording_open( + path: str | bytes | os.PathLike[str] | os.PathLike[bytes], + flags: int, + mode: int = 0o777, + *, + dir_fd: int | None = None, + ) -> int: + observed_flags.append(flags) + if dir_fd is None: + return real_open(path, flags, mode) + return real_open(path, flags, mode, dir_fd=dir_fd) + + monkeypatch.setattr(os, "open", recording_open) + + resolved = Pg8000CandidateServiceFileResolver(service_file)("analytics") + assert resolved["host"] == "db.example" + assert observed_flags + if hasattr(os, "O_NONBLOCK"): + assert observed_flags[0] & os.O_NONBLOCK + + +def test_candidate_service_file_rejects_non_regular_opened_descriptor( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Do not treat pipes or devices as caller-selected service-file authority.""" + service_file = tmp_path / "pg_service.conf" + service_file.write_text("[analytics]\nhost=db.example\n", encoding="utf-8") + real_fstat = os.fstat + + def fifo_fstat(fd: int) -> os.stat_result: + observed = real_fstat(fd) + values = list(observed) + values[0] = stat.S_IFIFO | 0o600 + return os.stat_result(values) + + monkeypatch.setattr(os, "fstat", fifo_fstat) + + with pytest.raises(Pg8000CandidateInvalidConninfoError): + Pg8000CandidateServiceFileResolver(service_file)("analytics") + + +def test_candidate_service_file_normalizes_descriptor_read_failure( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Normalize descriptor read failure without exposing path or content details.""" + service_file = tmp_path / "pg_service.conf" + service_file.write_text("[analytics]\nhost=db.example\n", encoding="utf-8") + + def failing_read(_fd: int, _size: int) -> bytes: + raise OSError("synthetic descriptor read failure") + + monkeypatch.setattr(os, "read", failing_read) + + with pytest.raises( + Pg8000CandidateInvalidConninfoError, + match="PostgreSQL connection selector is invalid", + ): + Pg8000CandidateServiceFileResolver(service_file)("analytics") + + +def test_candidate_service_file_normalizes_close_only_failure( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Fail closed when descriptor cleanup is the only failed operation.""" + service_file = tmp_path / "pg_service.conf" + service_file.write_text("[analytics]\nhost=db.example\n", encoding="utf-8") + real_close = os.close + + def failing_close(fd: int) -> None: + real_close(fd) + raise OSError("synthetic descriptor close failure") + + monkeypatch.setattr(os, "close", failing_close) + + with pytest.raises( + Pg8000CandidateInvalidConninfoError, + match="PostgreSQL connection selector is invalid", + ): + Pg8000CandidateServiceFileResolver(service_file)("analytics") + + +def test_candidate_service_file_preserves_primary_failure_when_close_also_fails( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Do not let a cleanup failure replace an existing descriptor failure.""" + service_file = tmp_path / "pg_service.conf" + service_file.write_text("[analytics]\nhost=db.example\n", encoding="utf-8") + real_close = os.close + + def failing_fstat(_fd: int) -> os.stat_result: + raise OSError("synthetic descriptor metadata failure") + + def failing_close(fd: int) -> None: + real_close(fd) + raise OSError("synthetic descriptor close failure") + + monkeypatch.setattr(os, "fstat", failing_fstat) + monkeypatch.setattr(os, "close", failing_close) + + with pytest.raises( + Pg8000CandidateInvalidConninfoError, + match="PostgreSQL connection selector is invalid", + ): + Pg8000CandidateServiceFileResolver(service_file)("analytics") + + +def test_candidate_service_file_rejects_metadata_drift_during_read( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Reject bytes if retained service-file metadata changes during inspection.""" + service_file = tmp_path / "pg_service.conf" + service_file.write_text("[analytics]\nhost=db.example\n", encoding="utf-8") + real_fstat = os.fstat + fstat_calls = 0 + + def drifting_fstat(fd: int) -> os.stat_result: + nonlocal fstat_calls + fstat_calls += 1 + observed = real_fstat(fd) + if fstat_calls == 1: + return observed + values = list(observed) + values[6] = observed.st_size + 1 + return os.stat_result(values) + + monkeypatch.setattr(os, "fstat", drifting_fstat) + + with pytest.raises( + Pg8000CandidateInvalidConninfoError, + match="PostgreSQL connection selector is invalid", + ): + Pg8000CandidateServiceFileResolver(service_file)("analytics") + assert fstat_calls >= 2 diff --git a/tests/test_pg8000_driver_candidate_adapter.py b/tests/test_pg8000_driver_candidate_adapter.py new file mode 100644 index 000000000..4ccb0e861 --- /dev/null +++ b/tests/test_pg8000_driver_candidate_adapter.py @@ -0,0 +1,322 @@ +"""Contract tests for the pg8000 DB-API candidate adapter boundary. + +These tests deliberately use small DB-API-shaped fakes instead of importing +pg8000. They pin the pg-llm-batch side of the candidate contract first, while +real pg8000 1.31.5 and PostgreSQL acceptance remains a separate mandatory gate +before any runtime dependency or default-driver change. +""" + +from __future__ import annotations + +from types import ModuleType + +import pytest + +from pg_llm_batch.pg8000_driver_candidate_adapter import ( + Pg8000CandidateAdapterError, + Pg8000CandidateConnectionAdapter, + Pg8000CandidateCursorAdapter, + validate_pg8000_dbapi_module, +) +from pg_llm_batch.postgres_driver_port import PostgresConnectionPort, PostgresCursorPort + + +class _FakeCursor: + """Record DB-API cursor calls while exposing configurable materialized rows.""" + + def __init__(self) -> None: + self.execute_calls: list[tuple[str, object | None]] = [] + self.executemany_calls: list[tuple[str, object]] = [] + self.fetchone_value: object | None = None + self.fetchmany_value: list[object] = [] + self.fetchall_value: list[object] = [] + self.rowcount: object = 0 + self.close_count = 0 + self.close_error: BaseException | None = None + self.enter_count = 0 + self.exit_args: tuple[object, object, object] | None = None + + def execute(self, query: str, params: object | None = None) -> _FakeCursor: + self.execute_calls.append((query, params)) + return self + + def executemany(self, query: str, params_seq: object) -> _FakeCursor: + self.executemany_calls.append((query, params_seq)) + return self + + def fetchone(self) -> object | None: + return self.fetchone_value + + def fetchmany(self, size: int) -> list[object]: + assert size > 0 + return self.fetchmany_value + + def fetchall(self) -> list[object]: + return self.fetchall_value + + def close(self) -> None: + self.close_count += 1 + if self.close_error is not None: + raise self.close_error + + def __enter__(self) -> _FakeCursor: + self.enter_count += 1 + return self + + def __exit__(self, exc_type: object, exc: object, traceback: object) -> bool: + self.exit_args = (exc_type, exc, traceback) + return False + + +class _FakeConnection: + """Record one connection's transaction and cursor activity.""" + + def __init__(self) -> None: + self.cursor_value = _FakeCursor() + self.autocommit: object = False + self.closed: object = False + self.commit_count = 0 + self.rollback_count = 0 + self.close_count = 0 + self.enter_count = 0 + self.exit_args: tuple[object, object, object] | None = None + + def cursor(self) -> _FakeCursor: + return self.cursor_value + + def commit(self) -> None: + self.commit_count += 1 + + def rollback(self) -> None: + self.rollback_count += 1 + + def close(self) -> None: + self.close_count += 1 + self.closed = True + + def __enter__(self) -> _FakeConnection: + self.enter_count += 1 + return self + + def __exit__(self, exc_type: object, exc: object, traceback: object) -> bool: + self.exit_args = (exc_type, exc, traceback) + return False + + +def _dbapi_module( + *, + apilevel: object = "2.0", + paramstyle: object = "format", + threadsafety: object = 1, +) -> ModuleType: + """Build one exact module-shaped DB-API authority for candidate contract tests.""" + module = ModuleType("pg8000.dbapi") + module.apilevel = apilevel + module.paramstyle = paramstyle + module.threadsafety = threadsafety + return module + + +def test_candidate_dbapi_module_requires_dbapi_2_and_format_parameter_style() -> None: + module = _dbapi_module() + + validate_pg8000_dbapi_module(module) + + module.paramstyle = "named" + with pytest.raises(Pg8000CandidateAdapterError, match="parameter style is incompatible"): + validate_pg8000_dbapi_module(module) + + module.paramstyle = "format" + module.apilevel = "1.0" + with pytest.raises(Pg8000CandidateAdapterError, match="API level is incompatible"): + validate_pg8000_dbapi_module(module) + + +def test_candidate_dbapi_module_requires_module_only_connection_thread_safety() -> None: + """Reject metadata that would misstate pg8000 connection-sharing semantics.""" + validate_pg8000_dbapi_module(_dbapi_module(threadsafety=1)) + + for invalid in (0, 2, 3, True, "1"): + with pytest.raises( + Pg8000CandidateAdapterError, + match="thread safety is incompatible", + ): + validate_pg8000_dbapi_module(_dbapi_module(threadsafety=invalid)) + + +def test_candidate_dbapi_module_rejects_shaped_or_behavior_bearing_metadata() -> None: + class _StringSubclass(str): + pass + + with pytest.raises(Pg8000CandidateAdapterError, match="module identity is invalid"): + validate_pg8000_dbapi_module(object()) + with pytest.raises(Pg8000CandidateAdapterError, match="API level is incompatible"): + validate_pg8000_dbapi_module(_dbapi_module(apilevel=_StringSubclass("2.0"))) + with pytest.raises(Pg8000CandidateAdapterError, match="parameter style is incompatible"): + validate_pg8000_dbapi_module(_dbapi_module(paramstyle=_StringSubclass("format"))) + + +def test_candidate_cursor_preserves_parameter_binding_and_wrapper_identity() -> None: + raw = _FakeCursor() + adapter = Pg8000CandidateCursorAdapter(raw) + params = ("tenant-a", 7) + + result = adapter.execute("SELECT %s, %s", params) + many_result = adapter.executemany("INSERT INTO t VALUES (%s)", [(1,), (2,)]) + + assert isinstance(adapter, PostgresCursorPort) + assert result is adapter + assert many_result is adapter + assert raw.execute_calls == [("SELECT %s, %s", params)] + assert raw.executemany_calls == [("INSERT INTO t VALUES (%s)", [(1,), (2,)])] + + +def test_candidate_cursor_normalizes_pg8000_list_rows_to_tuples() -> None: + raw = _FakeCursor() + adapter = Pg8000CandidateCursorAdapter(raw) + + assert adapter.fetchone() is None + + raw.fetchone_value = ["one", 1] + raw.fetchmany_value = [["two", 2], ("three", 3)] + raw.fetchall_value = [("four", 4), ["five", 5]] + + assert adapter.fetchone() == ("one", 1) + assert adapter.fetchmany(2) == [("two", 2), ("three", 3)] + assert adapter.fetchall() == [("four", 4), ("five", 5)] + + +def test_candidate_cursor_rejects_non_positional_rows_and_invalid_fetch_size() -> None: + raw = _FakeCursor() + adapter = Pg8000CandidateCursorAdapter(raw) + raw.fetchone_value = {"id": 1} + + with pytest.raises(Pg8000CandidateAdapterError, match="result row is invalid"): + adapter.fetchone() + with pytest.raises(Pg8000CandidateAdapterError, match="fetch size is invalid"): + adapter.fetchmany(0) + with pytest.raises(Pg8000CandidateAdapterError, match="fetch size is invalid"): + adapter.fetchmany(True) + + +def test_candidate_cursor_normalizes_unknown_row_count_and_rejects_bad_sentinels() -> None: + raw = _FakeCursor() + adapter = Pg8000CandidateCursorAdapter(raw) + + raw.rowcount = -1 + assert adapter.row_count() is None + raw.rowcount = 4 + assert adapter.row_count() == 4 + raw.rowcount = -2 + with pytest.raises(Pg8000CandidateAdapterError, match="row count is invalid"): + adapter.row_count() + raw.rowcount = "4" + with pytest.raises(Pg8000CandidateAdapterError, match="row count is invalid"): + adapter.row_count() + + +def test_candidate_cursor_context_owns_dbapi_cleanup_without_raw_context_dependency() -> None: + raw = _FakeCursor() + adapter = Pg8000CandidateCursorAdapter(raw) + error = RuntimeError("boom") + + assert adapter.__enter__() is adapter + assert raw.enter_count == 0 + assert adapter.__exit__(RuntimeError, error, None) is False + assert raw.close_count == 1 + assert raw.exit_args is None + + +def test_candidate_cursor_context_preserves_application_error_over_cleanup_failure() -> None: + """Cleanup failure must not replace the application error already in flight.""" + raw = _FakeCursor() + adapter = Pg8000CandidateCursorAdapter(raw) + application_error = ValueError("application") + raw.close_error = RuntimeError("cleanup") + + with pytest.raises(ValueError) as caught: + adapter.__exit__(ValueError, application_error, None) + + assert caught.value is application_error + assert raw.close_count == 1 + + +def test_candidate_cursor_context_propagates_cleanup_failure_without_application_error() -> None: + """A close-only failure remains visible when no earlier error has priority.""" + raw = _FakeCursor() + adapter = Pg8000CandidateCursorAdapter(raw) + cleanup_error = RuntimeError("cleanup") + raw.close_error = cleanup_error + + with pytest.raises(RuntimeError) as caught: + adapter.__exit__(None, None, None) + + assert caught.value is cleanup_error + assert raw.close_count == 1 + + +def test_candidate_connection_uses_one_raw_connection_for_execution_and_transactions() -> None: + raw = _FakeConnection() + adapter = Pg8000CandidateConnectionAdapter(raw) + + cursor = adapter.cursor() + direct_cursor = adapter.execute("SELECT %s", (9,)) + adapter.commit() + adapter.rollback() + + assert isinstance(adapter, PostgresConnectionPort) + assert isinstance(cursor, Pg8000CandidateCursorAdapter) + assert isinstance(direct_cursor, Pg8000CandidateCursorAdapter) + assert raw.cursor_value.execute_calls == [("SELECT %s", (9,))] + assert raw.commit_count == 1 + assert raw.rollback_count == 1 + + +def test_candidate_connection_validates_autocommit_and_owns_closed_state() -> None: + raw = _FakeConnection() + adapter = Pg8000CandidateConnectionAdapter(raw) + + adapter.set_autocommit(True) + assert raw.autocommit is True + assert adapter.is_closed() is False + + with pytest.raises(Pg8000CandidateAdapterError, match="autocommit is invalid"): + adapter.set_autocommit(1) # type: ignore[arg-type] + + raw.closed = 0 + assert adapter.is_closed() is False + del raw.closed + assert adapter.is_closed() is False + + adapter.close() + assert raw.close_count == 1 + assert adapter.is_closed() is True + + +def test_candidate_connection_context_commits_and_closes_without_raw_context_dependency() -> None: + raw = _FakeConnection() + adapter = Pg8000CandidateConnectionAdapter(raw) + + assert adapter.__enter__() is adapter + assert raw.enter_count == 0 + assert adapter.__exit__(None, None, None) is False + assert raw.commit_count == 1 + assert raw.rollback_count == 0 + assert raw.close_count == 1 + assert raw.exit_args is None + assert adapter.is_closed() is True + + +def test_candidate_connection_context_rolls_back_and_closes_on_exception() -> None: + raw = _FakeConnection() + adapter = Pg8000CandidateConnectionAdapter(raw) + error = ValueError("bad") + + assert adapter.__enter__() is adapter + assert raw.enter_count == 0 + assert adapter.__exit__(ValueError, error, None) is False + assert raw.commit_count == 0 + assert raw.rollback_count == 1 + assert raw.close_count == 1 + assert raw.exit_args is None + assert adapter.is_closed() is True diff --git a/tests/test_pg8000_driver_candidate_close_state.py b/tests/test_pg8000_driver_candidate_close_state.py new file mode 100644 index 000000000..b588f48df --- /dev/null +++ b/tests/test_pg8000_driver_candidate_close_state.py @@ -0,0 +1,38 @@ +"""Recovery-state regressions for the candidate pg8000 connection adapter.""" + +from __future__ import annotations + +import pytest + +from pg_llm_batch.pg8000_driver_candidate_adapter import ( + Pg8000CandidateConnectionAdapter, +) + + +class _ProtocolCloseFailureConnection: + """Model pg8000 closing its socket while protocol-level close reports failure.""" + + def __init__(self) -> None: + self.closed = False + self.close_count = 0 + + def close(self) -> None: + """Release the underlying capability, then report the protocol failure.""" + self.close_count += 1 + self.closed = True + raise RuntimeError("protocol close failed") + + +def test_candidate_marks_connection_closed_when_protocol_close_reports_failure() -> None: + """A released pg8000 socket must not remain reusable after close raises.""" + raw = _ProtocolCloseFailureConnection() + adapter = Pg8000CandidateConnectionAdapter(raw) + + assert adapter.is_closed() is False + + with pytest.raises(RuntimeError, match="protocol close failed"): + adapter.close() + + assert raw.close_count == 1 + assert raw.closed is True + assert adapter.is_closed() is True diff --git a/tests/test_pg8000_driver_candidate_error_classification.py b/tests/test_pg8000_driver_candidate_error_classification.py new file mode 100644 index 000000000..dffdbb4be --- /dev/null +++ b/tests/test_pg8000_driver_candidate_error_classification.py @@ -0,0 +1,85 @@ +"""Candidate-only PostgreSQL error-classification contract tests. + +These tests keep pg8000 out of the production dependency graph. The exact +candidate artifact is injected as a DB-API module so the classifier can prove +one PostgreSQL SQLSTATE without trusting shaped exception objects or message +text. Real PostgreSQL acceptance remains a separate CI smoke before candidate +promotion. +""" + +from __future__ import annotations + +from types import ModuleType + +import pytest + +from pg_llm_batch.pg8000_driver_candidate_errors import ( + Pg8000CandidateErrorEvidenceError, + is_pg8000_candidate_undefined_function, +) + + +class _ProgrammingError(Exception): + """Stand in for the exact candidate DB-API ProgrammingError class.""" + + +def _dbapi_module() -> ModuleType: + """Build one exact module-shaped DB-API authority for classifier tests.""" + module = ModuleType("pg8000.dbapi") + module.ProgrammingError = _ProgrammingError + return module + + +def test_candidate_classifies_only_exact_undefined_function_sqlstate() -> None: + module = _dbapi_module() + + assert is_pg8000_candidate_undefined_function( + _ProgrammingError({"S": "ERROR", "C": "42883", "M": "hidden"}), + dbapi_module=module, + ) is True + assert is_pg8000_candidate_undefined_function( + _ProgrammingError({"S": "ERROR", "C": "42P01", "M": "hidden"}), + dbapi_module=module, + ) is False + + +def test_candidate_classifier_rejects_untrusted_module_authority() -> None: + with pytest.raises( + Pg8000CandidateErrorEvidenceError, + match="DB-API module authority is invalid", + ): + is_pg8000_candidate_undefined_function( + _ProgrammingError({"C": "42883"}), + dbapi_module=object(), + ) + + +def test_candidate_classifier_does_not_execute_or_trust_shaped_payloads() -> None: + module = _dbapi_module() + + class _MappingLike: + def get(self, key: object) -> object: + raise AssertionError("mapping-like payload was evaluated") + + assert is_pg8000_candidate_undefined_function( + _ProgrammingError(_MappingLike()), + dbapi_module=module, + ) is False + assert is_pg8000_candidate_undefined_function( + RuntimeError({"C": "42883"}), + dbapi_module=module, + ) is False + + +def test_candidate_classifier_rejects_malformed_programming_error_authority() -> None: + module = _dbapi_module() + module.ProgrammingError = "ProgrammingError" + + with pytest.raises( + Pg8000CandidateErrorEvidenceError, + match="ProgrammingError authority is invalid", + ): + is_pg8000_candidate_undefined_function( + _ProgrammingError({"C": "42883"}), + dbapi_module=module, + ) diff --git a/tests/test_pg8000_driver_candidate_error_precedence.py b/tests/test_pg8000_driver_candidate_error_precedence.py new file mode 100644 index 000000000..45db4bfe3 --- /dev/null +++ b/tests/test_pg8000_driver_candidate_error_precedence.py @@ -0,0 +1,147 @@ +"""Regression tests for candidate PostgreSQL transaction cleanup precedence. + +The pg8000 anti-corruption adapter must attempt connection cleanup after a +transaction failure without letting a later close failure replace the earlier +commit or rollback failure. An application exception also remains primary when +rollback succeeds but later connection cleanup fails. Direct connection +execution must likewise close the internally created cursor when execution +fails, because the caller never receives that cursor and therefore cannot +release it. Candidate connection state remains terminal after a close attempt +because pg8000 1.31.5 releases its retained transport in the driver's ``finally`` +path even when protocol-level cleanup raises. +""" + +from __future__ import annotations + +import pytest + +from pg_llm_batch.pg8000_driver_candidate_adapter import ( + Pg8000CandidateConnectionAdapter, +) + + +class _TransactionAndCloseFailureConnection: + """Expose deterministic transaction and close failures for precedence tests.""" + + def __init__(self, *, fail_commit: bool) -> None: + self.fail_commit = fail_commit + self.commit_count = 0 + self.rollback_count = 0 + self.close_count = 0 + + def commit(self) -> None: + self.commit_count += 1 + if self.fail_commit: + raise RuntimeError("commit failed") + + def rollback(self) -> None: + self.rollback_count += 1 + if not self.fail_commit: + raise RuntimeError("rollback failed") + + def close(self) -> None: + self.close_count += 1 + raise OSError("close failed") + + +class _RollbackSuccessCloseFailureConnection: + """Succeed rollback but fail cleanup after an application exception.""" + + def __init__(self) -> None: + self.rollback_count = 0 + self.close_count = 0 + + def rollback(self) -> None: + self.rollback_count += 1 + + def close(self) -> None: + self.close_count += 1 + raise OSError("close failed") + + +class _ExecuteAndCloseFailureCursor: + """Fail direct execution and optionally fail the required cursor cleanup.""" + + def __init__(self, *, fail_close: bool) -> None: + self.fail_close = fail_close + self.execute_count = 0 + self.close_count = 0 + + def execute(self, query: str, params: object | None = None) -> None: + """Raise the primary execution failure after recording one attempt.""" + del query, params + self.execute_count += 1 + raise RuntimeError("execute failed") + + def close(self) -> None: + """Record cleanup and optionally expose a secondary cleanup failure.""" + self.close_count += 1 + if self.fail_close: + raise OSError("cursor close failed") + + +class _ExecuteFailureConnection: + """Return one retained failing cursor to the connection adapter.""" + + def __init__(self, cursor: _ExecuteAndCloseFailureCursor) -> None: + self.cursor_value = cursor + + def cursor(self) -> _ExecuteAndCloseFailureCursor: + """Return the exact cursor whose ownership transfers to direct execute.""" + return self.cursor_value + + +def test_candidate_context_preserves_commit_failure_when_close_also_fails() -> None: + raw = _TransactionAndCloseFailureConnection(fail_commit=True) + adapter = Pg8000CandidateConnectionAdapter(raw) + + with pytest.raises(RuntimeError, match="commit failed"): + adapter.__exit__(None, None, None) + + assert raw.commit_count == 1 + assert raw.rollback_count == 0 + assert raw.close_count == 1 + assert adapter.is_closed() is True + + +def test_candidate_context_preserves_rollback_failure_when_close_also_fails() -> None: + raw = _TransactionAndCloseFailureConnection(fail_commit=False) + adapter = Pg8000CandidateConnectionAdapter(raw) + application_error = ValueError("application failed") + + with pytest.raises(RuntimeError, match="rollback failed"): + adapter.__exit__(ValueError, application_error, None) + + assert raw.commit_count == 0 + assert raw.rollback_count == 1 + assert raw.close_count == 1 + assert adapter.is_closed() is True + + +def test_candidate_context_preserves_application_error_when_only_close_fails() -> None: + """Cleanup failure must not replace an application error after rollback.""" + raw = _RollbackSuccessCloseFailureConnection() + adapter = Pg8000CandidateConnectionAdapter(raw) + application_error = ValueError("application failed") + + with pytest.raises(ValueError) as caught: + adapter.__exit__(ValueError, application_error, None) + + assert caught.value is application_error + assert raw.rollback_count == 1 + assert raw.close_count == 1 + assert adapter.is_closed() is True + + +@pytest.mark.parametrize("fail_close", [False, True]) +def test_candidate_direct_execute_closes_cursor_and_preserves_primary_failure( + fail_close: bool, +) -> None: + raw_cursor = _ExecuteAndCloseFailureCursor(fail_close=fail_close) + adapter = Pg8000CandidateConnectionAdapter(_ExecuteFailureConnection(raw_cursor)) + + with pytest.raises(RuntimeError, match="execute failed"): + adapter.execute("SELECT %s", (1,)) + + assert raw_cursor.execute_count == 1 + assert raw_cursor.close_count == 1 diff --git a/tests/test_pg8000_driver_candidate_fetch_bound.py b/tests/test_pg8000_driver_candidate_fetch_bound.py new file mode 100644 index 000000000..ec228ba89 --- /dev/null +++ b/tests/test_pg8000_driver_candidate_fetch_bound.py @@ -0,0 +1,54 @@ +"""Regression contract for bounded pg8000 candidate fetches. + +The PostgreSQL driver port promises that ``fetchmany(size)`` returns at most the +requested row budget. A candidate that over-delivers rows or returns an unsized +iterable must fail at the anti-corruption boundary rather than expanding or +obscuring an application resource budget. +""" + +from __future__ import annotations + +from collections.abc import Iterator + +import pytest + +from pg_llm_batch.pg8000_driver_candidate_adapter import ( + Pg8000CandidateAdapterError, + Pg8000CandidateCursorAdapter, +) + + +class _OverDeliveringCursor: + """Return more rows than requested to model a non-conforming DB-API candidate.""" + + def fetchmany(self, size: int) -> list[list[int]]: + assert size == 1 + return [[1], [2]] + + +class _UnsizedCursor: + """Return an iterable whose cardinality cannot be checked before materialization.""" + + def fetchmany(self, size: int) -> Iterator[list[int]]: + assert size == 1 + yield [1] + + +def test_candidate_fetchmany_rejects_driver_overdelivery() -> None: + adapter = Pg8000CandidateCursorAdapter(_OverDeliveringCursor()) + + with pytest.raises( + Pg8000CandidateAdapterError, + match="fetch result exceeds requested size", + ): + adapter.fetchmany(1) + + +def test_candidate_fetchmany_rejects_unsized_driver_results() -> None: + adapter = Pg8000CandidateCursorAdapter(_UnsizedCursor()) + + with pytest.raises( + Pg8000CandidateAdapterError, + match="fetch result is invalid", + ): + adapter.fetchmany(1) diff --git a/tests/test_pg8000_driver_candidate_jsonb.py b/tests/test_pg8000_driver_candidate_jsonb.py new file mode 100644 index 000000000..611aedb53 --- /dev/null +++ b/tests/test_pg8000_driver_candidate_jsonb.py @@ -0,0 +1,50 @@ +"""Candidate JSONB adaptation contract for the permissive PostgreSQL driver lane. + +The pg8000 migration remains candidate-only until the exact artifact proves the +complete PostgreSQL, conninfo, RLS, recovery, packaging, SBOM, and provenance +contract. These tests pin only the JSONB parameter boundary needed by the +existing ``PostgresDriverPort.jsonb`` capability. +""" + +from __future__ import annotations + +import json +import math + +import pytest + +from pg_llm_batch.pg8000_driver_candidate_jsonb import ( + Pg8000CandidateJsonbError, + adapt_pg8000_jsonb, +) + + +def test_candidate_jsonb_serializes_exact_json_semantics_for_dbapi_binding() -> None: + """Return UTF-8 JSON text that pg8000 DB-API can bind to a JSONB cast.""" + payload = { + "request_id": "opaque-1", + "labels": ["한국어", "English", None], + "enabled": True, + "count": 3, + } + + adapted = adapt_pg8000_jsonb(payload) + + assert type(adapted) is str + assert adapted.encode("utf-8") + assert json.loads(adapted) == payload + assert payload["labels"] == ["한국어", "English", None] + + +@pytest.mark.parametrize("value", [math.nan, math.inf, -math.inf]) +def test_candidate_jsonb_rejects_non_finite_numbers(value: float) -> None: + """Reject values PostgreSQL JSONB cannot represent as standards-compliant JSON.""" + with pytest.raises(Pg8000CandidateJsonbError, match="JSONB value is invalid"): + adapt_pg8000_jsonb({"value": value}) + + +def test_candidate_jsonb_rejects_unencodable_or_non_json_values() -> None: + """Normalize invalid candidate payloads to one non-content-bearing error.""" + for value in ({"value": object()}, {"value": "\ud800"}): + with pytest.raises(Pg8000CandidateJsonbError, match="JSONB value is invalid"): + adapt_pg8000_jsonb(value) diff --git a/tests/test_pg8000_driver_candidate_thread_affinity.py b/tests/test_pg8000_driver_candidate_thread_affinity.py new file mode 100644 index 000000000..0f121decf --- /dev/null +++ b/tests/test_pg8000_driver_candidate_thread_affinity.py @@ -0,0 +1,133 @@ +"""Concurrency regressions for the candidate pg8000 DB-API boundary. + +pg8000 1.31.5 declares DB-API ``threadsafety == 1``: threads may share the +module, but not connections. The admitted candidate boundary must therefore fail +before raw driver access when a connection or cursor capability crosses its +creating thread. This is candidate-admission evidence only; it does not promote +pg8000 into the production dependency graph. +""" + +from __future__ import annotations + +from concurrent.futures import ThreadPoolExecutor +from typing import Callable + +import pytest + +from pg_llm_batch.pg8000_driver_candidate_adapter import Pg8000CandidateAdapterError +from pg_llm_batch.pg8000_thread_affine_candidate_adapter import ( + Pg8000ThreadAffineCandidateConnectionAdapter, + Pg8000ThreadAffineCandidateCursorAdapter, +) + + +class _RawCursor: + def __init__(self) -> None: + self.calls = 0 + self.rowcount = 1 + + def execute(self, _query: str, _params: object | None = None) -> None: + self.calls += 1 + + def executemany(self, _query: str, _params_seq: object) -> None: + self.calls += 1 + + def fetchone(self) -> tuple[int]: + self.calls += 1 + return (1,) + + def fetchmany(self, _size: int) -> list[tuple[int]]: + self.calls += 1 + return [(1,)] + + def fetchall(self) -> list[tuple[int]]: + self.calls += 1 + return [(1,)] + + def close(self) -> None: + self.calls += 1 + + +class _RawConnection: + def __init__(self) -> None: + self.calls = 0 + self._autocommit = False + + @property + def autocommit(self) -> bool: + return self._autocommit + + @autocommit.setter + def autocommit(self, value: bool) -> None: + self.calls += 1 + self._autocommit = value + + def cursor(self) -> _RawCursor: + self.calls += 1 + return _RawCursor() + + def commit(self) -> None: + self.calls += 1 + + def rollback(self) -> None: + self.calls += 1 + + def close(self) -> None: + self.calls += 1 + + +def _run_on_worker(operation: Callable[[], object]) -> object: + with ThreadPoolExecutor(max_workers=1) as executor: + return executor.submit(operation).result() + + +@pytest.mark.parametrize( + "operation", + ("cursor", "commit", "rollback", "set_autocommit", "close"), +) +def test_candidate_connection_rejects_cross_thread_driver_access(operation: str) -> None: + raw = _RawConnection() + adapter = Pg8000ThreadAffineCandidateConnectionAdapter(raw) + + callbacks: dict[str, Callable[[], object]] = { + "cursor": adapter.cursor, + "commit": adapter.commit, + "rollback": adapter.rollback, + "set_autocommit": lambda: adapter.set_autocommit(True), + "close": adapter.close, + } + + with pytest.raises( + Pg8000CandidateAdapterError, + match="must not be shared across threads", + ): + _run_on_worker(callbacks[operation]) + + assert raw.calls == 0 + + +@pytest.mark.parametrize( + "operation", + ("execute", "executemany", "fetchone", "fetchmany", "fetchall", "row_count", "close"), +) +def test_candidate_cursor_rejects_cross_thread_driver_access(operation: str) -> None: + raw = _RawCursor() + adapter = Pg8000ThreadAffineCandidateCursorAdapter(raw) + + callbacks: dict[str, Callable[[], object]] = { + "execute": lambda: adapter.execute("SELECT %s", (1,)), + "executemany": lambda: adapter.executemany("SELECT %s", [(1,)]), + "fetchone": adapter.fetchone, + "fetchmany": lambda: adapter.fetchmany(1), + "fetchall": adapter.fetchall, + "row_count": adapter.row_count, + "close": lambda: adapter.__exit__(None, None, None), + } + + with pytest.raises( + Pg8000CandidateAdapterError, + match="must not be shared across threads", + ): + _run_on_worker(callbacks[operation]) + + assert raw.calls == 0 diff --git a/tests/test_pg_tiktoken_runtime_authority.py b/tests/test_pg_tiktoken_runtime_authority.py index ab7b131ba..bf182ccff 100644 --- a/tests/test_pg_tiktoken_runtime_authority.py +++ b/tests/test_pg_tiktoken_runtime_authority.py @@ -3,7 +3,6 @@ from __future__ import annotations -from pg_llm_batch import token_counter as tc_mod from pg_llm_batch.token_counter import TokenCounter @@ -27,7 +26,7 @@ def fetchone(self) -> tuple[bool, bool, bool]: class _ProbeConnection: - """Minimal Psycopg connection double for a runtime capability probe.""" + """Minimal driver-port connection double for a runtime capability probe.""" def __init__(self) -> None: self.closed = False @@ -39,6 +38,12 @@ def __init__(self) -> None: def cursor(self) -> _ProbeCursor: return _ProbeCursor(self) + def set_autocommit(self, enabled: bool) -> None: + self.autocommit = enabled + + def is_closed(self) -> bool: + return self.closed + def commit(self) -> None: self.commit_calls += 1 @@ -47,24 +52,32 @@ def close(self) -> None: self.closed = True -class _ProbePsycopg: +class _ProbeDriver: """Return one inspectable connection without granting installation authority.""" def __init__(self) -> None: self.connection = _ProbeConnection() self.connect_calls = 0 - def connect(self, _dsn: str) -> _ProbeConnection: + def connect( + self, + _dsn: str, + *, + connect_timeout_seconds: int | None = None, + ) -> _ProbeConnection: + assert connect_timeout_seconds is None self.connect_calls += 1 return self.connection + def is_undefined_function(self, _error: BaseException) -> bool: + return False + -def test_runtime_pg_tiktoken_readiness_never_installs_extensions(monkeypatch) -> None: +def test_runtime_pg_tiktoken_readiness_never_installs_extensions() -> None: """Ordinary token counting must inspect capability without executing DDL.""" - driver = _ProbePsycopg() - monkeypatch.setattr(tc_mod, "psycopg", driver) + driver = _ProbeDriver() - counter = TokenCounter("postgresql://database") + counter = TokenCounter("postgresql://database", postgres_driver=driver) assert counter._pg_available is True assert driver.connect_calls == 1 diff --git a/tests/test_postgres_driver_candidate.py b/tests/test_postgres_driver_candidate.py new file mode 100644 index 000000000..a2891482c --- /dev/null +++ b/tests/test_postgres_driver_candidate.py @@ -0,0 +1,331 @@ +from __future__ import annotations + +import pytest + +from pg_llm_batch.postgres_driver_candidate import ( + REQUIRED_POSTGRES_DRIVER_CAPABILITIES, + REQUIRED_POSTGRES_DRIVER_PYTHON_VERSIONS, + PostgresDriverCandidateEvidence, + PostgresDriverCandidateEvidenceError, + evaluate_postgres_driver_candidate, +) + + +FULL_CAPABILITIES = frozenset(REQUIRED_POSTGRES_DRIVER_CAPABILITIES) +FULL_PYTHON_VERSIONS = tuple(sorted(REQUIRED_POSTGRES_DRIVER_PYTHON_VERSIONS)) +SOURCE_SHA = "a" * 40 +ARTIFACT_SHA256 = "b" * 64 +VULNERABILITY_REPORT_SHA256 = "c" * 64 +LICENSE_REPORT_SHA256 = "d" * 64 +CAPABILITY_REPORT_SHA256 = "e" * 64 + + +def _evidence(**overrides: object) -> PostgresDriverCandidateEvidence: + values: dict[str, object] = { + "package_name": "candidate-driver", + "package_version": "1.2.3", + "license_spdx": "BSD-3-Clause", + "license_report_sha256": LICENSE_REPORT_SHA256, + "python_versions": FULL_PYTHON_VERSIONS, + "source_commit_sha": SOURCE_SHA, + "artifact_sha256": ARTIFACT_SHA256, + "vulnerability_report_sha256": VULNERABILITY_REPORT_SHA256, + "capability_report_sha256": CAPABILITY_REPORT_SHA256, + "known_vulnerability_ids": (), + "capabilities": FULL_CAPABILITIES, + } + values.update(overrides) + return PostgresDriverCandidateEvidence(**values) # type: ignore[arg-type] + + +def test_complete_permissive_candidate_is_eligible_only_for_parity_validation() -> None: + decision = evaluate_postgres_driver_candidate(_evidence()) + + assert decision.eligible_for_parity_validation is True + assert decision.production_approved is False + assert decision.reasons == () + + +def test_candidate_requires_immutable_license_report_identity() -> None: + decision = evaluate_postgres_driver_candidate( + _evidence(license_report_sha256=LICENSE_REPORT_SHA256) + ) + + assert decision.eligible_for_parity_validation is True + assert decision.production_approved is False + + +@pytest.mark.parametrize( + "license_report_sha256", + ["d" * 63, "z" * 64, 64], +) +def test_candidate_rejects_malformed_license_report_identity( + license_report_sha256: object, +) -> None: + with pytest.raises(PostgresDriverCandidateEvidenceError, match="license report"): + _evidence(license_report_sha256=license_report_sha256) + + +@pytest.mark.parametrize( + "capability_report_sha256", + ["e" * 63, "z" * 64, 64], +) +def test_candidate_rejects_malformed_capability_report_identity( + capability_report_sha256: object, +) -> None: + with pytest.raises(PostgresDriverCandidateEvidenceError, match="capability report"): + _evidence(capability_report_sha256=capability_report_sha256) + + +def test_candidate_contract_covers_issue_322_type_and_parameter_parity() -> None: + assert { + "result_row_semantics", + "sql_parameter_style_adaptation", + "uuid_timestamp_adaptation", + } <= REQUIRED_POSTGRES_DRIVER_CAPABILITIES + + +def test_candidate_contract_preserves_each_supported_dsn_selector_family() -> None: + assert { + "conninfo_keyword_parse_render", + "conninfo_service_selector", + "conninfo_uri_parse_render", + } <= REQUIRED_POSTGRES_DRIVER_CAPABILITIES + + +def test_candidate_contract_requires_every_repository_ci_python_version() -> None: + assert REQUIRED_POSTGRES_DRIVER_PYTHON_VERSIONS == frozenset( + {"3.10", "3.11", "3.12", "3.13", "3.14"} + ) + + +def test_candidate_rejects_known_vulnerabilities_before_parity_validation() -> None: + decision = evaluate_postgres_driver_candidate( + _evidence(known_vulnerability_ids=("CVE-2025-61385",)) + ) + + assert decision.eligible_for_parity_validation is False + assert decision.production_approved is False + assert decision.reasons == ("known_vulnerability:CVE-2025-61385",) + + +@pytest.mark.parametrize( + ("license_spdx", "expected_reason"), + [ + ("LGPL-3.0-only", "license_not_approved"), + ("GPL-3.0-only", "license_not_approved"), + ("AGPL-3.0-only", "license_not_approved"), + ("LicenseRef-Proprietary", "license_not_approved"), + ], +) +def test_candidate_fails_closed_when_license_is_not_explicitly_permissive( + license_spdx: str, + expected_reason: str, +) -> None: + decision = evaluate_postgres_driver_candidate(_evidence(license_spdx=license_spdx)) + + assert decision.eligible_for_parity_validation is False + assert decision.production_approved is False + assert expected_reason in decision.reasons + + +@pytest.mark.parametrize( + "missing_version", ["3.10", "3.11", "3.12", "3.13", "3.14"] +) +def test_candidate_requires_every_repository_ci_python_version( + missing_version: str, +) -> None: + candidate_versions = tuple( + version for version in FULL_PYTHON_VERSIONS if version != missing_version + ) + decision = evaluate_postgres_driver_candidate( + _evidence(python_versions=candidate_versions) + ) + + assert decision.eligible_for_parity_validation is False + assert decision.reasons == (f"missing_python_version:{missing_version}",) + + +def test_candidate_reports_every_missing_runtime_capability_deterministically() -> None: + decision = evaluate_postgres_driver_candidate( + _evidence(capabilities=frozenset({"parameterized_sql", "jsonb"})) + ) + + expected_missing = sorted(FULL_CAPABILITIES - {"parameterized_sql", "jsonb"}) + assert decision.eligible_for_parity_validation is False + assert decision.reasons == tuple( + f"missing_capability:{capability}" for capability in expected_missing + ) + + +@pytest.mark.parametrize( + ("field_name", "value"), + [ + ("package_name", ""), + ("package_name", 7), + ("package_version", ""), + ("package_version", False), + ("license_spdx", ""), + ("license_spdx", object()), + ("python_versions", ()), + ("python_versions", ["3.14"]), + ("source_commit_sha", "a" * 39), + ("source_commit_sha", "g" * 40), + ("source_commit_sha", 40), + ("artifact_sha256", "b" * 63), + ("artifact_sha256", "z" * 64), + ("artifact_sha256", 64), + ("vulnerability_report_sha256", "c" * 63), + ("vulnerability_report_sha256", "z" * 64), + ("vulnerability_report_sha256", 64), + ("capability_report_sha256", "e" * 63), + ("capability_report_sha256", "z" * 64), + ("capability_report_sha256", 64), + ("known_vulnerability_ids", ["CVE-2025-61385"]), + ("capabilities", frozenset()), + ("capabilities", {"parameterized_sql"}), + ], +) +def test_candidate_rejects_incomplete_or_nonimmutable_evidence( + field_name: str, + value: object, +) -> None: + with pytest.raises(PostgresDriverCandidateEvidenceError): + _evidence(**{field_name: value}) + + +def test_candidate_rejects_ambiguous_python_version_tokens() -> None: + with pytest.raises(PostgresDriverCandidateEvidenceError, match="Python version"): + _evidence(python_versions=("3.14+",)) + + +def test_candidate_rejects_non_string_python_version_tokens() -> None: + with pytest.raises(PostgresDriverCandidateEvidenceError, match="Python version"): + _evidence(python_versions=("3.14", 314)) + + +def test_candidate_rejects_unknown_capability_names() -> None: + with pytest.raises(PostgresDriverCandidateEvidenceError, match="capability"): + _evidence(capabilities=FULL_CAPABILITIES | {"model_routing"}) + + +@pytest.mark.parametrize( + "vulnerability_id", + [ + "", + "CVE 2025 61385", + "CVE-2025-61385\nGHSA-wq2g-r956-j8cc", + "x" * 129, + ], +) +def test_candidate_rejects_malformed_vulnerability_identifiers( + vulnerability_id: str, +) -> None: + with pytest.raises(PostgresDriverCandidateEvidenceError, match="vulnerability"): + _evidence(known_vulnerability_ids=(vulnerability_id,)) + + +def test_candidate_rejects_duplicate_vulnerability_identifiers() -> None: + with pytest.raises(PostgresDriverCandidateEvidenceError, match="vulnerability"): + _evidence(known_vulnerability_ids=("CVE-2025-61385", "CVE-2025-61385")) + + +def test_candidate_rejects_unbounded_python_version_evidence() -> None: + versions = FULL_PYTHON_VERSIONS + tuple(f"4.{minor}" for minor in range(28)) + + with pytest.raises(PostgresDriverCandidateEvidenceError, match="Python version"): + _evidence(python_versions=versions) + + +def test_candidate_rejects_unbounded_vulnerability_evidence() -> None: + vulnerability_ids = tuple(f"CVE-2099-{index:04d}" for index in range(257)) + + with pytest.raises(PostgresDriverCandidateEvidenceError, match="vulnerability"): + _evidence(known_vulnerability_ids=vulnerability_ids) + + +@pytest.mark.parametrize( + ("field_name", "mutated_value"), + [ + ("python_versions", ["3.14"]), + ("known_vulnerability_ids", ["CVE-2025-61385"]), + ("capabilities", set(FULL_CAPABILITIES)), + ], +) +def test_candidate_evaluation_revalidates_post_construction_container_mutation( + field_name: str, + mutated_value: object, +) -> None: + evidence = _evidence() + object.__setattr__(evidence, field_name, mutated_value) + + with pytest.raises(PostgresDriverCandidateEvidenceError): + evaluate_postgres_driver_candidate(evidence) + + +def test_candidate_evaluation_normalizes_deleted_authority_field() -> None: + evidence = _evidence() + object.__delattr__(evidence, "capabilities") + + with pytest.raises(PostgresDriverCandidateEvidenceError): + evaluate_postgres_driver_candidate(evidence) + + +def test_candidate_evaluation_rejects_shaped_object_before_member_access() -> None: + class CandidateShapedObject: + @property + def license_spdx(self) -> str: + raise AssertionError("candidate-shaped object member was evaluated") + + with pytest.raises(PostgresDriverCandidateEvidenceError): + evaluate_postgres_driver_candidate(CandidateShapedObject()) # type: ignore[arg-type] + + +@pytest.mark.parametrize("field_name", ["package_name", "package_version", "license_spdx"]) +def test_candidate_rejects_surrounding_whitespace_in_identity_evidence( + field_name: str, +) -> None: + with pytest.raises(PostgresDriverCandidateEvidenceError): + _evidence(**{field_name: " candidate-driver "}) + + +@pytest.mark.parametrize( + "identity_value", + [ + "candidate\ndriver", + "candidate\tdriver", + "candidate\u00a0driver", + "candidate\x7fdriver", + ], +) +def test_candidate_rejects_embedded_whitespace_or_control_identity_evidence( + identity_value: str, +) -> None: + with pytest.raises(PostgresDriverCandidateEvidenceError): + _evidence(package_name=identity_value) + + +@pytest.mark.parametrize( + "identity_value", + [ + "candidate\u200bdriver", + "candidate\u202edriver", + "candidate\u2066driver", + ], +) +def test_candidate_rejects_unicode_format_controls_in_identity_evidence( + identity_value: str, +) -> None: + with pytest.raises(PostgresDriverCandidateEvidenceError): + _evidence(package_name=identity_value) + + +@pytest.mark.parametrize("field_name", ["package_name", "package_version", "license_spdx"]) +def test_candidate_rejects_unbounded_identity_evidence(field_name: str) -> None: + with pytest.raises(PostgresDriverCandidateEvidenceError): + _evidence(**{field_name: "x" * 257}) + + +def test_candidate_rejects_duplicate_python_version_evidence() -> None: + with pytest.raises(PostgresDriverCandidateEvidenceError, match="Python version"): + _evidence(python_versions=("3.14", "3.14")) diff --git a/tests/test_postgres_driver_candidate_capability_shape.py b/tests/test_postgres_driver_candidate_capability_shape.py new file mode 100644 index 000000000..ecd6946bf --- /dev/null +++ b/tests/test_postgres_driver_candidate_capability_shape.py @@ -0,0 +1,48 @@ +"""Capability-shape regressions for PostgreSQL driver candidate evidence.""" + +from __future__ import annotations + +from dataclasses import replace + +import pytest + +from pg_llm_batch.postgres_driver_candidate import ( + REQUIRED_POSTGRES_DRIVER_CAPABILITIES, + REQUIRED_POSTGRES_DRIVER_PYTHON_VERSIONS, + PostgresDriverCandidateEvidence, + PostgresDriverCandidateEvidenceError, +) + + +class _StringSubclass(str): + """Represent a shaped string that must not enter supply-chain evidence.""" + + +def _valid_evidence() -> PostgresDriverCandidateEvidence: + """Build one exact primitive candidate receipt for shape-validation tests.""" + return PostgresDriverCandidateEvidence( + package_name="pg8000", + package_version="1.31.5", + license_spdx="BSD-3-Clause", + license_report_sha256="1" * 64, + python_versions=tuple(sorted(REQUIRED_POSTGRES_DRIVER_PYTHON_VERSIONS)), + source_commit_sha="2" * 40, + artifact_sha256="3" * 64, + vulnerability_report_sha256="4" * 64, + capability_report_sha256="5" * 64, + known_vulnerability_ids=(), + capabilities=REQUIRED_POSTGRES_DRIVER_CAPABILITIES, + ) + + +def test_candidate_rejects_string_subclass_capability_before_set_comparison() -> None: + """Supply-chain capability evidence must contain exact built-in strings only.""" + capabilities = set(REQUIRED_POSTGRES_DRIVER_CAPABILITIES) + capabilities.remove("jsonb") + capabilities.add(_StringSubclass("jsonb")) + + with pytest.raises( + PostgresDriverCandidateEvidenceError, + match="PostgreSQL driver capability evidence is invalid", + ): + replace(_valid_evidence(), capabilities=frozenset(capabilities)) diff --git a/tests/test_postgres_driver_candidate_package_name.py b/tests/test_postgres_driver_candidate_package_name.py new file mode 100644 index 000000000..a31774b8e --- /dev/null +++ b/tests/test_postgres_driver_candidate_package_name.py @@ -0,0 +1,38 @@ +from __future__ import annotations + +import pytest + +from pg_llm_batch.postgres_driver_candidate import ( + REQUIRED_POSTGRES_DRIVER_CAPABILITIES, + REQUIRED_POSTGRES_DRIVER_PYTHON_VERSIONS, + PostgresDriverCandidateEvidence, + PostgresDriverCandidateEvidenceError, +) + + +@pytest.mark.parametrize( + "package_name", + [ + ".candidate-driver", + "candidate-driver-", + "candidate/driver", + "candidate@driver", + "candidaté-driver", + ], +) +def test_candidate_rejects_non_pypa_distribution_names(package_name: str) -> None: + """Supply-chain evidence must use a valid Python distribution project name.""" + with pytest.raises(PostgresDriverCandidateEvidenceError, match="package name"): + PostgresDriverCandidateEvidence( + package_name=package_name, + package_version="1.2.3", + license_spdx="BSD-3-Clause", + license_report_sha256="d" * 64, + python_versions=tuple(sorted(REQUIRED_POSTGRES_DRIVER_PYTHON_VERSIONS)), + source_commit_sha="a" * 40, + artifact_sha256="b" * 64, + vulnerability_report_sha256="c" * 64, + capability_report_sha256="e" * 64, + known_vulnerability_ids=(), + capabilities=frozenset(REQUIRED_POSTGRES_DRIVER_CAPABILITIES), + ) diff --git a/tests/test_postgres_driver_candidate_report_provenance.py b/tests/test_postgres_driver_candidate_report_provenance.py new file mode 100644 index 000000000..4a0fafa44 --- /dev/null +++ b/tests/test_postgres_driver_candidate_report_provenance.py @@ -0,0 +1,48 @@ +from __future__ import annotations + +from dataclasses import fields + +import pytest + +from pg_llm_batch.postgres_driver_candidate import ( + PostgresDriverCandidateEvidence, + PostgresDriverCandidateEvidenceError, +) + + +def test_candidate_evidence_has_explicit_schema_version() -> None: + """Acquisition receipts need a versioned schema before their meaning can evolve.""" + evidence_fields = {field.name for field in fields(PostgresDriverCandidateEvidence)} + + assert "evidence_schema_version" in evidence_fields + + +def test_candidate_evidence_binds_capability_report_identity() -> None: + """Capability claims need immutable report identity before parity admission.""" + evidence_fields = {field.name for field in fields(PostgresDriverCandidateEvidence)} + + assert "capability_report_sha256" in evidence_fields + + +def test_candidate_schema_version_rejects_equality_spoofing() -> None: + """Untrusted receipt objects cannot impersonate the current schema by equality.""" + + class PretendsToBeCurrent: + def __eq__(self, other: object) -> bool: + return True + + with pytest.raises(PostgresDriverCandidateEvidenceError, match="schema version"): + PostgresDriverCandidateEvidence( + package_name="candidate-driver", + package_version="1.2.3", + license_spdx="BSD-3-Clause", + license_report_sha256="d" * 64, + python_versions=("3.10", "3.11", "3.12", "3.13", "3.14"), + source_commit_sha="a" * 40, + artifact_sha256="b" * 64, + vulnerability_report_sha256="c" * 64, + capability_report_sha256="e" * 64, + known_vulnerability_ids=(), + capabilities=frozenset({"parameterized_sql"}), + evidence_schema_version=PretendsToBeCurrent(), # type: ignore[arg-type] + ) diff --git a/tests/test_postgres_driver_candidate_surrogate.py b/tests/test_postgres_driver_candidate_surrogate.py new file mode 100644 index 000000000..2e8660a4f --- /dev/null +++ b/tests/test_postgres_driver_candidate_surrogate.py @@ -0,0 +1,28 @@ +from __future__ import annotations + +import pytest + +from pg_llm_batch.postgres_driver_candidate import ( + REQUIRED_POSTGRES_DRIVER_CAPABILITIES, + REQUIRED_POSTGRES_DRIVER_PYTHON_VERSIONS, + PostgresDriverCandidateEvidence, + PostgresDriverCandidateEvidenceError, +) + + +def test_candidate_rejects_isolated_surrogate_with_domain_error() -> None: + """Malformed Unicode metadata must stay inside the candidate-evidence boundary.""" + with pytest.raises(PostgresDriverCandidateEvidenceError, match="package name"): + PostgresDriverCandidateEvidence( + package_name="candidate\ud800driver", + package_version="1.2.3", + license_spdx="BSD-3-Clause", + license_report_sha256="d" * 64, + python_versions=tuple(sorted(REQUIRED_POSTGRES_DRIVER_PYTHON_VERSIONS)), + source_commit_sha="a" * 40, + artifact_sha256="b" * 64, + vulnerability_report_sha256="c" * 64, + capability_report_sha256="e" * 64, + known_vulnerability_ids=(), + capabilities=frozenset(REQUIRED_POSTGRES_DRIVER_CAPABILITIES), + ) diff --git a/tests/test_postgres_driver_candidate_thread_capability.py b/tests/test_postgres_driver_candidate_thread_capability.py new file mode 100644 index 000000000..34e00d902 --- /dev/null +++ b/tests/test_postgres_driver_candidate_thread_capability.py @@ -0,0 +1,64 @@ +"""Commercial-candidate concurrency evidence must match DB-API connection ownership. + +A replacement driver can be permissively licensed and otherwise satisfy the SQL, +transaction, type, and conninfo contract while still declaring connections unsafe +to share across threads. Candidate admission therefore requires explicit evidence +for the connection thread-affinity policy instead of inferring concurrency safety +from module importability or successful single-thread PostgreSQL smokes. +""" + +from __future__ import annotations + +from pg_llm_batch.postgres_driver_candidate import ( + PostgresDriverCandidateEvidence, + evaluate_postgres_driver_candidate, +) + + +_LEGACY_CAPABILITIES_WITHOUT_THREAD_AFFINITY = frozenset( + { + "autocommit_state", + "connection_closed_state", + "connection_context", + "connection_context_commit_rollback", + "conninfo_keyword_parse_render", + "conninfo_service_selector", + "conninfo_uri_parse_render", + "cursor_context", + "finite_connect_timeout", + "invalid_conninfo_classification", + "jsonb", + "parameterized_sql", + "result_row_semantics", + "row_count", + "sql_parameter_style_adaptation", + "transaction_commit_rollback", + "undefined_function_classification", + "uuid_timestamp_adaptation", + } +) + + +def test_candidate_requires_explicit_connection_thread_affinity_evidence() -> None: + """Reject the former complete receipt when thread-ownership evidence is absent.""" + evidence = PostgresDriverCandidateEvidence( + package_name="candidate-driver", + package_version="1.2.3", + license_spdx="BSD-3-Clause", + license_report_sha256="d" * 64, + python_versions=("3.10", "3.11", "3.12", "3.13", "3.14"), + source_commit_sha="a" * 40, + artifact_sha256="b" * 64, + vulnerability_report_sha256="c" * 64, + capability_report_sha256="e" * 64, + known_vulnerability_ids=(), + capabilities=_LEGACY_CAPABILITIES_WITHOUT_THREAD_AFFINITY, + ) + + decision = evaluate_postgres_driver_candidate(evidence) + + assert decision.eligible_for_parity_validation is False + assert decision.production_approved is False + assert decision.reasons == ( + "missing_capability:connection_thread_affinity", + ) diff --git a/tests/test_postgres_driver_connection_context_contract.py b/tests/test_postgres_driver_connection_context_contract.py new file mode 100644 index 000000000..e15c8e1ec --- /dev/null +++ b/tests/test_postgres_driver_connection_context_contract.py @@ -0,0 +1,10 @@ +from __future__ import annotations + +from pg_llm_batch.postgres_driver_candidate import ( + REQUIRED_POSTGRES_DRIVER_CAPABILITIES, +) + + +def test_candidate_requires_transactional_connection_context_semantics() -> None: + """Reject drivers whose context manager closes without commit/rollback parity.""" + assert "connection_context_commit_rollback" in REQUIRED_POSTGRES_DRIVER_CAPABILITIES diff --git a/tests/test_postgres_driver_port.py b/tests/test_postgres_driver_port.py new file mode 100644 index 000000000..2b2a7fe68 --- /dev/null +++ b/tests/test_postgres_driver_port.py @@ -0,0 +1,235 @@ +from __future__ import annotations + +from collections.abc import Mapping +from typing import get_type_hints + +import pytest + +from pg_llm_batch.postgres_driver_port import ( + PostgresConnectionPort, + PostgresCursorPort, + PostgresDriverPort, +) + + +def test_cursor_port_covers_existing_database_interaction_surface() -> None: + assert PostgresCursorPort.__abstractmethods__ == { + "__enter__", + "__exit__", + "execute", + "executemany", + "fetchall", + "fetchmany", + "fetchone", + "row_count", + } + + +def test_cursor_port_declares_canonical_tuple_result_rows() -> None: + assert get_type_hints(PostgresCursorPort.fetchone)["return"] == ( + tuple[object, ...] | None + ) + assert get_type_hints(PostgresCursorPort.fetchmany)["return"] == list[ + tuple[object, ...] + ] + assert get_type_hints(PostgresCursorPort.fetchall)["return"] == list[ + tuple[object, ...] + ] + + +def test_connection_port_covers_transaction_and_cursor_lifecycle() -> None: + assert PostgresConnectionPort.__abstractmethods__ == { + "__enter__", + "__exit__", + "close", + "commit", + "cursor", + "execute", + "is_closed", + "rollback", + "set_autocommit", + } + + +def test_driver_port_covers_psycopg_replacement_capabilities_only() -> None: + assert PostgresDriverPort.__abstractmethods__ == { + "connect", + "is_invalid_conninfo", + "is_undefined_function", + "jsonb", + "make_conninfo", + "parse_conninfo", + } + assert not hasattr(PostgresDriverPort, "discover_model") + assert not hasattr(PostgresDriverPort, "route_model") + assert not hasattr(PostgresDriverPort, "select_provider") + + +class _Cursor(PostgresCursorPort): + def __init__(self) -> None: + self.executions: list[tuple[str, object | None]] = [] + self.affected_rows = 0 + + def execute(self, query: str, params: object | None = None) -> _Cursor: + self.executions.append((query, params)) + self.affected_rows = 1 + return self + + def executemany(self, query: str, params_seq: object) -> _Cursor: + self.executions.append((query, params_seq)) + self.affected_rows = 2 + return self + + def fetchone(self) -> tuple[object, ...] | None: + return ("row",) + + def fetchmany(self, size: int) -> list[tuple[object, ...]]: + return [("row",)] * size + + def fetchall(self) -> list[tuple[object, ...]]: + return [("row",)] + + def row_count(self) -> int: + return self.affected_rows + + def __enter__(self) -> _Cursor: + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + traceback: object | None, + ) -> None: + return None + + +class _Connection(PostgresConnectionPort): + def __init__(self) -> None: + self.cursor_instance = _Cursor() + self.committed = False + self.rolled_back = False + self.closed = False + self.autocommit = False + + def cursor(self) -> _Cursor: + return self.cursor_instance + + def execute(self, query: str, params: object | None = None) -> _Cursor: + return self.cursor_instance.execute(query, params) + + def commit(self) -> None: + self.committed = True + + def rollback(self) -> None: + self.rolled_back = True + + def set_autocommit(self, enabled: bool) -> None: + self.autocommit = enabled + + def is_closed(self) -> bool: + return self.closed + + def close(self) -> None: + self.closed = True + + def __enter__(self) -> _Connection: + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + traceback: object | None, + ) -> None: + self.close() + + +class _UndefinedFunctionError(Exception): + pass + + +class _InvalidConninfoError(Exception): + pass + + +class _Driver(PostgresDriverPort): + def connect( + self, + dsn: str, + *, + connect_timeout_seconds: int | None = None, + ) -> _Connection: + assert dsn == "service=pg_llm_batch" + assert connect_timeout_seconds == 5 + return _Connection() + + def parse_conninfo(self, dsn: str) -> Mapping[str, str]: + key, value = dsn.split("=", 1) + return {key: value} + + def make_conninfo(self, params: Mapping[str, str]) -> str: + return " ".join(f"{key}={value}" for key, value in sorted(params.items())) + + def jsonb(self, value: object) -> object: + return ("jsonb", value) + + def is_invalid_conninfo(self, error: BaseException) -> bool: + return isinstance(error, _InvalidConninfoError) + + def is_undefined_function(self, error: BaseException) -> bool: + return isinstance(error, _UndefinedFunctionError) + + +def test_complete_port_can_run_without_psycopg_types() -> None: + driver = _Driver() + + connection = driver.connect( + "service=pg_llm_batch", + connect_timeout_seconds=5, + ) + assert connection.is_closed() is False + connection.set_autocommit(True) + assert connection.autocommit is True + + with connection as active_connection: + with active_connection.cursor() as cursor: + cursor.execute("SELECT %s", ("tenant-a",)) + assert cursor.row_count() == 1 + cursor.executemany("SELECT %s", [("tenant-a",), ("tenant-b",)]) + assert cursor.row_count() == 2 + assert cursor.fetchone() == ("row",) + assert cursor.fetchmany(1) == [("row",)] + assert cursor.fetchall() == [("row",)] + active_connection.commit() + + assert connection.committed is True + assert connection.is_closed() is True + assert driver.parse_conninfo("service=pg_llm_batch") == { + "service": "pg_llm_batch" + } + assert driver.make_conninfo({"service": "pg_llm_batch"}) == ( + "service=pg_llm_batch" + ) + assert driver.jsonb({"request_id": "opaque"}) == ( + "jsonb", + {"request_id": "opaque"}, + ) + assert driver.is_invalid_conninfo(_InvalidConninfoError()) is True + assert driver.is_invalid_conninfo(RuntimeError()) is False + assert driver.is_undefined_function(_UndefinedFunctionError()) is True + assert driver.is_undefined_function(RuntimeError()) is False + + +def test_incomplete_driver_cannot_be_instantiated() -> None: + class _IncompleteDriver(PostgresDriverPort): + def connect( + self, + dsn: str, + *, + connect_timeout_seconds: int | None = None, + ) -> PostgresConnectionPort: + raise AssertionError("not called") + + with pytest.raises(TypeError): + _IncompleteDriver() diff --git a/tests/test_postgres_driver_remote_lifecycle.py b/tests/test_postgres_driver_remote_lifecycle.py new file mode 100644 index 000000000..6770ad2d6 --- /dev/null +++ b/tests/test_postgres_driver_remote_lifecycle.py @@ -0,0 +1,181 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Driver-port regressions for durable remote batch lifecycle persistence.""" + +from __future__ import annotations + +from datetime import datetime, timezone +from typing import Any + +import pytest + +from pg_llm_batch import db + + +class _Cursor: + """Expose the cursor behavior required by the lifecycle migration seam.""" + + def __init__(self, driver: _Driver) -> None: + self.driver = driver + + def __enter__(self) -> _Cursor: + return self + + def __exit__(self, *_exc: Any) -> None: + return None + + def execute(self, query: str, params: object | None = None) -> _Cursor: + self.driver.executions.append((query, params)) + return self + + def fetchone(self) -> tuple[object, ...] | None: + if not self.driver.rows: + return None + return self.driver.rows.pop(0) + + def row_count(self) -> int: + return self.driver.affected_rows + + +class _Connection: + """Retain one fake transaction and expose commit evidence.""" + + def __init__(self, driver: _Driver) -> None: + self.driver = driver + + def __enter__(self) -> _Connection: + return self + + def __exit__(self, *_exc: Any) -> None: + return None + + def cursor(self) -> _Cursor: + return _Cursor(self.driver) + + def commit(self) -> None: + self.driver.commits += 1 + + +class _Driver: + """Minimal driver-shaped test double that contains no Psycopg type.""" + + def __init__( + self, + *, + rows: list[tuple[object, ...]] | None = None, + affected_rows: int = 1, + ) -> None: + self.rows = list(rows or []) + self.affected_rows = affected_rows + self.executions: list[tuple[str, object | None]] = [] + self.connections: list[str] = [] + self.commits = 0 + + def connect(self, dsn: str) -> _Connection: + self.connections.append(dsn) + return _Connection(self) + + +def _persisted_remote_batch_row( + observed: datetime, +) -> tuple[object, ...]: + """Return one exact persisted lifecycle row used by stale-write regressions.""" + return ( + "standalone", + "primary", + "batch-1", + 1, + None, + "/v1/responses", + "in_progress", + None, + None, + 2, + 1, + 0, + {}, + observed, + observed, + None, + observed, + ) + + +def test_observation_order_reservation_uses_injected_driver_without_psycopg( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Driver migration must retain the default standalone tenant boundary.""" + driver = _Driver(rows=[(41,)]) + monkeypatch.setattr(db, "psycopg", None) + + order = db.reserve_remote_batch_observation_order( + "postgresql://x", + postgres_driver=driver, + ) + + assert order == 41 + assert driver.connections == ["postgresql://x"] + assert driver.executions == [ + ( + "SELECT set_config('pg_llm_batch.tenant_scope', %s, true)", + ("standalone",), + ), + ("SELECT nextval('llm_remote_batch_observation_sequence')", None), + ] + + +def test_stale_lifecycle_write_reads_persisted_state_through_injected_driver( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Port row-count semantics must preserve the stale-write recovery contract.""" + observed = datetime(2026, 9, 2, 12, 0, tzinfo=timezone.utc) + driver = _Driver( + rows=[_persisted_remote_batch_row(observed)], + affected_rows=0, + ) + monkeypatch.setattr(db, "psycopg", None) + + snapshot = db.persist_remote_batch_state( + "postgresql://x", + "primary", + { + "id": "batch-1", + "endpoint": "/v1/responses", + "status": "in_progress", + "request_counts": {"total": 2, "completed": 2, "failed": 0}, + }, + observation_order=2, + observed_at=observed, + postgres_driver=driver, + ) + + assert snapshot["observation_order"] == 1 + assert snapshot["completed_requests"] == 1 + assert any( + "SELECT tenant_scope" in query + for query, _params in driver.executions + ) + assert driver.commits == 1 + + +def test_remote_lifecycle_read_uses_injected_driver_without_psycopg( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Tenant-scoped reads must not reacquire the legacy driver implicitly.""" + observed = datetime(2026, 9, 2, 12, 0, tzinfo=timezone.utc) + driver = _Driver(rows=[_persisted_remote_batch_row(observed)]) + monkeypatch.setattr(db, "psycopg", None) + + snapshot = db.get_remote_batch_state( + "postgresql://x", + "primary", + "batch-1", + postgres_driver=driver, + ) + + assert snapshot is not None + assert snapshot["tenant_scope"] == "standalone" + assert snapshot["observation_order"] == 1 + assert driver.executions[0] == ( + "SELECT set_config('pg_llm_batch.tenant_scope', %s, true)", + ("standalone",), + ) diff --git a/tests/test_postgres_driver_review_contracts.py b/tests/test_postgres_driver_review_contracts.py new file mode 100644 index 000000000..6ac3039e4 --- /dev/null +++ b/tests/test_postgres_driver_review_contracts.py @@ -0,0 +1,167 @@ +from __future__ import annotations + +from datetime import datetime, timezone +from typing import Any, get_type_hints + +import pytest + +from pg_llm_batch import db +from pg_llm_batch.exceptions import ValidationError +from pg_llm_batch.postgres_driver_port import PostgresCursorPort +from pg_llm_batch.psycopg_driver_adapter import PsycopgCursorAdapter + + +class _RawCursor: + def __init__(self, rowcount: object) -> None: + self.rowcount = rowcount + + +class _Cursor: + def __init__(self, driver: _Driver) -> None: + self.driver = driver + + def __enter__(self) -> _Cursor: + return self + + def __exit__(self, *_exc: Any) -> None: + return None + + def execute(self, query: str, params: object | None = None) -> _Cursor: + self.driver.executions.append((query, params)) + return self + + def fetchone(self) -> tuple[object, ...] | None: + if not self.driver.rows: + return None + return self.driver.rows.pop(0) + + def row_count(self) -> int | None: + return self.driver.affected_rows + + +class _Connection: + def __init__(self, driver: _Driver) -> None: + self.driver = driver + + def __enter__(self) -> _Connection: + return self + + def __exit__(self, *_exc: Any) -> None: + return None + + def cursor(self) -> _Cursor: + return _Cursor(self.driver) + + def commit(self) -> None: + self.driver.commits += 1 + + +class _Driver: + def __init__( + self, + *, + rows: list[tuple[object, ...]] | None = None, + affected_rows: int | None = 1, + ) -> None: + self.rows = list(rows or []) + self.affected_rows = affected_rows + self.executions: list[tuple[str, object | None]] = [] + self.connections: list[str] = [] + self.commits = 0 + + def connect(self, dsn: str) -> _Connection: + self.connections.append(dsn) + return _Connection(self) + + +def _persisted_remote_batch_row(observed: datetime) -> tuple[object, ...]: + return ( + "standalone", + "primary", + "batch-1", + 1, + None, + "/v1/responses", + "in_progress", + None, + None, + 2, + 1, + 0, + {}, + observed, + observed, + None, + observed, + ) + + +def test_cursor_port_exposes_unknown_row_count_without_driver_sentinel() -> None: + """The provider-neutral cursor contract must represent unknown counts explicitly.""" + assert get_type_hints(PostgresCursorPort.row_count)["return"] == int | None + + +def test_psycopg_cursor_normalizes_unknown_row_count() -> None: + """Psycopg's -1 sentinel must not leak through the provider-neutral port.""" + assert PsycopgCursorAdapter(_RawCursor(-1)).row_count() is None + + +def test_observation_order_binds_validated_tenant_scope_before_sequence_io() -> None: + """A tenant reservation must establish transaction-local RLS identity first.""" + driver = _Driver(rows=[(41,)]) + + order = db.reserve_remote_batch_observation_order( + "postgresql://x", + tenant_scope="tenant-a", + postgres_driver=driver, + ) + + assert order == 41 + assert driver.executions == [ + ( + "SELECT set_config('pg_llm_batch.tenant_scope', %s, true)", + ("tenant-a",), + ), + ("SELECT nextval('llm_remote_batch_observation_sequence')", None), + ] + + +def test_observation_order_rejects_invalid_tenant_before_database_io() -> None: + """Untrusted tenant text must fail before a connection or sequence reservation.""" + driver = _Driver(rows=[(41,)]) + + with pytest.raises(ValidationError, match="tenant_scope"): + db.reserve_remote_batch_observation_order( + "postgresql://x", + tenant_scope="tenant scope", + postgres_driver=driver, + ) + + assert driver.connections == [] + assert driver.executions == [] + + +def test_unknown_remote_lifecycle_row_count_reads_persisted_state() -> None: + """Unknown affected-row evidence must not be guessed as a successful UPSERT.""" + observed = datetime(2026, 9, 3, 1, 0, tzinfo=timezone.utc) + driver = _Driver( + rows=[_persisted_remote_batch_row(observed)], + affected_rows=None, + ) + + snapshot = db.persist_remote_batch_state( + "postgresql://x", + "primary", + { + "id": "batch-1", + "endpoint": "/v1/responses", + "status": "in_progress", + "request_counts": {"total": 2, "completed": 2, "failed": 0}, + }, + observation_order=2, + observed_at=observed, + postgres_driver=driver, + ) + + assert snapshot["observation_order"] == 1 + assert snapshot["completed_requests"] == 1 diff --git a/tests/test_postgres_driver_runtime_selection.py b/tests/test_postgres_driver_runtime_selection.py new file mode 100644 index 000000000..b7240c0ab --- /dev/null +++ b/tests/test_postgres_driver_runtime_selection.py @@ -0,0 +1,125 @@ +"""Regressions for the retained PostgreSQL driver selection boundary. + +The commercial driver migration must leave concrete Psycopg authority in one +infrastructure adapter rather than importing the package from each bounded +context. These tests exercise the default connection path through a lazy runtime +selector while preserving explicit driver injection. +""" + +from __future__ import annotations + +from typing import Any + +import pg_llm_batch.checkpoint_store as checkpoint_store +import pg_llm_batch.config as config +import pg_llm_batch.db as db +import pg_llm_batch.health as health +import pg_llm_batch.orchestrator as orchestrator +import pg_llm_batch.token_counter as token_counter +from pg_llm_batch.postgres_driver_runtime import retained_postgres_driver + + +class _Connection: + """Represent one exact connection returned by the selected fake driver.""" + + +class _Driver: + """Capture default-driver connection attempts without a concrete client.""" + + def __init__(self) -> None: + self.dsns: list[str] = [] + self.connection_kwargs: list[dict[str, Any]] = [] + self.connection = _Connection() + + def connect(self, dsn: str, **kwargs: Any) -> _Connection: + """Record the exact DSN and connection options before returning the fake.""" + self.dsns.append(dsn) + self.connection_kwargs.append(dict(kwargs)) + return self.connection + + +def test_db_default_connection_uses_runtime_driver_selector(monkeypatch) -> None: + """Low-level DB helpers must not own a second concrete-driver import path.""" + driver = _Driver() + monkeypatch.setattr(db, "retained_postgres_driver", lambda: driver) + + connection = db._connect_database("postgresql://unit", None) + + assert connection is driver.connection + assert driver.dsns == ["postgresql://unit"] + assert driver.connection_kwargs == [{}] + + +def test_checkpoint_default_connection_uses_runtime_driver_selector(monkeypatch) -> None: + """Checkpoint persistence must share the same retained-driver authority.""" + driver = _Driver() + monkeypatch.setattr(checkpoint_store, "retained_postgres_driver", lambda: driver) + + connection = checkpoint_store._connect_postgres("postgresql://unit", None) + + assert connection is driver.connection + assert driver.dsns == ["postgresql://unit"] + assert driver.connection_kwargs == [{}] + + +def test_health_default_connection_uses_runtime_driver_selector(monkeypatch) -> None: + """Readiness must share the retained driver and preserve its finite timeout.""" + driver = _Driver() + monkeypatch.setattr(health, "retained_postgres_driver", lambda: driver) + + connection = health._connect_health_database("postgresql://unit", None) + + assert connection is driver.connection + assert driver.dsns == ["postgresql://unit"] + assert driver.connection_kwargs == [{"connect_timeout_seconds": 5}] + + +def test_config_default_connection_uses_runtime_driver_selector(monkeypatch) -> None: + """Configuration persistence must not retain a second concrete-client authority.""" + driver = _Driver() + monkeypatch.setattr(config, "retained_postgres_driver", lambda: driver) + + connection = config._connect_store_database( + "postgresql://unit", + None, + missing_dependency_message="driver unavailable", + ) + + assert connection is driver.connection + assert driver.dsns == ["postgresql://unit"] + assert driver.connection_kwargs == [{}] + + +def test_token_counter_default_driver_uses_runtime_selector(monkeypatch) -> None: + """Token counting must acquire its retained database capability centrally.""" + driver = _Driver() + monkeypatch.setattr(token_counter, "retained_postgres_driver", lambda: driver) + monkeypatch.setattr( + token_counter.TokenCounter, + "_ensure_pg_tiktoken", + lambda self: False, + ) + + counter = token_counter.TokenCounter("postgresql://unit") + + assert counter._postgres_driver is driver + + +def test_orchestrator_default_driver_uses_runtime_selector(monkeypatch) -> None: + """Batch assembly must not retain a direct concrete Psycopg authority path.""" + driver = _Driver() + monkeypatch.setattr(orchestrator, "retained_postgres_driver", lambda: driver) + + service = orchestrator.PostgresBatchOrchestrator("postgresql://unit") + + assert service._postgres_driver is driver + + +def test_runtime_selector_returns_postgres_driver_port() -> None: + """The retained selector must expose only the provider-neutral driver port.""" + driver = retained_postgres_driver() + + assert callable(driver.connect) + assert callable(driver.parse_conninfo) + assert callable(driver.make_conninfo) + assert callable(driver.jsonb) diff --git a/tests/test_psycopg_driver_adapter.py b/tests/test_psycopg_driver_adapter.py new file mode 100644 index 000000000..45a643002 --- /dev/null +++ b/tests/test_psycopg_driver_adapter.py @@ -0,0 +1,295 @@ +from __future__ import annotations + +from typing import Any + +import pytest +from psycopg import ProgrammingError +from psycopg.errors import UndefinedFunction +from psycopg.types.json import Jsonb + +from pg_llm_batch.postgres_driver_port import PostgresDriverPort +from pg_llm_batch.psycopg_driver_adapter import ( + PsycopgConnectionAdapter, + PsycopgCursorAdapter, + PsycopgDriverAdapter, + PsycopgDriverAdapterError, +) + + +class _RawCursor: + def __init__(self) -> None: + self.executions: list[tuple[str, object | None]] = [] + self.many_executions: list[tuple[str, object]] = [] + self.rows: list[object] = [["one"], ["two"]] + self.rowcount: object = 2 + self.entered = False + self.exited = False + + def execute(self, query: str, params: object | None = None) -> _RawCursor: + self.executions.append((query, params)) + return self + + def executemany(self, query: str, params_seq: object) -> None: + self.many_executions.append((query, params_seq)) + + def fetchone(self) -> object | None: + return self.rows[0] if self.rows else None + + def fetchmany(self, size: int) -> list[object]: + return list(self.rows[:size]) + + def fetchall(self) -> list[object]: + return list(self.rows) + + def __enter__(self) -> _RawCursor: + self.entered = True + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + traceback: object | None, + ) -> bool | None: + self.exited = True + return None + + +class _RawConnection: + def __init__(self) -> None: + self.autocommit = False + self.closed = False + self.cursor_value = _RawCursor() + self.commits = 0 + self.rollbacks = 0 + self.close_calls = 0 + self.entered = False + self.exited = False + + def cursor(self) -> _RawCursor: + return self.cursor_value + + def execute(self, query: str, params: object | None = None) -> _RawCursor: + return self.cursor_value.execute(query, params) + + def commit(self) -> None: + self.commits += 1 + + def rollback(self) -> None: + self.rollbacks += 1 + + def close(self) -> None: + self.close_calls += 1 + self.closed = True + + def __enter__(self) -> _RawConnection: + self.entered = True + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + traceback: object | None, + ) -> bool | None: + self.exited = True + return None + + +def test_psycopg_driver_is_provider_neutral_port_implementation() -> None: + assert isinstance(PsycopgDriverAdapter(), PostgresDriverPort) + + +def test_cursor_adapter_preserves_parameter_and_result_authority() -> None: + raw = _RawCursor() + cursor = PsycopgCursorAdapter(raw) + + assert cursor.execute("SELECT %s", ("tenant-a",)) is cursor + assert raw.executions == [("SELECT %s", ("tenant-a",))] + + batch = [(1,), (2,)] + assert cursor.executemany("INSERT INTO example VALUES (%s)", batch) is cursor + assert raw.many_executions == [("INSERT INTO example VALUES (%s)", batch)] + + assert cursor.fetchone() == ("one",) + assert cursor.fetchmany(1) == [("one",)] + assert cursor.fetchall() == [("one",), ("two",)] + assert cursor.row_count() == 2 + + +def test_cursor_adapter_rejects_none_inside_materialized_result_page() -> None: + raw = _RawCursor() + raw.rows = [["one"], None] + cursor = PsycopgCursorAdapter(raw) + + with pytest.raises(PsycopgDriverAdapterError, match="result row"): + cursor.fetchmany(2) + with pytest.raises(PsycopgDriverAdapterError, match="result row"): + cursor.fetchall() + + +@pytest.mark.parametrize("invalid_size", [True, 0, -1, 1.5]) +def test_cursor_adapter_rejects_non_positive_or_non_integer_fetch_budget( + invalid_size: object, +) -> None: + raw = _RawCursor() + + with pytest.raises(PsycopgDriverAdapterError, match="fetch size"): + PsycopgCursorAdapter(raw).fetchmany(invalid_size) # type: ignore[arg-type] + + +def test_cursor_adapter_rejects_driver_overdelivery_beyond_fetch_budget() -> None: + """Fail closed if a concrete driver violates the port's finite fetch budget.""" + + class _OverdeliveringCursor(_RawCursor): + def fetchmany(self, size: int) -> list[object]: + assert size == 1 + return list(self.rows) + + raw = _OverdeliveringCursor() + + with pytest.raises(PsycopgDriverAdapterError, match="exceeds requested size"): + PsycopgCursorAdapter(raw).fetchmany(1) + + +def test_cursor_adapter_rejects_non_integer_row_count() -> None: + raw = _RawCursor() + raw.rowcount = True + + with pytest.raises(PsycopgDriverAdapterError, match="row count"): + PsycopgCursorAdapter(raw).row_count() + + +def test_cursor_adapter_preserves_context_manager_boundary() -> None: + raw = _RawCursor() + cursor = PsycopgCursorAdapter(raw) + + with cursor as entered: + assert entered is cursor + + assert raw.entered is True + assert raw.exited is True + + +def test_connection_adapter_preserves_transaction_and_session_semantics() -> None: + raw = _RawConnection() + connection = PsycopgConnectionAdapter(raw) + + cursor = connection.cursor() + assert isinstance(cursor, PsycopgCursorAdapter) + assert cursor.execute("SELECT %s", (1,)) is cursor + + direct = connection.execute("SELECT %s", (2,)) + assert isinstance(direct, PsycopgCursorAdapter) + assert raw.cursor_value.executions[-1] == ("SELECT %s", (2,)) + + connection.commit() + connection.rollback() + assert raw.commits == 1 + assert raw.rollbacks == 1 + + connection.set_autocommit(True) + assert raw.autocommit is True + assert connection.is_closed() is False + + connection.close() + assert raw.close_calls == 1 + assert connection.is_closed() is True + + +def test_connection_adapter_rejects_non_boolean_autocommit() -> None: + raw = _RawConnection() + + with pytest.raises(PsycopgDriverAdapterError, match="autocommit"): + PsycopgConnectionAdapter(raw).set_autocommit(1) # type: ignore[arg-type] + + assert raw.autocommit is False + + +def test_connection_adapter_preserves_context_manager_boundary() -> None: + raw = _RawConnection() + connection = PsycopgConnectionAdapter(raw) + + with connection as entered: + assert entered is connection + + assert raw.entered is True + assert raw.exited is True + + +def test_driver_uses_psycopg_conninfo_and_jsonb_contracts() -> None: + driver = PsycopgDriverAdapter() + + parsed = driver.parse_conninfo("host=localhost dbname='batch db' application_name=pg-llm-batch") + assert parsed == { + "host": "localhost", + "dbname": "batch db", + "application_name": "pg-llm-batch", + } + + rendered = driver.make_conninfo(parsed) + assert driver.parse_conninfo(rendered) == parsed + + adapted = driver.jsonb({"count": 1}) + assert isinstance(adapted, Jsonb) + assert adapted.obj == {"count": 1} + + +def test_driver_classifies_only_adapter_owned_conninfo_grammar_failures() -> None: + driver = PsycopgDriverAdapter() + + with pytest.raises(PsycopgDriverAdapterError) as invalid_conninfo: + driver.parse_conninfo("host='unterminated") + + assert driver.is_invalid_conninfo(invalid_conninfo.value) is True + assert driver.is_invalid_conninfo(ProgrammingError("syntax error")) is False + assert driver.is_invalid_conninfo(RuntimeError("invalid conninfo")) is False + + +def test_driver_classifies_only_psycopg_undefined_function() -> None: + driver = PsycopgDriverAdapter() + + assert driver.is_undefined_function(UndefinedFunction("missing")) is True + assert driver.is_undefined_function(RuntimeError("missing")) is False + + +def test_driver_connect_preserves_exact_timeout_and_wraps_connection( + monkeypatch: pytest.MonkeyPatch, +) -> None: + raw = _RawConnection() + calls: list[tuple[str, dict[str, Any]]] = [] + + def fake_connect(dsn: str, **kwargs: Any) -> _RawConnection: + calls.append((dsn, kwargs)) + return raw + + monkeypatch.setattr("pg_llm_batch.psycopg_driver_adapter.psycopg.connect", fake_connect) + + connection = PsycopgDriverAdapter().connect( + "host=localhost dbname=batch", + connect_timeout_seconds=7, + ) + + assert isinstance(connection, PsycopgConnectionAdapter) + assert calls == [("host=localhost dbname=batch", {"connect_timeout": 7})] + + +def test_driver_connect_fails_closed_on_invalid_timeout_before_driver_call( + monkeypatch: pytest.MonkeyPatch, +) -> None: + called = False + + def fake_connect(*args: object, **kwargs: object) -> None: + nonlocal called + called = True + + monkeypatch.setattr("pg_llm_batch.psycopg_driver_adapter.psycopg.connect", fake_connect) + + for invalid in (True, 0, -1, 1.5): + with pytest.raises(PsycopgDriverAdapterError, match="timeout"): + PsycopgDriverAdapter().connect( + "host=localhost dbname=batch", + connect_timeout_seconds=invalid, # type: ignore[arg-type] + ) + + assert called is False \ No newline at end of file diff --git a/tests/test_token_counter.py b/tests/test_token_counter.py index a8774a768..406a4d98f 100644 --- a/tests/test_token_counter.py +++ b/tests/test_token_counter.py @@ -3,22 +3,27 @@ from __future__ import annotations +from threading import RLock + import pytest -from pg_llm_batch import token_counter as tc_mod from pg_llm_batch import db as db_mod +from pg_llm_batch import token_counter as tc_mod from pg_llm_batch.exceptions import TokenLimitExceededError, ValidationError from pg_llm_batch.models import BatchRequest +from pg_llm_batch.postgres_driver_runtime import PostgresDriverUnavailableError from pg_llm_batch.token_counter import BatchAccumulator, TokenCounter from tests.conftest import FakePsycopg +from tests.fake_postgres_driver_port import FakePsycopgDriverPort @pytest.fixture() def fake_pg(monkeypatch): + """Bind the legacy in-memory SQL fake through the production driver port.""" fake = FakePsycopg() - monkeypatch.setattr(tc_mod, "psycopg", fake) - monkeypatch.setattr(tc_mod, "UndefinedFunction", fake.errors.UndefinedFunction) - monkeypatch.setattr(db_mod, "psycopg", fake) + driver = FakePsycopgDriverPort(fake) + monkeypatch.setattr(tc_mod, "retained_postgres_driver", lambda: driver) + monkeypatch.setattr(db_mod, "retained_postgres_driver", lambda: driver) return fake @@ -45,7 +50,10 @@ def test_db_tokenizer_lookup_prefers_mapping(fake_pg, monkeypatch): monkeypatch.setattr( tc_mod, "get_model_metadata", - lambda dsn, model: {"mode": "chat", "tokenizer_model": "o200k_base"}, + lambda dsn, model, *, postgres_driver=None: { + "mode": "chat", + "tokenizer_model": "o200k_base", + }, ) counter = TokenCounter("postgresql://x") assert counter.get_tiktoken_name("some-deployment") == "o200k_base" @@ -79,7 +87,7 @@ def test_split_oversized_batch(fake_pg): requests = [ BatchRequest(user_prompt="a b", model="m"), # 2 BatchRequest(user_prompt="c d", model="m"), # 2 -> new batch - BatchRequest(user_prompt="e", model="m"), # 1 -> fits with prev + BatchRequest(user_prompt="e", model="m"), # 1 -> fits with prev ] batches = counter.split_oversized_batch(requests) assert len(batches) == 2 @@ -185,12 +193,12 @@ def test_count_tokens_fails_closed_when_extension_disappears(fake_pg, monkeypatc counter.count_tokens("hello", "gpt-4o") assert counter._pg_available is False - monkeypatch.setattr(tc_mod, "psycopg", None) - unavailable = TokenCounter("postgresql://x") - with pytest.raises(RuntimeError, match="requires pg_tiktoken"): - unavailable.count_tokens("hello", "gpt-4o") - with pytest.raises(RuntimeError, match="integration is unavailable"): - unavailable._count_tokens_postgres("hello", "gpt-4o") + def unavailable_driver(): + raise PostgresDriverUnavailableError("Retained PostgreSQL driver is unavailable") + + monkeypatch.setattr(tc_mod, "retained_postgres_driver", unavailable_driver) + with pytest.raises(PostgresDriverUnavailableError, match="driver is unavailable"): + TokenCounter("postgresql://x") def test_pg_tiktoken_probe_failure_closes_connection(monkeypatch): @@ -202,7 +210,8 @@ def close(self): counter = object.__new__(TokenCounter) counter._pg_conn = Connection() - monkeypatch.setattr(tc_mod, "psycopg", object()) + counter._pg_connection_lock = RLock() + counter._postgres_driver = object() monkeypatch.setattr( counter, "_get_pg_conn", @@ -210,12 +219,6 @@ def close(self): ) assert counter._ensure_pg_tiktoken() is False assert counter._pg_conn is None - monkeypatch.setattr(tc_mod, "psycopg", None) - assert counter._ensure_pg_tiktoken() is False - - monkeypatch.setattr(tc_mod, "psycopg", object()) - counter._pg_conn = None - assert counter._ensure_pg_tiktoken() is False def test_postgres_count_falls_back_to_encode(fake_pg): @@ -240,11 +243,12 @@ def fetchone(self): cursor = Cursor() class Connection: - closed = False - def cursor(self): return cursor + def is_closed(self): + return False + counter = TokenCounter("postgresql://x") counter._pg_conn = Connection() assert counter._count_tokens_postgres("one two three four", "gpt-4o") == 4 @@ -270,14 +274,15 @@ def fetchone(self): return None class Connection: - closed = False - def __init__(self, cursor): self._cursor = cursor def cursor(self): return self._cursor + def is_closed(self): + return False + counter = TokenCounter("postgresql://x") counter._pg_conn = Connection(Cursor(fallback=False)) assert counter._count_tokens_postgres("text", "model") == 0 diff --git a/tests/test_token_counter_driver_port.py b/tests/test_token_counter_driver_port.py new file mode 100644 index 000000000..acfcf4b4c --- /dev/null +++ b/tests/test_token_counter_driver_port.py @@ -0,0 +1,232 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Driver-port regressions for PostgreSQL token-counting migration.""" + +from __future__ import annotations + +from concurrent.futures import ThreadPoolExecutor +from threading import Barrier, Lock +import time +from typing import Any + +import pytest + +import pg_llm_batch.token_counter as token_counter_module +from pg_llm_batch.token_counter import TokenCounter + + +class _UndefinedFunctionError(RuntimeError): + """Represent one driver-classified undefined PostgreSQL function.""" + + +class _OtherDriverError(RuntimeError): + """Represent a database failure that must not disable pg_tiktoken availability.""" + + +class _Cursor: + """Return deterministic pg_tiktoken probe and count rows.""" + + def __init__(self, driver: _Driver) -> None: + self.driver = driver + + def __enter__(self) -> _Cursor: + return self + + def __exit__(self, *_exc: Any) -> None: + return None + + def execute(self, query: str, params: object | None = None) -> _Cursor: + self.driver.executions.append((query, params)) + if "tiktoken_count" in query and "to_regprocedure" not in query: + self.driver.enter_count_execution() + try: + if self.driver.execution_delay_seconds: + time.sleep(self.driver.execution_delay_seconds) + if self.driver.primary_error is not None: + error = self.driver.primary_error + self.driver.primary_error = None + raise error + self.driver.rows.append((7,)) + finally: + self.driver.leave_count_execution() + elif "tiktoken_encode" in query and "to_regprocedure" not in query: + self.driver.rows.append((9,)) + return self + + def fetchone(self) -> tuple[object, ...] | None: + if not self.driver.rows: + return None + return self.driver.rows.pop(0) + + +class _Connection: + """Expose the connection lifecycle required by TokenCounter.""" + + def __init__(self, driver: _Driver) -> None: + self.driver = driver + self.closed = False + self.autocommit_values: list[bool] = [] + + def cursor(self) -> _Cursor: + return _Cursor(self.driver) + + def set_autocommit(self, enabled: bool) -> None: + self.autocommit_values.append(enabled) + + def is_closed(self) -> bool: + return self.closed + + def close(self) -> None: + self.closed = True + + +class _Driver: + """Minimal concrete-driver-free port fake for token counting.""" + + def __init__( + self, + *, + primary_error: BaseException | None = None, + execution_delay_seconds: float = 0.0, + ) -> None: + self.primary_error = primary_error + self.execution_delay_seconds = execution_delay_seconds + self.executions: list[tuple[str, object | None]] = [] + self.rows: list[tuple[object, ...]] = [(True, True, True)] + self.connections: list[_Connection] = [] + self.dsn_values: list[str] = [] + self._execution_lock = Lock() + self.active_count_executions = 0 + self.max_active_count_executions = 0 + + def connect( + self, + dsn: str, + *, + connect_timeout_seconds: int | None = None, + ) -> _Connection: + assert connect_timeout_seconds is None + self.dsn_values.append(dsn) + connection = _Connection(self) + self.connections.append(connection) + return connection + + def is_undefined_function(self, error: BaseException) -> bool: + return isinstance(error, _UndefinedFunctionError) + + def enter_count_execution(self) -> None: + """Record concurrent use of the shared token-counting connection.""" + with self._execution_lock: + self.active_count_executions += 1 + self.max_active_count_executions = max( + self.max_active_count_executions, + self.active_count_executions, + ) + + def leave_count_execution(self) -> None: + """Release one deterministic concurrent-execution observation.""" + with self._execution_lock: + self.active_count_executions -= 1 + + +def _deny_default_driver(monkeypatch: pytest.MonkeyPatch) -> None: + """Fail if explicit token-counter injection silently reacquires the runtime default.""" + + def fail_default_driver(): + raise AssertionError("default PostgreSQL runtime driver was reached") + + monkeypatch.setattr( + token_counter_module, + "retained_postgres_driver", + fail_default_driver, + ) + + +def test_token_counter_uses_injected_driver_without_default_driver( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A replacement candidate must exercise pg_tiktoken through its own port.""" + driver = _Driver() + _deny_default_driver(monkeypatch) + metadata_calls: list[tuple[str, str, object]] = [] + + def _metadata(dsn: str, model: str, *, postgres_driver: object = None) -> dict[str, str]: + metadata_calls.append((dsn, model, postgres_driver)) + return {"tokenizer_model": "o200k_base"} + + monkeypatch.setattr(token_counter_module, "get_model_metadata", _metadata) + + counter = TokenCounter("postgresql://x", postgres_driver=driver) + + assert counter.count_tokens("hello", "model-a") == 7 + assert driver.dsn_values == ["postgresql://x"] + assert driver.connections[0].autocommit_values == [True] + assert metadata_calls == [("postgresql://x", "model-a", driver)] + + +def test_token_counter_serializes_shared_driver_connection_use( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A DB-API level-1 candidate must never receive concurrent connection calls.""" + driver = _Driver(execution_delay_seconds=0.03) + _deny_default_driver(monkeypatch) + monkeypatch.setattr( + token_counter_module, + "get_model_metadata", + lambda _dsn, _model, *, postgres_driver=None: {"tokenizer_model": "o200k_base"}, + ) + counter = TokenCounter("postgresql://x", postgres_driver=driver) + counter.get_encoder("model-a") + start = Barrier(4) + + def _count_one(index: int) -> int: + start.wait() + return counter.count_tokens(f"hello-{index}", "model-a") + + with ThreadPoolExecutor(max_workers=4) as executor: + results = list(executor.map(_count_one, range(4))) + + assert results == [7, 7, 7, 7] + assert driver.max_active_count_executions == 1 + + +def test_token_counter_uses_driver_error_classification_for_encode_fallback( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Undefined-function fallback must depend only on the driver-port classifier.""" + driver = _Driver(primary_error=_UndefinedFunctionError("undefined function")) + _deny_default_driver(monkeypatch) + monkeypatch.setattr( + token_counter_module, + "get_model_metadata", + lambda _dsn, _model, *, postgres_driver=None: {"tokenizer_model": "o200k_base"}, + ) + + counter = TokenCounter("postgresql://x", postgres_driver=driver) + + assert counter.count_tokens("hello", "model-a") == 9 + assert any("tiktoken_encode" in query for query, _params in driver.executions) + + +def test_non_undefined_driver_error_discards_cached_connection_before_retry( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A transient DB failure must retry on a fresh connection without disabling pg_tiktoken.""" + driver = _Driver(primary_error=_OtherDriverError("temporary database failure")) + _deny_default_driver(monkeypatch) + monkeypatch.setattr( + token_counter_module, + "get_model_metadata", + lambda _dsn, _model, *, postgres_driver=None: {"tokenizer_model": "o200k_base"}, + ) + + counter = TokenCounter("postgresql://x", postgres_driver=driver) + + with pytest.raises(RuntimeError, match="Token counting requires pg_tiktoken"): + counter.count_tokens("first", "model-a") + + assert len(driver.connections) == 1 + assert driver.connections[0].closed is True + assert counter.count_tokens("second", "model-a") == 7 + assert len(driver.connections) == 2 + assert driver.connections[1].autocommit_values == [True] + assert not any("tiktoken_encode" in query for query, _params in driver.executions) diff --git a/tests/test_token_counter_undefined_function_cleanup.py b/tests/test_token_counter_undefined_function_cleanup.py new file mode 100644 index 000000000..a74f17c83 --- /dev/null +++ b/tests/test_token_counter_undefined_function_cleanup.py @@ -0,0 +1,115 @@ +"""Recovery regressions for terminal pg_tiktoken capability loss. + +A replacement PostgreSQL driver may classify both supported pg_tiktoken entry +points as undefined at runtime, for example after an extension rollback or a +misrouted connection. Once the counter marks the capability unavailable it +must also release the cached session; retaining an unusable connection would +leak a database resource that the counter intentionally never retries. +""" + +from __future__ import annotations + +from typing import Any + +import pytest + +import pg_llm_batch.token_counter as token_counter_module +from pg_llm_batch.token_counter import TokenCounter + + +class _UndefinedFunctionError(RuntimeError): + """Represent one driver-classified PostgreSQL undefined-function error.""" + + +class _Cursor: + """Expose a healthy capability probe followed by two missing functions.""" + + def __init__(self, driver: _Driver) -> None: + self._driver = driver + + def __enter__(self) -> _Cursor: + return self + + def __exit__(self, *_exc: Any) -> None: + return None + + def execute(self, query: str, params: object | None = None) -> _Cursor: + if "to_regprocedure" in query: + self._driver.rows.append((True, True, True)) + return self + if "tiktoken_count" in query or "tiktoken_encode" in query: + raise _UndefinedFunctionError("pg_tiktoken function unavailable") + raise AssertionError(f"unexpected SQL in cleanup regression: {query!r}") + + def fetchone(self) -> tuple[object, ...] | None: + if not self._driver.rows: + return None + return self._driver.rows.pop(0) + + +class _Connection: + """Track whether terminal capability loss releases the retained session.""" + + def __init__(self, driver: _Driver) -> None: + self._driver = driver + self.closed = False + + def cursor(self) -> _Cursor: + return _Cursor(self._driver) + + def set_autocommit(self, enabled: bool) -> None: + assert enabled is True + + def is_closed(self) -> bool: + return self.closed + + def close(self) -> None: + self.closed = True + + +class _Driver: + """Minimal driver port for terminal undefined-function recovery evidence.""" + + def __init__(self) -> None: + self.rows: list[tuple[object, ...]] = [] + self.connections: list[_Connection] = [] + + def connect( + self, + dsn: str, + *, + connect_timeout_seconds: int | None = None, + ) -> _Connection: + assert dsn == "postgresql://candidate" + assert connect_timeout_seconds is None + connection = _Connection(self) + self.connections.append(connection) + return connection + + def is_undefined_function(self, error: BaseException) -> bool: + return isinstance(error, _UndefinedFunctionError) + + +def test_terminal_undefined_function_failure_closes_cached_driver_session( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Both missing pg_tiktoken entry points must disable and release the session.""" + driver = _Driver() + monkeypatch.setattr(token_counter_module, "psycopg", None) + monkeypatch.setattr( + token_counter_module, + "get_model_metadata", + lambda _dsn, _model, *, postgres_driver=None: { + "tokenizer_model": "o200k_base" + }, + ) + + counter = TokenCounter("postgresql://candidate", postgres_driver=driver) + + with pytest.raises(RuntimeError, match="Token counting requires pg_tiktoken"): + counter.count_tokens("hello", "model-a") + + assert len(driver.connections) == 1 + assert driver.connections[0].closed is True + assert counter._pg_conn is None + assert counter._pg_available is False diff --git a/tests/test_workflow_contracts.py b/tests/test_workflow_contracts.py index 6c0315f24..445aee8d0 100644 --- a/tests/test_workflow_contracts.py +++ b/tests/test_workflow_contracts.py @@ -151,6 +151,72 @@ def test_ci_workflow_enforces_supported_versions_and_quality_gates() -> None: _assert_external_actions_are_pinned(workflow) +def test_ci_pg8000_candidate_parity_is_immutable_and_queue_conservative() -> None: + """Keep replacement-driver proof exact without creating another runner lane.""" + workflow = _read(".github/workflows/ci.yml") + project = _read("pyproject.toml") + current_setup_uv = ( + "astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d" + ) + stale_setup_uv = ( + "astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9" + ) + + assert "pg8000-candidate-python314:" not in workflow + assert "pg8000==1.31.5" in workflow + assert ( + "0af2c1926b153307639868d2ee5cef6cd3a7d07448e12736989b10e1d491e201" + in workflow + ) + assert "tests/smoke_pg8000_candidate_postgres.py" in workflow + assert workflow.count(current_setup_uv) == 3 + assert stale_setup_uv not in workflow + assert "pg8000-candidate-ci-password" not in workflow + assert "secrets.token_urlsafe(32)" in workflow + assert "::add-mask::$candidate_password" in workflow + assert "PG_LLM_BATCH_POSTGRES_PASSWORD=$candidate_password" not in workflow + assert "PG8000_CANDIDATE_PASSWORD_FILE" in workflow + assert "Tear down candidate PostgreSQL runtime" in workflow + assert '"pg8000' not in project.casefold() + + +def test_ci_pg8000_candidate_pins_and_hashes_full_dependency_closure() -> None: + """Candidate proof must not resolve mutable transitive wheels at install time.""" + workflow = _read(".github/workflows/ci.yml") + + exact_artifacts = { + "pg8000==1.31.5": ( + "0af2c1926b153307639868d2ee5cef6cd3a7d07448e12736989b10e1d491e201" + ), + "python-dateutil==2.9.0.post0": ( + "a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427" + ), + "scramp==1.4.17": ( + "a4e3fd2e8169461a28a13777a166d3da94274454f0714a7d3023fee124474ac8" + ), + "asn1crypto==1.5.1": ( + "db4e40728b728508912cbb3d44f19ce188f218e9eba635821bb4b68564f8fd67" + ), + "six==1.17.0": ( + "4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274" + ), + } + for requirement, digest in exact_artifacts.items(): + assert requirement in workflow + assert digest in workflow + + assert "pip download --no-deps --only-binary=:all:" in workflow + assert "uv pip install --python .venv/bin/python --no-deps" in workflow + assert "/tmp/pg8000-candidate/pg8000-1.31.5-py3-none-any.whl" in workflow + assert ( + "/tmp/pg8000-candidate/python_dateutil-2.9.0.post0-py2.py3-none-any.whl" + in workflow + ) + assert "/tmp/pg8000-candidate/scramp-1.4.17-py3-none-any.whl" in workflow + assert "/tmp/pg8000-candidate/asn1crypto-1.5.1-py2.py3-none-any.whl" in workflow + assert "/tmp/pg8000-candidate/six-1.17.0-py2.py3-none-any.whl" in workflow + + def test_workflow_step_field_matching_ignores_comments_and_unrelated_values() -> None: """Comment or nested text must not masquerade as workflow step fields.""" decoy = """ - name: Decoy diff --git a/tools/verify_candidate_wheel_licenses.py b/tools/verify_candidate_wheel_licenses.py new file mode 100644 index 000000000..514b3fa34 --- /dev/null +++ b/tools/verify_candidate_wheel_licenses.py @@ -0,0 +1,206 @@ +#!/usr/bin/env python3 +"""Verify license metadata for the exact pg8000 candidate wheel closure. + +Artifact hashes prove which wheels were downloaded, but they do not prove that +every transitive package satisfies the repository's commercial inbound-license +policy. This verifier reads METADATA directly from the already hash-verified +wheels, requires the exact reviewed package/version set, rejects GPL-family +metadata, and requires positive permissive-license evidence for every wheel. +It never imports or executes candidate package code. +""" + +from __future__ import annotations + +from email.parser import Parser +from pathlib import Path +import re +import sys +import zipfile + + +_MAX_METADATA_BYTES = 524_288 +_CANONICAL_NAME_SEPARATOR = re.compile(r"[-_.]+") +_GPL_FAMILY = re.compile(r"(?:^|[^a-z])(agpl|lgpl|gpl)(?:[^a-z]|$)") + +_EXPECTED_WHEELS = { + "pg8000-1.31.5-py3-none-any.whl": ( + "pg8000", + "1.31.5", + ("bsd",), + ), + "python_dateutil-2.9.0.post0-py2.py3-none-any.whl": ( + "python-dateutil", + "2.9.0.post0", + ("apache software license", "bsd license"), + ), + "scramp-1.4.17-py3-none-any.whl": ( + "scramp", + "1.4.17", + ("mit-0", "mit no attribution", "mit no attribution license"), + ), + "asn1crypto-1.5.1-py2.py3-none-any.whl": ( + "asn1crypto", + "1.5.1", + ("mit",), + ), + "six-1.17.0-py2.py3-none-any.whl": ( + "six", + "1.17.0", + ("mit",), + ), +} + + +class CandidateWheelLicenseError(RuntimeError): + """Reject candidate metadata that cannot support commercial admission.""" + + +def _canonical_distribution_name(value: str) -> str: + """Apply the package-name normalization used for exact identity comparison.""" + return _CANONICAL_NAME_SEPARATOR.sub("-", value).casefold() + + +def _read_metadata(wheel_path: Path) -> str: + """Read one bounded wheel METADATA member without extracting package code.""" + try: + with zipfile.ZipFile(wheel_path) as archive: + metadata_members = [ + info + for info in archive.infolist() + if info.filename.endswith(".dist-info/METADATA") + ] + if len(metadata_members) != 1: + raise CandidateWheelLicenseError( + "candidate wheel metadata identity is invalid" + ) + metadata_member = metadata_members[0] + if metadata_member.file_size > _MAX_METADATA_BYTES: + raise CandidateWheelLicenseError( + "candidate wheel metadata exceeds the bounded evidence size" + ) + raw_metadata = archive.read(metadata_member) + except (OSError, zipfile.BadZipFile, RuntimeError): + raise CandidateWheelLicenseError("candidate wheel could not be inspected") from None + + if len(raw_metadata) > _MAX_METADATA_BYTES: + raise CandidateWheelLicenseError( + "candidate wheel metadata exceeds the bounded evidence size" + ) + try: + return raw_metadata.decode("utf-8") + except UnicodeDecodeError: + raise CandidateWheelLicenseError("candidate wheel metadata is not UTF-8") from None + + +def _license_evidence(metadata_text: str) -> tuple[str, ...]: + """Collect declared license fields and classifiers as separate evidence lines.""" + message = Parser().parsestr(metadata_text) + evidence_values: list[str] = [] + for header in ("License-Expression", "License"): + value = message.get(header) + if value: + evidence_values.append(value.casefold()) + evidence_values.extend( + classifier.casefold() + for classifier in message.get_all("Classifier", []) + if classifier.casefold().startswith("license ::") + ) + return tuple(evidence_values) + + +def _contains_marker(evidence_lines: tuple[str, ...], marker: str) -> bool: + """Match one approved marker on non-alphanumeric boundaries only. + + License metadata is untrusted decision input. Substring matching would accept + an unrelated word such as ``limited`` for the reviewed ``MIT`` marker. The + boundary check still accepts SPDX identifiers and classifier phrases while + preventing a permissive token from being smuggled inside another word. + """ + marker_pattern = re.compile( + rf"(? None: + """Validate exact package identity and positive/negative license evidence.""" + metadata_text = _read_metadata(wheel_path) + message = Parser().parsestr(metadata_text) + package_name = message.get("Name") + package_version = message.get("Version") + if ( + type(package_name) is not str + or _canonical_distribution_name(package_name) != expected_name + or type(package_version) is not str + or package_version != expected_version + ): + raise CandidateWheelLicenseError("candidate wheel package identity is invalid") + + license_evidence = _license_evidence(metadata_text) + if not license_evidence: + raise CandidateWheelLicenseError( + "candidate wheel lacks approved license evidence" + ) + joined_evidence = "\n".join(license_evidence) + if ( + _GPL_FAMILY.search(joined_evidence) is not None + or "gnu general public license" in joined_evidence + or "gnu lesser general public license" in joined_evidence + or "gnu affero general public license" in joined_evidence + ): + raise CandidateWheelLicenseError( + "candidate wheel contains a disallowed license" + ) + if not any( + _contains_marker(license_evidence, marker) for marker in approved_markers + ): + raise CandidateWheelLicenseError( + "candidate wheel lacks approved license evidence" + ) + + +def verify_candidate_wheel_licenses(directory: Path) -> None: + """Verify the exact immutable pg8000 candidate closure in ``directory``. + + The directory is expected to contain only the five wheel artifacts that the + CI digest gate pins and later installs with ``--no-deps``. Requiring the same + exact set prevents an unreviewed extra artifact from entering license evidence + without a corresponding digest and policy decision. + """ + if not isinstance(directory, Path) or not directory.is_dir(): + raise CandidateWheelLicenseError("candidate wheel directory is invalid") + wheel_names = {path.name for path in directory.glob("*.whl")} + if wheel_names != set(_EXPECTED_WHEELS): + raise CandidateWheelLicenseError("candidate wheel set is invalid") + + for filename, (package_name, package_version, approved_markers) in sorted( + _EXPECTED_WHEELS.items() + ): + _verify_one_wheel( + directory / filename, + expected_name=package_name, + expected_version=package_version, + approved_markers=approved_markers, + ) + + +def main(argv: list[str] | None = None) -> int: + """Run license admission against one exact candidate wheel directory.""" + arguments = sys.argv[1:] if argv is None else argv + if len(arguments) != 1: + raise SystemExit("usage: verify_candidate_wheel_licenses.py WHEEL_DIRECTORY") + try: + verify_candidate_wheel_licenses(Path(arguments[0])) + except CandidateWheelLicenseError as exc: + raise SystemExit(str(exc)) from None + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())