quarb_sqlite/

lib.rs

//! SQLite adapter for the Quarb query engine.
//!
//! A thin catalog driver over the shared relational model
//! (`quarb-relational`): it introspects the schema via PRAGMAs
//! (`table_info`, `foreign_key_list`), streams every user table,
//! maps SQLite's storage classes to values (`NULL` → null,
//! `INTEGER`/`REAL` → numbers, text as itself, blobs as a size
//! placeholder), and hands the result to [`RelationalModel`].
//!
//! The arbor mapping, the foreign-key reference machinery (`-->`
//! chains, `->`/`<-` crosslinks, `<--` reverse resolution, the
//! table-naming hint), and the metadata surface are documented on
//! the shared model.

use quarb::{AstAdapter, NodeId, Value};
use quarb_relational::{RelationalModel, RowSpec, TableSpec};
use rusqlite::Connection;
use rusqlite::types::ValueRef;

/// An error loading a database.
#[derive(Debug, thiserror::Error)]
pub enum SqliteError {
    #[error("sqlite: {0}")]
    Sqlite(#[from] rusqlite::Error),
    #[error("pushdown plan: {0}")]
    Plan(String),
}

/// A SQLite database, materialized as an arbor.
pub struct SqliteAdapter {
    model: RelationalModel,
}

fn to_value(v: ValueRef<'_>) -> Value {
    match v {
        ValueRef::Null => Value::Null,
        ValueRef::Integer(n) => Value::Int(n),
        ValueRef::Real(f) => Value::Float(f),
        ValueRef::Text(t) => Value::Str(String::from_utf8_lossy(t).into_owned()),
        ValueRef::Blob(b) => Value::Str(format!("<blob {} bytes>", b.len())),
    }
}

/// Quote an identifier (table or column name) for interpolation into
/// SQL: wrap in double quotes, doubling any embedded quote, so
/// reserved words and special characters survive.
fn quote_ident(name: &str) -> String {
    quarb_relational::quote_double(name)
}

/// Introspect the schema: every user table's and view's spec, via
/// PRAGMAs. Views are containers like tables — a single-column view
/// (the `SELECT DISTINCT` dimension idiom) is named by that column,
/// the way a single-column primary key names a table's rows.
/// `refs` is the declared-references document: each `(field,
/// container)` pair synthesizes a foreign-key edge from every
/// matching column (bare `column`, or `table.column` to scope) to
/// the target container's key, unless the schema already declares
/// one there.
fn specs(conn: &Connection, refs: &[(String, String)]) -> Result<Vec<TableSpec>, SqliteError> {
    let named: Vec<(String, bool)> = conn
        .prepare(
            "SELECT name, type = 'view' FROM sqlite_master \
             WHERE type IN ('table', 'view') \
             AND name NOT LIKE 'sqlite_%' ORDER BY name",
        )?
        .query_map([], |r| Ok((r.get(0)?, r.get(1)?)))?
        .collect::<Result<_, _>>()?;

    let mut out = Vec::new();
    for (name, is_view) in named {
        // Columns and the primary key (a single-column pk names the
        // rows; a composite or absent pk falls back to rowid).
        let mut columns = Vec::new();
        let mut json_columns = Vec::new();
        let mut pk_cols: Vec<(i64, usize)> = Vec::new();
        let mut column_types: Vec<String> = Vec::new();
        {
            let mut stmt = conn.prepare(&format!("PRAGMA table_info({})", quote_ident(&name)))?;
            let mut rows = stmt.query([])?;
            while let Some(r) = rows.next()? {
                let col: String = r.get(1)?;
                let declared: String = r.get(2)?;
                let pk: i64 = r.get(5)?;
                if pk > 0 {
                    pk_cols.push((pk, columns.len()));
                }
                // a column declared JSON or JSONB declares its
                // cells' notation (spec: Grafting)
                if declared.trim().eq_ignore_ascii_case("json")
                    || declared.trim().eq_ignore_ascii_case("jsonb")
                {
                    json_columns.push(columns.len());
                }
                columns.push(col);
                column_types.push(declared);
            }
        }
        pk_cols.sort();
        // A single-column primary key names the rows. A view has no
        // primary key ever; when it projects exactly one column —
        // the `SELECT DISTINCT` dimension idiom — that column is
        // the row's identity, so it names the rows (and keys `-->`
        // resolution into the view).
        let pk = (pk_cols.len() == 1)
            .then(|| pk_cols[0].1)
            .or_else(|| (is_view && columns.len() == 1).then_some(0));

        // Declared foreign keys (an omitted target column means the
        // target's primary key).
        let mut fks = Vec::new();
        {
            let mut stmt =
                conn.prepare(&format!("PRAGMA foreign_key_list({})", quote_ident(&name)))?;
            let mut rows = stmt.query([])?;
            while let Some(r) = rows.next()? {
                let target: String = r.get(2)?;
                let from: String = r.get(3)?;
                let to: Option<String> = r.get(4)?;
                if let Some(idx) = columns.iter().position(|c| *c == from) {
                    fks.push((idx, target, to.unwrap_or_default()));
                }
            }
        }
        // Declared client-side references: synthesize an edge to the
        // target container's key for every matching column the
        // schema does not already cover.
        for (field, container) in refs {
            let col = match field.split_once('.') {
                Some((table, col)) if table == name => col,
                Some(_) => continue,
                None => field.as_str(),
            };
            if let Some(idx) = columns.iter().position(|c| c == col)
                && !fks.iter().any(|(i, _, _)| *i == idx)
                // A container's key column pointing at the container
                // itself is the identity, not an edge.
                && !(*container == name && pk == Some(idx))
            {
                fks.push((idx, container.clone(), String::new()));
            }
        }
        out.push(TableSpec {
            name,
            columns,
            pk,
            fks,
            json_columns,
            column_types,
        });
    }
    Ok(out)
}

/// Stream one table's rows, optionally filtered (partial
/// pushdown: the engine re-applies the pushed predicates).
fn fetch_rows_where(
    conn: &Connection,
    spec: &TableSpec,
    where_sql: Option<&str>,
) -> Result<Vec<RowSpec>, SqliteError> {
    let cols = spec
        .columns
        .iter()
        .map(|c| quote_ident(c))
        .collect::<Vec<_>>()
        .join(", ");
    let table = quote_ident(&spec.name);
    let filter = match where_sql {
        Some(w) => format!(" WHERE {w}"),
        None => String::new(),
    };
    // Prefer SQLite's rowid: stable row identity, and the row's edge
    // name for tables without a single-column primary key. WITHOUT
    // ROWID tables have no rowid column, so `SELECT rowid` fails to
    // prepare; fall back to a positional ordinal (below), ordered by
    // the key for determinism.
    let with_rowid = format!("SELECT rowid, {cols} FROM {table}{filter} ORDER BY rowid");
    if let Ok(mut stmt) = conn.prepare(&with_rowid) {
        let mut rows = stmt.query([])?;
        let mut out = Vec::new();
        while let Some(r) = rows.next()? {
            let rowid: i64 = r.get(0)?;
            let values: Vec<Value> = (0..spec.columns.len())
                .map(|i| to_value(r.get_ref(i + 1).expect("column in range")))
                .collect();
            out.push(RowSpec {
                rowid: Some(rowid),
                values,
            });
        }
        return Ok(out);
    }
    // WITHOUT ROWID: no rowid column to select or order by. Order by
    // the single-column primary key (which names the rows) when
    // present, else by every column so the ordinal keys are stable.
    let order = match spec.pk {
        Some(i) => quote_ident(&spec.columns[i]),
        None => cols.clone(),
    };
    let mut stmt = conn.prepare(&format!(
        "SELECT {cols} FROM {table}{filter} ORDER BY {order}"
    ))?;
    let mut rows = stmt.query([])?;
    let mut out = Vec::new();
    let mut rowid = 0i64;
    while let Some(r) = rows.next()? {
        rowid += 1;
        let values: Vec<Value> = (0..spec.columns.len())
            .map(|i| to_value(r.get_ref(i).expect("column in range")))
            .collect();
        out.push(RowSpec {
            rowid: Some(rowid),
            values,
        });
    }
    Ok(out)
}

/// The rows of a table whose column — `col`, or the row key (the
/// primary key, else `rowid`) — holds one of `values`, in rowid
/// order: what a hop reads instead of the table. Each value is
/// bound in its own storage class and, where it reads as a number
/// or is one, in the other too, so the comparison finds what the
/// engine's text equality would under any column affinity; the
/// statement can only return more, never less, and the model
/// matches by text afterwards. A WITHOUT ROWID table is declined
/// (its rows carry positional ids, which a keyed read cannot
/// assign).
fn fetch_rows_keyed(
    conn: &Connection,
    spec: &TableSpec,
    filter: Option<&str>,
    col: Option<usize>,
    values: &[Value],
) -> Result<Vec<RowSpec>, String> {
    use rusqlite::types::Value as Sql;
    let column = match col.or(spec.pk) {
        Some(i) => quote_ident(&spec.columns[i]),
        None => "rowid".to_string(),
    };
    let mut binds: Vec<Sql> = Vec::new();
    for v in values {
        match v {
            Value::Int(i) => {
                binds.push(Sql::Integer(*i));
                binds.push(Sql::Text(i.to_string()));
            }
            Value::Float(f) => {
                binds.push(Sql::Real(*f));
                binds.push(Sql::Text(v.to_string()));
            }
            Value::Bool(b) => binds.push(Sql::Integer(*b as i64)),
            Value::Null => {}
            other => {
                let text = other.to_string();
                if let Ok(i) = text.parse::<i64>() {
                    binds.push(Sql::Integer(i));
                } else if let Ok(f) = text.parse::<f64>() {
                    binds.push(Sql::Real(f));
                }
                binds.push(Sql::Text(text));
            }
        }
    }
    if binds.is_empty() {
        return Ok(Vec::new());
    }
    let cols = spec
        .columns
        .iter()
        .map(|c| quote_ident(c))
        .collect::<Vec<_>>()
        .join(", ");
    let marks = vec!["?"; binds.len()].join(", ");
    let filter = match filter {
        Some(w) => format!("({w}) AND "),
        None => String::new(),
    };
    let sql = format!(
        "SELECT rowid, {cols} FROM {} WHERE {filter}{column} IN ({marks}) ORDER BY rowid",
        quote_ident(&spec.name)
    );
    let mut stmt = conn.prepare(&sql).map_err(|e| e.to_string())?;
    let mut rows = stmt
        .query(rusqlite::params_from_iter(binds.iter()))
        .map_err(|e| e.to_string())?;
    let mut out = Vec::new();
    while let Some(r) = rows.next().map_err(|e| e.to_string())? {
        let rowid: i64 = r.get(0).map_err(|e| e.to_string())?;
        let values: Vec<Value> = (0..spec.columns.len())
            .map(|i| to_value(r.get_ref(i + 1).expect("column in range")))
            .collect();
        out.push(RowSpec {
            rowid: Some(rowid),
            values,
        });
    }
    Ok(out)
}

/// `ANALYZE`'s figure for one value of `column`: the average rows
/// per distinct key of an index that opens with the column
/// (`sqlite_stat1.stat`'s second number). An error when no index
/// opens with it or the statistics are missing.
fn index_rows_per_key(conn: &Connection, table: &str, column: &str) -> Result<u64, String> {
    let mut list = conn
        .prepare(&format!("PRAGMA index_list({})", quote_ident(table)))
        .map_err(|e| e.to_string())?;
    let names: Vec<String> = list
        .query_map([], |r| r.get::<_, String>(1))
        .map_err(|e| e.to_string())?
        .filter_map(Result::ok)
        .collect();
    for idx in names {
        let first: Option<String> = conn
            .query_row(
                &format!("PRAGMA index_info({})", quote_ident(&idx)),
                [],
                |r| r.get(2),
            )
            .ok();
        if first.as_deref() != Some(column) {
            continue;
        }
        let stat: String = conn
            .query_row(
                "SELECT stat FROM sqlite_stat1 WHERE tbl = ?1 AND idx = ?2",
                [table, idx.as_str()],
                |r| r.get(0),
            )
            .map_err(|e| e.to_string())?;
        return stat
            .split_whitespace()
            .nth(1)
            .and_then(|t| t.parse().ok())
            .ok_or_else(|| "no per-key figure in sqlite_stat1".to_string());
    }
    Err(format!("no index opens with {column}"))
}

impl SqliteAdapter {
    /// Open a database file (read-only): the catalog now, each
    /// table's rows on first touch (the connection stays owned by
    /// the adapter).
    pub fn open(path: &std::path::Path) -> Result<Self, SqliteError> {
        Self::open_impl(path, None, &[])
    }

