Skip to content
bmsuissePublic

About

Tools for developing with Postgres, mostly around git integration. Also MS SQL if you really need...

Resources

Stars

0 stars

Watchers

0 watching

Forks

Latest commit

 

History

215 Commits

Folders and files

Repository files navigation

pgdevkit

A helper for developing with Postgres.

pgdb compare

Compare a directory of SQL scripts (see the database-in-source layout convention) against a live database and report differences:

pgdb compare --url postgresql://user:pass@host:port/db path/to/database/

Entra ID auth (Azure Postgres / Databricks Lakebase)

Pass --entra-user <identity> to pgdb compare to authenticate with an Entra ID token instead of a static password. Which token flow is used is auto-detected from the database hostname:

  • Azure Database for PostgreSQL (*.postgres.database.azure.com, *.postgres.cosmos.azure.com) — the default: fetches a token via DefaultAzureCredential and uses it directly as the password. Requires the azure extra: pip install pgdevkit[azure].
  • Databricks Lakebase (*.database.azuredatabricks.net, *.database.cloud.databricks.com) — fetches a Databricks-scoped Entra token, then exchanges it for a short-lived Postgres credential via the Databricks workspace API. Also requires --databricks-workspace-host and --databricks-instance:
pgdb compare --url postgresql://instance-abc.database.azuredatabricks.net:5432/databricks_postgres \
  --entra-user alice@example.com \
  --databricks-workspace-host https://adb-123456789.azuredatabricks.net \
  --databricks-instance myinstance \
  path/to/database/

(--url's own user/password, if any, are discarded and replaced — --entra-user plus the fetched token become the connection's actual credentials.)

MSSQL

pgdb compare/pgdb fetch-missing default to Postgres. Pass --dialect mssql to compare against a SQL Server database instead:

pgdb compare --dialect mssql --url "Server=host,1433;Database=db;UID=user;PWD=pass" path/to/database/

Requires the mssql extra: pip install pgdevkit[mssql] (pulls in mssql-python, which bundles its own driver — no system ODBC driver install needed). MSSQL has no composite type or native enum equivalent, so those areas of a database/ tree don't have a direct equivalent on this backend — see docs/database-layout.md.

pgdb compare --dialect mssql --entra-user <identity> appends Authentication=ActiveDirectoryDefault to the connection string so mssql-python gets an Entra ID token via DefaultAzureCredential (also install the azure extra). The identity is whatever that credential chain resolves; the flag's value only switches Entra auth on. Current Azure SQL/SQL Server (2025+) does have a native json column type, which parses/introspects/diffs like any other column type; see "pgdevkit.db — helpers for application code" below for how JSON values are handled on the CRUD side (write-side serialization only, no auto-parsing on read — mssql-python doesn't distinguish json columns from nvarchar).

Area and schema filtering

Two independent, composable ways to narrow which files a command touches: area is an explicit opt-in tag; schema is derived automatically from each file's own SQL.

Area tagging

Any migration file or database/ code file can declare one or more areas by starting with a -- area: comment:

-- area: billing
CREATE TABLE billing.invoices (id int primary key);

A file can declare more than one area, either comma-separated on one line (-- area: billing, reporting) or across several -- area: lines — the declared areas union. The directive is only recognized in the file's leading comment block (blank lines and -- comments at the very top, stopping at the first real statement); a -- area: comment later in the file doesn't count. A file with no directive is untagged, and untagged files are treated as shared/common.

Schema filtering

No tag needed — schema membership is parsed straight out of the SQL itself: every schema-qualified (or default-schema, when unqualified) table/view/ function/index reference across every statement in the file, DDL or DML alike, plus any CREATE SCHEMA name. A file whose schema(s) can't be determined (unparseable content, or no table/schema reference in it at all) is treated the same as an untagged file — always kept.

Options

pgdb compare, pgdb migrate check, pgdb migrate apply, pgdb testdb up, and pgdb testdb reset all accept:

  • --area NAME (repeatable) — restrict to files declaring one of the given areas, plus every untagged file (untagged files always stay in scope).
  • --exclude-area NAME (repeatable) — drop files declaring one of the given areas; untagged files are never dropped by this.
  • --schema NAME (repeatable) — restrict to files referencing one of the given schemas, plus every file with no detectable schema reference.
  • --exclude-schema NAME (repeatable) — drop files referencing one of the given schemas; files with no detectable reference are never dropped.

