From e5ce3b91dc005241dc7d1dcb47a9eacb593a9c9a Mon Sep 17 00:00:00 2001 From: Joe Bylund Date: Mon, 17 Aug 2026 11:11:58 -0400 Subject: [PATCH] Declare a primary key, and let psql print it Table takes primary_key, one column or a tuple of them, normalised to a tuple so nothing downstream branches on the spelling and validated at construction against the table's own columns -- a key naming a column the table does not have is knowable there and should raise on the line that wrote it. A tuple rather than a set because a composite key's order is part of it: it is what conkey records and what psql prints. `\d commits` now ends with Indexes: "commits_pkey" PRIMARY KEY, btree (sha) which is byte-for-byte what PostgreSQL 18 prints for the same declaration. Getting there needed four things that are each invisible on their own: - pg_class.relhasindex, which psql gates the entire footer on. Left False, the rows below are built and never asked for. - a pg_class row for the *index*, relkind 'i'. psql reads the index's name as c2.relname off a second pg_class in the same query, so without it the join drops every index it was told about. - pg_get_indexdef, which cannot be the constant "" the other catalog functions are: psql echoes everything after " USING " straight into the footer, so that substring is where `btree (sha)` comes from. It is rewritten into a reference to a column carrying the rendered text, exactly as format_type already is, and pg_get_constraintdef gets the same treatment for #129's benefit. - pg_index and pg_constraint leaving catalog_rewrite's _ALWAYS_EMPTY_TABLES. Safe because psql reads both in a JOIN rather than in one of the correlated subqueries that substitution exists for. A key's columns are reported not null through both attnotnull and information_schema.columns.is_nullable, which had to be done in both places -- they answer one question and a client reads whichever it reads. PostgreSQL 17 also records them as contype='n' rows in pg_constraint; pg_mimic reports server_version 16, so it does not. Index and constraint OIDs come from their own ranges rather than continuing the table numbering, so declaring a key cannot renumber a table -- table OIDs are positional and a client may have introspected them already. There is a test for that specifically. Identifiers in a rendered definition are quoted only when they have to be, which is what pg_get_constraintdef does: btree (sha) and btree ("userId"). examples/git_sql.py declares the four keys that are actually sound for git data: commits.sha, files.path, branches.name and commit_files (sha, path). Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG.md | 10 +++ docs/session-api.md | 12 ++++ docs/whats-handled.md | 7 +- examples/git_sql.py | 10 ++- pg_mimic/catalog.py | 129 +++++++++++++++++++++++++++++++++--- pg_mimic/catalog_data.py | 8 +++ pg_mimic/catalog_rewrite.py | 39 +++++++++-- pg_mimic/declared.py | 32 +++++++++ tests/test_catalog.py | 18 +++++ tests/test_declared.py | 31 +++++++++ tests/test_pg_catalog.py | 87 +++++++++++++++++++++++- tests/test_psql.py | 75 ++++++++++++++++++++- 12 files changed, 434 insertions(+), 24 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6b85af5..203aa03 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,16 @@ what shipped rather than what was written down at the time. ## Unreleased ### Added +- A `Table` can declare a **primary key** — `Table("commits", {...}, primary_key="sha")`, or a + tuple for a composite one — and psql's `\d` prints the footer it prints for a real + table: `"commits_pkey" PRIMARY KEY, btree (sha)`. Its columns are reported `not null` + by both `pg_catalog.pg_attribute.attnotnull` and + `information_schema.columns.is_nullable`, the key appears in `pg_index` and + `pg_constraint` with a real `pg_get_indexdef()` and `pg_get_constraintdef()`, and + `\di` lists it. Nothing is enforced: pg_mimic stores no rows of its own, so a declared + key is what the catalog reports rather than a uniqueness check. Declaring the key does + not renumber any table's OID. (#127) + - `Schema` and `Table`: what a `Session.schema()` declares, as objects rather than a nested dict. `Schema([Table("users", {"id": "integer"})])` says the same thing `{"users": {"id": "integer"}}` did, and the dict is still accepted — but a `Table` has diff --git a/docs/session-api.md b/docs/session-api.md index f17390c..90240ed 100644 --- a/docs/session-api.md +++ b/docs/session-api.md @@ -59,6 +59,18 @@ class MySession(Session): ) ``` +A `Table` can declare a primary key, which is what psql's `\d` reads to print its `Indexes:` +footer. One column or several, in key order: + +```python +Table("commits", {"sha": "text", "author": "text"}, primary_key="sha") +Table("commit_files", {"sha": "text", "path": "text"}, primary_key=("sha", "path")) +``` + +Its columns are then reported `not null`, as they are in Postgres, by both `attnotnull` and +`information_schema.columns.is_nullable`. Nothing is enforced — pg_mimic stores no rows of its own, +so a declared key is what the catalog reports rather than a uniqueness check anything applies. + The nested `{table: {column: type_name}}` dict this returned before is still accepted, so an existing session needs no change. A `Schema` is what pg_mimic normalises either into, and it is what `TableSession.schema()` hands back. diff --git a/docs/whats-handled.md b/docs/whats-handled.md index 7ba6150..93c0ab9 100644 --- a/docs/whats-handled.md +++ b/docs/whats-handled.md @@ -57,8 +57,11 @@ before it ever reaches your `Session`, so real clients/ORMs/`psql` work without - `information_schema.tables` and `information_schema.columns` at PostgreSQL's full width — all 12 columns and all 44, in the server's own order, so `SELECT *` lines up positionally and a query for `column_default`, `udt_name` or `numeric_precision` answers about the columns that exist instead of - returning nothing. Plus the slice of `pg_catalog` psql's `\dt`, `\d ` and `\dn` read — all - built from your `Session.schema()` declaration — see [the session API](./session-api.md#declaring-a-schema). + returning nothing. Plus the slice of `pg_catalog` psql's `\dt`, `\d
`, `\di` and `\dn` read — + all built from your `Session.schema()` declaration — see + [the session API](./session-api.md#declaring-a-schema). A declared primary key gets the + `Indexes:` footer `\d` prints for one, reports its columns as `not null`, and appears in + `pg_index`, `pg_constraint` and `\di`. - Multi-statement simple-query batches (`"BEGIN; INSERT ...; COMMIT;"` sent as one `'Q'` message, e.g. by `psql -f script.sql`) are split into individual statements (via sqlglot), each getting its own `RowDescription`/`DataRow*`/`CommandComplete` — not silently merged or truncated. diff --git a/examples/git_sql.py b/examples/git_sql.py index 4e77c76..b52a34a 100644 --- a/examples/git_sql.py +++ b/examples/git_sql.py @@ -173,13 +173,19 @@ def _interval_delta(count: str, unit: str) -> timedelta: "deletions": "integer", "files_changed": "integer", }, + primary_key="sha", ), Table( "commit_files", {"sha": "text", "path": "text", "insertions": "integer", "deletions": "integer"}, + primary_key=("sha", "path"), + ), + Table("files", {"path": "text", "ext": "text", "size_bytes": "integer", "lines": "integer"}, primary_key="path"), + Table( + "branches", + {"name": "text", "is_head": "boolean", "upstream": "text", "last_commit_at": "timestamp"}, + primary_key="name", ), - Table("files", {"path": "text", "ext": "text", "size_bytes": "integer", "lines": "integer"}), - Table("branches", {"name": "text", "is_head": "boolean", "upstream": "text", "last_commit_at": "timestamp"}), ] ) diff --git a/pg_mimic/catalog.py b/pg_mimic/catalog.py index 36b4a23..4e4456f 100644 --- a/pg_mimic/catalog.py +++ b/pg_mimic/catalog.py @@ -33,7 +33,7 @@ from .arrays import ARRAY_OID, is_array_oid from .catalog_data import INFORMATION_SCHEMA_SCHEMA, PG_CATALOG_SCHEMA from .catalog_rewrite import rewrite_for_executor -from .declared import Schema, resolve +from .declared import Schema, Table, resolve from .describe import oid_for_declared_type, result_columns from .errors import FEATURE_NOT_SUPPORTED, UNDEFINED_COLUMN, PgError from .results import ResultColumn @@ -69,8 +69,15 @@ _PUBLIC_OID = 2200 _OWNER_OID = 10 _HEAP_AM_OID = 2 +_BTREE_AM_OID = 403 # Postgres' own oid for btree, which psql prints by name # Real Postgres starts user objects here, and psql only cares that they're distinct. _FIRST_USER_OID = 16384 +# Indexes and constraints get their own OID ranges rather than continuing the table +# numbering, so that declaring a primary key cannot renumber a *table* -- table OIDs +# are positional and a client may have introspected them already. The gaps are wide +# enough that a schema would have to declare 100000 tables to collide. +_FIRST_INDEX_OID = _FIRST_USER_OID + 100_000 +_FIRST_CONSTRAINT_OID = _FIRST_USER_OID + 200_000 _CATALOG_NAME = "pg_mimic" @@ -192,9 +199,12 @@ def _build_information_schema(declared: Schema) -> _CatalogInput: table_name=table_name, column_name=col_name, ordinal_position=i, - # A declared schema names types, not constraints, so nothing here - # is NOT NULL -- as pg_catalog already says with attnotnull. - is_nullable="YES", + # A primary key's columns are implicitly NOT NULL, which is the one + # thing a declared schema settles here. Everything else is nullable: + # nothing declares a NOT NULL of its own yet. pg_catalog says the + # same through attnotnull, and the two must not disagree -- a client + # reading one and a client reading the other are asking one question. + is_nullable="NO" if col_name in table.primary_key else "YES", data_type=col_type, udt_catalog=_CATALOG_NAME, # The underlying type, named as pg_type names it -- `int4`, @@ -301,10 +311,32 @@ def _typname_for(rows: list[dict], oid: int) -> str: return "unknown" +_PLAIN_IDENTIFIER = re.compile(r"^[a-z_][a-z0-9_$]*$") + + +def _rendered_identifier(name: str) -> str: + """A name as Postgres would print it in a constraint or index definition. + + Quoted only when it has to be, which is what `pg_get_constraintdef` does -- an + unquoted `sha` and a quoted `"userId"`. Reserved words are not consulted: the list + is long, it changes between releases, and a name that needs quoting for *that* + reason renders as bare here rather than wrongly quoted somewhere else. + """ + return name if _PLAIN_IDENTIFIER.match(name) else '"' + name.replace('"', '""') + '"' + + +def _key_positions(table: Table) -> list[int]: + """A key's columns as 1-based attnums, in key order -- Postgres' conkey and indkey.""" + positions = {column: position for position, column in enumerate(table.columns, start=1)} + return [positions[column] for column in table.primary_key] + + def _build_pg_catalog(declared: Schema, database: str = "postgres") -> _CatalogInput: """The slice of pg_catalog psql's \\d family reads, from what the session declares.""" class_rows = [] attribute_rows = [] + index_rows: list[dict] = [] + constraint_rows: list[dict] = [] for index, (table_name, table) in enumerate(declared.tables.items()): table_oid = _FIRST_USER_OID + index class_rows.append( @@ -317,7 +349,9 @@ def _build_pg_catalog(declared: Schema, database: str = "postgres") -> _CatalogI "relam": _HEAP_AM_OID, "reltuples": -1, # never analysed "relpages": 0, - "relhasindex": False, + # psql gates the whole `Indexes:` footer on this, so a declared primary + # key that does not set it builds rows nothing ever asks for. + "relhasindex": bool(table.primary_key), "relpersistence": "p", "reltablespace": 0, "relispartition": False, @@ -338,7 +372,11 @@ def _build_pg_catalog(declared: Schema, database: str = "postgres") -> _CatalogI "attname": col_name, "atttypid": oid_for_declared_type(col_type), "attnum": position, - "attnotnull": False, + # A primary key's columns are implicitly NOT NULL in Postgres, and + # psql's Nullable column reads exactly this. PostgreSQL 17 also + # records them as contype='n' rows in pg_constraint; pg_mimic reports + # server_version 16, so it does not. + "attnotnull": col_name in table.primary_key, "atthasdef": False, "attisdropped": False, "attidentity": "", @@ -353,6 +391,76 @@ def _build_pg_catalog(declared: Schema, database: str = "postgres") -> _CatalogI } ) + if table.primary_key: + index_oid = _FIRST_INDEX_OID + index + constraint_oid = _FIRST_CONSTRAINT_OID + index + key_name = f"{table_name}_pkey" + key_columns = ", ".join(_rendered_identifier(column) for column in table.primary_key) + # An index is a relation, so it needs its own pg_class row: psql reads the + # index's *name* as c2.relname off a second pg_class in the same query, and + # without this row the join drops every index it was told about. + class_rows.append( + { + "oid": index_oid, + "relname": key_name, + "relnamespace": _PUBLIC_OID, + "relkind": "i", + "relowner": _OWNER_OID, + "relam": _BTREE_AM_OID, + "reltuples": -1, + "relpages": 0, + "relhasindex": False, + "relpersistence": "p", + "reltablespace": 0, + "relispartition": False, + "relrowsecurity": False, + "relforcerowsecurity": False, + "relreplident": "n", + "reloftype": 0, + "relhastriggers": False, + "relchecks": 0, + "relhasrules": False, + "reltoastrelid": 0, + } + ) + index_rows.append( + { + # psql echoes everything after " USING " verbatim into the footer, + # which is where `btree (sha)` comes from -- so the access method and + # the column list have to be in this string, not derived from relam. + "indexdef": f"CREATE UNIQUE INDEX {key_name} ON {_rendered_identifier(table_name)} USING btree ({key_columns})", + "indexrelid": index_oid, + "indrelid": table_oid, + "indisprimary": True, + "indisunique": True, + "indisclustered": False, + "indisvalid": True, + "indisreplident": False, + "indkey": " ".join(str(position) for position in _key_positions(table)), + } + ) + constraint_rows.append( + { + "condef": f"PRIMARY KEY ({key_columns})", + "oid": constraint_oid, + "conname": key_name, + "conrelid": table_oid, + "contype": "p", + "conparentid": 0, + "conindid": index_oid, + "confrelid": 0, + "condeferrable": False, + "condeferred": False, + "convalidated": True, + "connoinherit": False, + "conislocal": True, + "coninhcount": 0, + # Postgres' int2[] of attnums, as its array literal -- the column is + # declared TEXT here, and a Python list has no wire encoding. + "conkey": "{" + ",".join(str(position) for position in _key_positions(table)) + "}", + } + ) + tables = { "pg_catalog": { "pg_namespace": [ @@ -361,13 +469,16 @@ def _build_pg_catalog(declared: Schema, database: str = "postgres") -> _CatalogI {"oid": 13000, "nspname": "information_schema", "nspowner": _OWNER_OID}, ], "pg_class": class_rows, - "pg_am": [{"oid": _HEAP_AM_OID, "amname": "heap", "amhandler": "heap_tableam_handler", "amtype": "t"}], + "pg_am": [ + {"oid": _HEAP_AM_OID, "amname": "heap", "amhandler": "heap_tableam_handler", "amtype": "t"}, + {"oid": _BTREE_AM_OID, "amname": "btree", "amhandler": "bthandler", "amtype": "i"}, + ], "pg_attribute": attribute_rows, "pg_type": _pg_type_rows(), "pg_attrdef": [], "pg_collation": [], - "pg_constraint": [], - "pg_index": [], + "pg_constraint": constraint_rows, + "pg_index": index_rows, "pg_inherits": [], "pg_rewrite": [], "pg_trigger": [], diff --git a/pg_mimic/catalog_data.py b/pg_mimic/catalog_data.py index 1c452f8..3b9da53 100644 --- a/pg_mimic/catalog_data.py +++ b/pg_mimic/catalog_data.py @@ -78,6 +78,9 @@ # than returning the no rows it should. "pg_attrdef": {"adrelid": "INT", "adnum": "INT", "adbin": "TEXT"}, "pg_constraint": { + # condef is pg_get_constraintdef()'s answer, carried on the row -- see the + # note on pg_index.indexdef. + "condef": "TEXT", "oid": "INT", "conname": "TEXT", "conrelid": "INT", @@ -94,6 +97,11 @@ "conkey": "TEXT", }, "pg_index": { + # indexdef carries what pg_get_indexdef() would render, for the same reason + # attformattype carries format_type()'s answer: it depends on the row, so it + # cannot be a constant like the other catalog functions. catalog_rewrite + # turns the call into a reference to this column. + "indexdef": "TEXT", "indexrelid": "INT", "indrelid": "INT", "indisprimary": "BOOLEAN", diff --git a/pg_mimic/catalog_rewrite.py b/pg_mimic/catalog_rewrite.py index 93c6d29..a842a73 100644 --- a/pg_mimic/catalog_rewrite.py +++ b/pg_mimic/catalog_rewrite.py @@ -49,9 +49,9 @@ def _regexp_like(this: Any, expression: Any, *_: Any) -> Any: "pg_table_is_visible": lambda connection: exp.true(), "pg_type_is_visible": lambda connection: exp.true(), "pg_function_is_visible": lambda connection: exp.true(), - # No indexes, constraints or defaults to describe. - "pg_get_indexdef": lambda connection: exp.Literal.string(""), - "pg_get_constraintdef": lambda connection: exp.Literal.string(""), + # pg_get_indexdef and pg_get_constraintdef are *not* here: their answer depends on + # the row, so they are rewritten into a reference to the column carrying it, the + # same way format_type is. No column defaults to describe. "pg_get_expr": lambda connection: exp.Literal.string(""), "pg_get_partkeydef": lambda connection: exp.null(), "pg_get_viewdef": lambda connection: exp.Literal.string(""), @@ -67,13 +67,17 @@ def _regexp_like(this: Any, expression: Any, *_: Any) -> Any: _REGEX_OPERATORS = {"~", "!~", "~*", "!~*"} -# Catalog tables pg_mimic never has rows for: there are no column defaults, no -# non-default collations, and no indexes or constraints to describe. +# Catalog functions whose answer depends on the row rather than being constant, and the +# column each one's answer travels on. See the rewrite in rewrite_for_executor. +_ROW_DEPENDENT_FUNCTIONS = {"pg_get_indexdef": "indexdef", "pg_get_constraintdef": "condef"} + +# Catalog tables pg_mimic never has rows for: there are no column defaults and no +# non-default collations. pg_index and pg_constraint left this set in #127 -- a declared +# primary key puts rows in both, and psql reads them in a JOIN rather than in one of the +# correlated subqueries this substitution exists for. _ALWAYS_EMPTY_TABLES = { "pg_attrdef", "pg_collation", - "pg_constraint", - "pg_index", "pg_inherits", "pg_rewrite", "pg_trigger", @@ -159,6 +163,27 @@ def rewrite_for_executor(connection: Connection, expr: exp.Expression) -> exp.Ex parent = node.parent (parent if isinstance(parent, exp.Dot) else node).replace(column) + # pg_get_indexdef(i.indexrelid, 0, true) and pg_get_constraintdef(con.oid, true) + # depend on the row for the same reason format_type does, so they get the same + # treatment: the rendered text travels on the row and the call becomes a reference + # to it. psql echoes everything after " USING " in the indexdef straight into the + # `Indexes:` footer, which is where `btree (sha)` comes from -- a constant "" there + # would render an index with no columns. + for node in list(expr.find_all(exp.Anonymous)): + rendered = _ROW_DEPENDENT_FUNCTIONS.get(str(node.this).lower()) + if rendered is None or not node.expressions: + continue + first = node.expressions[0] + if not isinstance(first, exp.Column): + continue + # The argument says which relation the OID came from, so the replacement carries + # the same qualifier -- `con.oid` becomes `con.condef`. An unqualified `oid` + # becomes an unqualified `condef`, which qualify() then resolves, since only + # pg_constraint has that column. + column = exp.column(rendered, table=first.table) + parent = node.parent + (parent if isinstance(parent, exp.Dot) else node).replace(column) + # psql reads column defaults and collations with correlated scalar subqueries, # which sqlglot's executor cannot run at all. Both select from tables that are # always empty here -- pg_mimic models neither -- so the answer is NULL either diff --git a/pg_mimic/declared.py b/pg_mimic/declared.py index 0fdcce4..39609f5 100644 --- a/pg_mimic/declared.py +++ b/pg_mimic/declared.py @@ -48,10 +48,22 @@ class Table: A table with no columns at all is legal, because Postgres allows `CREATE TABLE footable ()`, permits queries over it, and keeps its row count. + + `primary_key` names this table's own columns, in key order, and is normalised to a + tuple so nothing downstream has to branch on the spelling:: + + Table("commits", {"sha": "text", ...}, primary_key="sha") + Table("commit_files", {"sha": "text", "path": "text"}, primary_key=("sha", "path")) + + A tuple rather than a set because a composite key's order is part of it -- it is + what `pg_constraint.conkey` records and what psql prints. pg_mimic enforces + nothing: nothing here stores rows, so a declared key is what the catalog reports, + not a uniqueness check anything applies. """ name: str columns: Mapping[str, str] + primary_key: str | Sequence[str] = () def __post_init__(self) -> None: if not isinstance(self.name, str) or not self.name: @@ -72,6 +84,26 @@ def __post_init__(self) -> None: # included. This is where that check goes, once #134 makes the ordinary # spellings resolve so validating them stops rejecting them. object.__setattr__(self, "columns", MappingProxyType(columns)) + object.__setattr__(self, "primary_key", _key_columns(self.name, "primary key", self.primary_key, columns)) + + +def _key_columns(table: str, kind: str, declared: str | Sequence[str], columns: Mapping[str, str]) -> tuple[str, ...]: + """A key spec as a tuple of this table's own columns, in the order it was given. + + A bare string is one column, which is the common case and not worth making callers + spell as a one-tuple. Checked against `columns` here rather than at resolve time, + because a key naming a column the table does not have is knowable at construction + and so should raise on the line that wrote it. Whether a *reference target* is + covered by a key is the other kind of question and cannot be answered until every + table is in -- see https://github.com/jbylund/pg_mimic/issues/129. + """ + key = (declared,) if isinstance(declared, str) else tuple(declared) + for column in key: + if column not in columns: + raise ValueError(f"{table!r} declares a {kind} on {column!r}, which is not one of its columns") + if len(set(key)) != len(key): + raise ValueError(f"{table!r} names a column twice in its {kind}: {key}") + return key class Schema: diff --git a/tests/test_catalog.py b/tests/test_catalog.py index 0083e28..051f981 100644 --- a/tests/test_catalog.py +++ b/tests/test_catalog.py @@ -4,6 +4,8 @@ import pytest +from pg_mimic import Schema, Table + def test_information_schema_tables(conn, mock_session): async def schema(): @@ -626,3 +628,19 @@ async def schema(): assert cur.fetchall() == [] cur.execute("SELECT relname FROM pg_catalog.pg_class") assert cur.fetchall() == [] + + +async def test_a_primary_keys_columns_are_reported_not_nullable(conn, mock_session): + """information_schema.columns.is_nullable and pg_catalog's attnotnull answer one + question, so they must not disagree -- a client reads whichever it reads.""" + + async def schema(): + return Schema([Table("commits", {"sha": "text", "author": "text"}, primary_key="sha")]) + + mock_session.schema = schema + + with conn.cursor() as cur: + cur.execute("SELECT column_name, is_nullable FROM information_schema.columns WHERE table_name = 'commits'") + assert cur.fetchall() == [("sha", "NO"), ("author", "YES")] + cur.execute("SELECT attname, attnotnull FROM pg_catalog.pg_attribute ORDER BY attnum") + assert cur.fetchall() == [("sha", True), ("author", False)] diff --git a/tests/test_declared.py b/tests/test_declared.py index e1df39c..928bbdc 100644 --- a/tests/test_declared.py +++ b/tests/test_declared.py @@ -133,3 +133,34 @@ def test_resolve_refuses_what_is_not_a_schema_at_all(declared): def test_a_table_refuses_a_name_or_type_that_is_not_a_non_empty_string(name, columns): with pytest.raises(ValueError): Table(name, columns) + + +# --- primary keys -------------------------------------------------------------------- + + +def test_a_primary_key_is_normalised_to_a_tuple(): + """A bare string is one column, which is the common case and not worth making a + caller spell as a one-tuple. Storing only the tuple means nothing downstream has to + branch on which spelling it was given.""" + assert Table("t", {"a": "text"}, primary_key="a").primary_key == ("a",) + assert Table("t", {"a": "text", "b": "text"}, primary_key=("a", "b")).primary_key == ("a", "b") + assert Table("t", {"a": "text", "b": "text"}, primary_key=["b", "a"]).primary_key == ("b", "a") + + +def test_a_table_declares_no_primary_key_by_default(): + assert Table("t", {"a": "text"}).primary_key == () + + +def test_a_composite_key_keeps_the_order_it_was_declared_in(): + """Not a set: the order is part of the key, and it is what conkey and psql report.""" + assert Table("t", {"a": "text", "b": "text"}, primary_key=("b", "a")).primary_key == ("b", "a") + + +def test_a_primary_key_must_name_the_tables_own_columns(): + with pytest.raises(ValueError, match="not one of its columns"): + Table("t", {"a": "text"}, primary_key="b") + + +def test_a_primary_key_cannot_name_a_column_twice(): + with pytest.raises(ValueError, match="twice in its primary key"): + Table("t", {"a": "text"}, primary_key=("a", "a")) diff --git a/tests/test_pg_catalog.py b/tests/test_pg_catalog.py index db39333..5003271 100644 --- a/tests/test_pg_catalog.py +++ b/tests/test_pg_catalog.py @@ -12,11 +12,12 @@ from __future__ import annotations +import pytest import sqlglot from pg_mimic.catalog import _build_pg_catalog from pg_mimic.catalog_rewrite import rewrite_for_executor -from pg_mimic.declared import resolve +from pg_mimic.declared import Schema, Table, resolve SCHEMA = {"users": {"id": "integer", "name": "text"}, "orders": {"id": "bigint", "total": "numeric"}} @@ -151,3 +152,87 @@ def test_pg_type_covers_the_types_pg_mimic_can_encode(): assert by_oid[TEXT]["typname"] == "text" assert by_oid[ARRAY_OID[INT8]]["typelem"] == INT8 # what marks it an array assert by_oid[ARRAY_OID[TEXT]]["typname"] == "_text" + + +# --- a declared primary key ---------------------------------------------------------- + +_PK_SCHEMA = Schema( + [ + Table("commits", {"sha": "text", "author": "text"}, primary_key="sha"), + Table("commit_files", {"sha": "text", "path": "text", "adds": "integer"}, primary_key=("sha", "path")), + Table("files", {"path": "text"}), + Table("odd", {"userId": "integer"}, primary_key="userId"), + ] +) + +# psql's index-footer query for `\d
`, verbatim except the OID literal. The +# LEFT JOIN to pg_constraint is why an index needs a matching constraint row, and +# `c2.relname` is why the index needs a pg_class row of its own. +_INDEX_FOOTER = """ +SELECT c2.relname, i.indisprimary, i.indisunique, i.indisclustered, i.indisvalid, + pg_catalog.pg_get_indexdef(i.indexrelid, 0, true), + pg_catalog.pg_get_constraintdef(con.oid, true), contype, condeferrable, condeferred, + i.indisreplident, c2.reltablespace +FROM pg_catalog.pg_class c, pg_catalog.pg_class c2, pg_catalog.pg_index i + LEFT JOIN pg_catalog.pg_constraint con + ON (conrelid = i.indrelid AND conindid = i.indexrelid AND contype IN ('p','u','x')) +WHERE c.oid = '{oid}' AND c.oid = i.indrelid AND i.indexrelid = c2.oid +ORDER BY i.indisprimary DESC, c2.relname +""" + + +def _oid_of(table_name): + rows = _run(f"SELECT oid FROM pg_catalog.pg_class WHERE relname = '{table_name}' AND relkind = 'r'", _PK_SCHEMA) + return rows[0][0] + + +@pytest.mark.parametrize( + argnames=["table_name", "index_name", "indexed"], + argvalues=[ + ["commits", "commits_pkey", "(sha)"], + ["commit_files", "commit_files_pkey", "(sha, path)"], + ["odd", "odd_pkey", '("userId")'], + ], + ids=["one column", "composite, in key order", "a name that needs quoting"], +) +def test_psqls_index_footer_query_answers_for_a_declared_primary_key(table_name, index_name, indexed): + """The whole point of the row-carried indexdef: psql echoes everything after + " USING " into the footer, so that substring is what renders `btree (sha)`.""" + rows = _run(_INDEX_FOOTER.format(oid=_oid_of(table_name)), _PK_SCHEMA) + assert len(rows) == 1 + relname, indisprimary, indisunique = rows[0][0], rows[0][1], rows[0][2] + indexdef, condef, contype = rows[0][5], rows[0][6], rows[0][7] + assert (relname, indisprimary, indisunique) == (index_name, True, True) + assert indexdef.endswith(f" USING btree {indexed}") + assert (condef, contype) == (f"PRIMARY KEY {indexed}", "p") + + +def test_a_table_with_no_primary_key_has_no_index_footer(): + assert _run(_INDEX_FOOTER.format(oid=_oid_of("files")), _PK_SCHEMA) == [] + + +def test_declaring_a_primary_key_does_not_renumber_the_tables(): + """Table OIDs are positional and a client may have introspected them, so index and + constraint OIDs come from their own ranges rather than continuing the table run.""" + keyed = _run("SELECT relname, oid FROM pg_catalog.pg_class WHERE relkind = 'r' ORDER BY oid", _PK_SCHEMA) + unkeyed = _run( + "SELECT relname, oid FROM pg_catalog.pg_class WHERE relkind = 'r' ORDER BY oid", + Schema([Table(name, dict(table.columns)) for name, table in _PK_SCHEMA.tables.items()]), + ) + assert keyed == unkeyed + indexes = _run("SELECT oid FROM pg_catalog.pg_class WHERE relkind = 'i'", _PK_SCHEMA) + assert not {oid for (oid,) in indexes} & {oid for _, oid in keyed} + + +def test_a_primary_keys_columns_are_not_null_and_the_others_are(): + rows = _run( + "SELECT a.attname, a.attnotnull FROM pg_catalog.pg_attribute a " + f"WHERE a.attrelid = '{_oid_of('commit_files')}' ORDER BY a.attnum", + _PK_SCHEMA, + ) + assert rows == [("sha", True), ("path", True), ("adds", False)] + + +def test_conkey_is_the_key_columns_attnums_in_key_order(): + rows = _run("SELECT conname, conkey FROM pg_catalog.pg_constraint ORDER BY conname", _PK_SCHEMA) + assert dict(rows) == {"commit_files_pkey": "{1,2}", "commits_pkey": "{1}", "odd_pkey": "{1}"} diff --git a/tests/test_psql.py b/tests/test_psql.py index aed3156..d06159c 100644 --- a/tests/test_psql.py +++ b/tests/test_psql.py @@ -17,7 +17,7 @@ import pytest -from pg_mimic import ResultColumn, Session +from pg_mimic import ResultColumn, Schema, Session, Table from pg_mimic.testing import serve_in_thread psql_required = pytest.mark.skipif(shutil.which("psql") is None, reason="psql is not installed") @@ -39,8 +39,8 @@ async def query(self, sql, params): yield ("row",) -def _psql(command): - with serve_in_thread(SchemaSession) as server: +def _psql(command, session=None): + with serve_in_thread(session or SchemaSession) as server: return subprocess.run( ["psql", server.dsn(user="test", dbname="test"), "-c", command], capture_output=True, @@ -137,3 +137,72 @@ async def schema(): "SELECT pg_catalog.format_type(a.atttypid, a.atttypmod) FROM pg_catalog.pg_attribute AS a WHERE a.attname = 'id'" ) assert cur.fetchall() == [("integer",)] + + +# --- a declared primary key ---------------------------------------------------------- + +_PK_SCHEMA = Schema( + [ + Table("commits", {"sha": "text", "author": "text"}, primary_key="sha"), + Table("commit_files", {"sha": "text", "path": "text"}, primary_key=("sha", "path")), + Table("files", {"path": "text"}), + Table("odd", {"userId": "integer"}, primary_key="userId"), + ] +) + + +class PrimaryKeySession(SchemaSession): + async def schema(self): + return _PK_SCHEMA + + +@psql_required +@pytest.mark.parametrize( + argnames=["table_name", "footer"], + argvalues=[ + ["commits", '"commits_pkey" PRIMARY KEY, btree (sha)'], + ["commit_files", '"commit_files_pkey" PRIMARY KEY, btree (sha, path)'], + ["odd", '"odd_pkey" PRIMARY KEY, btree ("userId")'], + ], + ids=["one column", "composite", "a name that needs quoting"], +) +def test_describe_table_shows_a_declared_primary_key(table_name, footer): + """Byte-for-byte what PostgreSQL 18 prints for the same declaration, which is the + assertion that matters: psql composes this footer from four columns across three + catalog tables, and any one of them missing renders it wrong or not at all.""" + result = _psql(f"\\d {table_name}", PrimaryKeySession) + assert result.stderr.strip() == "", result.stderr + assert "Indexes:" in result.stdout + assert footer in result.stdout + + +@psql_required +def test_a_table_with_no_primary_key_has_no_indexes_footer(): + result = _psql("\\d files", PrimaryKeySession) + assert result.stderr.strip() == "", result.stderr + assert "path" in result.stdout + assert "Indexes:" not in result.stdout + + +@psql_required +def test_a_primary_keys_columns_are_reported_not_null(): + result = _psql("\\d commits", PrimaryKeySession) + assert result.stderr.strip() == "", result.stderr + assert "not null" in result.stdout + + +@psql_required +def test_list_indexes_shows_the_key_index(): + result = _psql("\\di", PrimaryKeySession) + assert result.stderr.strip() == "", result.stderr + assert "commits_pkey" in result.stdout and "index" in result.stdout + + +@psql_required +def test_listing_tables_does_not_list_the_key_index(): + """`\\dt` filters on relkind, and an index is a relation -- so the pg_class row an + index needs for its name must not leak into the table list.""" + result = _psql("\\dt", PrimaryKeySession) + assert result.stderr.strip() == "", result.stderr + assert "commits" in result.stdout + assert "pkey" not in result.stdout