    /// [`open`], with a declared-references document: each `(field,
    /// container)` pair supplies the edge the schema does not
    /// declare (see [`quarb_relational::parse_refs`]).
    pub fn open_with_refs(
        path: &std::path::Path,
        refs: &[(String, String)],
    ) -> Result<Self, SqliteError> {
        Self::open_impl(path, None, refs)
    }

    /// [`open`], with one table's fetch filtered by a WHERE clause
    /// (partial pushdown; the engine re-applies the predicates).
    pub fn open_filtered(
        path: &std::path::Path,
        table: &str,
        where_sql: &str,
    ) -> Result<Self, SqliteError> {
        Self::open_impl(path, Some((table.to_string(), where_sql.to_string())), &[])
    }

    /// [`open_filtered`], with a declared-references document.
    pub fn open_filtered_with_refs(
        path: &std::path::Path,
        table: &str,
        where_sql: &str,
        refs: &[(String, String)],
    ) -> Result<Self, SqliteError> {
        Self::open_impl(path, Some((table.to_string(), where_sql.to_string())), refs)
    }

    fn open_impl(
        path: &std::path::Path,
        filter: Option<(String, String)>,
        refs: &[(String, String)],
    ) -> Result<Self, SqliteError> {
        let conn = std::rc::Rc::new(Connection::open_with_flags(
            path,
            rusqlite::OpenFlags::SQLITE_OPEN_READ_ONLY,
        )?);
        let specs = specs(&conn, refs)?;
        let (whole, keyed, stats, listed) = (conn.clone(), conn.clone(), conn.clone(), conn);
        let mut model = RelationalModel::lazy(
            specs,
            Box::new(move |_, spec| {
                fetch_rows_where(&whole, spec, None).map_err(|e| e.to_string())
            }),
        )
        // A hop reads the rows it lands on, by key.
        .with_key_fetcher(Box::new(move |_, spec, col, values| {
            fetch_rows_keyed(&keyed, spec, None, col, values)
        }))
        // `ANALYZE`'s statistics, when the file carries them: the
        // table's row count as the planner knows it. A table with
        // no index has one row with a null `idx`; an indexed table
        // has one row per index, each opening with the table's row
        // count. Without statistics the estimate is unknown — a
        // count would be the scan.
        .with_estimator(std::rc::Rc::new(move |_, spec, equals| {
            let stat: String = stats
                .query_row(
                    "SELECT stat FROM sqlite_stat1 WHERE tbl = ?1 \
                     ORDER BY idx IS NULL DESC LIMIT 1",
                    [spec.name.as_str()],
                    |r| r.get(0),
                )
                .map_err(|e| e.to_string())?;
            let n: u64 = stat
                .split_whitespace()
                .next()
                .and_then(|t| t.parse().ok())
                .ok_or("no row count in sqlite_stat1")?;
            // An equality on a column some index opens with: the
            // index's statistics say how many rows one value of it
            // selects on average, the planner's own figure. The
            // smallest over the equalities bounds the read.
            let mut per_key: Option<u64> = None;
            for (col, _) in equals {
                let column = &spec.columns[*col];
                if let Ok(k) = index_rows_per_key(&stats, &spec.name, column) {
                    per_key = Some(per_key.map_or(k, |p| p.min(k)));
                }
            }
            Ok(quarb::Estimate {
                nodes: match (equals.is_empty(), per_key) {
                    (true, _) => (n, n),
                    (false, Some(k)) => (0, k.min(n)),
                    (false, None) => (0, n),
                },
                bytes: None,
                requests: Some(0),
                latency: quarb::Latency::Local,
                confidence: quarb::Confidence::Guess,
            })
        }));
        // The pushed-down WHERE filters the table's listing only:
        // hops into it go by key, unfiltered, so they stay exact.
        if let Some((table, where_sql)) = filter {
            model = model.with_filtered_listing(
                &table,
                Box::new(move |_, spec| {
                    fetch_rows_where(&listed, spec, Some(&where_sql)).map_err(|e| e.to_string())
                }),
            );
        }
        Ok(SqliteAdapter { model })
    }