All four can be combined — a file must pass every filter it's subject to (an area match doesn't excuse a schema mismatch, and vice versa), and a file matching both an included and an excluded value on the same axis is excluded. Passing none of them applies no filtering (the default, unchanged behavior).

pgdb migrate apply path/to/database/_migration_scripts --url ... --area billing
pgdb compare path/to/database/ --url ... --exclude-area reporting
pgdb migrate check path/to/database/_migration_scripts --url ... --schema billing --exclude-schema reporting
pgdb testdb up --schema billing

compare's default report (no --report-extra-db) only checks that the filtered scripts exist correctly in the DB, so it composes safely with area and schema filtering. Passing --report-extra-db together with either kind of filter also reports every DB object outside the filtered area(s)/ schema(s) as "missing in scripts" — since the live database has no concept of areas, and isn't itself filtered by --schema either — only the scripts side is filtered — so treat that combination's "missing in scripts" results with that in mind (the CLI prints a warning when you combine them).

pgdb fetch-missing deliberately has no --area/--exclude-area (or --schema/--exclude-schema): it diffs the full database against scripts to find genuinely untracked objects, so narrowing the scripts side would make every object tracked under a different area/schema look "missing" too — and --write would then reconstruct a duplicate file for something that already exists.

pgdevkit.areas exposes the tag-filtering logic for scripting: parse_areas/file_areas read a file's declared areas, and area_allowed/filter_by_area apply the only/exclude semantics above. pgdevkit.schemas exposes the equivalent for schema filtering: sql_schemas/file_schemas detect a file's referenced schemas, and schema_allowed/filter_by_schema apply the same only/exclude semantics. pgdevkit.migrate.list_migration_files/pending_migrations and pgdevkit.parser.parse_directory take both pairs of keyword arguments (areas/exclude_areas and schemas/exclude_schemas); pgdevkit.fetch_missing.find_missing_objects takes neither, for the reason above.

Environment-tagged files (<name>.<env>.sql)

A file whose name ends .<env>.sql (e.g. grants.prod.sql, seed.staging.sql) is only in scope when targeting that environment; a plain <name>.sql file is untagged and always in scope, regardless of environment. .init.sql (see docs/database-layout.md) is reserved and is never treated as an environment tag.

  • pgdb testdb up/pgdb testdb reset accept --env (default local_test) — so an untagged grants.sql always applies, but grants.prod.sql is skipped unless run with --env prod.
  • pgdb migrate check/pgdb migrate apply accept --env too, but it's optional with no default: omit it and every file is a candidate regardless of its tag (unchanged, today's behavior); pass it to restrict to files tagged for that environment plus untagged ones.
pgdb testdb up --env prod   # apply prod-tagged files too, against the local test container
pgdb migrate apply path/to/database/_migration_scripts --url ... --env prod

pgdevkit.envtag exposes the same logic for scripting: file_env reads a file's tag, env_allowed applies the filtering semantics above, and strip_env_suffix returns a tagged file's logical name (e.g. grants.prod.sql -> "grants").

pgdb testdb

Manages a single shared, Podman-backed Postgres container for local tests across all your projects — no more one-container-per-project-per-worktree. Isolation between projects and worktrees is per-database, inside one container.

Add to pyproject.toml:

[tool.pgdevkit]
name = "myproject"        # optional; defaults to the repo directory name
database_dir = "database" # optional; defaults to "database"

Add to conftest.py:

import os
import pytest
from pgdevkit.testdb import ensure_testdb

@pytest.fixture(scope="session", autouse=True)
def ensure_test_postgres():
    for k, v in ensure_testdb().items():
        os.environ[k] = v

CLI: pgdb testdb up|reset|run-sql|status|shell|clean|list-orphaned. up/reset accept --env (default local_test) — see "Environment-tagged files" above.

up/reset accept --area/--exclude-area and --schema/--exclude-schema (see "Area and schema filtering" above) to scope which database/ files get applied — e.g. pgdb testdb up --schema billing for a test DB with only the billing schema's tables/views/functions, without waiting on the rest of the project's schema to apply. ensure_testdb/reset_testdb take the same keyword arguments when called from Python (e.g. from a pytest fixture).

Every git worktree/branch of a project gets its own database, named after project_name + branch (see pgdevkit.testdb.naming.workspace_db_name). Removing a worktree (or deleting its directory without git worktree remove) doesn't drop its database — pgdb testdb list-orphaned lists this project's databases whose worktree no longer exists, and pgdb testdb clean --orphaned drops them (as opposed to --all, which drops every database of this project regardless of whether its worktree is still live). The same is available from Python as pgdevkit.testdb.find_orphaned_dbs() and pgdevkit.testdb.clean_testdb(orphaned=True).

