diff --git a/Cargo.lock b/Cargo.lock index 192f953..a424869 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1691,7 +1691,7 @@ dependencies = [ [[package]] name = "datafusion_pg_catalog" version = "0.1.0" -source = "git+https://github.com/ybrs/pg_catalog?branch=main#319178ba798740e8e3b988ffb59f83757a5f98a4" +source = "git+https://github.com/ybrs/pg_catalog?rev=4ee6e83b86a3444dac1aa61f07e930d4222b6316#4ee6e83b86a3444dac1aa61f07e930d4222b6316" dependencies = [ "anyhow", "arrow", diff --git a/Cargo.toml b/Cargo.toml index 81f9e5e..f2ea11e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -37,7 +37,8 @@ sqlparser = "0.61.0" chrono = "0.4.41" log = "0.4" datafusion = "53.1.0" -datafusion_pg_catalog = { git = "https://github.com/ybrs/pg_catalog", branch = "main", package = "datafusion_pg_catalog" } +datafusion_pg_catalog = { git = "https://github.com/ybrs/pg_catalog", rev = "4ee6e83b86a3444dac1aa61f07e930d4222b6316", package = "datafusion_pg_catalog" } +# Local dev override (uncomment to build against the in-tree ../pg_catalog): # datafusion_pg_catalog = { path = "../pg_catalog" } env_logger = "0.11.8" diff --git a/README.md b/README.md index 56b744e..f76949f 100644 --- a/README.md +++ b/README.md @@ -245,6 +245,34 @@ For example server.start(catalog_emulation=True) ``` +#### Lazy (callback-driven) catalog + +The `register_*` calls above snapshot the catalog at startup. For a **live** +source, install a lazy catalog instead: supply one source object and Riffq pulls +catalog metadata from it on every `pg_catalog` scan, so tables created after +startup show up automatically — nothing is cached. + +```python +class MyCatalog: + def databases(self, callback): + callback([{"oid": 16384, "name": "appdb"}]) + def schemas(self, database, callback): + callback([{"oid": 16385, "name": "public"}]) + def relations(self, database, schema, callback): + callback([{"oid": 20001, "reltype_oid": 30001, "name": "users", + "kind": "table", "has_index": False}]) + def columns(self, database, schema, relation, callback): + callback([{"name": "id", "type_oid": 23, "nullable": False}]) # 23 = int4 + +server.set_lazy_catalog(MyCatalog()) # replaces the eager register_* calls +server.start(catalog_emulation=True) +``` + +You own the OIDs (stable + unique); `type_oid` is a `pg_type` OID. See +[docs/catalog.md](docs/catalog.md) for the full contract and +[`example/lazy_catalog.py`](example/lazy_catalog.py) for a runnable example; +[Teleduck](teleduck/) uses this path against a live DuckDB connection. + --- diff --git a/docs/catalog.md b/docs/catalog.md index fcafd32..28ec0a6 100644 --- a/docs/catalog.md +++ b/docs/catalog.md @@ -78,6 +78,86 @@ for schema_name, table_name in list_tables_from_duckdb(): server.start(catalog_emulation=True) ``` +## Lazy (callback-driven) catalog + +The `register_*` calls above take a **snapshot at startup**: if your underlying +data changes (a table is created, dropped, or altered), the emulated catalog +goes stale until you re-register. + +For a live source, install a **lazy catalog** instead. You supply one source +object and Riffq pulls catalog metadata from it on *every* `pg_catalog` / +`information_schema` scan, so the catalog always reflects the current state — +tables created after startup included. Nothing is cached. + +```python +server.set_lazy_catalog(source) # instead of register_database/schema/table +server.start(catalog_emulation=True) +``` + +The `source` object implements four methods. Each receives a `callback` and +invokes it with a list of row dicts (mirroring the Rust `LazyCatalogSource` +trait one method per catalog level): + +```python +class MyCatalog: + def databases(self, callback): + # -> pg_database + callback([{"oid": 16384, "name": "appdb"}]) # "datdba"? optional + + def schemas(self, database, callback): + # -> pg_namespace + callback([{"oid": 16385, "name": "public"}]) # "owner_oid"? optional + + def relations(self, database, schema, callback): + # -> pg_class (+ pg_type rowtype) + callback([{ + "oid": 20001, "reltype_oid": 30001, "name": "users", + "kind": "table", # "table" | "view" | "materialized_view" + # all optional, default off / NULL: + "owner_oid": 10, # -> pg_tables.tableowner (omit if no ownership) + "has_index": True, # -> pg_tables.hasindexes + "has_rules": False, "has_triggers": False, "row_security": False, + }]) + + def columns(self, database, schema, relation, callback): + # -> pg_attribute (+ information_schema.columns) + callback([ + {"name": "id", "type_oid": 23, "nullable": False}, # 23 = int4 + {"name": "name", "type_oid": 25, "nullable": True}, # 25 = text + ]) + +server.set_lazy_catalog(MyCatalog()) +server.start(catalog_emulation=True) +``` + +### Rules of the contract + +- **You own the OIDs.** Riffq writes them through verbatim. They must be stable + across calls (so catalog joins like `pg_class.oid = pg_attribute.attrelid` + resolve) and unique among your objects. Keep them **above the built-in range** + so they don't collide with built-in rows on OID joins — use `16384` + (PostgreSQL's first normal object id) as a safe base, as the example and + Teleduck do. A common trick is `16384 + stable_hash(name)` (see the example + below). +- **`type_oid` is a `pg_type` OID** you choose, e.g. `23` int4, `20` int8, + `25` text, `16` bool, `701` float8, `1700` numeric, `1082` date, `1114` + timestamp. This keeps Riffq independent of your engine's type system. +- **Merged with built-ins.** Your rows are merged with the built-in system rows + (so `int4`, `pg_class`, etc. still resolve). A user object whose name collides + with a built-in **replaces** it; two of your objects with the same identity + (e.g. two tables of the same name in one schema) is an **error** surfaced to the + client. +- **Errors propagate.** An exception raised in any source method is returned to + the SQL client as an error — it never fails silently. +- **Opt-in & exclusive.** When a lazy source is set, the eager + `register_database`/`register_schema`/`register_table` calls are ignored. + Requires `start(catalog_emulation=True)`. + +A complete runnable example backed by an in-memory dict (no external engine) is +in [`example/lazy_catalog.py`](https://github.com/ybrs/riffq/blob/main/example/lazy_catalog.py), +and [Teleduck](https://github.com/ybrs/riffq/tree/main/teleduck) uses this path +against a live DuckDB connection. + ## Examples For a minimal end-to-end example that registers a database, schema, and table and asserts they appear via `pg_catalog`, see: @@ -221,5 +301,15 @@ RiffqServer.register_table(database_name: str, schema_name: str, table_name: str, columns: list[dict]) -> None: ``` +### Install a lazy (callback-driven) catalog source. + +Pull catalog metadata from `source` on every scan instead of snapshotting it. +See [Lazy (callback-driven) catalog](#lazy-callback-driven-catalog) for the +source object's contract. Mutually exclusive with the eager `register_*` calls. + +```python +RiffqServer.set_lazy_catalog(source) -> None: +``` + [pg_catalog_rs]: https://github.com/ybrs/pg_catalog diff --git a/example/lazy_catalog.py b/example/lazy_catalog.py new file mode 100644 index 0000000..96b18b5 --- /dev/null +++ b/example/lazy_catalog.py @@ -0,0 +1,143 @@ +"""Lazy (callback-driven) catalog example. + +Riffq can answer PostgreSQL catalog queries (``pg_catalog`` / ``information_schema``) +from a *source object* that is consulted on every scan, instead of a snapshot +registered up front. This example backs that source with a plain in-memory dict +-- no external engine -- so you can see the whole contract in one file. + +Run it:: + + python example/lazy_catalog.py + +then connect with psql and watch the catalog stay in sync as you create tables:: + + psql -h 127.0.0.1 -p 5444 -U user -d appdb + + appdb=> SELECT relname FROM pg_class WHERE relname='orders'; -- 0 rows + appdb=> CREATE TABLE orders(id INT, total INT); -- handled below + appdb=> SELECT relname FROM pg_class WHERE relname='orders'; -- now 1 row + appdb=> \\d orders +""" + +import logging + +import pyarrow as pa +import riffq + +logging.basicConfig(level=logging.INFO) + +# The "engine": a live, in-memory catalog. In a real app this would be DuckDB, +# PostgreSQL, a config file, a remote service -- anything. The lazy source below +# reads whatever is here *at query time*, so mutating it is reflected instantly. +CATALOG = { + "appdb": { + "public": { + # table_name -> list of (column_name, pg_type_oid, nullable) + "users": [("id", 23, False), ("name", 25, True)], # 23=int4, 25=text + } + } +} + +# pg_type OIDs the example understands, keyed by a coarse type name. +TYPE_OIDS = {"int": 23, "bigint": 20, "text": 25, "bool": 16, "float": 701} + + +def stable_oid(salt: str, *parts: str) -> int: + """Derive a stable, built-in-clear OID from a name. + + The same inputs always return the same OID, so ``pg_class.oid`` and + ``pg_attribute.attrelid`` agree across scans and catalog joins resolve. + Distinct object classes use distinct salts to avoid collisions, and the + result is kept well above the built-in OID range. + """ + h = 5381 + for ch in (salt + "\x00" + "\x00".join(parts)): + h = (h * 33 + ord(ch)) & 0x7FFFFFFF + return 16384 + (h % 2_000_000_000) + + +class DictCatalogSource: + """A lazy catalog source over the in-memory ``CATALOG`` dict. + + Each method receives a ``callback`` and invokes it with a list of row dicts, + mirroring Riffq's Rust ``LazyCatalogSource`` trait one method per level. + """ + + def databases(self, callback): + callback( + [{"oid": stable_oid("db", name), "name": name} for name in CATALOG] + ) + + def schemas(self, database, callback): + callback( + [ + {"oid": stable_oid("ns", database, schema), "name": schema} + for schema in CATALOG.get(database, {}) + ] + ) + + def relations(self, database, schema, callback): + tables = CATALOG.get(database, {}).get(schema, {}) + callback( + [ + { + "oid": stable_oid("rel", database, schema, name), + "reltype_oid": stable_oid("type", database, schema, name), + "name": name, + "kind": "table", + } + for name in tables + ] + ) + + def columns(self, database, schema, relation, callback): + cols = CATALOG.get(database, {}).get(schema, {}).get(relation, []) + callback( + [ + {"name": col, "type_oid": type_oid, "nullable": nullable} + for (col, type_oid, nullable) in cols + ] + ) + + +class Connection(riffq.BaseConnection): + """Data path. Catalog queries are served by the lazy source above; everything + else lands here. We implement a toy ``CREATE TABLE`` so you can watch the + catalog update live, and echo a single row for any other query.""" + + def handle_auth(self, user, password, host, database=None, callback=callable): + # Accept any credentials in this example. + return callback(True) + + def handle_query(self, sql, callback=callable, **kwargs): + text = sql.strip().rstrip(";") + low = text.lower() + + if low.startswith("create table "): + # "create table public.orders(id int, total int)" (schema optional) + head, _, body = text[len("create table "):].partition("(") + qualified = head.strip() + schema, _, name = qualified.rpartition(".") + schema = schema or "public" + cols = [] + for part in body.rstrip(")").split(","): + bits = part.split() + if len(bits) >= 2: + cols.append((bits[0].strip('"'), TYPE_OIDS.get(bits[1].lower(), 25), True)) + CATALOG.setdefault("appdb", {}).setdefault(schema, {})[name] = cols + return callback("CREATE TABLE", is_tag=True) + + # Any other statement: return a single dummy row so clients are happy. + batch = self.arrow_batch([pa.array([1])], ["?column?"]) + return self.send_reader(batch, callback) + + +def main(): + server = riffq.RiffqServer("127.0.0.1:5444", connection_cls=Connection) + # Install the lazy source instead of register_database/schema/table. + server.set_lazy_catalog(DictCatalogSource()) + server.start(catalog_emulation=True) + + +if __name__ == "__main__": + main() diff --git a/pysrc/riffq/connection.py b/pysrc/riffq/connection.py index 1dc0741..b80da6a 100644 --- a/pysrc/riffq/connection.py +++ b/pysrc/riffq/connection.py @@ -163,6 +163,43 @@ def set_tls(self, crt:str, key:str): self._server.set_tls(crt, key) + def set_lazy_catalog(self, source: Any) -> None: + """Install a lazy (callback-driven) catalog source for catalog emulation. + + Instead of eagerly pre-registering every database/schema/table with + `register_database`/`register_schema`/`register_table`, supply one + ``source`` object and the server pulls catalog metadata from it on every + ``pg_catalog`` / ``information_schema`` scan — so the catalog always + reflects the source's live state. + + ``source`` must expose four methods, each of which receives a ``callback`` + and invokes it with a list of row dicts: + + - ``databases(callback)`` -> ``[{"oid": int, "name": str, "datdba"?: int}]`` + - ``schemas(database, callback)`` -> ``[{"oid": int, "name": str, "owner_oid"?: int}]`` + - ``relations(database, schema, callback)`` -> + ``[{"oid": int, "reltype_oid": int, "name": str, + "kind"?: "table"|"view"|"materialized_view", + "owner_oid"?: int, "has_index"?: bool, "has_rules"?: bool, + "has_triggers"?: bool, "row_security"?: bool}]`` — the optional flags + populate ``pg_tables`` (``tableowner`` via ``owner_oid``, ``hasindexes``, + etc.); omit ``owner_oid`` for backends without ownership (it stays blank) + - ``columns(database, schema, relation, callback)`` -> + ``[{"name": str, "type_oid": int, "nullable": bool}]`` + + OIDs are supplied by the source and written through verbatim; they must be + stable across calls (so catalog joins resolve) and unique among the + source's objects. A duplicate object (same name in the same scope) is an + error; a source object whose name collides with a built-in replaces it. + + When a lazy source is set, the eager ``register_*`` registrations are + ignored. Requires ``start(catalog_emulation=True)``. + + Args: + source: The lazy catalog source object described above. + """ + self._server.set_lazy_catalog(source) + def register_database(self, database_name: str) -> None: """Register a logical database for catalog emulation. diff --git a/src/lib.rs b/src/lib.rs index 61b2454..1972826 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -46,11 +46,20 @@ use datafusion::execution::context::SessionContext; use datafusion_pg_catalog::{ dispatch_query, get_base_session_context, + get_base_session_context_with_lazy_catalog, register_user_database, register_schema, register_user_tables, ColumnDef, + ColumnSpec, + DatabaseDef, + LazyCatalogOptions, + LazyCatalogSource, + RelationDef, + RelationKind, + SchemaDef, }; +use datafusion::error::{DataFusionError, Result as DFResult}; use postgres_types::FromSql; use chrono::{DateTime, Duration, NaiveDate}; @@ -1702,6 +1711,278 @@ fn setup_tls(cert_path: &str, key_path: &str) -> Result { Ok(TlsAcceptor::from(Arc::new(config))) } +/// Convert a Python exception raised by a lazy catalog source into a +/// `DataFusionError`, so the failure propagates to the SQL client instead of +/// being swallowed (the lazy catalog contract forbids failing silently). +fn py_to_df(e: PyErr) -> DataFusionError { + DataFusionError::Execution(format!("lazy catalog source error: {e}")) +} + +/// The synchronous callback object handed to a Python lazy-catalog method. The +/// Python source calls it with the list of rows it produced; we capture that +/// list so the surrounding Rust method can marshal it. Mirrors the +/// `&mut dyn FnMut(Vec<...>)` callback of the Rust `LazyCatalogSource` trait. +#[pyclass] +struct CatalogCallback { + rows: Arc>>>, +} + +#[pymethods] +impl CatalogCallback { + /// Record the rows the Python source passed in. Expected to be invoked once, + /// synchronously, before the calling method returns. + fn __call__(&self, rows: Py) { + *self.rows.lock().unwrap() = Some(rows); + } +} + +/// Read a required integer field from a row dict. +fn req_i32(d: &Bound<'_, PyDict>, key: &str) -> DFResult { + match d.get_item(key).map_err(py_to_df)? { + Some(v) => v.extract::().map_err(py_to_df), + None => Err(DataFusionError::Execution(format!( + "lazy catalog row is missing required field '{key}'" + ))), + } +} + +/// Read a required string field from a row dict. +fn req_str(d: &Bound<'_, PyDict>, key: &str) -> DFResult { + match d.get_item(key).map_err(py_to_df)? { + Some(v) => v.extract::().map_err(py_to_df), + None => Err(DataFusionError::Execution(format!( + "lazy catalog row is missing required field '{key}'" + ))), + } +} + +/// Read a required boolean field from a row dict. +fn req_bool(d: &Bound<'_, PyDict>, key: &str) -> DFResult { + match d.get_item(key).map_err(py_to_df)? { + Some(v) => v.extract::().map_err(py_to_df), + None => Err(DataFusionError::Execution(format!( + "lazy catalog row is missing required field '{key}'" + ))), + } +} + +/// Read an optional integer field from a row dict (absent or `None` -> `None`). +fn opt_i32(d: &Bound<'_, PyDict>, key: &str) -> DFResult> { + match d.get_item(key).map_err(py_to_df)? { + Some(v) if !v.is_none() => Ok(Some(v.extract::().map_err(py_to_df)?)), + _ => Ok(None), + } +} + +/// Read an optional string field from a row dict (absent or `None` -> `None`). +fn opt_str(d: &Bound<'_, PyDict>, key: &str) -> DFResult> { + match d.get_item(key).map_err(py_to_df)? { + Some(v) if !v.is_none() => Ok(Some(v.extract::().map_err(py_to_df)?)), + _ => Ok(None), + } +} + +/// Read an optional boolean field from a row dict, defaulting to `default` when +/// absent or `None`. +fn opt_bool_or(d: &Bound<'_, PyDict>, key: &str, default: bool) -> DFResult { + match d.get_item(key).map_err(py_to_df)? { + Some(v) if !v.is_none() => v.extract::().map_err(py_to_df), + _ => Ok(default), + } +} + +/// Downcast the Python value a source returned into a list of row dicts. +fn row_dicts<'py>( + method: &str, + list: &Bound<'py, PyAny>, +) -> DFResult>> { + let list: &Bound<'py, PyList> = list.downcast().map_err(|e| { + DataFusionError::Execution(format!("{method}() must pass a list of dicts: {e}")) + })?; + let mut out = Vec::with_capacity(list.len()); + for item in list.iter() { + let d: Bound<'py, PyDict> = item.downcast_into().map_err(|e| { + DataFusionError::Execution(format!("{method}() rows must be dicts: {e}")) + })?; + out.push(d); + } + Ok(out) +} + +/// Parse `databases()` rows: `{oid, name, [datdba]}` -> [`DatabaseDef`]. +fn parse_databases(list: &Bound<'_, PyAny>) -> DFResult> { + row_dicts("databases", list)? + .iter() + .map(|d| { + let oid = req_i32(d, "oid")?; + let name = req_str(d, "name")?; + let datdba = opt_i32(d, "datdba")?.unwrap_or(10); + Ok(DatabaseDef::new(oid, name, datdba)) + }) + .collect() +} + +/// Parse `schemas()` rows: `{oid, name, [owner_oid]}` -> [`SchemaDef`]. +fn parse_schemas(list: &Bound<'_, PyAny>) -> DFResult> { + row_dicts("schemas", list)? + .iter() + .map(|d| { + Ok(SchemaDef { + oid: req_i32(d, "oid")?, + name: req_str(d, "name")?, + owner_oid: opt_i32(d, "owner_oid")?, + }) + }) + .collect() +} + +/// Parse `relations()` rows: `{oid, reltype_oid, name, [kind]}` -> [`RelationDef`]. +fn parse_relations(list: &Bound<'_, PyAny>) -> DFResult> { + row_dicts("relations", list)? + .iter() + .map(|d| { + let kind = match opt_str(d, "kind")?.as_deref() { + Some("table") | None => RelationKind::Table, + Some("view") => RelationKind::View, + Some("materialized_view") | Some("matview") => RelationKind::MaterializedView, + Some(other) => { + return Err(DataFusionError::Execution(format!( + "unknown relation kind '{other}' (use table/view/materialized_view)" + ))) + } + }; + Ok(RelationDef { + oid: req_i32(d, "oid")?, + reltype_oid: req_i32(d, "reltype_oid")?, + name: req_str(d, "name")?, + kind, + owner_oid: opt_i32(d, "owner_oid")?, + has_index: opt_bool_or(d, "has_index", false)?, + has_rules: opt_bool_or(d, "has_rules", false)?, + has_triggers: opt_bool_or(d, "has_triggers", false)?, + row_security: opt_bool_or(d, "row_security", false)?, + }) + }) + .collect() +} + +/// Parse `columns()` rows: `{name, type_oid, nullable}` -> [`ColumnSpec`]. +fn parse_columns(list: &Bound<'_, PyAny>) -> DFResult> { + row_dicts("columns", list)? + .iter() + .map(|d| { + Ok(ColumnSpec::new( + req_str(d, "name")?, + req_i32(d, "type_oid")?, + req_bool(d, "nullable")?, + )) + }) + .collect() +} + +/// A [`LazyCatalogSource`] backed by a Python object whose `databases`, +/// `schemas`, `relations`, and `columns` methods each accept a callback and +/// invoke it with a list of row dicts. Each trait method acquires the GIL, hands +/// Python a [`CatalogCallback`], and marshals the captured rows into the +/// pg_catalog definition types. Errors raised in Python surface as +/// `DataFusionError` to the SQL client. +struct PyLazyCatalogSource { + obj: Py, +} + +impl PyLazyCatalogSource { + /// Call `method` on the Python source with `str_args` followed by a fresh + /// callback, returning whatever list the callback captured (or `None` if the + /// source never invoked it). + fn pull( + &self, + py: Python<'_>, + method: &str, + str_args: &[&str], + ) -> DFResult>> { + let cell: Arc>>> = Arc::new(Mutex::new(None)); + let wrapper = Py::new( + py, + CatalogCallback { + rows: cell.clone(), + }, + ) + .map_err(py_to_df)?; + + let mut items: Vec> = Vec::with_capacity(str_args.len() + 1); + for s in str_args { + items.push((*s).into_py_any(py).map_err(py_to_df)?); + } + items.push(wrapper.into_py_any(py).map_err(py_to_df)?); + let args = PyTuple::new(py, items).map_err(py_to_df)?; + + self.obj + .bind(py) + .call_method1(method, args) + .map_err(py_to_df)?; + + let captured = cell.lock().unwrap().take(); + Ok(captured) + } +} + +impl LazyCatalogSource for PyLazyCatalogSource { + fn databases(&self, callback: &mut dyn FnMut(Vec)) -> DFResult<()> { + let defs = Python::attach(|py| -> DFResult> { + match self.pull(py, "databases", &[])? { + Some(list) => parse_databases(list.bind(py)), + None => Ok(Vec::new()), + } + })?; + callback(defs); + Ok(()) + } + + fn schemas(&self, database: &str, callback: &mut dyn FnMut(Vec)) -> DFResult<()> { + let defs = Python::attach(|py| -> DFResult> { + match self.pull(py, "schemas", &[database])? { + Some(list) => parse_schemas(list.bind(py)), + None => Ok(Vec::new()), + } + })?; + callback(defs); + Ok(()) + } + + fn relations( + &self, + database: &str, + schema: &str, + callback: &mut dyn FnMut(Vec), + ) -> DFResult<()> { + let defs = Python::attach(|py| -> DFResult> { + match self.pull(py, "relations", &[database, schema])? { + Some(list) => parse_relations(list.bind(py)), + None => Ok(Vec::new()), + } + })?; + callback(defs); + Ok(()) + } + + fn columns( + &self, + database: &str, + schema: &str, + relation: &str, + callback: &mut dyn FnMut(Vec), + ) -> DFResult<()> { + let defs = Python::attach(|py| -> DFResult> { + match self.pull(py, "columns", &[database, schema, relation])? { + Some(list) => parse_columns(list.bind(py)), + None => Ok(Vec::new()), + } + })?; + callback(defs); + Ok(()) + } +} + #[pyclass] pub struct Server { addr: String, @@ -1714,6 +1995,7 @@ pub struct Server { databases: Vec, schemas: Vec<(String, String)>, tables: Vec<(String, String, String, Vec>)>, + lazy_catalog_source: Arc>>>, } #[pymethods] @@ -1731,6 +2013,7 @@ impl Server { databases: Vec::new(), schemas: Vec::new(), tables: Vec::new(), + lazy_catalog_source: Arc::new(Mutex::new(None)), } } @@ -1766,6 +2049,19 @@ impl Server { } } + /// Install a lazy catalog source. `source` is a Python object whose + /// `databases(callback)`, `schemas(database, callback)`, + /// `relations(database, schema, callback)`, and + /// `columns(database, schema, relation, callback)` methods each call their + /// `callback` with a list of row dicts. When set, the base catalog context is + /// built with the lazy providers (so `pg_catalog`/`information_schema` reflect + /// the source live on every scan) and the eager `register_database`/ + /// `register_schema`/`register_table` registrations are skipped. Requires + /// `start(catalog_emulation=True)` for the catalog queries to be routed here. + fn set_lazy_catalog(&mut self, _py: Python, source: Py) { + *self.lazy_catalog_source.lock().unwrap() = Some(source); + } + fn register_database(&mut self, database_name: String) { self.databases.push(database_name); } @@ -1840,8 +2136,33 @@ impl Server { rt.block_on(async move { let py_worker = Arc::new(PythonWorker::new(query_cb, connect_cb, disconnect_cb, auth_cb)); let mut ctx_map: HashMap> = HashMap::new(); - - if self.databases.is_empty() { + + // Clone the Python lazy-catalog source out of the server, if one was set. + let lazy_source = self + .lazy_catalog_source + .lock() + .unwrap() + .as_ref() + .map(|o| Python::attach(|py| o.clone_ref(py))); + + if let Some(obj) = lazy_source { + // Lazy path: a single catalog context whose pg_catalog / + // information_schema tables (and views) are sourced from Python on + // every scan. Built BEFORE the views are created so they bind to + // the lazy providers. Eager register_* is skipped below. + let source: Arc = Arc::new(PyLazyCatalogSource { obj }); + let (raw_ctx, _) = get_base_session_context_with_lazy_catalog( + None, + "datafusion".to_string(), + "public".to_string(), + None, + source, + LazyCatalogOptions::all(), + ) + .await + .unwrap(); + ctx_map.insert("datafusion".to_string(), Arc::new(raw_ctx)); + } else if self.databases.is_empty() { if catalog_emulation { let (raw_ctx, _) = get_base_session_context(None, "datafusion".to_string(), "public".to_string(), None).await.unwrap(); ctx_map.insert("datafusion".to_string(), Arc::new(raw_ctx)); @@ -1851,28 +2172,32 @@ impl Server { } } else { for db in &self.databases { - let (raw_ctx, _) = get_base_session_context(None, - db.to_string(), + let (raw_ctx, _) = get_base_session_context(None, + db.to_string(), "main".to_string(), None).await.unwrap(); ctx_map.insert(db.clone(), Arc::new(raw_ctx)); } } - - for ctx in ctx_map.values() { - for db in &self.databases { - register_user_database(ctx, db).await.unwrap(); + + // The eager registrations are mutually exclusive with the lazy source: + // when a source is installed it is authoritative for user objects. + if self.lazy_catalog_source.lock().unwrap().is_none() { + for ctx in ctx_map.values() { + for db in &self.databases { + register_user_database(ctx, db).await.unwrap(); + } } - } - for (db, schema) in &self.schemas { - if let Some(c) = ctx_map.get(db) { - register_schema(c, db, schema).await.unwrap(); + for (db, schema) in &self.schemas { + if let Some(c) = ctx_map.get(db) { + register_schema(c, db, schema).await.unwrap(); + } } - } - for (db, schema, table, cols) in &self.tables { - if let Some(c) = ctx_map.get(db) { - register_user_tables(c, db, schema, table, cols.clone()).await.unwrap(); + for (db, schema, table, cols) in &self.tables { + if let Some(c) = ctx_map.get(db) { + register_user_tables(c, db, schema, table, cols.clone()).await.unwrap(); + } } } diff --git a/teleduck/README.md b/teleduck/README.md index 6b06b28..a43362c 100644 --- a/teleduck/README.md +++ b/teleduck/README.md @@ -10,6 +10,19 @@ teleduck mydata.db and connect with any PostgreSQL client. +## How it works + +Teleduck runs a `riffq` server in catalog-emulation mode. Data queries are +executed against DuckDB; catalog queries (`pg_catalog`, `information_schema`) are +answered by `riffq`'s **lazy catalog**: Teleduck registers a single source +(`DuckdbCatalogSource`) via `server.set_lazy_catalog(...)`, and `riffq` reads the +live DuckDB schema from it on every catalog scan. As a result the emulated +catalog always reflects the current database — a table or index created **after** +Teleduck started shows up immediately, with no re-registration. + +DuckDB has no table ownership, so `pg_tables.tableowner` is intentionally left +blank; index presence (`hasindexes`) is reported from `duckdb_indexes()`. + ## Command line options Teleduck serves the PostgreSQL protocol over TLS by default. The command diff --git a/teleduck/TODO.md b/teleduck/TODO.md new file mode 100644 index 0000000..ed184ee --- /dev/null +++ b/teleduck/TODO.md @@ -0,0 +1,54 @@ +# Teleduck TODO + +## ⚠️ Column type mapping — `duckdb_type_to_oid` returns wrong pg_type OIDs + +`duckdb_type_to_oid` in `src/teleduck/server.py` maps a DuckDB declared type to a +PostgreSQL `pg_type` OID for `pg_attribute.atttypid` / `information_schema.columns`. +Two cases are wrong (verified empirically against DuckDB 1.5.3) because they rely +on naive substring checks. These should be fixed together with the broader +column-type-mapping work (see `../../pg_catalog/TODO.md`). + +### 1. `TIMESTAMP WITH TIME ZONE` → reported as `timestamp` (1114) instead of `timestamptz` (1184) + +DuckDB's `information_schema.columns.data_type` returns the **spelled-out** string, +not `TIMESTAMPTZ`: + +``` +CREATE TABLE t(a TIMESTAMPTZ); +SELECT data_type FROM information_schema.columns WHERE table_name='t'; +-- => 'TIMESTAMP WITH TIME ZONE' + +duckdb_type_to_oid("TIMESTAMP WITH TIME ZONE"): + "TIMESTAMP" in dt -> True + return 1184 if "TZ" in dt else 1114 + "TZ" in "TIMESTAMP WITH TIME ZONE" -> False # <-- the bug + => 1114 (timestamp, WRONG; should be 1184 timestamptz) +``` + +Fix: also test `"TIME ZONE" in dt` (or normalize the type string first). + +### 2. `INTERVAL` → reported as `int4` (23) + +`"INTERVAL".upper()` contains the substring `"INT"`, and the broad `if "INT" in dt` +arm matches before any interval handling: + +``` +CREATE TABLE t(b INTERVAL); +SELECT data_type FROM information_schema.columns WHERE table_name='t'; +-- => 'INTERVAL' + +duckdb_type_to_oid("INTERVAL"): + "INT" in "INTERVAL" -> True + => 23 (int4, WRONG; should be 1186 interval) +``` + +Fix: handle `INTERVAL` (and guard against other `*INT*`-containing names like +`POINT`) before the bare `"INT" in dt` fallback. + +### General + +The whole `duckdb_type_to_oid` table is best-effort and only covers common scalar +types; unrecognized types fall back to `text` (25). When the dedicated +type-mapping branch lands, add a test that round-trips representative DuckDB types +(timestamptz, interval, decimal, the integer family, bool, float/double, date, +time) through `pg_attribute`/`information_schema.columns` and asserts the OIDs. diff --git a/teleduck/src/teleduck.egg-info/SOURCES.txt b/teleduck/src/teleduck.egg-info/SOURCES.txt index f694a17..00af2bd 100644 --- a/teleduck/src/teleduck.egg-info/SOURCES.txt +++ b/teleduck/src/teleduck.egg-info/SOURCES.txt @@ -14,5 +14,6 @@ src/teleduck/certs/server.key tests/test_checkpoint_on_shutdown.py tests/test_cli.py tests/test_duckdb_catalog.py +tests/test_lazy_catalog.py tests/test_read_only.py tests/test_sql_init.py \ No newline at end of file diff --git a/teleduck/src/teleduck/server.py b/teleduck/src/teleduck/server.py index 77a8276..41f7964 100644 --- a/teleduck/src/teleduck/server.py +++ b/teleduck/src/teleduck/server.py @@ -24,6 +24,147 @@ def map_type(data_type: str) -> str: return "datetime" return "str" + +def duckdb_type_to_oid(data_type: str) -> int: + """Map a DuckDB declared type to a PostgreSQL ``pg_type`` OID. + + Used to fill ``pg_attribute.atttypid`` / ``information_schema.columns`` for + the lazy catalog. Only the common scalar types are distinguished; anything + unrecognized falls back to ``text`` (25), which is the safest default for + client type introspection. + """ + dt = data_type.upper() + if "BIGINT" in dt or "INT8" in dt or "HUGEINT" in dt or "LONG" in dt: + return 20 # int8 + if "SMALLINT" in dt or "INT2" in dt or "TINYINT" in dt: + return 21 # int2 + if "INT" in dt: + return 23 # int4 + if "BOOL" in dt: + return 16 # bool + if "DOUBLE" in dt or "FLOAT8" in dt: + return 701 # float8 + if "REAL" in dt or "FLOAT" in dt: + return 700 # float4 + if "DECIMAL" in dt or "NUMERIC" in dt: + return 1700 # numeric + if "TIMESTAMP" in dt or "DATETIME" in dt: + return 1184 if "TZ" in dt else 1114 # timestamptz / timestamp + if dt == "DATE": + return 1082 # date + if "TIME" in dt: + return 1083 # time + return 25 # text / varchar / everything else + + +def _stable_oid(salt: str, *parts: str) -> int: + """Derive a stable, built-in-clear OID from a namespace `salt` and `parts`. + + The same inputs always yield the same OID (so ``pg_class.oid`` and + ``pg_attribute.attrelid`` agree across scans and joins resolve), distinct + object classes use distinct salts to avoid collisions, and the result sits + well above the built-in OID range and inside the signed-32-bit range that + the catalog row types use. + """ + key = "\x00".join((salt,) + parts) + h = int(hashlib.sha1(key.encode("utf-8")).hexdigest()[:8], 16) + return 16384 + (h % 2_000_000_000) + + +class DuckdbCatalogSource: + """A lazy ``pg_catalog`` source backed by a live DuckDB connection. + + Each method queries DuckDB's own catalog on demand and hands the rows to the + ``callback``, so the emulated ``pg_catalog`` / ``information_schema`` always + reflects DuckDB's current schema -- including tables created after the server + started. Mirrors the Rust ``LazyCatalogSource`` trait one method per level. + + A fresh ``cursor()`` is used per call because the underlying connection is + shared with the query path and DuckDB connections are not concurrency-safe; + cursors give an independent, GIL-serialized handle. + """ + + def __init__(self, con): + self._con = con + + def databases(self, callback): + rows = self._con.cursor().execute( + "SELECT database_name FROM duckdb_databases() WHERE internal = false" + ).fetchall() + callback( + [ + {"oid": _stable_oid("db", name), "name": name} + for (name,) in rows + if name not in ("system", "temp") + ] + ) + + def schemas(self, database, callback): + rows = self._con.cursor().execute( + "SELECT DISTINCT table_schema FROM information_schema.tables " + "WHERE table_catalog = ? " + "AND table_schema NOT IN ('pg_catalog', 'information_schema')", + (database,), + ).fetchall() + callback( + [ + {"oid": _stable_oid("ns", database, schema), "name": schema} + for (schema,) in rows + ] + ) + + def relations(self, database, schema, callback): + rows = self._con.cursor().execute( + "SELECT table_name, table_type FROM information_schema.tables " + "WHERE table_catalog = ? AND table_schema = ?", + (database, schema), + ).fetchall() + # Tables that carry at least one index, so pg_tables.hasindexes is true. + # DuckDB has no triggers/rules/row-level-security, so those flags stay + # false (their default), which is truthful rather than blank. DuckDB also + # has no table ownership, so "owner_oid" is intentionally omitted and + # pg_tables.tableowner is left blank. + indexed = { + name + for (name,) in self._con.cursor().execute( + "SELECT DISTINCT table_name FROM duckdb_indexes() " + "WHERE database_name = ? AND schema_name = ?", + (database, schema), + ).fetchall() + } + out = [] + for table_name, table_type in rows: + kind = "view" if (table_type or "").upper().startswith("VIEW") else "table" + out.append( + { + "oid": _stable_oid("rel", database, schema, table_name), + "reltype_oid": _stable_oid("type", database, schema, table_name), + "name": table_name, + "kind": kind, + "has_index": table_name in indexed, + } + ) + callback(out) + + def columns(self, database, schema, relation, callback): + rows = self._con.cursor().execute( + "SELECT column_name, data_type, is_nullable FROM information_schema.columns " + "WHERE table_catalog = ? AND table_schema = ? AND table_name = ? " + "ORDER BY ordinal_position", + (database, schema, relation), + ).fetchall() + callback( + [ + { + "name": col_name, + "type_oid": duckdb_type_to_oid(data_type), + "nullable": str(is_nullable).upper() == "YES", + } + for (col_name, data_type, is_nullable) in rows + ] + ) + + class Connection(riffq.BaseConnection): def _handle_query(self, sql, callback, **kwargs): cur = duckdb_con.cursor() @@ -213,46 +354,12 @@ def _checkpoint_and_close(): server.set_tls(cert_path, key_path) - def register_schemas_and_tables_in_database(database_name): - if database_name in ("system", "temp"): - return - - # duckdb_con.execute(f"use {database_name}") - - tbls = duckdb_con.execute( - "SELECT table_schema, table_name FROM information_schema.tables " - "WHERE table_schema NOT IN ('pg_catalog','information_schema')" \ - "and table_catalog = ?", (database_name,) - ).fetchall() + # Drive pg_catalog lazily from the live DuckDB connection: every catalog + # scan re-reads DuckDB's schema, so tables created after startup show up + # without any re-registration. (Replaces the previous eager walk that + # snapshotted databases/schemas/tables once at boot.) + server.set_lazy_catalog(DuckdbCatalogSource(duckdb_con)) - for schema_name, table_name in tbls: - server._server.register_schema(database_name, schema_name) - cols_info = duckdb_con.execute( - "SELECT column_name, data_type, is_nullable FROM information_schema.columns " - "WHERE table_schema=? AND table_name=?", - (schema_name, table_name), - ).fetchall() - columns = [] - for col_name, data_type, is_nullable in cols_info: - columns.append( - { - col_name: { - "type": map_type(data_type), - "nullable": is_nullable.upper() == "YES", - } - } - ) - server._server.register_table(database_name, schema_name, table_name, columns) - - - databases = duckdb_con.execute( - "SELECT database_name, path, type FROM duckdb_databases() where internal=false" - ).fetchall() - - for database_name, path, type in databases: - server._server.register_database(database_name) - register_schemas_and_tables_in_database(database_name) - # riffq catches SIGINT/SIGTERM inside its tokio runtime and invokes this # before start() returns. a python signal.signal handler cannot be used # here: start() parks the main thread in rust, so a python-level handler diff --git a/teleduck/tests/test_lazy_catalog.py b/teleduck/tests/test_lazy_catalog.py new file mode 100644 index 0000000..b4cec54 --- /dev/null +++ b/teleduck/tests/test_lazy_catalog.py @@ -0,0 +1,105 @@ +import multiprocessing +import socket +import time +import tempfile +from pathlib import Path +import psycopg +import unittest +import duckdb +from server_readiness import wait_for_catalog, stop_server + + +def _run_server(db_file: str, port: int): + import sys + repo_root = Path(__file__).resolve().parents[1] + sys.path.insert(0, str(repo_root / "src")) + for mod in ["teleduck.server", "teleduck"]: + if mod in sys.modules: + del sys.modules[mod] + from teleduck.server import run_server + run_server(db_file, port, read_only=False) + + +class DuckDbLazyCatalogTest(unittest.TestCase): + """Proves the catalog is lazy: a table created after the server started is + visible in pg_catalog without any re-registration.""" + + @classmethod + def setUpClass(cls): + cls.port = 55491 + fd, cls.db_file = tempfile.mkstemp(suffix=".db") + Path(cls.db_file).unlink() # remove so DuckDB can create it + with duckdb.connect(cls.db_file) as con: + con.execute("CREATE TABLE users(id INTEGER, name VARCHAR)") + cls.database_name = con.execute( + "SELECT database_name FROM duckdb_databases() WHERE internal=false" + ).fetchall()[0][0] + + cls.proc = multiprocessing.Process( + target=_run_server, args=(cls.db_file, cls.port), daemon=True + ) + cls.proc.start() + start = time.time() + while time.time() - start < 60: + with socket.socket() as sock: + if sock.connect_ex(("127.0.0.1", cls.port)) == 0: + break + time.sleep(0.1) + else: + stop_server(cls.proc) + raise RuntimeError("Server did not start") + + wait_for_catalog( + cls.port, + "db", + f"SELECT datname FROM pg_catalog.pg_database WHERE datname='{cls.database_name}'", + cls.database_name, + ) + + @classmethod + def tearDownClass(cls): + stop_server(cls.proc) + Path(cls.db_file).unlink(missing_ok=True) + + def test_table_created_after_startup_is_visible(self): + conn = psycopg.connect( + f"postgresql://user:123@127.0.0.1:{self.port}/db", autocommit=True + ) + with conn.cursor() as cur: + # The table does not exist yet ... + cur.execute("SELECT count(*) FROM pg_catalog.pg_class WHERE relname='late_arrivals'") + self.assertEqual(cur.fetchone()[0], 0) + + # ... create it through the (DuckDB-backed) data path ... + cur.execute("CREATE TABLE late_arrivals(id INTEGER, note VARCHAR)") + + # ... and the lazy catalog reflects it on the next scan. + cur.execute("SELECT relname FROM pg_catalog.pg_class WHERE relname='late_arrivals'") + self.assertEqual(cur.fetchone()[0], "late_arrivals") + + # Its columns flow through information_schema too. + cur.execute( + "SELECT column_name FROM information_schema.columns " + "WHERE table_name='late_arrivals' ORDER BY ordinal_position" + ) + self.assertEqual([r[0] for r in cur.fetchall()], ["id", "note"]) + + # pg_tables flags are populated (not blank) for the user table, and + # reflect reality: no index yet, and DuckDB has no triggers. + cur.execute( + "SELECT hasindexes, hastriggers FROM pg_catalog.pg_tables " + "WHERE tablename='late_arrivals'" + ) + self.assertEqual(cur.fetchone(), (False, False)) + + # Create an index; hasindexes flips true on the next (lazy) scan. + cur.execute("CREATE INDEX late_idx ON late_arrivals(id)") + cur.execute( + "SELECT hasindexes FROM pg_catalog.pg_tables WHERE tablename='late_arrivals'" + ) + self.assertEqual(cur.fetchone()[0], True) + conn.close() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_lazy_catalog.py b/tests/test_lazy_catalog.py new file mode 100644 index 0000000..aa95147 --- /dev/null +++ b/tests/test_lazy_catalog.py @@ -0,0 +1,235 @@ +import multiprocessing +import socket +import time +import psycopg +import unittest +from helpers import stop_server + + +def _run_server(port: int): + import riffq + from riffq.helpers import to_arrow + + # In-process catalog state. The lazy source reads it on every scan, so a + # table appended here (by handle_query, below) appears in pg_class on the + # very next catalog query -- without any re-registration. + state = {"tables": ["users"]} + + def rel_oid(name): + # Stable OID per table name so pg_class.oid and pg_attribute.attrelid + # agree across scans and the joins resolve. + return 20000 + (abs(hash(name)) % 5000) + + class FakeSource: + """A lazy catalog source over the in-process `state` dict.""" + + def databases(self, callback): + callback([{"oid": 16384, "name": "appdb"}]) + + def schemas(self, database, callback): + if database == "appdb": + callback([{"oid": 16385, "name": "public"}]) + + def relations(self, database, schema, callback): + if database == "appdb" and schema == "public": + callback( + [ + { + "oid": rel_oid(name), + "reltype_oid": rel_oid(name) + 100000, + "name": name, + "kind": "table", + } + for name in state["tables"] + ] + ) + + def columns(self, database, schema, relation, callback): + callback( + [ + {"name": "id", "type_oid": 23, "nullable": False}, + {"name": "name", "type_oid": 25, "nullable": True}, + ] + ) + + def handle_query(sql, callback, **kwargs): + s = sql.strip().lower() + # Emulate a data backend: a CREATE TABLE mutates the live catalog state. + if s.startswith("create table "): + rest = s[len("create table "):].strip() + name = rest.split("(")[0].split()[0] + if name not in state["tables"]: + state["tables"].append(name) + callback(to_arrow([{"name": "status", "type": "str"}], [["OK"]])) + return + callback(to_arrow([{"name": "val", "type": "int"}], [[1]])) + + server = riffq.Server(f"127.0.0.1:{port}") + server.set_lazy_catalog(FakeSource()) + server.on_query(handle_query) + server.start(catalog_emulation=True) + + +class LazyCatalogTest(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.port = 55490 + cls.proc = multiprocessing.Process( + target=_run_server, args=(cls.port,), daemon=True + ) + cls.proc.start() + start = time.time() + while time.time() - start < 10: + with socket.socket() as sock: + if sock.connect_ex(("127.0.0.1", cls.port)) == 0: + break + time.sleep(0.1) + else: + stop_server(cls.proc) + raise RuntimeError("Server did not start") + + @classmethod + def tearDownClass(cls): + stop_server(cls.proc) + + def _conn(self): + return psycopg.connect( + f"postgresql://user@127.0.0.1:{self.port}/db", autocommit=True + ) + + def test_lazy_objects_and_builtins(self): + with self._conn() as conn, conn.cursor() as cur: + # The lazy database and its built-in neighbours both show up. + cur.execute("SELECT datname FROM pg_catalog.pg_database WHERE datname='appdb'") + self.assertEqual(cur.fetchone()[0], "appdb") + cur.execute("SELECT datname FROM pg_catalog.pg_database WHERE datname='postgres'") + self.assertEqual(cur.fetchone()[0], "postgres") + + cur.execute("SELECT nspname FROM pg_catalog.pg_namespace WHERE nspname='public'") + self.assertIsNotNone(cur.fetchone()) + + cur.execute("SELECT relname FROM pg_catalog.pg_class WHERE relname='users'") + self.assertEqual(cur.fetchone()[0], "users") + + def test_lazy_join_resolves(self): + # pg_class join pg_attribute over the lazy 'users' relation. + with self._conn() as conn, conn.cursor() as cur: + cur.execute( + "SELECT a.attname FROM pg_catalog.pg_class c " + "JOIN pg_catalog.pg_attribute a ON a.attrelid = c.oid " + "WHERE c.relname = 'users' ORDER BY a.attnum" + ) + self.assertEqual([r[0] for r in cur.fetchall()], ["id", "name"]) + + def test_lazy_reflects_table_created_after_startup(self): + with self._conn() as conn, conn.cursor() as cur: + cur.execute("SELECT count(*) FROM pg_catalog.pg_class WHERE relname='orders'") + self.assertEqual(cur.fetchone()[0], 0) + + # Create the table through the data path; the source now reports it. + cur.execute("CREATE TABLE orders(id int)") + + cur.execute("SELECT count(*) FROM pg_catalog.pg_class WHERE relname='orders'") + self.assertEqual( + cur.fetchone()[0], 1, "lazy catalog must reflect the new table" + ) + + +def _run_faulty_server(port: int, mode: str): + import riffq + from riffq.helpers import to_arrow + + class FaultySource: + """databases/schemas are fine so the server is healthy; relations() is + broken in the way selected by `mode`, to exercise error propagation.""" + + def databases(self, callback): + callback([{"oid": 16384, "name": "appdb"}]) + + def schemas(self, database, callback): + if database == "appdb": + callback([{"oid": 16385, "name": "public"}]) + + def relations(self, database, schema, callback): + if mode == "raise": + raise ValueError("boom from source") + if mode == "missing_field": + # 'oid' is required; omit it. + callback([{"reltype_oid": 30001, "name": "broken", "kind": "table"}]) + if mode == "bad_kind": + callback([{"oid": 20001, "reltype_oid": 30001, "name": "broken", + "kind": "nonsense"}]) + + def columns(self, database, schema, relation, callback): + callback([]) + + def handle_query(sql, callback, **kwargs): + callback(to_arrow([{"name": "v", "type": "int"}], [[1]])) + + server = riffq.Server(f"127.0.0.1:{port}") + server.set_lazy_catalog(FaultySource()) + server.on_query(handle_query) + server.start(catalog_emulation=True) + + +class LazyCatalogErrorPathTest(unittest.TestCase): + """The bridge must surface source errors to the SQL client, never silently + return an empty catalog.""" + + PORT_BASE = 55470 + MODES = {"raise": 0, "missing_field": 1, "bad_kind": 2} + + @classmethod + def setUpClass(cls): + cls.procs = {} + for mode, off in cls.MODES.items(): + port = cls.PORT_BASE + off + proc = multiprocessing.Process( + target=_run_faulty_server, args=(port, mode), daemon=True + ) + proc.start() + cls.procs[mode] = (proc, port) + start = time.time() + while time.time() - start < 10: + with socket.socket() as sock: + if sock.connect_ex(("127.0.0.1", port)) == 0: + break + time.sleep(0.1) + else: + stop_server(proc) + raise RuntimeError(f"server for mode {mode} did not start") + + @classmethod + def tearDownClass(cls): + for proc, _ in cls.procs.values(): + stop_server(proc) + + def _expect_error(self, mode, needle): + _, port = self.procs[mode] + conn = psycopg.connect( + f"postgresql://user@127.0.0.1:{port}/db", autocommit=True + ) + try: + with conn.cursor() as cur: + with self.assertRaises(psycopg.Error) as ctx: + cur.execute("SELECT relname FROM pg_catalog.pg_class") + cur.fetchall() + self.assertIn(needle, str(ctx.exception).lower()) + finally: + conn.close() + + def test_exception_in_source_propagates(self): + # A Python exception in a source method becomes a query error. + self._expect_error("raise", "boom from source") + + def test_missing_required_field_errors(self): + # A row missing the required 'oid' is a hard error, not a dropped row. + self._expect_error("missing_field", "oid") + + def test_unknown_relation_kind_errors(self): + # An unrecognized 'kind' string is rejected. + self._expect_error("bad_kind", "kind") + + +if __name__ == "__main__": + unittest.main()