    /// Open a database held entirely in memory — the bytes of a
    /// `.db` file that never touched a filesystem (the browser's
    /// uploaded files). Deserializes read-only, then materializes
    /// eagerly, as [`SqliteAdapter::load`] does.
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, SqliteError> {
        let mut conn = Connection::open_in_memory()?;
        conn.deserialize_read_exact(rusqlite::MAIN_DB, bytes, bytes.len(), true)?;
        Self::load(&conn)
    }

    /// Materialize every user table of an open connection, eagerly
    /// (in-memory databases, tests).
    pub fn load(conn: &Connection) -> Result<Self, SqliteError> {
        Self::load_with_refs(conn, &[])
    }

    /// [`load`], with a declared-references document.
    pub fn load_with_refs(
        conn: &Connection,
        refs: &[(String, String)],
    ) -> Result<Self, SqliteError> {
        let mut input = Vec::new();
        for spec in specs(conn, refs)? {
            let rows = fetch_rows_where(conn, &spec, None)?;
            input.push((spec, rows));
        }
        Ok(SqliteAdapter {
            model: RelationalModel::build(input),
        })
    }

    /// A human-readable locator: `/table/key` for rows.
    /// Surface a filtered fetch's error now (see
    /// `RelationalModel::prefetch`): the partial-pushdown caller
    /// falls back to the scan instead of answering empty.
    pub fn prefetch(&self, table: &str) -> Result<(), String> {
        self.model.prefetch(table)
    }

