1use std::collections::HashMap;
2use std::path::PathBuf;
3use std::{fs, io};
4
5use sofie_core::sql_lexer::{Lexed, Token, lex};
6use sofie_model::Module;
7
8use change::Change;
9
10mod change;
11
12#[derive(Debug, Default)]
14pub struct Schema {
15 checks: Vec<Check>,
16 indexes: Vec<Index>,
17 columns: Vec<Column>,
18}
19
20#[derive(Debug)]
21struct Index {
22 table: String,
23 name: String,
24 columns: Vec<String>,
25 unique: bool,
26}
27
28#[derive(Debug)]
29struct Column {
30 table: String,
31 name: String,
32 ty: Vec<Token>,
33 nullable: bool,
34}
35
36const AFTER_TYPE: [&str; 11] = [
38 "not",
39 "null",
40 "default",
41 "check",
42 "unique",
43 "primary",
44 "references",
45 "constraint",
46 "generated",
47 "collate",
48 "using",
49];
50
51#[derive(Debug)]
52struct Check {
53 table: String,
54 name: String,
55 column: Option<String>,
56 tokens: Vec<Token>,
57}
58
59impl Check {
60 fn covers(&self, table: &str, column: &str) -> bool {
61 self.table == table
62 && (self.column.as_deref() == Some(column)
63 || self.tokens.iter().any(|token| is_named(token, column)))
64 }
65
66 fn literals(&self) -> Vec<&str> {
67 self.tokens
68 .iter()
69 .filter_map(|token| match token {
70 Token::Text(text) => Some(text.as_str()),
71 _ => None,
72 })
73 .collect()
74 }
75}
76
77impl Schema {
78 pub fn read(&mut self, migration: &Lexed) {
79 for statement in migration.tokens.split(|token| *token == Token::Symbol(';')) {
80 self.statement(statement);
81 }
82 }
83
84 pub(super) fn column_type(&self, table: &str, column: &str) -> Option<&[Token]> {
85 self.columns
86 .iter()
87 .find(|known| known.table == table && known.name == column)
88 .map(|known| known.ty.as_slice())
89 }
90
91 #[must_use]
94 pub fn check_literals(&self, table: &str, column: &str) -> Option<Vec<Vec<&str>>> {
95 let covering: Vec<Vec<&str>> = self
96 .checks
97 .iter()
98 .filter(|check| check.covers(table, column))
99 .map(Check::literals)
100 .collect();
101 (!covering.is_empty()).then_some(covering)
102 }
103
104 #[must_use]
105 pub fn knows_table(&self, table: &str) -> bool {
106 self.columns.iter().any(|known| known.table == table)
107 }
108
109 #[must_use]
110 pub fn has_column(&self, table: &str, column: &str) -> bool {
111 self.column_type(table, column).is_some()
112 }
113
114 #[must_use]
115 pub fn columns(&self, table: &str) -> Vec<&str> {
116 self.columns
117 .iter()
118 .filter(|column| column.table == table)
119 .map(|column| column.name.as_str())
120 .collect()
121 }
122
123 #[must_use]
124 pub fn is_unique(&self, table: &str, column: &str) -> bool {
125 self.is_unique_over(table, &[column])
126 }
127
128 #[must_use]
129 pub fn is_unique_over(&self, table: &str, columns: &[&str]) -> bool {
130 self.indexes
131 .iter()
132 .any(|index| index.unique && index.table == table && index.columns == columns)
133 }
134
135 #[must_use]
137 pub fn has_type(&self, table: &str, column: &str, sql_type: &str) -> Option<bool> {
138 let migrated = self.column_type(table, column)?;
139 Some(type_family(migrated) == type_family(&lex(sql_type).tokens))
140 }
141
142 #[must_use]
143 pub fn is_nullable(&self, table: &str, column: &str) -> Option<bool> {
144 self.columns
145 .iter()
146 .find(|known| known.table == table && known.name == column)
147 .map(|known| known.nullable)
148 }
149
150 #[must_use]
152 pub fn is_indexed(&self, table: &str, column: &str) -> bool {
153 self.indexes.iter().any(|index| {
154 index.table == table && index.columns.first().is_some_and(|first| first == column)
155 })
156 }
157
158 #[must_use]
159 pub fn has_index(&self, name: &str) -> bool {
160 self.indexes.iter().any(|index| index.name == name)
161 }
162
163 pub(super) fn statement(&mut self, tokens: &[Token]) {
164 Change { schema: self }.statement(tokens);
165 }
166}
167
168pub(super) fn type_of(definition: &[Token]) -> &[Token] {
170 let mut depth = 0_usize;
171 let end = definition
172 .iter()
173 .position(|token| {
174 match token {
175 Token::Symbol('(') => depth += 1,
176 Token::Symbol(')') => depth = depth.saturating_sub(1),
177 _ => {}
178 }
179 depth == 0 && AFTER_TYPE.contains(&token.keyword())
180 })
181 .unwrap_or(definition.len());
182 &definition[..end]
183}
184
185fn type_family(ty: &[Token]) -> String {
186 let mut depth = 0_usize;
187 let mut parts = Vec::new();
188 for token in ty {
189 match token {
190 Token::Symbol('(') => depth += 1,
191 Token::Symbol(')') => depth = depth.saturating_sub(1),
192 Token::Symbol(symbol) if depth == 0 => parts.push(symbol.to_string()),
193 _ if depth == 0 => parts.push(word(token).to_owned()),
194 _ => {}
195 }
196 }
197 let name = parts.join(" ");
198 let family = match name.as_str() {
199 "int8" | "bigserial" | "serial8" => "bigint",
200 "int" | "int4" | "serial" | "serial4" => "integer",
201 "bool" => "boolean",
202 "varchar" | "character varying" | "char" | "character" | "bpchar" | "citext" => "text",
203 "timestamp with time zone" => "timestamptz",
204 "timestamp without time zone" => "timestamp",
205 "float8" | "float" => "double precision",
206 _ => return name,
207 };
208 family.to_owned()
209}
210
211pub(super) fn word(token: &Token) -> &str {
213 match token {
214 Token::Word(word) | Token::Quoted(word) => word,
215 _ => "",
216 }
217}
218
219pub(super) fn is_identifier(token: &Token) -> bool {
220 matches!(token, Token::Word(_) | Token::Quoted(_))
221}
222
223fn is_named(token: &Token, name: &str) -> bool {
224 is_identifier(token) && word(token) == name
225}
226
227pub(super) fn skip_words<'t>(tokens: &'t [Token], words: &[&str]) -> &'t [Token] {
228 let skipped = tokens
229 .iter()
230 .take_while(|token| words.contains(&token.keyword()))
231 .count();
232 &tokens[skipped..]
233}
234
235pub(super) fn table_name(tokens: &[Token]) -> (String, &[Token]) {
236 let tokens = skip_words(tokens, &["if", "not", "exists", "only"]);
237 let mut end = 1;
238 while tokens.get(end) == Some(&Token::Symbol('.')) {
239 end += 2;
240 }
241 let name = tokens
242 .get(..end.min(tokens.len()))
243 .and_then(<[Token]>::last)
244 .map_or("", word);
245 (name.to_owned(), tokens.get(end..).unwrap_or_default())
246}
247
248pub(super) fn enclosed(tokens: &[Token]) -> &[Token] {
249 let mut depth = 0_usize;
250 for (index, token) in tokens.iter().enumerate() {
251 match token {
252 Token::Symbol('(') => depth += 1,
253 Token::Symbol(')') if depth == 1 => return &tokens[1..index],
254 Token::Symbol(')') => depth = depth.saturating_sub(1),
255 _ => {}
256 }
257 }
258 tokens.get(1..).unwrap_or_default()
259}
260
261pub(super) fn split_top_level(tokens: &[Token]) -> Vec<&[Token]> {
262 let mut parts = Vec::new();
263 let (mut depth, mut start) = (0_usize, 0);
264 for (index, token) in tokens.iter().enumerate() {
265 match token {
266 Token::Symbol('(') => depth += 1,
267 Token::Symbol(')') => depth = depth.saturating_sub(1),
268 Token::Symbol(',') if depth == 0 => {
269 parts.push(&tokens[start..index]);
270 start = index + 1;
271 }
272 _ => {}
273 }
274 }
275 parts.push(&tokens[start..]);
276 parts
277}
278
279fn identifiers(tokens: &[Token]) -> Vec<String> {
281 split_top_level(tokens)
282 .into_iter()
283 .filter_map(|element| {
284 if element.get(1) != Some(&Token::Symbol('(')) {
285 return element.first().map(word).map(str::to_owned);
286 }
287 let read: Vec<&str> = element
288 .iter()
289 .enumerate()
290 .filter(|(index, token)| {
291 is_identifier(token) && element.get(index + 1) != Some(&Token::Symbol('('))
292 })
293 .map(|(_, token)| word(token))
294 .collect();
295 match read.as_slice() {
296 [column] => Some((*column).to_owned()),
297 _ => None,
298 }
299 })
300 .collect()
301}
302
303pub(super) struct Migrations(HashMap<PathBuf, io::Result<Lexed>>);
304
305impl Migrations {
306 pub(super) fn read(modules: &[Module]) -> Self {
307 let lexed = modules
308 .iter()
309 .flat_map(|module| &module.migrations)
310 .map(|file| (file.clone(), fs::read_to_string(file).map(|sql| lex(&sql))))
311 .collect();
312 Self(lexed)
313 }
314
315 pub(super) fn of<'m>(
317 &'m self,
318 module: &'m Module,
319 ) -> impl Iterator<Item = (&'m PathBuf, &'m io::Result<Lexed>)> {
320 module
321 .migrations
322 .iter()
323 .filter_map(|file| Some((file, self.0.get(file)?)))
324 }
325}
326
327#[cfg(test)]
328mod tests {
329 use super::*;
330
331 fn read(sql: &str) -> Schema {
332 let mut constraints = Schema::default();
333 constraints.read(&lex(sql));
334 constraints
335 }
336
337 #[test]
338 fn finds_checks_in_tables_columns_and_later_migrations() {
339 let constraints = read(
340 "CREATE TABLE product (\n id uuid PRIMARY KEY,\n name text NOT NULL CHECK (char_length(name) <= 120),\n\
341 -- a comment; with a semicolon\n price bigint,\n CONSTRAINT positive CHECK (price > 0)\n);\n\
342 ALTER TABLE \"product\" ADD COLUMN status text, ADD CHECK (status IN ('DRAFT', 'LIVE'));",
343 );
344
345 assert_eq!(
346 constraints.check_literals("product", "name"),
347 Some(vec![vec![]])
348 );
349 assert!(constraints.check_literals("product", "price").is_some());
350 assert_eq!(
351 constraints.check_literals("product", "status"),
352 Some(vec![vec!["DRAFT", "LIVE"]])
353 );
354 assert_eq!(constraints.check_literals("product", "id"), None);
355 assert_eq!(constraints.check_literals("order", "price"), None);
356 }
357
358 #[test]
359 fn finds_unique_columns_indexes_and_constraints() {
360 let constraints = read(
361 "CREATE TABLE catalog.product (sku text UNIQUE, slug text, code text, a int, b int, UNIQUE (a, b));\n\
362 CREATE UNIQUE INDEX IF NOT EXISTS product_slug ON ONLY product (slug) WHERE deleted_at IS NULL;\n\
363 ALTER TABLE product ADD CONSTRAINT product_code UNIQUE (code);",
364 );
365
366 for column in ["sku", "slug", "code"] {
367 assert!(constraints.is_unique("product", column), "{column}");
368 }
369 assert!(!constraints.is_unique("product", "a"));
370 }
371
372 #[test]
373 fn each_check_lists_its_own_literals_and_a_dropped_one_lists_none() {
374 let constraints = read(
375 "CREATE TABLE line (status text CHECK (status <> 'CLOSED'), kind text, \
376 CONSTRAINT kinds CHECK (kind IN ('PLAIN')));\n\
377 ALTER TABLE line ADD CHECK (status IN ('OPEN'));\n\
378 ALTER TABLE line DROP CONSTRAINT kinds, ADD CONSTRAINT kinds2 CHECK (kind IN ('GIFT'));\n\
379 ALTER TABLE line DROP CONSTRAINT IF EXISTS line_status_check;",
380 );
381
382 assert_eq!(
383 constraints.check_literals("line", "status"),
384 Some(vec![vec!["OPEN"]])
385 );
386 assert_eq!(
387 constraints.check_literals("line", "kind"),
388 Some(vec![vec!["GIFT"]])
389 );
390 }
391
392 #[test]
393 fn a_renamed_column_keeps_its_checks_and_uniques_and_a_dropped_one_loses_them() {
394 let constraints = read(
395 "CREATE TABLE line (state text CHECK (state IN ('OPEN')), code text UNIQUE, note text CHECK (note <> ''));\n\
396 ALTER TABLE line RENAME COLUMN state TO status;\n\
397 ALTER TABLE line RENAME code TO sku;\n\
398 ALTER TABLE line DROP COLUMN note;",
399 );
400
401 assert!(constraints.has_column("line", "status"));
402 assert!(!constraints.has_column("line", "state"));
403 assert_eq!(
404 constraints.check_literals("line", "status"),
405 Some(vec![vec!["OPEN"]])
406 );
407 assert!(constraints.is_unique("line", "sku"));
408 assert!(!constraints.has_column("line", "note"));
409 assert_eq!(constraints.check_literals("line", "note"), None);
410 }
411
412 #[test]
413 fn an_index_serves_its_leading_column_until_the_column_is_dropped() {
414 let schema = read(
415 "CREATE TABLE line (id uuid, account uuid, kind text, note text);\n\
416 CREATE INDEX CONCURRENTLY IF NOT EXISTS line_account ON line (account, kind);\n\
417 CREATE INDEX line_note ON line (note);\n\
418 ALTER TABLE line DROP COLUMN note;",
419 );
420
421 assert!(schema.is_indexed("line", "account"));
422 assert!(!schema.is_indexed("line", "kind"));
423 assert!(!schema.is_indexed("line", "note"));
424 assert_eq!(schema.columns("line"), ["id", "account", "kind"]);
425 }
426
427 #[test]
428 fn a_dropped_index_or_table_is_gone_and_a_renamed_table_keeps_its_schema() {
429 let schema = read(
430 "CREATE TABLE line (id uuid, account uuid, code text UNIQUE, kind text CHECK (kind IN ('A')));\n\
431 CREATE INDEX line_account ON line (account);\n\
432 CREATE TABLE note (id uuid);\n\
433 DROP INDEX CONCURRENTLY IF EXISTS public.line_account;\n\
434 ALTER TABLE line RENAME TO entry;\n\
435 DROP TABLE IF EXISTS note CASCADE;",
436 );
437
438 assert!(!schema.is_indexed("entry", "account"));
439 assert!(!schema.has_index("line_account"));
440 assert!(!schema.knows_table("line"));
441 assert_eq!(schema.columns("entry"), ["id", "account", "code", "kind"]);
442 assert!(schema.is_unique("entry", "code"));
443 assert!(schema.has_index("line_code_key"));
444 assert_eq!(
445 schema.check_literals("entry", "kind"),
446 Some(vec![vec!["A"]])
447 );
448 assert!(!schema.knows_table("note"));
449 }
450
451 #[test]
452 fn a_primary_key_serves_a_lookup_by_its_leading_column() {
453 let schema = read(
454 "CREATE TABLE member (account uuid, team uuid, PRIMARY KEY (account, team));\n\
455 CREATE TABLE seat (number int PRIMARY KEY, CONSTRAINT seat_row UNIQUE (number));\n\
456 ALTER TABLE seat DROP CONSTRAINT seat_row;",
457 );
458
459 assert!(schema.is_indexed("member", "account"));
460 assert!(!schema.is_indexed("member", "team"));
461 assert!(schema.is_unique("seat", "number"));
462 assert!(!schema.has_index("seat_row"));
463 }
464
465 #[test]
466 fn a_column_s_type_is_read_under_its_aliases_and_its_nullability_follows_alters() {
467 let schema = read(
468 "CREATE TABLE line (id uuid PRIMARY KEY, version int8 NOT NULL, note varchar(20), \
469 at timestamp(3) with time zone, count integer, owner uuid, PRIMARY KEY (owner));\n\
470 ALTER TABLE line ALTER COLUMN note SET NOT NULL, ALTER COLUMN version DROP NOT NULL;",
471 );
472
473 assert_eq!(schema.has_type("line", "version", "bigint"), Some(true));
474 assert_eq!(schema.has_type("line", "note", "text"), Some(true));
475 assert_eq!(schema.has_type("line", "at", "timestamptz"), Some(true));
476 assert_eq!(schema.has_type("line", "count", "text"), Some(false));
477 assert_eq!(schema.has_type("line", "gone", "text"), None);
478 assert_eq!(schema.is_nullable("line", "id"), Some(false));
479 assert_eq!(schema.is_nullable("line", "note"), Some(false));
480 assert_eq!(schema.is_nullable("line", "version"), Some(true));
481 assert_eq!(schema.is_nullable("line", "count"), Some(true));
482 assert_eq!(schema.is_nullable("line", "owner"), Some(false));
483 }
484
485 #[test]
486 fn a_quoted_column_named_like_a_keyword_is_a_column() {
487 let constraints = read(
488 "CREATE TABLE line (id uuid);\nALTER TABLE line ADD COLUMN \"check\" text, ADD \"unique\" text;",
489 );
490
491 assert!(constraints.has_column("line", "check"));
492 assert!(constraints.has_column("line", "unique"));
493 assert_eq!(constraints.check_literals("line", "id"), None);
494 }
495
496 #[test]
497 fn a_unique_expression_of_one_column_keeps_that_column_unique() {
498 let constraints = read(
499 "CREATE TABLE account (email text, name text);\n\
500 CREATE UNIQUE INDEX account_email ON account (lower(email));",
501 );
502
503 assert!(constraints.is_unique("account", "email"));
504 assert!(!constraints.is_unique("account", "lower"));
505 }
506}