If your project's own test setup also creates a sibling database per worktree (e.g. <main_db>_myservice for a mock service used only by that project's tests), add its literal suffix so orphan detection knows it belongs to a live worktree too:

[tool.pgdevkit]
extra_db_suffixes = ["_myservice"]

pgdevkit.testdb.workspace_db_names(project_root=None) is the single-workspace analog of find_orphaned_dbs(): it returns the exact set of DB names (main + any extra_db_suffixes) owned by the branch currently checked out at project_root, without touching Postgres or diffing against other worktrees. Useful for a caller that's about to remove one specific worktree and wants to know exactly which DB(s) go with it.

Container connection defaults (localhost:54322, postgres/testpwd) can be overridden with PGDEVKIT_TESTDB_HOST, PGDEVKIT_TESTDB_PORT, PGDEVKIT_TESTDB_USER, PGDEVKIT_TESTDB_PASSWORD, and the image with PGDEVKIT_TESTDB_IMAGE. Before touching the Docker API, pgdevkit first checks (with a short timeout) whether Postgres is already reachable at that address and skips container management if so. Set PGDEVKIT_SKIP_CONTAINER=1 to always assume it's already there and skip that check too.

Docker-level resource knobs are unset by default (Docker/Podman's own defaults apply) and can be overridden with PGDEVKIT_TESTDB_SHM_SIZE (e.g. 256m, useful when Docker's 64MB default shared-memory size is too small for parallel query workers), PGDEVKIT_TESTDB_MEM_LIMIT (e.g. 1g), PGDEVKIT_TESTDB_CPUS (e.g. 2 or 0.5), and PGDEVKIT_TESTDB_TMPFS — docker-CLI --tmpfs syntax (/path:options), semicolon-separated for multiple mounts, e.g. /var/lib/postgresql/data:size=512m.

All of the above (including _HOST/_PORT/_USER/_PASSWORD/_IMAGE) only take effect when pgdevkit actually creates the container, not when it finds and restarts an existing stopped one with the same name — remove the existing container (docker rm -f pgdevkit-postgres) first if you've changed any of these on a machine that already has one.