    pub fn locator(&self, node: NodeId) -> String {
        self.model.locator(node)
    }
}

/// Execute pushed-down SQL directly (read-only): the column names
/// and rows, ordered by `order_table`'s key when one is given (the
/// pushdown contract: row order must match the adapter's document
/// order, which is that key).
/// Whether `cols` covers the primary key or a non-partial UNIQUE
/// index of `table` — the soundness condition for executing a
/// witness-JOIN pushdown (see `raw_query`).
fn unique_key(conn: &Connection, table: &str, cols: &[String]) -> Result<bool, SqliteError> {
    use std::collections::HashSet;
    let want: HashSet<&str> = cols.iter().map(|s| s.as_str()).collect();
    if want.is_empty() {
        return Ok(false);
    }
    // Primary key (pk ordinal > 0 in table_info).
    let mut stmt = conn.prepare(&format!("PRAGMA table_info({})", quote_ident(table)))?;
    let mut pk: HashSet<String> = HashSet::new();
    let mut rows = stmt.query([])?;
    while let Some(r) = rows.next()? {
        let name: String = r.get(1)?;
        let ord: i64 = r.get(5)?;
        if ord > 0 {
            pk.insert(name);
        }
    }
    if !pk.is_empty() && pk.iter().all(|c| want.contains(c.as_str())) {
        return Ok(true);
    }
    // Non-partial UNIQUE indexes.
    let mut stmt = conn.prepare(&format!("PRAGMA index_list({})", quote_ident(table)))?;
    let mut uniques: Vec<String> = Vec::new();
    let mut rows = stmt.query([])?;
    while let Some(r) = rows.next()? {
        let name: String = r.get(1)?;
        let is_unique: i64 = r.get(2)?;
        let partial: i64 = r.get(4).unwrap_or(0);
        if is_unique == 1 && partial == 0 {
            uniques.push(name);
        }
    }
    for idx in uniques {
        let mut stmt = conn.prepare(&format!("PRAGMA index_info({})", quote_ident(&idx)))?;
        let mut cols_of: Vec<String> = Vec::new();
        let mut rows = stmt.query([])?;
        while let Some(r) = rows.next()? {
            cols_of.push(r.get(2)?);
        }
        if !cols_of.is_empty() && cols_of.iter().all(|c| want.contains(c.as_str())) {
            return Ok(true);
        }
    }
    Ok(false)
}

/// The catalog for a pushed plan's column-type assumptions
/// (`quarb_sql::ColumnTypes::hold`): every column of `tables`, with
/// whether SQLite gives it a numeric affinity. A column of INTEGER,
/// REAL or NUMERIC affinity holds numbers and compares with a
/// numeric literal as Quarb's scan does; a TEXT or untyped column
/// compares by storage class (`'150'` is not `150`, and any text
/// is greater than any number), which the scan's comparison by
/// value is not.
pub fn column_catalog(
    path: &std::path::Path,
    tables: &[&str],
) -> Result<Vec<(String, bool, bool)>, SqliteError> {
    let conn = Connection::open_with_flags(path, rusqlite::OpenFlags::SQLITE_OPEN_READ_ONLY)?;
    let mut catalog = Vec::new();
    for table in tables {
        let mut stmt = conn.prepare(&format!("PRAGMA table_info({})", quote_ident(table)))?;
        let rows = stmt.query_map([], |r| Ok((r.get::<_, String>(1)?, r.get::<_, String>(2)?)))?;
        for row in rows {
            let (name, declared) = row?;
            // numeric affinity, and whether the column is declared
            // json (the graft's declaration, kept by SQLite as the
            // type name)
            let json = matches!(
                declared.trim().to_ascii_lowercase().as_str(),
                "json" | "jsonb"
            );
            catalog.push((name, numeric_affinity(&declared), json));
        }
    }
    Ok(catalog)
}

/// SQLite's affinity rules (its "Datatypes" document, section 3.1),
/// reduced to the question asked: INTEGER, REAL or NUMERIC.
fn numeric_affinity(declared: &str) -> bool {
    let d = declared.to_ascii_uppercase();
    if d.contains("INT") {
        true
    } else if d.contains("CHAR") || d.contains("CLOB") || d.contains("TEXT") {
        false
    } else {
        !(d.contains("BLOB") || d.trim().is_empty())
    }
}

pub fn raw_query(
    path: &std::path::Path,
    sql: &str,
    order_table: Option<&str>,
    join_left: Option<(&str, &[String])>,
) -> Result<(Vec<String>, Vec<Vec<Value>>), SqliteError> {
    let conn = Connection::open_with_flags(path, rusqlite::OpenFlags::SQLITE_OPEN_READ_ONLY)?;
    // A witness-JOIN plan is only sound when the ON binds the left
    // table by a unique key (else SQL multiplies rows where
    // Quarb's existential binding does not). Verify against the
    // catalog; refusing sends the caller back to the scan.
    if let Some((table, cols)) = join_left
        && !unique_key(&conn, table, cols)?
    {
        return Err(SqliteError::Plan(format!(
            "join ON does not bind {table} by a unique key; \
             the SQL JOIN would multiply rows"
        )));
    }
    let sql = match order_table {
        Some(t) => {
            let key = specs(&conn, &[])?
                .into_iter()
                .find(|s| s.name == t)
                .and_then(|s| s.pk.map(|i| s.columns[i].clone()))
                .unwrap_or_else(|| "rowid".to_string());
            format!("{sql} ORDER BY {t}.{key}")
        }
        None => sql.to_string(),
    };
    // --explain prints what ran: the ORDER BY key is a
    // catalog lookup made here, so the final statement
    // exists nowhere earlier.
    quarb_relational::record_executed(&sql);
    let mut stmt = conn.prepare(&sql)?;
    let cols: Vec<String> = stmt.column_names().iter().map(|c| c.to_string()).collect();
    let n = cols.len();
    let mut rows = stmt.query([])?;
    let mut out = Vec::new();
    while let Some(r) = rows.next()? {
        out.push(
            (0..n)
                .map(|i| to_value(r.get_ref(i).expect("column in range")))
                .collect(),
        );
    }
    Ok((cols, out))
}

impl AstAdapter for SqliteAdapter {
    // Every trait method, through the relational model: the search
    // fast paths, the plan and its estimate included.
    quarb::forward_ast_adapter!(model);
}