Skip to main content

coven_database/
schema_introspection.rs

1use std::collections::BTreeMap;
2
3use rusqlite::{Connection, OptionalExtension};
4
5#[derive(Debug, thiserror::Error)]
6pub enum CreateTableSchemaError {
7    #[error("read CREATE TABLE schema for {table:?} failed: {source}")]
8    Read {
9        table: String,
10        #[source]
11        source: rusqlite::Error,
12    },
13    #[error("no CREATE TABLE schema for {0}")]
14    Missing(String),
15    #[error("bad CREATE TABLE SQL for {table}: {sql}")]
16    Malformed { table: String, sql: String },
17}
18
19/// `CREATE TABLE` text for `table` from `sqlite_master`.
20pub(crate) fn create_table_sql(
21    conn: &Connection,
22    table: &str,
23) -> Result<String, CreateTableSchemaError> {
24    let create = conn
25        .query_row(
26            "SELECT sql FROM sqlite_master WHERE type='table' AND name = ?1",
27            [table],
28            |row| row.get::<_, Option<String>>(0),
29        )
30        .optional()
31        .map_err(|source| CreateTableSchemaError::Read {
32            table: table.to_string(),
33            source,
34        })?;
35    create
36        .flatten()
37        .ok_or_else(|| CreateTableSchemaError::Missing(table.to_string()))
38}
39
40/// Qualify a `CREATE TABLE <name> ...` statement so it builds the table inside
41/// the attached schema `alias`, replacing only the table-name token.
42pub fn rewrite_create_into_schema(
43    create: &str,
44    table: &str,
45    alias: &str,
46) -> Result<String, CreateTableSchemaError> {
47    let Some((name_start, name_end, parsed_table)) = create_table_name_token(create) else {
48        return Err(CreateTableSchemaError::Malformed {
49            table: table.to_string(),
50            sql: create.to_string(),
51        });
52    };
53    if parsed_table != table {
54        return Err(CreateTableSchemaError::Malformed {
55            table: table.to_string(),
56            sql: create.to_string(),
57        });
58    }
59
60    let qualified = format!("{alias}.{}", quote_ident(table));
61    let mut out = String::with_capacity(create.len() + qualified.len());
62    out.push_str(&create[..name_start]);
63    out.push_str(&qualified);
64    out.push_str(&create[name_end..]);
65    Ok(out)
66}
67
68fn create_table_name_token(create: &str) -> Option<(usize, usize, String)> {
69    let mut pos = consume_keyword_ws(create, skip_ascii_ws(create, 0), "CREATE")?;
70    pos = consume_keyword_ws(create, pos, "TABLE")?;
71
72    if keyword_at(create, pos, "IF") {
73        pos = consume_keyword_ws(create, pos, "IF")?;
74        pos = consume_keyword_ws(create, pos, "NOT")?;
75        pos = consume_keyword_ws(create, pos, "EXISTS")?;
76    }
77
78    parse_identifier_token(create, pos)
79}
80
81fn skip_ascii_ws(sql: &str, mut pos: usize) -> usize {
82    while sql.as_bytes().get(pos).is_some_and(u8::is_ascii_whitespace) {
83        pos += 1;
84    }
85    pos
86}
87
88fn keyword_at(sql: &str, pos: usize, keyword: &str) -> bool {
89    let Some(end) = pos.checked_add(keyword.len()) else {
90        return false;
91    };
92    sql.get(pos..end)
93        .is_some_and(|token| token.eq_ignore_ascii_case(keyword))
94        && sql.as_bytes().get(end).is_some_and(u8::is_ascii_whitespace)
95}
96
97fn consume_keyword(sql: &str, pos: usize, keyword: &str) -> Option<usize> {
98    keyword_at(sql, pos, keyword).then_some(pos + keyword.len())
99}
100
101fn consume_keyword_ws(sql: &str, pos: usize, keyword: &str) -> Option<usize> {
102    Some(skip_ascii_ws(sql, consume_keyword(sql, pos, keyword)?))
103}
104
105fn parse_identifier_token(sql: &str, pos: usize) -> Option<(usize, usize, String)> {
106    match sql.as_bytes().get(pos).copied()? {
107        b'"' => parse_delimited_identifier(sql, pos, b'"'),
108        _ => parse_bare_identifier(sql, pos),
109    }
110}
111
112fn parse_delimited_identifier(
113    sql: &str,
114    start: usize,
115    delimiter: u8,
116) -> Option<(usize, usize, String)> {
117    let bytes = sql.as_bytes();
118    let mut pos = start + 1;
119    let mut out = String::new();
120    while pos < bytes.len() {
121        if bytes[pos] == delimiter {
122            if bytes.get(pos + 1).copied() == Some(delimiter) {
123                out.push(delimiter as char);
124                pos += 2;
125            } else {
126                return Some((start, pos + 1, out));
127            }
128        } else {
129            let ch = sql[pos..].chars().next()?;
130            out.push(ch);
131            pos += ch.len_utf8();
132        }
133    }
134    None
135}
136
137fn parse_bare_identifier(sql: &str, start: usize) -> Option<(usize, usize, String)> {
138    let mut pos = start;
139    while pos < sql.len() {
140        let b = sql.as_bytes()[pos];
141        if b.is_ascii_whitespace() || b == b'(' {
142            break;
143        }
144        let ch = sql[pos..].chars().next()?;
145        pos += ch.len_utf8();
146    }
147    (pos > start).then(|| (start, pos, sql[start..pos].to_string()))
148}
149
150/// Column names of `table`, in declared order, via `PRAGMA table_info`. The
151/// index of a name here is the index SQLite session changesets report for that
152/// column.
153pub(crate) fn table_columns(conn: &Connection, table: &str) -> rusqlite::Result<Vec<String>> {
154    let sql = format!("PRAGMA table_info({})", quote_ident(table));
155    let mut stmt = conn.prepare(&sql)?;
156    let columns = stmt
157        .query_map([], |row| row.get::<_, String>(1))?
158        .collect::<Result<Vec<_>, _>>()?;
159    Ok(columns)
160}
161
162/// Quote an SQL identifier (table/column name), doubling any embedded quote, so
163/// a trusted-but-unbindable name interpolates safely. Identifiers cannot be
164/// passed as bound parameters; this is the safe interpolation path for them.
165pub fn quote_ident(ident: &str) -> String {
166    format!("\"{}\"", ident.replace('"', "\"\""))
167}
168
169#[derive(Debug, thiserror::Error)]
170pub enum ForeignKeySchemaError {
171    #[error(transparent)]
172    Sqlite(#[from] rusqlite::Error),
173    #[error("foreign key on {child_table:?} is malformed: {reason}")]
174    Malformed { child_table: String, reason: String },
175}
176
177#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
178pub struct ForeignKeyColumn {
179    pub child: String,
180    pub parent: String,
181}
182
183#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
184pub struct ForeignKeyEdge {
185    pub parent_table: String,
186    pub columns: Vec<ForeignKeyColumn>,
187    pub on_update: String,
188    pub on_delete: String,
189    pub match_clause: String,
190}
191
192struct ForeignKeyRow {
193    sequence: i64,
194    parent_table: String,
195    child_column: String,
196    parent_column: Option<String>,
197    on_update: String,
198    on_delete: String,
199    match_clause: String,
200}
201
202/// Every outgoing foreign key on `child_table`, grouped by constraint and
203/// sorted by its complete parent/column/action shape. SQLite's constraint ids
204/// and PRAGMA listing order are deliberately discarded.
205pub(crate) fn foreign_key_edges(
206    conn: &Connection,
207    child_table: &str,
208) -> Result<Vec<ForeignKeyEdge>, ForeignKeySchemaError> {
209    let sql = format!("PRAGMA foreign_key_list({})", quote_ident(child_table));
210    let mut statement = conn.prepare(&sql)?;
211    let rows = statement.query_map([], |row| {
212        Ok((
213            row.get::<_, i64>(0)?,
214            ForeignKeyRow {
215                sequence: row.get(1)?,
216                parent_table: row.get(2)?,
217                child_column: row.get(3)?,
218                parent_column: row.get(4)?,
219                on_update: row.get::<_, String>(5)?.to_ascii_uppercase(),
220                on_delete: row.get::<_, String>(6)?.to_ascii_uppercase(),
221                match_clause: row.get::<_, String>(7)?.to_ascii_uppercase(),
222            },
223        ))
224    })?;
225    let mut grouped: BTreeMap<i64, Vec<ForeignKeyRow>> = BTreeMap::new();
226    for row in rows {
227        let (id, row) = row?;
228        grouped.entry(id).or_default().push(row);
229    }
230
231    let mut edges = Vec::with_capacity(grouped.len());
232    for mut rows in grouped.into_values() {
233        rows.sort_by_key(|row| row.sequence);
234        let first = rows
235            .first()
236            .ok_or_else(|| ForeignKeySchemaError::Malformed {
237                child_table: child_table.to_string(),
238                reason: "constraint has no columns".to_string(),
239            })?;
240        if rows.iter().any(|row| {
241            row.parent_table != first.parent_table
242                || row.on_update != first.on_update
243                || row.on_delete != first.on_delete
244                || row.match_clause != first.match_clause
245        }) {
246            return Err(ForeignKeySchemaError::Malformed {
247                child_table: child_table.to_string(),
248                reason: "one constraint reports inconsistent parent or actions".to_string(),
249            });
250        }
251        let omitted_parent_columns = rows.iter().all(|row| row.parent_column.is_none());
252        if !omitted_parent_columns && rows.iter().any(|row| row.parent_column.is_none()) {
253            return Err(ForeignKeySchemaError::Malformed {
254                child_table: child_table.to_string(),
255                reason: "one constraint mixes named and omitted parent columns".to_string(),
256            });
257        }
258        let inferred_parent_columns = if omitted_parent_columns {
259            primary_key_columns(conn, &first.parent_table)?
260        } else {
261            Vec::new()
262        };
263        if omitted_parent_columns && inferred_parent_columns.len() != rows.len() {
264            return Err(ForeignKeySchemaError::Malformed {
265                child_table: child_table.to_string(),
266                reason: format!(
267                    "{} child columns reference {} primary-key columns",
268                    rows.len(),
269                    inferred_parent_columns.len(),
270                ),
271            });
272        }
273        let columns = rows
274            .iter()
275            .enumerate()
276            .map(|(position, row)| ForeignKeyColumn {
277                child: row.child_column.clone(),
278                parent: row
279                    .parent_column
280                    .clone()
281                    .unwrap_or_else(|| inferred_parent_columns[position].clone()),
282            })
283            .collect();
284        edges.push(ForeignKeyEdge {
285            parent_table: first.parent_table.clone(),
286            columns,
287            on_update: first.on_update.clone(),
288            on_delete: first.on_delete.clone(),
289            match_clause: first.match_clause.clone(),
290        });
291    }
292    edges.sort();
293    Ok(edges)
294}
295
296fn primary_key_columns(
297    conn: &Connection,
298    table: &str,
299) -> Result<Vec<String>, ForeignKeySchemaError> {
300    let sql = format!("PRAGMA table_info({})", quote_ident(table));
301    let mut statement = conn.prepare(&sql)?;
302    let rows = statement.query_map([], |row| {
303        Ok((row.get::<_, i64>(5)?, row.get::<_, String>(1)?))
304    })?;
305    let mut columns = rows
306        .collect::<rusqlite::Result<Vec<_>>>()?
307        .into_iter()
308        .filter(|(rank, _)| *rank > 0)
309        .collect::<Vec<_>>();
310    columns.sort_by_key(|(rank, _)| *rank);
311    Ok(columns.into_iter().map(|(_, name)| name).collect())
312}
313
314#[cfg(test)]
315mod tests {
316    use super::*;
317
318    #[test]
319    fn create_table_rewrite_qualifies_table_token() {
320        let cases = [
321            (
322                "CREATE TABLE nodes (id TEXT PRIMARY KEY)",
323                "CREATE TABLE empty.\"nodes\" (id TEXT PRIMARY KEY)",
324            ),
325            (
326                "CREATE TABLE \"nodes\" (id TEXT PRIMARY KEY)",
327                "CREATE TABLE empty.\"nodes\" (id TEXT PRIMARY KEY)",
328            ),
329            (
330                "CREATE TABLE IF NOT EXISTS nodes (id TEXT PRIMARY KEY)",
331                "CREATE TABLE IF NOT EXISTS empty.\"nodes\" (id TEXT PRIMARY KEY)",
332            ),
333            (
334                "CREATE TABLE nodes (id TEXT PRIMARY KEY, parent_id TEXT REFERENCES \"nodes\" (id))",
335                "CREATE TABLE empty.\"nodes\" (id TEXT PRIMARY KEY, parent_id TEXT REFERENCES \"nodes\" (id))",
336            ),
337        ];
338
339        for (create, expected) in cases {
340            let rewritten = rewrite_create_into_schema(create, "nodes", "empty").expect("rewrite");
341            assert_eq!(rewritten, expected);
342        }
343    }
344
345    #[test]
346    fn create_table_rewrite_rejects_mismatched_table_token() {
347        let err = rewrite_create_into_schema(
348            "CREATE TABLE other_nodes (id TEXT PRIMARY KEY)",
349            "nodes",
350            "empty",
351        )
352        .expect_err("mismatched table token must fail");
353        assert!(
354            matches!(
355                err,
356                CreateTableSchemaError::Malformed { ref table, ref sql }
357                    if table == "nodes"
358                        && sql == "CREATE TABLE other_nodes (id TEXT PRIMARY KEY)"
359            ),
360            "unexpected error: {err}"
361        );
362    }
363}