Container management goes through the Docker API (the docker package, docker.from_env(), falling back to Podman's rootful/rootless socket) — it works against a real Docker daemon or Podman transparently, no CLI binary required either way.

To point at a local Postgres install instead of the container — useful when neither is available, or you'd rather use peer authentication as the current OS user — set PGDEVKIT_TESTDB_HOST to the unix socket directory (e.g. /var/run/postgresql) and PGDEVKIT_TESTDB_PASSWORD="". The role named by PGDEVKIT_TESTDB_USER must exist and match your OS user (CREATE ROLE <user> SUPERUSER LOGIN;) and pg_hba.conf must allow peer auth for local connections (Debian/Ubuntu Postgres ships this by default).

MSSQL

Add engine = "mssql" to [tool.pgdevkit] (or set PGDEVKIT_TESTDB_ENGINE=mssql for a one-off run) to manage a shared SQL Server container instead of Postgres — same one-container-per-machine, one-database-per-workspace model. Requires the mssql extra (see above).

Container defaults (localhost:14330, sa/a generated complexity-valid password) can be overridden with PGDEVKIT_TESTDB_MSSQL_HOST, _PORT, _USER, _PASSWORD, _IMAGE, _MEMORY_LIMIT_MB. The container only bootstraps the sa login — additional logins are a known limitation.

The same Docker-level resource knobs as the Postgres container are available here too: PGDEVKIT_TESTDB_MSSQL_SHM_SIZE, PGDEVKIT_TESTDB_MSSQL_MEM_LIMIT (a container-level cgroup limit, distinct from _MEMORY_LIMIT_MB above which only tunes SQL Server's own internal memory management), PGDEVKIT_TESTDB_MSSQL_CPUS, and PGDEVKIT_TESTDB_MSSQL_TMPFS. As with the Postgres container, these only take effect when pgdevkit creates the container, not when restarting an existing stopped one (docker rm -f pgdevkit-mssql first if you've changed any of them).

pgdb testdb shell execs into sqlcmd (an external prerequisite, the same category as psql for the Postgres path) rather than a Python REPL.

pgdb migrate

Applies numbered, forward-only SQL migration files from a directory to a live Postgres or MSSQL database, tracking each one in a schema.table (default public.schema_migrations, dbo.schema_migrations for MSSQL) so repeat runs only apply what's pending.

pgdb migrate check  path/to/database/_migration_scripts --url postgresql://user:pass@host:port/db
pgdb migrate apply  path/to/database/_migration_scripts --url postgresql://user:pass@host:port/db

MSSQL

Pass --dialect mssql (needs the mssql extra) to both check and apply, with --url as an ODBC connection string:

pgdb migrate apply path/to/database/_migration_scripts --dialect mssql \
    --url "Server=host,1433;Database=db;UID=user;PWD=pass"
  • Migration files are T-SQL: they are split into batches on standalone GO lines (needed for e.g. CREATE VIEW, which must be first in its batch), and all batches of one file run in a single transaction that is rolled back if any batch fails (including a failure in a later statement of a multi-statement batch). GO <count> is not honored: the batch runs once.
  • The tracking table needs filename nvarchar(450) primary key, applied_at datetimeoffset not null default sysdatetimeoffset(), applied_by nvarchar(128) not null default suser_sname(). As on Postgres, a migration creating it can bootstrap it. Override its name with mssql_migrations_table in [tool.pgdevkit] (the Postgres migrations_table key is deliberately ignored for MSSQL) or with --tracking-table.
  • --ask auto-detects "already done" for single-statement GO batches that create a table, view or schema, or add a column; anything else is asked about.
  • Post-apply verification checks every CREATE TABLE target via OBJECT_ID.
  • --entra-user appends Authentication=ActiveDirectoryDefault to the connection string, so mssql-python fetches an Entra ID token through azure-identity's DefaultAzureCredential (install pgdevkit[mssql,azure]). The identity is whatever that credential chain resolves — the flag's value only switches Entra auth on. It can't be combined with an Authentication= already in the connection string.
  • pgdevkit.migrate.missing_privileges is not available for MSSQL.

For Postgres, --entra-user works the same as pgdb compare (see above). The tracking table needs filename text primary key, applied_at timestamptz not null default now(), applied_by text not null default current_user (a migration file that creates it, in the same directory, is the usual way to bootstrap it — inserting into a not-yet-existing tracking table is tolerated so that migration can still run).

The tracking table defaults to public.schema_migrations. Override it per-project in pyproject.toml:

[tool.pgdevkit]
migrations_table = "myschema.migrations"

or per-invocation with --tracking-table, which takes precedence over the pyproject.toml value.

--ask prints each pending file and asks yes/no/already-done/quit before running it. Answering yes queues the file on a background worker and moves straight to the next prompt — you can keep reviewing while earlier files are still executing, instead of waiting on each one before seeing the next. A tqdm progress bar tracks the queue; migrations still run one at a time, in file order. Without --ask, apply queues every pending file up front and shows the same progress bar. --file <name> applies a single file (still through the same verify-and-track path) instead of walking all pending ones. Pass --yes to skip the "about to run migrations against ..." confirmation prompt (e.g. in CI).

After each file's DDL is applied, apply re-checks that every CREATE TABLE statement's target actually exists (via to_regclass) before recording the file as applied — catching a migration that silently rolled back. That check parses each statement with sqlglot and only falls back to a regex (run against comment-stripped SQL) for statements sqlglot's postgres dialect can't parse, so a CREATE TABLE mentioned only in a -- comment is never mistaken for a real one.

pgdevkit.migrate is also usable directly as a library — list_migration_files, applied_migrations, pending_migrations, and apply_migration are the same functions the CLI calls, so a project can script around them without shelling out. The database-touching ones take dialect="postgres" | "mssql".

pgdevkit.db — helpers for application code

Install with the db extra: pip install pgdevkit[db].

  • TableModel (formerly PostgresTableModel, still importable under that name) — a pydantic.BaseModel base class for models that map 1:1 to a table row, for either engine. Implement get_table_name() (returns (schema, table)) and get_primary_key() on each model.
  • PgPool — an async connection pool keyed off {env_prefix}HOST/PORT/DB/USER/PASSWORD env vars. Call await pool.open() once at startup, then use async with pool.connection() as con:. Pass entra_user to authenticate via Entra ID instead of a static password — same host-based auto-detection as pgdb compare's --entra-user. For Lakebase hosts, also set the {env_prefix}DATABRICKS_WORKSPACE_HOST and {env_prefix}DATABRICKS_INSTANCE env vars. Values (password, dsn_params) are quoted properly, so spaces, backslashes and quotes are fine. min_size, timeout (how long to wait for a free connection before PoolTimeout), configure (async hook run on every new connection, e.g. to set session timeouts or application_name) and check are passed through to the psycopg pool; max_size defaults to 40 and min_size to 4 (or max_size if smaller).
  • fetch_all(query, params=None, *, model=None, row_mapper=None, con=None, pool=None, cancel=None, statement_timeout=None) — run one custom query and get every row back as dicts, as validated Pydantic model instances, or through a row_mapper. Pass con to run on a connection you already hold (it is borrowed: never committed, closed or released); otherwise a connection is taken from pool or the pool registered with set_default_pool() (a ConnectionSource such as PgPool, or a plain callable returning an async connection context manager). Set the cancel asyncio.Event to abort the query on the server (QueryCanceled; to stream a result as JSON from FastAPI see PostgresJsonResponse, which watches for the disconnect itself, and for buffered endpoints see CancelOnDisconnectRoute); cancelling the awaiting task cancels it there too, and CancelledError propagates as usual. statement_timeout=<seconds> lets Postgres abort the query after that long (QueryCanceled, "canceling statement due to statement timeout"); it is applied per call with SET LOCAL semantics, so it never outlives the call (a borrowed con gets its previous setting back); see Statement timeout for what it does to a timeout the connection already has. query must be a literal string, a sqlglot expression, a psycopg.sql composable or a t-string — never an f-string (a sqlglot expression is only as safe as the strings it was built from: pass user values as exp.Placeholder + params). A literal % in a sqlglot expression (LIKE 'a%', the modulo operator) works with or without params: it is escaped for psycopg when params is given (plain strings and psycopg.sql composables are passed on as is, so there you write %% yourself when you pass params).
  • fetch_one, fetch_scalar, execute — siblings of fetch_all with the same con/pool/cancel/ statement_timeout keywords and the same query types (including t-strings, which take no params), so there is no need for hand-written cursor code. fetch_one(query, params, *, model=None, row_mapper=None) returns the first row (dict, model or row_mapper result) or None, fetch_scalar(query, params) the first column of the first row or None (await fetch_scalar("SELECT name FROM widget WHERE id = %(id)s", {"id": 1})), and execute(query, params) the row count of a write without RETURNING (-1 for statements without one, like DDL). Without con, execute runs on a pooled connection that commits when it returns (true for PgPool/psycopg_pool; a custom connection source decides for itself); on your con nothing is committed.
  • readonly_transaction(con) — async context manager that runs its block in a read-only transaction: async with pool.connection() as con, readonly_transaction(con): await fetch_all(sql, con=con). Postgres rejects writes with ReadOnlySqlTransaction; the transaction is committed when the block ends normally and rolled back on error, and read_only is always put back afterwards (normally to None), so a connection that goes back to a shared pool is never left read-only. con must not be inside a transaction already (a fresh pooled connection is fine). It is a safety net, not a sandbox for untrusted SQL: text without bound parameters uses the simple query protocol, where COMMIT; INSERT ... ends the read-only transaction, TEMP tables stay writable and SET statements stick to the connection. For LLM-generated SQL also connect as a role that only has SELECT privileges, and consider statement_timeout=at_most(...).
  • ComplexHelper — handles composite, enum and JSONB columns for the CRUD functions: pass complex_helper=ComplexHelper(con) (optionally with normalizers), or complex_helper="auto" for a default one, to any pg_* function that takes the argument. Nothing is registered on the connection: psycopg's type registration is global to a connection (and so to the next user of a pooled one: a plain SELECT of an enum would return psycopg's auto-numbered Enum member instead of the label, #54). The helpers therefore register the enum/composite adapters only on the cursor that runs a write statement (pg_insert, pg_update/pg_update_dict, pg_upsert/pg_upsert_dict, the *_many variants), and only for the columns where a dict/list value had to be converted. Reads (pg_retrieve*) select such columns as to_jsonb(...) and need no registration. After any pg_* call the connection reads enums as label strings and composites as text (e.g. '(1,2)'), as before the call. RETURNING rows of pg_insert/pg_update_dict/pg_upsert_dict (and so pg_update/pg_upsert) select such columns as to_jsonb(...) too, so with a helper they are plain Python values (a composite as a dict, an enum as its label) whichever columns were written; without one it is a plain RETURNING *. If you call recursive_convert yourself, register the types on the cursor that sends the values, with the info from load_all_complex_types/load_complex_type: await helper.register_on(cur, info), then execute the converted value on that same cur (passing the connection instead would restore psycopg's global behaviour). What it learns from the catalog (types and column types) is cached per connection, so a helper created per call queries the catalog only the first time a connection is used (the SELECT/INSERT itself is then the only query). The cache cannot see later DDL, nor DDL that was rolled back after it had been cached: call helper.clear_cache() after CREATE TYPE/ALTER TYPE/column changes or rollbacks of those (typical for tests that create types inside a rolled-back transaction on a reused connection). ComplexHelper(con, cache=False) gives the helper a private cache and the pre-cache behaviour.
  • CRUD functions — pg_retrieve, pg_retrieve_many (equality filters, plus optional where= as a t-string or a psycopg.sql composable with params=, order_by= as a column name, a list of names or (name, "asc"|"desc") pairs, or a composable, and limit=; prefer a t-string, since in a composable a literal % must be written %% whenever filters or params are present), pg_insert, pg_insert_many, pg_update, pg_update_dict, pg_upsert, pg_upsert_dict, pg_upsert_many, pg_upsert_many_dict, pg_delete, pg_delete_dict — typed (TableModel-based) or dict-based CRUD against a table, built on psycopg for safe identifier/value handling. The mssql extra provides an mssql_*-prefixed mirror of the same functions in pgdevkit.db.mssql_crud, built on mssql-python (MERGE-based upsert, OUTPUT instead of RETURNING) — MSSQL has no composite/enum equivalent, so complex_helper is always None on that path. It does have a native json column type on current versions (and the older NVARCHAR(MAX)-plus-OPENJSON() convention works on any version), but mssql-python has no auto-serialization for dict/list parameter values (binding one raises TypeError) and no way to distinguish a json column from nvarchar on fetch — so every mssql_* write function serializes dict/list values to JSON text automatically (db.mssql_sql.json_encode_value), while reads always come back as plain str; deserialize with json.loads() yourself if you need the parsed value back.
  • SqlLoader — loads and caches .sql files from {root}/<topic>/<name>.sql, for keeping hand-written queries out of Python source.
from pgdevkit.db import PgPool, PostgresTableModel, pg_retrieve, pg_upsert

class Widget(PostgresTableModel):
    id: int
    name: str

    @staticmethod
    def get_table_name() -> tuple[str, str]:
        return ("public", "widget")

    @staticmethod
    def get_primary_key() -> list[str]:
        return ["id"]

pool = PgPool(env_prefix="POSTGRES_")
await pool.open()
async with pool.connection() as con:
    widget = await pg_retrieve(con, Widget, {"id": 1})
    await pg_upsert(con, Widget(id=1, name="thing"), Widget)

pgdevkit.fastapi — streaming JSON responses and cancel-on-disconnect

Requires the fastapi extra (pip install pgdevkit[fastapi], FastAPI >= 0.118).

PostgresJsonResponse runs a query and streams its rows to the client as a JSON array. Postgres builds the JSON (row_to_json), so there is no row-by-row Python serialization and no pydantic round trip, and memory use stays flat for big results. If the client goes away, the query is cancelled in Postgres instead of running to completion. (fetch_all(..., cancel=event) cancels a buffered query too; for fetch_all-based endpoints put the route on CancelOnDisconnectRoute (below) instead of wiring the event yourself.)

from pgdevkit.db import PgPool, set_default_pool
from pgdevkit.fastapi import PostgresJsonResponse

pool = PgPool(env_prefix="APP_POSTGRES_")
set_default_pool(pool)  # once at startup, shared with fetch_all

@router.get("/articles", responses={200: {"model": list[ArticleOut]}})
async def articles(lng: str) -> PostgresJsonResponse:
    return PostgresJsonResponse(
        "select id, name from articles where lng = %(lng)s order by name",
        {"lng": lng},
    )

The responses= entry keeps the OpenAPI schema (and a generated frontend client) typed; response_model is ignored for a Response, and the rows are not validated against the model.

  • Connections work as in fetch_all: the pool registered with set_default_pool(), an explicit pool= (anything with .connection(), e.g. PgPool), or a con= you opened yourself. A pooled connection is acquired when streaming starts and released when it ends. A con must stay open until the response has been sent, i.e. come from a Depends with yield (FastAPI >= 0.118); async with pool.connection() as conn: return PostgresJsonResponse(q, con=conn) closes it too early. Prefer the pool.
  • Query — a literal string, a psycopg.sql composable, a sqlglot expression (any exp.Expr, so Select, Union, Subquery and other exp.Query values are fine) or a t-string, with values bound through params (%(name)s; a t-string carries its own values and takes no params). Passing only literals is enforced by the type checker, not at runtime: the query is wrapped as a subquery, never parsed, so never build it from user input. For a data-modifying CTE, or custom JSON, pass query_produces_json=True and return one text column per row (a json column or NULL fails). A trailing ; is fine in a string, a ; followed by a comment is not.
  • Options — batch_size (rows per chunk, >= 1, default 1000; don't take it from user input), statement_timeout (seconds, see below), status_code, headers (e.g. {"Cache-Control": "max-age=3600"}).

Cancellation

The response watches for http.disconnect itself (Starlette only does that for ASGI spec < 2.4), also while the query is still running and nothing has been sent yet. It then cancels the query with cancel_safe(), shielded from the cancellation that is tearing the request down, and closes the connection if it was left mid-query: a pool opens a fresh one instead of reusing it. Only a query that is still running is cancelled, so a late cancel can never hit the next statement on a reused connection. A frontend AbortController / closed tab therefore frees the database.

With con=, a disconnect aborts the query and your transaction, unlike fetch_all: the connection is afterwards either left in an aborted transaction (idle, with autocommit) or closed, depending on whether the stream was waiting for Postgres or for the client. Don't use it afterwards, and don't count on earlier writes in that transaction.

Connections are held while the client reads, so a client that reads very slowly pins one (statement_timeout does not help, see below). With PgPool (40 connections by default; max_size=) that means a few slow readers, or many parallel streams of big results, can use up the pool and make every other request wait for a connection. Enforce a write timeout in the reverse proxy and size the pool accordingly. Disconnect behaviour can only be tested over a real socket (httpx.ASGITransport never sends http.disconnect); see tests/test_pg_json_granian.py.

Cancelling other endpoints when the client leaves

PostgresJsonResponse is the only thing that notices a disconnect by itself. ASGI servers such as Granian do not cancel a request handler when the client goes away, so a handler built on fetch_all (or a raw cursor, Meilisearch, ...) runs to the end and its result is thrown away, which defeats a frontend AbortController. Put such read endpoints on a router with CancelOnDisconnectRoute: it runs the handler as a task and cancels it on http.disconnect (in the request's own task, with an asyncio.timeout as the cancel scope, as PostgresJsonResponse does), and psycopg/pgdevkit then cancel the running statement on the server (see fetch_all's cancel).

from fastapi import APIRouter
from pgdevkit.db import fetch_all
from pgdevkit.fastapi import CancelOnDisconnectRoute

cancellable = APIRouter(route_class=CancelOnDisconnectRoute)  # or router.add_api_route(..., route_class_override=...)

@cancellable.post("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/customers/overview")  # a POST that only reads
async def overview(payload: Filter) -> list[Row]:
    return await fetch_all("select ...", {"q": payload.q})  # aborted in Postgres when the client leaves

app.include_router(cancellable)
  • Opt-in, reads only. A cancelled write is rolled back at best and half-applied at worst. Don't select routes by HTTP method: search and grid endpoints are often POSTs.
  • Everything inside the route is cancelled, dependencies included (authentication, audit logging, counters), and the client gets no response, so BackgroundTasks registered by the handler never run. Don't rely on side effects that must always happen in an opted-in route: shield them or do them in a middleware. A Depends with yield sees an HTTPException (status 499) after the cancellation, so a transaction rolls back instead of committing.
  • The client that left gets a 499 nobody reads; the middleware stack sees a normal response (no No response returned).
  • The request body is not buffered: the handler reads it from a Request whose receive() is fed by a reader task that also spots the disconnect. That reader runs at most a few chunks ahead of the handler, so a route that never reads a large body can't pin it in memory, but a disconnect is then only noticed once the body is consumed.
  • Shared work must survive one caller's cancellation. Await shared tasks and futures (single-flight caches, "refresh once, many wait" loaders) through asyncio.shield; otherwise one aborted request cancels the load every other request is waiting on.
  • A returned StreamingResponse is not covered (Starlette watches for the disconnect there only on servers with ASGI spec < 2.4; PostgresJsonResponse always does). Threads (asyncio.to_thread) and external services (Databricks statements, HTTP calls) keep running; only the awaiting stops.
  • Like the above, disconnects only show up over a real socket: see tests/test_pg_json_granian.py for a Granian test.

Statement timeout

statement_timeout=<seconds> (optional, also on fetch_all) lets Postgres abort the query after that long. It is set with SET LOCAL semantics for this one query: transaction-scoped, so it can't leak into the pool and works behind PgBouncer's transaction pooling, and a con you passed gets its previous setting back. Before the first byte the client gets a 504 {"error": "Query timed out"} (the key is error_key, see Errors); after that the response is cut short like any other error. It needs a connection that is not in autocommit mode. Without it, the connection's own statement_timeout applies.

A plain value replaces whatever timeout the connection already has, in both directions: with a database-level safety net of 600 s, a per-call statement_timeout=120 shortens it for that call, but statement_timeout=900 loosens it. To make a per-call value only ever tighten it, wrap it: statement_timeout=at_most(120) (from pgdevkit.db import at_most) uses min(current, 120), where a current timeout of 0 (none) counts as unlimited. at_most works anywhere statement_timeout does, PostgresJsonResponse included.

It limits query execution, nothing else: Postgres defers a cancel while it is blocked writing to a client that has stopped reading, so a stalled reader keeps its connection regardless (a backend was observed staying active in ClientWrite well past a 1 s timeout). For that, use a write timeout in the reverse proxy.

Errors

An error before the first byte becomes a JSON {"error": ...} response (500, or the status and headers of an HTTPException raised by a custom pool); the text is hidden unless you subclass and set expose_errors = True (it contains SQL fragments and parameter values: never in production):

class AppJsonResponse(PostgresJsonResponse):
    expose_errors = IS_DEV

FastAPI's own error handlers answer with {"detail": ...} instead. If a frontend already handles that shape, set error_key = "detail" in the subclass and these responses use it too.

After the first byte the status line is gone: the error is logged and re-raised, and the response ends without the closing ], so clients fail to parse it rather than accepting a silently shortened array. (Granian, which our apps run on, ends such a response cleanly with the truncated body.)

Notes

Timestamps: Postgres' row_to_json writes a timestamptz as 2026-01-01T10:00:00+00:00 (with the session time zone's offset, e.g. +01:00), where Pydantic/FastAPI serialize UTC values as ...Z. Both are valid ISO-8601 and parse to the same instant, but code that compares or displays the raw strings sees a difference; normalize in the frontend, or in the query (to_char(...)).

