From 0539a86fc5e433d431fdbf3c55f24ea7eec4b017 Mon Sep 17 00:00:00 2001 From: Joe Bylund Date: Mon, 17 Aug 2026 10:52:05 -0400 Subject: [PATCH] Declare tables as objects, so constraints have somewhere to live Session.schema() returned {table_name: {column_name: type_str}}, which describes columns and nothing else -- so there was nowhere to say that commits.sha is a primary key or that commit_files.sha references it, and therefore nowhere for pg_constraint and pg_index to get rows from. This adds the objects that have somewhere to put them, and deliberately nothing else: no constraint can be declared yet, and no behaviour changes but the one noted below. pg_mimic/declared.py is a leaf -- catalog, copy and tables all import it and it imports nothing back, because #92 records the one real import cycle this package has and this is not joining it. resolve() is the only thing that reads schema()'s return value, since it is the single place that knows every shape a session may hand back: a Schema, a sequence of Table, the historical nested dict, or None. Two invariants are now written down and tested rather than implicit. Table order decides table OIDs (_FIRST_USER_OID + index) and column order is ordinal_position and attnum, so sorting either mapping renumbers something a client may already have introspected. And a Table is frozen with its columns copied into a read-only mapping, because session_factory runs once per client: a module-level Table is shared by every connection, and frozen=True alone would still allow table.columns["sha"] = "integer". Schema holds dict[str, Table] built from Table.name rather than accepting a mapping, so a name and its key cannot disagree and a duplicate raises. resolve() refuses Mapping[str, Table] for the same reason -- it is the one shape where a table could end up with two names -- and everything reaches the tables through the single read-only .tables accessor, so Schema grows no partial collection protocol that could be mistaken for the old dict. The one deliberate behaviour change: TableSession.schema() returns a Schema. Not a simplification -- the real one is retiring tables.py's private _Table, which needs #134 first -- but pg_mimic's own session should not hand back the legacy shape once Schema is the canonical one. Schema.column_types() is the same declaration in the old shape. Session.schema() also defaults to an empty Schema rather than None, which deletes the branch whose mis-binding once crashed every catalog query for a session that did not override it; resolve() still tolerates None because an override may return one, and one in the suite does. Also here because the refactor already touched every call site: catalog.py used "schema" for three different things, and _build_information_schema took one sense and returned another, so `schema, tables = _build_...()` read as though the first contained the second when it was really sqlglot's schema= and tables=. That pair is now a named _CatalogInput(column_types, rows). examples/git_sql.py is converted as the first real consumer -- an API with no caller is untested design -- which surfaced the entanglement #135 will fix: its one dict served as both the declaration and sqlglot's schema= argument. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG.md | 19 +++++ docs/session-api.md | 33 +++++-- docs/whats-handled.md | 2 +- examples/git_sql.py | 77 ++++++++--------- pg_mimic/__init__.py | 6 ++ pg_mimic/catalog.py | 70 +++++++++------ pg_mimic/copy.py | 22 ++--- pg_mimic/declared.py | 166 ++++++++++++++++++++++++++++++++++++ pg_mimic/session.py | 18 +++- pg_mimic/tables.py | 19 +++-- tests/test_catalog.py | 5 +- tests/test_declared.py | 135 +++++++++++++++++++++++++++++ tests/test_pg_catalog.py | 7 +- tests/test_table_session.py | 10 ++- 14 files changed, 494 insertions(+), 95 deletions(-) create mode 100644 pg_mimic/declared.py create mode 100644 tests/test_declared.py diff --git a/CHANGELOG.md b/CHANGELOG.md index ce5a3de..6bd11ca 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/docs/session-api.md b/docs/session-api.md index 7798c98..f17390c 100644 --- a/docs/session-api.md +++ b/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 diff --git a/docs/whats-handled.md b/docs/whats-handled.md index 2503bd0..7ba6150 100644 --- a/docs/whats-handled.md +++ b/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 ` 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. diff --git a/examples/git_sql.py b/examples/git_sql.py index a66e1dc..9bb7ae9 100644 --- a/examples/git_sql.py +++ b/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 diff --git a/pg_mimic/__init__.py b/pg_mimic/__init__.py index d63aa2c..0e24d62 100644 --- a/pg_mimic/__init__.py +++ b/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", diff --git a/pg_mimic/catalog.py b/pg_mimic/catalog.py index c2e5759..36b4a23 100644 --- a/pg_mimic/catalog.py +++ b/pg_mimic/catalog.py @@ -22,7 +22,7 @@ import logging import re -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, NamedTuple from sqlglot import exp from sqlglot.errors import OptimizeError @@ -33,6 +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 .describe import oid_for_declared_type, result_columns from .errors import FEATURE_NOT_SUPPORTED, UNDEFINED_COLUMN, PgError from .results import ResultColumn @@ -138,8 +139,22 @@ def _information_schema_row(view: str, **derived) -> dict: return row -def _build_information_schema(user_schema: dict) -> tuple[dict, dict]: - """user_schema: {table_name: {col_name: type_str}} (Session.schema()'s shape). +class _CatalogInput(NamedTuple): + """sqlglot's two arguments for a catalog query, which are not a `Schema`. + + `column_types` is the executor's `schema=` -- the column types of the *catalog's + own* tables, `{"pg_class": {"oid": "INT", ...}}` -- and `rows` is its `tables=`. + Named because "schema" already meant three things in this module: what the session + declares, this type map, and a Postgres namespace. `schema, tables = _build_...()` + read as though the first contained the second. + """ + + column_types: dict + rows: dict + + +def _build_information_schema(declared: Schema) -> _CatalogInput: + """The information_schema views, from what the session declares. The columns filled in below are the ones a declared schema genuinely settles, plus the handful Postgres answers with a constant for any ordinary table -- @@ -151,7 +166,7 @@ def _build_information_schema(user_schema: dict) -> tuple[dict, dict]: type_rows = pg_type_by_oid() tables_rows = [] columns_rows = [] - for table_name, cols in user_schema.items(): + for table_name, table in declared.tables.items(): tables_rows.append( _information_schema_row( "tables", @@ -167,7 +182,7 @@ def _build_information_schema(user_schema: dict) -> tuple[dict, dict]: is_typed="NO", ) ) - for i, (col_name, col_type) in enumerate(cols.items(), start=1): + for i, (col_name, col_type) in enumerate(table.columns.items(), start=1): oid = oid_for_declared_type(col_type) columns_rows.append( _information_schema_row( @@ -203,8 +218,8 @@ def _build_information_schema(user_schema: dict) -> tuple[dict, dict]: ) ) - tables = {"information_schema": {"tables": tables_rows, "columns": columns_rows}} - return INFORMATION_SCHEMA_SCHEMA, tables + rows = {"information_schema": {"tables": tables_rows, "columns": columns_rows}} + return _CatalogInput(INFORMATION_SCHEMA_SCHEMA, rows) def _pg_type_rows() -> list[dict]: @@ -286,11 +301,11 @@ def _typname_for(rows: list[dict], oid: int) -> str: return "unknown" -def _build_pg_catalog(user_schema: dict, database: str = "postgres") -> tuple[dict, dict]: - """The slice of pg_catalog psql's \\d family reads, from Session.schema().""" +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 = [] - for index, (table_name, cols) in enumerate(user_schema.items()): + for index, (table_name, table) in enumerate(declared.tables.items()): table_oid = _FIRST_USER_OID + index class_rows.append( { @@ -316,7 +331,7 @@ def _build_pg_catalog(user_schema: dict, database: str = "postgres") -> tuple[di "reltoastrelid": 0, } ) - for position, (col_name, col_type) in enumerate(cols.items(), start=1): + for position, (col_name, col_type) in enumerate(table.columns.items(), start=1): attribute_rows.append( { "attrelid": table_oid, @@ -396,7 +411,7 @@ def _build_pg_catalog(user_schema: dict, database: str = "postgres") -> tuple[di ], } } - return PG_CATALOG_SCHEMA, tables + return _CatalogInput(PG_CATALOG_SCHEMA, tables) # --- the statements the middleware hands back ---------------------------------------- @@ -461,7 +476,7 @@ def _declared_columns(expr: exp.Expression, schema: dict) -> list[ResultColumn] return None -def _as_statement(expr: exp.Expression, schema: dict, tables: dict, *, strict: bool) -> Statement | None: +def _as_statement(expr: exp.Expression, catalog: _CatalogInput, *, strict: bool) -> Statement | None: """Run a catalog query, or decide what its failure means. See #39. `strict` says whether an executor failure is this connection's problem or the @@ -469,7 +484,7 @@ def _as_statement(expr: exp.Expression, schema: dict, tables: dict, *, strict: b themselves, and False for pg_catalog, which is overwhelmingly psql's own SQL. """ try: - result = sqlglot_execute(expr, schema=schema, tables=tables, dialect="postgres") + result = sqlglot_execute(expr, schema=catalog.column_types, tables=catalog.rows, dialect="postgres") except OptimizeError as error: # A column the catalog does not model. On information_schema that is the # client's own query and Postgres answers 42703; no rows is the same lie the @@ -512,30 +527,31 @@ def _as_statement(expr: exp.Expression, schema: dict, tables: dict, *, strict: b sql = expr.sql(dialect="postgres") rows = [tuple(row) for row in result.rows] - declared = _declared_columns(expr, schema) + declared = _declared_columns(expr, catalog.column_types) if declared is not None and len(declared) == len(result.columns): return StaticStatement(sql, declared, rows) return statement_from_rows(sql, result.columns, rows) -async def _user_schema(connection: Connection) -> dict: - """What the session declares, or nothing -- never None. +async def _declared_schema(connection: Connection) -> Schema: + """What the session declares, as a Schema -- never None. - `Session.schema()` returns None by default, and the `or {}` guarding that used - to bind to the else branch instead of the whole conditional, so a session that - did not override it crashed every catalog query with `'NoneType' object has no - attribute 'items'`. Nothing caught it because the sessions in this suite either - declare a schema or are TableSession, which always does. + The base `Session.schema()` used to return None, and the `or {}` guarding that + bound to the else branch instead of the whole conditional, so a session that did + not override it crashed every catalog query with `'NoneType' object has no + attribute 'items'`. That branch is gone: the default is an empty `Schema` and + `resolve()` absorbs every other shape, None included, since a session's own + override may still return one. """ schema_fn = getattr(connection.session, "schema", None) - return ((await schema_fn()) if schema_fn is not None else {}) or {} + return resolve(await schema_fn()) if schema_fn is not None else Schema() async def information_schema_statement(connection: Connection, expr: exp.Select) -> Statement | None: - schema, tables = _build_information_schema(await _user_schema(connection)) - return _as_statement(expr, schema, tables, strict=True) + catalog = _build_information_schema(await _declared_schema(connection)) + return _as_statement(expr, catalog, strict=True) async def pg_catalog_statement(connection: Connection, expr: exp.Select) -> Statement | None: - schema, tables = _build_pg_catalog(await _user_schema(connection), connection.state.database) - return _as_statement(rewrite_for_executor(connection, expr), schema, tables, strict=False) + catalog = _build_pg_catalog(await _declared_schema(connection), connection.state.database) + return _as_statement(rewrite_for_executor(connection, expr), catalog, strict=False) diff --git a/pg_mimic/copy.py b/pg_mimic/copy.py index 51649f0..32852dc 100644 --- a/pg_mimic/copy.py +++ b/pg_mimic/copy.py @@ -31,6 +31,7 @@ from dataclasses import dataclass from typing import Any, AsyncIterator, Iterator, Sequence +from .declared import Schema, resolve from .errors import ( FEATURE_NOT_SUPPORTED, INTERNAL_ERROR, @@ -788,18 +789,19 @@ async def header_names(self) -> list[str]: if self.parsed.columns: return list(self.parsed.columns) schema_fn = getattr(self._session, "schema", None) - schema = (await schema_fn()) if schema_fn is not None and self.parsed.table is not None else None - schema = schema or {} - # Exact first: a quoted `COPY "People"` names the schema() key as spelled. + declared = resolve(await schema_fn()) if schema_fn is not None and self.parsed.table is not None else Schema() + # Exact first: a quoted `COPY "People"` names the declared table as spelled. # Only an unquoted name -- already folded to lower case by - # _unquote_identifier -- falls back to matching a key case-insensitively, + # _unquote_identifier -- falls back to matching a name case-insensitively, # which is the direction Postgres folds in. - columns = schema.get(self.parsed.table) - if not columns: - folded = {str(name).lower(): value for name, value in schema.items()} - columns = folded.get(self.parsed.table or "") - if columns: - return [str(name) for name in columns] + table = declared.tables.get(self.parsed.table or "") + if table is None: + folded = {name.lower(): value for name, value in declared.tables.items()} + table = folded.get(self.parsed.table or "") + # A table declared with no columns at all is legal (see pg_mimic.declared) and + # still cannot name a header, so it falls through to the refusal below. + if table is not None and table.columns: + return list(table.columns) raise PgError( FEATURE_NOT_SUPPORTED, "COPY TO STDOUT with HEADER needs the column names: list them in the statement " diff --git a/pg_mimic/declared.py b/pg_mimic/declared.py new file mode 100644 index 0000000..0fdcce4 --- /dev/null +++ b/pg_mimic/declared.py @@ -0,0 +1,166 @@ +"""What a session declares about its tables: `Table`, `Schema`, and `resolve()`. + +`Session.schema()` used to return `{table_name: {column_name: type_str}}`, which +describes columns and nothing else -- so there was nowhere to say that `commits.sha` +is a primary key or that `commit_files.sha` references it, and therefore nowhere for +`pg_constraint` and `pg_index` to get rows from. These two objects are that +somewhere. See https://github.com/jbylund/pg_mimic/issues/126. + +A leaf on purpose. `catalog`, `copy` and `tables` all import this, so it imports +nothing from the package in return -- #92 records the one real import cycle +pg_mimic already has and this is not joining it. + +`resolve()` is the only thing that should read `schema()`'s return value, because it +is the one place that knows every shape a session is allowed to hand back: a +`Schema`, a sequence of `Table`, the historical nested dict, or None. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType + +__all__ = ["Schema", "Table", "resolve"] + + +@dataclass(frozen=True) +class Table: + """One table as the catalog should describe it. + + `columns` maps a column name to the SQL spelling of its type -- `{"sha": "text"}` + -- which is exactly what `Session.schema()` has always declared, and what + `pg_mimic.describe.oid_for_declared_type` reads. + + **Column order is load-bearing.** It is `information_schema.ordinal_position` and + `pg_attribute.attnum`, both taken straight from the order this mapping iterates + in, so re-ordering it renumbers the columns of a table the client may already have + introspected. + + Frozen because `session_factory` runs once per client: a module-level `Table` is + shared by every connection, and mutable metadata there lets one connection alter + another's catalog. `frozen=True` alone only blocks rebinding an attribute, not + `table.columns["sha"] = "integer"`, so the mapping is copied into a read-only view + on the way in. + + Not promised to be hashable -- the generated `__hash__` hashes the field tuple and + a mapping is not hashable, so hashing raises. Nothing needs it. + + A table with no columns at all is legal, because Postgres allows + `CREATE TABLE footable ()`, permits queries over it, and keeps its row count. + """ + + name: str + columns: Mapping[str, str] + + def __post_init__(self) -> None: + if not isinstance(self.name, str) or not self.name: + raise ValueError(f"a table name must be a non-empty string, not {self.name!r}") + # Names are the identifier *as written*, the way a CREATE TABLE reads one: + # pg_mimic folds the *query* (see tables.py's normalize_identifiers call) and + # matches it against declared names untouched, which is why {"Users": ...} + # answers `FROM "Users"` and not `FROM Users`. Everything lives in `public`, + # so a name carries no schema qualifier. + columns = dict(self.columns) + for column, type_name in columns.items(): + if not isinstance(column, str) or not column: + raise ValueError(f"a column name must be a non-empty string, not {column!r} in table {self.name!r}") + if not isinstance(type_name, str) or not type_name: + raise ValueError(f"{self.name!r}.{column} declares {type_name!r}, which is not a type name") + # TODO(#134): the spelling itself is not checked yet, and an unrecognised + # one silently resolves to text -- `varchar(255)` and `numeric(10,2)` + # 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)) + + +class Schema: + """The tables a session declares. + + Indexed by name for the lookups psql's `\\d` family needs -- its SQL is full of + `WHERE c.oid = '16384'` -- built from `Table.name` rather than accepting a mapping, + so a name and its key cannot disagree and a duplicate raises instead of quietly + winning. + + **Table order is load-bearing**, for the same reason column order is: a table's OID + is its position (`_FIRST_USER_OID + index` in `pg_mimic.catalog`). Insertion order + is what `tables` iterates in, and sorting it renumbers every table. + + Mutating a `Schema` after it has been served is undefined. Assembly happens once, + before serving; `tables` is exposed read-only so that the constraint-declaring + methods in https://github.com/jbylund/pg_mimic/issues/130 can be the only way in, + which also keeps an explicit immutable snapshot additive if one is ever wanted. + """ + + __slots__ = ("_tables",) + + def __init__(self, tables: Sequence[Table] = ()): + by_name: dict[str, Table] = {} + for table in tables: + if not isinstance(table, Table): + raise TypeError(f"a Schema holds Table objects, not {type(table).__name__}") + if table.name in by_name: + raise ValueError(f"two tables are both named {table.name!r}") + by_name[table.name] = table + self._tables = by_name + + @property + def tables(self) -> Mapping[str, Table]: + return MappingProxyType(self._tables) + + def column_types(self) -> dict[str, dict[str, str]]: + """The declared spellings in the historical nested shape. + + For the callers that still want `{table: {column: type}}` -- sqlglot's own + `schema=` argument takes that shape, and `examples/git_sql.py` passes it + straight through. Deliberately *not* sqlglot's typed view of a schema, which + quotes its names and uses sqlglot's spellings and is + https://github.com/jbylund/pg_mimic/issues/135. + """ + return {name: dict(table.columns) for name, table in self._tables.items()} + + def __repr__(self) -> str: + return f"Schema({list(self._tables)!r})" + + def __eq__(self, other: object) -> bool: + if not isinstance(other, Schema): + return NotImplemented + # Order-sensitive, because it decides OIDs: two schemas that declare the same + # tables in a different order are not interchangeable. + return list(self._tables.items()) == list(other._tables.items()) + + +Declared = Schema | Sequence[Table] | Mapping[str, Mapping[str, str]] | None + + +def resolve(declared: Declared) -> Schema: + """Whatever `Session.schema()` returned, as a `Schema`. + + None becomes an empty one. The base `Session.schema()` no longer returns None -- + it returns an empty `Schema`, because None and `{}` produced identical catalogs and + the branch guarding the difference is what once crashed every catalog query -- but + a session's own override may still return it, and one in the suite does. + + A `Mapping[str, Table]` is refused rather than accepted. It is the one shape where + the key and `Table.name` could disagree, and picking a winner is worse than saying + so. + """ + if declared is None: + return Schema() + if isinstance(declared, Schema): + return declared + if isinstance(declared, Mapping): + tables = [] + for name, columns in declared.items(): + if isinstance(columns, Table): + raise TypeError( + "schema() may return a Schema or a sequence of Table, but not a mapping of " + f"name to Table -- {name!r} would have two names. Pass Schema([...]) instead." + ) + tables.append(Table(str(name), columns)) + return Schema(tables) + if isinstance(declared, (str, bytes)): + raise TypeError(f"schema() returned {declared!r}, which is not a schema") + if isinstance(declared, Sequence): + return Schema(declared) + raise TypeError(f"schema() returned {type(declared).__name__}, which is not a Schema, a sequence of Table, or a mapping") diff --git a/pg_mimic/session.py b/pg_mimic/session.py index 58ad675..9bfad40 100644 --- a/pg_mimic/session.py +++ b/pg_mimic/session.py @@ -25,6 +25,7 @@ from abc import ABC, abstractmethod from typing import Any, AsyncIterable, AsyncIterator, Awaitable, Callable, Iterable, Sequence +from .declared import Declared, Schema from .results import ResultColumn from .state import SessionState, SettingValue from .types import TEXT @@ -363,10 +364,21 @@ async def discard(self, kind: str) -> None: Does nothing by default. """ - async def schema(self) -> dict | None: + async def schema(self) -> Declared: """Optional: describe your tables for information_schema/pg_catalog - emulation. See pg_mimic.catalog.""" - return None + emulation. See pg_mimic.catalog and pg_mimic.declared. + + Return a `Schema` of `Table`s, a sequence of `Table`, or the nested + `{table: {column: type_name}}` dict this has always accepted. + + An empty `Schema` by default, rather than None: the two produced identical + catalogs, and the branch that guarded the difference is what once crashed + every catalog query for a session that did not override this. #136 covers + whether "I declare no tables" is worth distinguishing from "I make no claim + about my tables" -- if it ever is, it wants a named sentinel rather than a + resurrected None. + """ + return Schema() async def copy_in(self, sql: str, rows: AsyncIterator[Row]) -> int | None: """Optional: handle `COPY ... FROM STDIN`. diff --git a/pg_mimic/tables.py b/pg_mimic/tables.py index 16b4004..1b36230 100644 --- a/pg_mimic/tables.py +++ b/pg_mimic/tables.py @@ -61,13 +61,15 @@ from sqlglot.planner import Plan from sqlglot.schema import ensure_schema +from .analysis import AnalyzedQuery, annotate +from .arrays import element_oid_of, is_array_oid + # The declared-schema half of describe(), shared with any other session that # declares one (see pg_mimic.describe, and #88 for the copy that had diverged). # _PG_NAME and _pg_type_name are re-exported by being imported here: they read as # TableSession's own vocabulary from the outside, and the catalog's round-trip test # takes them off this module. -from .analysis import AnalyzedQuery, annotate -from .arrays import element_oid_of, is_array_oid +from .declared import Schema, Table from .describe import ( _PG_NAME, _SQLGLOT_NAME, @@ -240,7 +242,7 @@ def __init__( raise ValueError(f"columns declares table(s) {unknown} that are not in tables ({sorted(tables)})") declared = {name: _declare_table(name, rows, declared_columns.get(name)) for name, rows in tables.items()} - self._schema = {name: table.pg_types for name, table in declared.items()} + self._schema = Schema([Table(name, table.pg_types) for name, table in declared.items()]) # Quoted, like the column names inside them: sqlglot folds a bare name it is # handed, so `{"Users": ...}` would otherwise declare a table called `users`. self._rows = {_quoted(name): table.rows for name, table in declared.items()} @@ -283,7 +285,14 @@ async def query(self, sql: str, params: list[Any]) -> list[Row]: rows = [row[: plan.visible_columns] for row in rows] # drop the added DISTINCT ON keys return rows - async def schema(self) -> dict: + async def schema(self) -> Schema: + """The declared tables, as a Schema -- see pg_mimic.declared. + + A `Schema` rather than the nested dict this returned until #126. Not a + simplification in itself, but pg_mimic's own session should not hand back the + legacy shape once `Schema` is the canonical one. A session of your own may + still return the dict; `resolve()` accepts it. + """ return self._schema async def prepare(self, sql: str, param_oids: list[int | None]) -> Statement: @@ -395,7 +404,7 @@ def _reject_unknown_tables(self, expression: exp.Expression) -> None: # Against `_schema`, the one view of the declaration still keyed by the # names as written -- which, the tree having been folded already, is # what `table.name` now is. - if table.name not in self._schema and table.name not in cte_names: + if table.name not in self._schema.tables and table.name not in cte_names: raise PgError(UNDEFINED_TABLE, f'relation "{table.name}" does not exist') diff --git a/tests/test_catalog.py b/tests/test_catalog.py index dc55a35..0083e28 100644 --- a/tests/test_catalog.py +++ b/tests/test_catalog.py @@ -567,10 +567,11 @@ def test_every_declared_catalog_column_has_a_value_in_every_row(): So declaring a column and populating it are one step, and this is the guard. """ from pg_mimic.catalog import _build_information_schema, _build_pg_catalog + from pg_mimic.declared import resolve for schema, tables in ( - _build_pg_catalog({"t": {"id": "integer", "name": "text"}}, "testdb"), - _build_information_schema({"t": {"id": "integer"}}), + _build_pg_catalog(resolve({"t": {"id": "integer", "name": "text"}}), "testdb"), + _build_information_schema(resolve({"t": {"id": "integer"}})), ): namespace = next(iter(schema)) for table, columns in schema[namespace].items(): diff --git a/tests/test_declared.py b/tests/test_declared.py new file mode 100644 index 0000000..e1df39c --- /dev/null +++ b/tests/test_declared.py @@ -0,0 +1,135 @@ +"""`Table`, `Schema` and `resolve()` -- what a session declares about its tables. + +The interesting assertions here are not that the containers hold what they were +given. They are the properties the rest of the package quietly depends on and that +nothing else would notice breaking: + +- **order**, which decides OIDs and `ordinal_position`, so sorting either mapping + renumbers something a client may already have introspected +- **immutability**, because a module-level `Table` is shared by every connection +- **which shapes `resolve()` refuses**, since the one it must refuse is the one where + a table could end up with two different names + +See https://github.com/jbylund/pg_mimic/issues/126. +""" + +from __future__ import annotations + +import pytest + +from pg_mimic import Schema, Table +from pg_mimic.declared import resolve + + +def test_columns_are_read_only_even_though_the_dataclass_is_frozen(): + """`frozen=True` blocks rebinding the attribute, not mutating the mapping it holds. + + Load-bearing rather than tidy: `session_factory` runs once per client, so a + module-level Table is shared by every connection, and a writable mapping there + lets one connection alter another's catalog. + """ + table = Table("commits", {"sha": "text"}) + + with pytest.raises(TypeError): + table.columns["sha"] = "integer" + with pytest.raises(Exception): + table.name = "other" + + assert dict(table.columns) == {"sha": "text"} + + +def test_a_table_does_not_alias_the_mapping_it_was_given(): + declared = {"sha": "text"} + table = Table("commits", declared) + declared["injected"] = "text" + assert dict(table.columns) == {"sha": "text"} + + +def test_column_order_is_kept_because_it_is_ordinal_position(): + columns = {"z": "text", "a": "integer", "m": "date"} + assert list(Table("t", columns).columns) == ["z", "a", "m"] + + +def test_table_order_is_kept_because_it_decides_oids(): + schema = Schema([Table("z", {}), Table("a", {}), Table("m", {})]) + assert list(schema.tables) == ["z", "a", "m"] + + +def test_a_table_with_no_columns_is_legal(): + """Postgres allows `CREATE TABLE footable ()`, queries it, and keeps its row count.""" + assert dict(Table("footable", {}).columns) == {} + assert list(Schema([Table("footable", {})]).tables) == ["footable"] + + +def test_a_schema_refuses_two_tables_with_one_name(): + with pytest.raises(ValueError, match="both named 'commits'"): + Schema([Table("commits", {"a": "text"}), Table("commits", {"b": "text"})]) + + +def test_a_schemas_mapping_cannot_be_mutated_through_the_accessor(): + schema = Schema([Table("commits", {"sha": "text"})]) + with pytest.raises(TypeError): + schema.tables["files"] = Table("files", {}) + assert list(schema.tables) == ["commits"] + + +def test_schema_equality_is_order_sensitive(): + """Two schemas declaring the same tables in a different order are not + interchangeable, because the order is what assigns their OIDs.""" + one, two = Table("one", {}), Table("two", {}) + assert Schema([one, two]) == Schema([one, two]) + assert Schema([one, two]) != Schema([two, one]) + + +def test_column_types_is_the_declaration_in_the_historical_shape(): + schema = Schema([Table("commits", {"sha": "text"}), Table("files", {"path": "text", "lines": "integer"})]) + assert schema.column_types() == {"commits": {"sha": "text"}, "files": {"path": "text", "lines": "integer"}} + + +def test_resolve_accepts_the_nested_dict_a_session_has_always_returned(): + schema = resolve({"commits": {"sha": "text"}, "files": {"path": "text"}}) + assert list(schema.tables) == ["commits", "files"] + assert dict(schema.tables["commits"].columns) == {"sha": "text"} + + +def test_resolve_accepts_a_sequence_of_tables(): + assert list(resolve([Table("commits", {"sha": "text"})]).tables) == ["commits"] + + +def test_resolve_passes_a_schema_through_untouched(): + schema = Schema([Table("commits", {"sha": "text"})]) + assert resolve(schema) is schema + + +def test_resolve_turns_none_into_an_empty_schema(): + """The base `Session.schema()` no longer returns None, but an override may, and + `test_catalog.py::test_a_session_that_declares_no_schema_can_still_be_introspected` + is a session that does.""" + assert resolve(None).tables == {} + + +def test_resolve_refuses_a_mapping_of_name_to_table(): + """The one shape where a table would have two names -- the key and `Table.name` -- + so it is refused rather than having a winner picked for it.""" + with pytest.raises(TypeError, match="two names"): + resolve({"commits": Table("comits", {"sha": "text"})}) + + +@pytest.mark.parametrize( + argnames=["declared"], + argvalues=[["commits"], [42], [object()]], + ids=["a string", "a number", "some other object"], +) +def test_resolve_refuses_what_is_not_a_schema_at_all(declared): + with pytest.raises(TypeError): + resolve(declared) + + +@pytest.mark.parametrize( + argnames=["name", "columns"], + argvalues=[["", {}], [None, {}], ["t", {"": "text"}], ["t", {"a": ""}], ["t", {"a": None}]], + ids=["empty table name", "table name is not a string", "empty column name", "empty type name", "type is not a string"], +) +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) diff --git a/tests/test_pg_catalog.py b/tests/test_pg_catalog.py index d66fa44..db39333 100644 --- a/tests/test_pg_catalog.py +++ b/tests/test_pg_catalog.py @@ -16,6 +16,7 @@ from pg_mimic.catalog import _build_pg_catalog from pg_mimic.catalog_rewrite import rewrite_for_executor +from pg_mimic.declared import resolve SCHEMA = {"users": {"id": "integer", "name": "text"}, "orders": {"id": "bigint", "total": "numeric"}} @@ -25,7 +26,7 @@ class FakeConnection: def _run(sql, schema=None): - sqlglot_schema, tables = _build_pg_catalog(SCHEMA if schema is None else schema) + sqlglot_schema, tables = _build_pg_catalog(resolve(SCHEMA if schema is None else schema)) from sqlglot.executor import execute expr = rewrite_for_executor(FakeConnection(), sqlglot.parse_one(sql, dialect="postgres")) @@ -134,7 +135,7 @@ def test_declared_types_map_to_real_oids(): rather than silently defaulting to text like an unrecognised spelling.""" from pg_mimic.types import INT8, NUMERIC - _schema, tables = _build_pg_catalog(SCHEMA) + _schema, tables = _build_pg_catalog(resolve(SCHEMA)) orders_oid = next(r["oid"] for r in tables["pg_catalog"]["pg_class"] if r["relname"] == "orders") attrs = {r["attname"]: r["atttypid"] for r in tables["pg_catalog"]["pg_attribute"] if r["attrelid"] == orders_oid} assert attrs == {"id": INT8, "total": NUMERIC} @@ -145,7 +146,7 @@ def test_pg_type_covers_the_types_pg_mimic_can_encode(): describe anything pg_mimic can put on the wire.""" from pg_mimic import ARRAY_OID, INT8, TEXT - _schema, tables = _build_pg_catalog(SCHEMA) + _schema, tables = _build_pg_catalog(resolve(SCHEMA)) by_oid = {r["oid"]: r for r in tables["pg_catalog"]["pg_type"]} assert by_oid[TEXT]["typname"] == "text" assert by_oid[ARRAY_OID[INT8]]["typelem"] == INT8 # what marks it an array diff --git a/tests/test_table_session.py b/tests/test_table_session.py index ba89451..77c1b4b 100644 --- a/tests/test_table_session.py +++ b/tests/test_table_session.py @@ -641,13 +641,21 @@ def test_the_catalog_reports_the_declared_names(mixed_conn): async def test_schema_is_derived_from_the_tables(): - assert await _session().schema() == { + """A Schema since #126, where this asserted the nested dict it returned before. + + `column_types()` is the same declaration in the old shape, so this stays an + assertion about the derived *types* rather than about the container. + """ + schema = await _session().schema() + assert schema.column_types() == { "users": {"id": "bigint", "name": "text", "joined": "date", "tags": "text[]"}, "orders": {"id": "bigint", "user_id": "bigint", "total": "numeric"}, "events": {"id": "bigint", "at": "timestamp"}, "notes": {"id": "bigint", "body": "text"}, "docs": {"id": "bigint", "body": "jsonb"}, } + # Order is load-bearing -- it is what decides each table's OID. + assert list(schema.tables) == ["users", "orders", "events", "notes", "docs"] # --- what the executor answers wrongly, repaired --------------------------------------