Declare tables as objects, so constraints have somewhere to live

CHANGELOG.md

@@ -10,6 +10,17 @@ what shipped rather than what was written down at the time.
 ## Unreleased
 
 ### Added
+- `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
+  somewhere to put a primary key or a foreign key, which a dict of column types does
+  not, and which is what `pg_constraint` and `pg_index` need before psql's `\d` can
+  report either. Declaring them comes next (#127 onwards); this is the shape they hang
+  off. A `Table` is immutable, because one declared at module scope is shared by every
+  connection, and both the table order and each table's column order are meaningful —
+  they decide table OIDs and `ordinal_position`. A table with no columns is legal, as it
+  is in Postgres. (#126)
+
 - `information_schema.tables` and `information_schema.columns` are served at
   PostgreSQL's full width: 12 columns and 44, up from 4 and 7, in the server's own
   `ordinal_position` order so `SELECT *` lines up positionally. A column pg_mimic
@@ -92,6 +103,14 @@ what shipped rather than what was written down at the time.
   it exists for. It now reports `self.port`. (#106)
 
 ### Changed
+- `TableSession.schema()` returns a `Schema` rather than the nested
+  `{table: {column: type_name}}` dict. `Schema.column_types()` is the same declaration in
+  the old shape for anything that wants it. A session of your own may keep returning the
+  dict — that is accepted and will stay accepted. `Session.schema()` also now defaults to
+  an empty `Schema` instead of `None`: the two produced identical catalogs, and the branch
+  guarding the difference is what once crashed every catalog query for a session that did
+  not override it. (#126)
+
 - **Breaking.** `SET` checks whether the parameter exists and whether a session may
   change it, where before it accepted any name at all. A name that is not a
   parameter is `42704 undefined_object`; a parameter outside a session's reach is

docs/session-api.md

@@ -42,11 +42,34 @@ silently reshaped, since Postgres has no wire representation for one.
 
 ## Declaring a schema
 
-`schema()` returns `{table: {column: type_name}}`, where each type name is a SQL spelling Postgres
-itself uses — `"integer"`, `"character varying"`, `"text[]"`. `information_schema` and the
-`pg_catalog` slice psql's `\d` reads are built out of it (see
-[What's handled automatically](./whats-handled.md)), and so is `TableSession`'s answer to
-every `Describe`.
+`schema()` returns a `Schema` of `Table`s. Each `Table` maps a column name to a SQL spelling
+Postgres itself uses — `"integer"`, `"character varying"`, `"text[]"`:
+
+```python
+from pg_mimic import Schema, Session, Table
+
+
+class MySession(Session):
+    async def schema(self):
+        return Schema(
+            [
+                Table("users", {"id": "integer", "name": "text"}),
+                Table("orders", {"id": "integer", "user_id": "integer", "total": "numeric"}),
+            ]
+        )
+```
+
+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.
+
+Two properties are worth knowing because nothing else will tell you: **the order of the tables and
+of each table's columns is meaningful** — it decides table OIDs and `ordinal_position` — and a
+`Table` is immutable, since one declared at module scope is shared by every connection.
+
+`information_schema` and the `pg_catalog` slice psql's `\d` reads are built out of the declaration
+(see [What's handled automatically](./whats-handled.md)), and so is `TableSession`'s answer to
+every `Describe`. A table with no columns is legal, as it is in Postgres.
 
 `oid_for_declared_type()` is how one of those names becomes an OID, and the reason a session that
 declares a schema needs no type table of its own. Its neighbour `oid_for_type()` answers the same

docs/whats-handled.md

@@ -58,7 +58,7 @@ before it ever reaches your `Session`, so real clients/ORMs/`psql` work without
   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()`.
+  built from your `Session.schema()` declaration — see [the session API](./session-api.md#declaring-a-schema).
 - 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

@@ -22,7 +22,7 @@
 Two parts are worth reading for their own sake.
 
 describe() never runs the query. Column shape comes from sqlglot's type
-annotator against SCHEMA and is exact even through aggregates and arithmetic
+annotator against the declared SCHEMA and is exact even through aggregates and arithmetic
 -- count(*) is int8, insertions / 2.0 is float8 -- so Parse/Describe is
 answered without touching git. examples/dbapi_proxy.py has to execute a SELECT
 to learn its own column types; a declared schema buys the difference. The
@@ -59,7 +59,9 @@
     INT8,
     NUMERIC,
     TIMESTAMP,
+    Schema,
     Session,
+    Table,
     oid_for_declared_type,
 )
 from pg_mimic.analysis import AnalyzedQuery
@@ -155,38 +157,37 @@ def _interval_delta(count: str, unit: str) -> timedelta:
 # (pg_mimic.describe), so neither can drift from the library's. This file used to
 # carry its own copy of both, and did drift: it described `select 3000000000` as
 # int4 long after #40 was fixed for TableSession.