NaN/Infinity: JSON has no such numbers, so row_to_json writes a float/numeric NaN or Infinity as the string "NaN"/"Infinity". A frontend expecting numbers must handle those strings, or the query should map them (nullif(x, 'NaN')).

Behind Starlette's BaseHTTPMiddleware a failure after the first byte has been reported to end as a clean 200 with a truncated body: the missing closing ] is then the only signal (clients fail to parse it), the status cannot carry it.

It does no authorization: scope the query yourself. The first byte is sent once batch_size rows (or the whole result) are available, so lower batch_size for earlier output; Postgres itself also sends rows in 8 kB buffers, so a slow query with tiny rows delivers late either way.

Releasing

Bump version in pyproject.toml as part of your PR, same as any other change. Once that PR merges to main and the Python Test workflow passes for that commit, .github/workflows/auto-release.yml automatically tags it vX.Y.Z, cuts a GitHub Release (skipping if that version was already released, e.g. a merge that didn't touch the version), and dispatches python-publish.yml to publish it to PyPI — no manual release step, and no extra secret to configure. Two non-obvious GitHub Actions quirks shaped this (see the comments at the top of auto-release.yml for the full reasoning, since both were hit and confirmed the hard way):

  • A release created with the default GITHUB_TOKEN does not trigger other workflows' release: published listeners (an anti-recursion safeguard) — workflow_dispatch is the documented exception, so auto-release.yml dispatches python-publish.yml directly (gh workflow run) instead of relying on the release to cascade into it.
  • python-publish.yml deliberately stays a plain, directly-triggered top-level workflow rather than something auto-release.yml calls via workflow_call: PyPI's OIDC trusted publishing does not support reusable/called workflows and silently rejects the token in that shape.

workflow_dispatch (or an actual GitHub UI release) on python-publish.yml still works as a manual fallback if you ever need to re-publish a version without going through auto-release.yml.

About

Tools for developing with Postgres, mostly around git integration. Also MS SQL if you really need...

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages