Declare a primary key, and let psql print it

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

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.

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 <table>` 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 <table>`, `\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.

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"}),
     ]
 )
 

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": [],

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",

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