Skip to main content

sofie_check/rules/
migration.rs

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/// What the migrations so far leave in the database: columns with their types, checks and indexes.
13#[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
36/// Words ending a column's type in its definition.
37const 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    /// The literals of each CHECK on the column, one list per CHECK, since
92    /// only a single CHECK listing every value admits them all.
93    #[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    /// Under any of Postgres's names for the type: `int8` is a `bigint`, `varchar(20)` a `text`.
136    #[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    /// An index leading with the column serves a lookup by it.
151    #[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
168/// A column definition's type: `varchar(20)`, `character varying`, `numeric(10, 2)`.
169pub(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
211/// An identifier's name, quoted or not, or a keyword.
212pub(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
279/// A unique expression of one column, such as `lower(email)`, keeps that column's values distinct too.
280fn 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    /// In the module's order, which later migrations build on.
316    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}