quarb_postgres/
lib.rs

1//! PostgreSQL adapter for the Quarb query engine.
2//!
3//! A thin catalog driver over the shared relational model
4//! (`quarb-relational`): it connects with `tokio-postgres` (on an
5//! internal current-thread runtime, so the public surface stays
6//! synchronous), introspects the `public` schema through the
7//! system catalogs (tables, columns, primary keys, foreign
8//! keys), streams every table, and hands the result to
9//! [`RelationalModel`].
10//!
11//! Type mapping: booleans, the integer family, and the float
12//! family arrive natively; text arrives as itself; every other
13//! type (numeric, dates, timestamps, uuid, json, arrays, …) is
14//! selected with a `::text` cast, where Quarb's numeric reading
15//! and comparisons take over — the same posture as the CSV
16//! adapter. `NULL` is null.
17//!
18//! The arbor mapping and the foreign-key reference machinery (`~>`
19//! chains, `->`/`<-` crosslinks, `<~` reverse resolution, the
20//! table-naming hint) are documented on the shared model. The
21//! database is materialized eagerly at connect and opened
22//! read-only in spirit (the adapter only ever issues SELECTs).
23
24use quarb::{AstAdapter, NodeId, Value};
25use quarb_relational::{RelationalModel, RowSpec, TableSpec};
26use tokio_postgres::types::Type;
27use tokio_postgres::{Client, NoTls, Row};
28
29/// An error connecting to or loading a database.
30#[derive(Debug, thiserror::Error)]
31pub enum PostgresError {
32    #[error("postgres: {0}")]
33    Postgres(#[from] tokio_postgres::Error),
34    #[error("pushdown plan: {0}")]
35    Plan(String),
36    #[error("postgres runtime: {0}")]
37    Runtime(#[from] std::io::Error),
38}
39
40/// A PostgreSQL database, materialized as an arbor.
41pub struct PostgresAdapter {
42    model: RelationalModel,
43}
44
45/// Whether `data_type` (as `information_schema` spells it) arrives
46/// natively; anything else is selected with a `::text` cast.
47fn native_type(data_type: &str) -> bool {
48    matches!(
49        data_type,
50        "boolean"
51            | "smallint"
52            | "integer"
53            | "bigint"
54            | "real"
55            | "double precision"
56            | "text"
57            | "character varying"
58            | "character"
59    )
60}
61
62fn cell(row: &Row, i: usize) -> Result<Value, tokio_postgres::Error> {
63    let ty = row.columns()[i].type_();
64    let value = match *ty {
65        Type::BOOL => row
66            .get::<_, Option<bool>>(i)
67            .map_or(Value::Null, Value::Bool),
68        Type::INT2 => row
69            .get::<_, Option<i16>>(i)
70            .map_or(Value::Null, |n| Value::Int(n as i64)),
71        Type::INT4 => row
72            .get::<_, Option<i32>>(i)
73            .map_or(Value::Null, |n| Value::Int(n as i64)),
74        Type::INT8 => row.get::<_, Option<i64>>(i).map_or(Value::Null, Value::Int),
75        Type::FLOAT4 => row
76            .get::<_, Option<f32>>(i)
77            .map_or(Value::Null, |f| Value::Float(f as f64)),
78        Type::FLOAT8 => row
79            .get::<_, Option<f64>>(i)
80            .map_or(Value::Null, Value::Float),
81        // Non-native types are `::text`-cast on the adapter's own fetch
82        // path, so this arm sees a real `text` column there. A full
83        // pushdown (`raw_query`) can hand us a raw numeric/date/
84        // timestamp/uuid/json column with no cast, though: decoding
85        // that as `String` would panic, so use `try_get` and surface
86        // the type error to the caller (qua then falls back to the scan
87        // path) rather than aborting the process.
88        _ => row
89            .try_get::<_, Option<String>>(i)?
90            .map_or(Value::Null, Value::Str),
91    };
92    Ok(value)
93}
94
95impl PostgresAdapter {
96    /// Connect and materialize the `public` schema. `config` is a
97    /// `tokio-postgres` connection string — URL form
98    /// (`postgres://user@host:port/db`) or keyword form
99    /// (`host=/run/postgresql user=me dbname=db`).
100    pub fn connect(config: &str) -> Result<Self, PostgresError> {
101        Self::connect_impl(config, None)
102    }
103
104    /// [`connect`], with one table's fetch filtered by a WHERE
105    /// clause (partial pushdown; the engine re-applies the
106    /// predicates).
107    pub fn connect_filtered(
108        config: &str,
109        table: &str,
110        where_sql: &str,
111    ) -> Result<Self, PostgresError> {
112        Self::connect_impl(config, Some((table.to_string(), where_sql.to_string())))
113    }
114
115    fn connect_impl(config: &str, filter: Option<(String, String)>) -> Result<Self, PostgresError> {
116        let rt = tokio::runtime::Builder::new_current_thread()
117            .enable_all()
118            .build()?;
119        // Introspect the catalog now; keep the runtime, client, and
120        // connection driver alive inside the fetcher so each table's
121        // rows can stream on first touch.
122        let (client, specs, types) = rt.block_on(async {
123            let (client, connection) = tokio_postgres::connect(config, NoTls).await?;
124            tokio::spawn(connection);
125            let (specs, types) = Self::introspect(&client).await?;
126            Ok::<_, tokio_postgres::Error>((client, specs, types))
127        })?;
128        let model = RelationalModel::lazy(
129            specs,
130            Box::new(move |t, spec| {
131                let w = filter
132                    .as_ref()
133                    .filter(|(tn, _)| *tn == spec.name)
134                    .map(|(_, w)| w.as_str());
135                rt.block_on(Self::fetch_rows(&client, spec, &types[t], w))
136                    .map_err(|e| e.to_string())
137            }),
138        );
139        Ok(PostgresAdapter { model })
140    }
141
142    /// The catalog: every public table's spec, plus its columns'
143    /// catalog types (which decide the `::text` casts at fetch).
144    async fn introspect(
145        client: &Client,
146    ) -> Result<(Vec<TableSpec>, Vec<Vec<String>>), tokio_postgres::Error> {
147        let names: Vec<String> = client
148            .query(
149                "SELECT table_name FROM information_schema.tables \
150                 WHERE table_schema = 'public' AND table_type = 'BASE TABLE' \
151                 ORDER BY table_name",
152                &[],
153            )
154            .await?
155            .iter()
156            .map(|r| r.get(0))
157            .collect();
158
159        let mut specs = Vec::new();
160        let mut all_types = Vec::new();
161        for name in names {
162            // Columns, in ordinal order, with their catalog types.
163            let cols = client
164                .query(
165                    "SELECT column_name, data_type \
166                     FROM information_schema.columns \
167                     WHERE table_schema = 'public' AND table_name = $1 \
168                     ORDER BY ordinal_position",
169                    &[&name],
170                )
171                .await?;
172            let columns: Vec<String> = cols.iter().map(|r| r.get(0)).collect();
173            let types: Vec<String> = cols.iter().map(|r| r.get(1)).collect();
174
175            // The primary key (single-column keys name the rows).
176            let pk_rows = client
177                .query(
178                    "SELECT kcu.column_name \
179                     FROM information_schema.table_constraints tc \
180                     JOIN information_schema.key_column_usage kcu \
181                       ON tc.constraint_name = kcu.constraint_name \
182                      AND tc.table_schema = kcu.table_schema \
183                      AND tc.table_name = kcu.table_name \
184                     WHERE tc.table_schema = 'public' AND tc.table_name = $1 \
185                       AND tc.constraint_type = 'PRIMARY KEY' \
186                     ORDER BY kcu.ordinal_position",
187                    &[&name],
188                )
189                .await?;
190            let pk = (pk_rows.len() == 1).then(|| {
191                let col: String = pk_rows[0].get(0);
192                columns.iter().position(|c| *c == col)
193            });
194            let pk = pk.flatten();
195
196            // Declared foreign keys. Read straight from `pg_catalog`:
197            // `conkey`/`confkey` are parallel column-number arrays, so
198            // unnesting them together (`WITH ORDINALITY` keeps them in
199            // lockstep) pairs each referencing column with its own
200            // referenced column — a composite key stays paired instead
201            // of cross-producting. Keying on `conrelid` (the owning
202            // table) also avoids the information_schema hazard that
203            // constraint names collide across tables.
204            let fk_rows = client
205                .query(
206                    "SELECT a.attname, cf.relname, af.attname \
207                     FROM pg_catalog.pg_constraint con \
208                     JOIN pg_catalog.pg_class c ON c.oid = con.conrelid \
209                     JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \
210                     JOIN pg_catalog.pg_class cf ON cf.oid = con.confrelid \
211                     JOIN LATERAL unnest(con.conkey, con.confkey) \
212                          WITH ORDINALITY AS k(conkey, confkey, ord) ON true \
213                     JOIN pg_catalog.pg_attribute a \
214                       ON a.attrelid = con.conrelid AND a.attnum = k.conkey \
215                     JOIN pg_catalog.pg_attribute af \
216                       ON af.attrelid = con.confrelid AND af.attnum = k.confkey \
217                     WHERE con.contype = 'f' AND n.nspname = 'public' \
218                       AND c.relname = $1 \
219                     ORDER BY con.oid, k.ord",
220                    &[&name],
221                )
222                .await?;
223            let mut fks = Vec::new();
224            for r in &fk_rows {
225                let from: String = r.get(0);
226                let target: String = r.get(1);
227                let to: String = r.get(2);
228                if let Some(idx) = columns.iter().position(|c| *c == from) {
229                    fks.push((idx, target, to));
230                }
231            }
232
233            specs.push(TableSpec {
234                name,
235                columns,
236                pk,
237                fks,
238            });
239            all_types.push(types);
240        }
241        Ok((specs, all_types))
242    }
243
244    /// Stream one table's rows: native types bare, the rest cast to
245    /// text; row order follows the key when one names rows.
246    async fn fetch_rows(
247        client: &Client,
248        spec: &TableSpec,
249        types: &[String],
250        where_sql: Option<&str>,
251    ) -> Result<Vec<RowSpec>, tokio_postgres::Error> {
252        let select: Vec<String> = spec
253            .columns
254            .iter()
255            .zip(types)
256            .map(|(c, t)| {
257                if native_type(t) {
258                    format!("\"{c}\"")
259                } else {
260                    format!("\"{c}\"::text")
261                }
262            })
263            .collect();
264        // Qualify with the table name so a non-native pk (selected as
265        // `"pk"::text`, whose output column is also named `pk`) sorts by
266        // the raw column, not the text-cast alias — lexicographic vs.
267        // numeric — keeping this in step with `raw_query`'s qualified
268        // `ORDER BY "t"."pk"`.
269        let order = match spec.pk {
270            Some(i) => format!(" ORDER BY \"{}\".\"{}\"", spec.name, spec.columns[i]),
271            None => String::new(),
272        };
273        let filter = match where_sql {
274            Some(w) => format!(" WHERE {w}"),
275            None => String::new(),
276        };
277        let rows = client
278            .query(
279                &format!(
280                    "SELECT {} FROM \"{}\"{filter}{order}",
281                    select.join(", "),
282                    spec.name
283                ),
284                &[],
285            )
286            .await?;
287        rows.iter()
288            .enumerate()
289            .map(|(i, r)| {
290                let values = (0..spec.columns.len())
291                    .map(|c| cell(r, c))
292                    .collect::<Result<Vec<_>, _>>()?;
293                Ok(RowSpec {
294                    rowid: i as i64 + 1,
295                    values,
296                })
297            })
298            .collect()
299    }
300
301    /// A human-readable locator: `/table/key` for rows.
302    pub fn locator(&self, node: NodeId) -> String {
303        self.model.locator(node)
304    }
305}
306
307/// Execute pushed-down SQL directly: the column names and rows,
308/// ordered by `order_table`'s key when one is given (the pushdown
309/// contract: row order must match the adapter's document order).
310pub fn raw_query(
311    config: &str,
312    sql: &str,
313    order_table: Option<&str>,
314    join_left: Option<(&str, &[String])>,
315) -> Result<(Vec<String>, Vec<Vec<Value>>), PostgresError> {
316    // Witness-JOIN plans carry a uniqueness obligation this
317    // driver does not yet verify against its catalog; decline
318    // so the caller falls back to the (sound) scan.
319    if join_left.is_some() {
320        return Err(PostgresError::Plan(
321            "witness-JOIN uniqueness not verified by this driver".into(),
322        ));
323    }
324
325    let rt = tokio::runtime::Builder::new_current_thread()
326        .enable_all()
327        .build()?;
328    rt.block_on(async {
329        let (client, connection) = tokio_postgres::connect(config, NoTls).await?;
330        tokio::spawn(connection);
331        let sql = match order_table {
332            Some(t) => {
333                let (specs, _) = PostgresAdapter::introspect(&client).await?;
334                let key = specs
335                    .into_iter()
336                    .find(|s| s.name == t)
337                    .and_then(|s| s.pk.map(|i| s.columns[i].clone()));
338                match key {
339                    Some(k) => format!("{sql} ORDER BY \"{t}\".\"{k}\""),
340                    None => sql.to_string(),
341                }
342            }
343            None => sql.to_string(),
344        };
345        let rows = client.query(&sql, &[]).await?;
346        let cols: Vec<String> = rows
347            .first()
348            .map(|r| r.columns().iter().map(|c| c.name().to_string()).collect())
349            .unwrap_or_default();
350        let out = rows
351            .iter()
352            .map(|r| {
353                (0..r.columns().len())
354                    .map(|i| cell(r, i))
355                    .collect::<Result<Vec<_>, _>>()
356            })
357            .collect::<Result<Vec<_>, _>>()?;
358        Ok((cols, out))
359    })
360}
361
362impl AstAdapter for PostgresAdapter {
363    fn root(&self) -> NodeId {
364        self.model.root()
365    }
366    fn children(&self, node: NodeId) -> Vec<NodeId> {
367        self.model.children(node)
368    }
369    fn name(&self, node: NodeId) -> Option<String> {
370        self.model.name(node)
371    }
372    fn parent(&self, node: NodeId) -> Option<NodeId> {
373        self.model.parent(node)
374    }
375    fn property(&self, node: NodeId, name: &str) -> Option<Value> {
376        self.model.property(node, name)
377    }
378    fn default_value(&self, node: NodeId) -> Option<Value> {
379        self.model.default_value(node)
380    }
381    fn metadata(&self, node: NodeId, key: &str) -> Option<Value> {
382        self.model.metadata(node, key)
383    }
384    fn resolve(&self, node: NodeId, property: &str, hint: Option<&str>) -> Option<NodeId> {
385        self.model.resolve(node, property, hint)
386    }
387    fn links(&self, node: NodeId) -> Vec<(String, NodeId)> {
388        self.model.links(node)
389    }
390    fn backlinks(&self, node: NodeId) -> Vec<(String, NodeId)> {
391        self.model.backlinks(node)
392    }
393}
quarb_postgres/lib.rs

quarb_postgres/
lib.rs