-SCHEMA = {
-    "commits": {
-        "sha": "text",
-        "author_name": "text",
-        "author_email": "text",
-        "committer_name": "text",
-        "committed_at": "timestamp",
-        "authored_at": "timestamp",
-        "subject": "text",
-        "insertions": "integer",
-        "deletions": "integer",
-        "files_changed": "integer",
-    },
-    "commit_files": {
-        "sha": "text",
-        "path": "text",
-        "insertions": "integer",
-        "deletions": "integer",
-    },
-    "files": {
-        "path": "text",
-        "ext": "text",
-        "size_bytes": "integer",
-        "lines": "integer",
-    },
-    "branches": {
-        "name": "text",
-        "is_head": "boolean",
-        "upstream": "text",
-        "last_commit_at": "timestamp",
-    },
-}
+SCHEMA = Schema(
+    [
+        Table(
+            "commits",
+            {
+                "sha": "text",
+                "author_name": "text",
+                "author_email": "text",
+                "committer_name": "text",
+                "committed_at": "timestamp",
+                "authored_at": "timestamp",
+                "subject": "text",
+                "insertions": "integer",
+                "deletions": "integer",
+                "files_changed": "integer",
+            },
+        ),
+        Table(
+            "commit_files",
+            {"sha": "text", "path": "text", "insertions": "integer", "deletions": "integer"},
+        ),
+        Table("files", {"path": "text", "ext": "text", "size_bytes": "integer", "lines": "integer"}),
+        Table("branches", {"name": "text", "is_head": "boolean", "upstream": "text", "last_commit_at": "timestamp"}),
+    ]
+)
+
+# sqlglot's `schema=` wants the nested dict, and so does the one place below that reads
+# a column's declared type. Derived once here rather than at each call site, and it is
+# the *declared* spellings rather than sqlglot's own view of them -- see
+# https://github.com/jbylund/pg_mimic/issues/135.
+COLUMN_TYPES = SCHEMA.column_types()
 
 
 # --- pushdown ------------------------------------------------------------------------
@@ -570,7 +571,7 @@ def _infer_param_oids(self, sql: str, declared: list[int | None]) -> list[int |
             # away. No `$n`, nothing to infer.
             return declared
         try:
-            expr = AnalyzedQuery(sqlglot.parse_one(sql, dialect="postgres"), schema=SCHEMA).qualified()
+            expr = AnalyzedQuery(sqlglot.parse_one(sql, dialect="postgres"), schema=COLUMN_TYPES).qualified()
         except Exception:
             return declared  # unparseable, or not about our tables -- leave it alone
         aliases = _alias_map(expr)
@@ -611,7 +612,7 @@ def _analyzed(self, sql: str) -> AnalyzedQuery | None:
         expr = sqlglot.parse_one(sql, dialect="postgres")
         if not isinstance(expr, exp.Select):
             return None
-        return AnalyzedQuery(expr, schema=SCHEMA)
+        return AnalyzedQuery(expr, schema=COLUMN_TYPES)
 
     def _analyze(self, sql: str, params: list | None) -> exp.Expression:
         """Qualify every column to a table, then substitute bound parameters.
@@ -622,14 +623,14 @@ def _analyze(self, sql: str, params: list | None) -> exp.Expression:
         expr = sqlglot.parse_one(sql, dialect="postgres")
         if not isinstance(expr, exp.Select):
             return expr
-        qualified = AnalyzedQuery(expr, schema=SCHEMA).qualified()
+        qualified = AnalyzedQuery(expr, schema=COLUMN_TYPES).qualified()
         _substitute_params(qualified, params)
         return qualified
 
     def _execute(self, expr: exp.Select):
         tables = self._tables(expr)  # a PgError from git or an unknown table stands as-is
         try:
-            return sqlglot_execute(expr, schema=SCHEMA, tables=tables, dialect="postgres")
+            return sqlglot_execute(expr, schema=COLUMN_TYPES, tables=tables, dialect="postgres")
         except Exception as error:
             # The executor supports a good deal less SQL than Postgres. Say so, rather
             # than returning an empty result the client reads as "no matching rows".
@@ -691,7 +692,7 @@ def _oid_of_compared_column(node: exp.Parameter, aliases: dict[str, str]) -> int
         return None
     for side in (parent.this, parent.expression):
         if isinstance(side, exp.Column):
-            declared = SCHEMA.get(aliases.get(side.table, ""), {}).get(side.name)
+            declared = COLUMN_TYPES.get(aliases.get(side.table, ""), {}).get(side.name)
             return oid_for_declared_type(declared) if declared else None
     return None
 

pg_mimic/__init__.py

@@ -29,6 +29,10 @@
     TrustAuthPlugin,
 )
 
+# Table and Schema are what a Session.schema() declares. Public because declaring a
+# primary key or a foreign key is something a user does, not something pg_mimic infers.
+from .declared import Schema, Table
+
 # The two ways a column's type is named on the way to an OID, exported together
 # because they are one pair: `oid_for_declared_type("text[]")` takes the SQL
 # spelling a Session.schema() declares, `oid_for_type(list[str])` the Python type
@@ -88,6 +92,8 @@
     "BaseSession",
     "Session",
     "TableSession",
+    "Schema",
+    "Table",
     "Statement",
     "StaticStatement",
     "Portal",