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