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
19pub(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
40pub 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
150pub(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
162pub 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
202pub(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}