1use crate::migrator::{Migration, MigratorError};
39use drizzle_types::Dialect;
40
41#[derive(Debug, Clone, Copy, PartialEq, Eq)]
43pub enum ObjectKind {
44 Table,
46 Index,
48 View,
50 Trigger,
52 Enum,
54}
55
56#[derive(Debug, Clone)]
58pub struct CatalogObject {
59 pub kind: ObjectKind,
61 pub schema: Option<String>,
63 pub name: String,
65 pub sql: Option<String>,
68 pub members: Vec<String>,
70 pub unique: bool,
72}
73
74#[derive(Debug, Clone, Default)]
76pub struct Catalog {
77 pub objects: Vec<CatalogObject>,
79}
80
81impl Catalog {
82 #[must_use]
84 pub const fn new() -> Self {
85 Self {
86 objects: Vec::new(),
87 }
88 }
89
90 pub fn push(&mut self, object: CatalogObject) {
92 self.objects.push(object);
93 }
94
95 #[must_use]
100 pub fn find(
101 &self,
102 kind: ObjectKind,
103 schema: Option<&str>,
104 name: &str,
105 ) -> Option<&CatalogObject> {
106 let same_name =
107 |object: &&CatalogObject| object.kind == kind && object.name.eq_ignore_ascii_case(name);
108
109 if let Some(schema) = schema {
110 return self.objects.iter().find(|object| {
111 same_name(object)
112 && object
113 .schema
114 .as_deref()
115 .is_some_and(|live| live.eq_ignore_ascii_case(schema))
116 });
117 }
118
119 if let Some(object) = self.objects.iter().find(|object| {
120 same_name(object)
121 && object
122 .schema
123 .as_deref()
124 .is_none_or(|live| live.eq_ignore_ascii_case("public"))
125 }) {
126 return Some(object);
127 }
128
129 let mut matches = self.objects.iter().filter(same_name);
130 let first = matches.next()?;
131 matches.next().is_none().then_some(first)
132 }
133}
134
135#[derive(Debug, Clone, PartialEq, Eq)]
137pub enum StatementTarget {
138 Table {
140 schema: Option<String>,
142 name: String,
144 columns: Vec<String>,
146 },
147 Index {
149 schema: Option<String>,
151 name: String,
153 unique: bool,
155 columns: Vec<String>,
157 },
158 View {
160 schema: Option<String>,
162 name: String,
164 },
165 Enum {
167 schema: Option<String>,
169 name: String,
171 values: Vec<String>,
173 },
174 Unclassified,
176}
177
178impl StatementTarget {
179 #[must_use]
181 pub const fn is_classified(&self) -> bool {
182 !matches!(self, Self::Unclassified)
183 }
184}
185
186#[derive(Debug, Clone, PartialEq, Eq)]
188pub enum Disposition {
189 Skip {
191 reason: String,
193 },
194 Execute,
196 Unresolvable {
198 reason: String,
200 },
201}
202
203#[derive(Debug, Clone)]
205pub struct Step {
206 pub index: usize,
208 pub sql: String,
210 pub disposition: Disposition,
212}
213
214#[derive(Debug, Clone)]
216pub struct Plan {
217 pub tag: String,
219 pub steps: Vec<Step>,
221}
222
223impl Plan {
224 #[must_use]
226 pub fn unresolvable(&self) -> Vec<&Step> {
227 self.steps
228 .iter()
229 .filter(|step| matches!(step.disposition, Disposition::Unresolvable { .. }))
230 .collect()
231 }
232
233 #[must_use]
235 pub fn skipped_count(&self) -> usize {
236 self.steps
237 .iter()
238 .filter(|step| matches!(step.disposition, Disposition::Skip { .. }))
239 .count()
240 }
241
242 #[must_use]
244 pub fn is_resolvable(&self) -> bool {
245 self.unresolvable().is_empty()
246 }
247
248 pub fn into_executable(self, table_ident: &str) -> Result<Vec<String>, MigratorError> {
255 let blockers = self.unresolvable();
256 if !blockers.is_empty() {
257 let mut message = format!(
258 "cannot repair migration `{}`: {} statement(s) could not be reconciled against \
259 the live schema.\n\
260 Repair only skips CREATE TABLE / CREATE [UNIQUE] INDEX / CREATE VIEW / \
261 CREATE TYPE ... AS ENUM statements that introspection proves are already \
262 present; anything else inside the interrupted region must be resolved by hand.",
263 self.tag,
264 blockers.len()
265 );
266 for step in blockers {
267 let reason = match &step.disposition {
268 Disposition::Unresolvable { reason } => reason.as_str(),
269 _ => unreachable!("filtered to unresolvable"),
270 };
271 message.push_str(&format!(
272 "\n statement {}: {reason}\n {}",
273 step.index + 1,
274 single_line(&step.sql)
275 ));
276 }
277 message.push_str(&format!(
278 "\nResolve those statements manually, then mark the migration complete:\n \
279 UPDATE {table_ident} SET \"applied_at\" = CURRENT_TIMESTAMP WHERE \"name\" = '{}';\n\
280 …or discard the marker and re-run from scratch:\n \
281 DELETE FROM {table_ident} WHERE \"name\" = '{}';",
282 self.tag.replace('\'', "''"),
283 self.tag.replace('\'', "''"),
284 ));
285 return Err(MigratorError::UnrepairableMigration(message));
286 }
287
288 Ok(self
289 .steps
290 .into_iter()
291 .filter(|step| step.disposition == Disposition::Execute)
292 .map(|step| step.sql)
293 .collect())
294 }
295}
296
297#[must_use]
301pub fn plan(dialect: Dialect, migration: &Migration, catalog: &Catalog) -> Plan {
302 let mut steps = Vec::new();
303 let mut in_applied_prefix = true;
306
307 for (index, sql) in migration.statements().iter().enumerate() {
308 if sql.trim().is_empty() {
309 continue;
310 }
311
312 let target = classify_statement(dialect, sql);
313 let disposition = match &target {
314 StatementTarget::Unclassified => {
315 if in_applied_prefix {
316 Disposition::Unresolvable {
317 reason: "not a provable CREATE statement, and it may already have run \
318 before the interruption"
319 .to_string(),
320 }
321 } else {
322 Disposition::Execute
323 }
324 }
325 _ => match reconcile(dialect, &target, sql, catalog) {
326 Reconciled::Present => Disposition::Skip {
327 reason: describe_present(&target),
328 },
329 Reconciled::Absent => {
330 in_applied_prefix = false;
331 Disposition::Execute
332 }
333 Reconciled::Conflict(reason) => Disposition::Unresolvable { reason },
334 },
335 };
336
337 steps.push(Step {
338 index,
339 sql: sql.clone(),
340 disposition,
341 });
342 }
343
344 Plan {
345 tag: migration.tag().to_string(),
346 steps,
347 }
348}
349
350enum Reconciled {
351 Present,
353 Absent,
355 Conflict(String),
357}
358
359fn reconcile(
360 dialect: Dialect,
361 target: &StatementTarget,
362 sql: &str,
363 catalog: &Catalog,
364) -> Reconciled {
365 match target {
366 StatementTarget::Table {
367 schema,
368 name,
369 columns,
370 } => {
371 let Some(object) = catalog.find(ObjectKind::Table, schema.as_deref(), name) else {
372 return Reconciled::Absent;
373 };
374 if object_matches_sql(dialect, object, sql) || members_match(&object.members, columns) {
375 Reconciled::Present
376 } else {
377 Reconciled::Conflict(format!(
378 "table `{name}` already exists but its columns ({}) do not match the ones this \
379 statement creates ({})",
380 join_or_unknown(&object.members),
381 join_or_unknown(columns)
382 ))
383 }
384 }
385 StatementTarget::Index {
386 schema,
387 name,
388 unique,
389 columns,
390 } => {
391 let Some(object) = catalog.find(ObjectKind::Index, schema.as_deref(), name) else {
392 return Reconciled::Absent;
393 };
394 let have_evidence = object.sql.is_some() || !object.members.is_empty();
398 if have_evidence && object.unique != *unique {
399 return Reconciled::Conflict(format!(
400 "index `{name}` already exists but its uniqueness differs (live: {}, \
401 statement: {})",
402 object.unique, unique
403 ));
404 }
405 if object_matches_sql(dialect, object, sql) || members_match(&object.members, columns) {
406 Reconciled::Present
407 } else {
408 Reconciled::Conflict(format!(
409 "index `{name}` already exists but covers ({}) instead of ({})",
410 join_or_unknown(&object.members),
411 join_or_unknown(columns)
412 ))
413 }
414 }
415 StatementTarget::View { schema, name } => {
416 let Some(object) = catalog.find(ObjectKind::View, schema.as_deref(), name) else {
417 return Reconciled::Absent;
418 };
419 if object_matches_sql(dialect, object, sql) {
420 Reconciled::Present
421 } else {
422 Reconciled::Conflict(format!(
423 "view `{name}` already exists but its definition differs from this statement"
424 ))
425 }
426 }
427 StatementTarget::Enum {
428 schema,
429 name,
430 values,
431 } => {
432 let Some(object) = catalog.find(ObjectKind::Enum, schema.as_deref(), name) else {
433 return Reconciled::Absent;
434 };
435 if object.members == *values {
436 Reconciled::Present
437 } else {
438 Reconciled::Conflict(format!(
439 "enum type `{name}` already exists with labels ({}) instead of ({})",
440 join_or_unknown(&object.members),
441 join_or_unknown(values)
442 ))
443 }
444 }
445 StatementTarget::Unclassified => Reconciled::Conflict("unclassified statement".to_string()),
446 }
447}
448
449fn object_matches_sql(dialect: Dialect, object: &CatalogObject, sql: &str) -> bool {
450 object.sql.as_deref().is_some_and(|live_sql| {
451 canonical_tokens(dialect, live_sql) == canonical_tokens(dialect, sql)
452 })
453}
454
455fn members_match(live: &[String], target: &[String]) -> bool {
456 !live.is_empty()
457 && live.len() == target.len()
458 && live
459 .iter()
460 .zip(target)
461 .all(|(a, b)| a.eq_ignore_ascii_case(b))
462}
463
464fn join_or_unknown(values: &[String]) -> String {
465 if values.is_empty() {
466 "unknown".to_string()
467 } else {
468 values.join(", ")
469 }
470}
471
472fn describe_present(target: &StatementTarget) -> String {
473 match target {
474 StatementTarget::Table { name, .. } => {
475 format!("table `{name}` already exists with a matching definition")
476 }
477 StatementTarget::Index { name, .. } => {
478 format!("index `{name}` already exists with a matching definition")
479 }
480 StatementTarget::View { name, .. } => {
481 format!("view `{name}` already exists with a matching definition")
482 }
483 StatementTarget::Enum { name, .. } => {
484 format!("enum type `{name}` already exists with matching labels")
485 }
486 StatementTarget::Unclassified => "already applied".to_string(),
487 }
488}
489
490fn single_line(sql: &str) -> String {
491 let collapsed = sql.split_whitespace().collect::<Vec<_>>().join(" ");
492 if collapsed.chars().count() > 160 {
493 let truncated: String = collapsed.chars().take(157).collect();
494 format!("{truncated}...")
495 } else {
496 collapsed
497 }
498}
499
500#[derive(Debug, Clone, PartialEq, Eq)]
505enum Tok {
506 Word(String),
508 Quoted(String),
510 Str(String),
512 Punct(char),
514}
515
516impl Tok {
517 fn ident(&self) -> Option<&str> {
518 match self {
519 Self::Word(value) | Self::Quoted(value) => Some(value),
520 _ => None,
521 }
522 }
523
524 fn is_word(&self, keyword: &str) -> bool {
525 matches!(self, Self::Word(value) if value == keyword)
526 }
527
528 fn is_punct(&self, character: char) -> bool {
529 matches!(self, Self::Punct(value) if *value == character)
530 }
531
532 fn canonical(&self) -> String {
533 match self {
534 Self::Word(value) => value.clone(),
535 Self::Quoted(value) => value.to_lowercase(),
536 Self::Str(value) => format!("'{value}'"),
537 Self::Punct(value) => value.to_string(),
538 }
539 }
540}
541
542fn tokenize(sql: &str) -> Vec<Tok> {
548 let bytes: Vec<char> = sql.chars().collect();
549 let mut tokens = Vec::new();
550 let mut index = 0;
551
552 while index < bytes.len() {
553 let ch = bytes[index];
554
555 if ch.is_whitespace() {
556 index += 1;
557 continue;
558 }
559
560 if ch == '-' && bytes.get(index + 1) == Some(&'-') {
562 while index < bytes.len() && bytes[index] != '\n' {
563 index += 1;
564 }
565 continue;
566 }
567 if ch == '/' && bytes.get(index + 1) == Some(&'*') {
568 index += 2;
569 while index < bytes.len() {
570 if bytes[index] == '*' && bytes.get(index + 1) == Some(&'/') {
571 index += 2;
572 break;
573 }
574 index += 1;
575 }
576 continue;
577 }
578
579 if let Some((closing, is_string)) = match ch {
581 '"' => Some(('"', false)),
582 '`' => Some(('`', false)),
583 '[' => Some((']', false)),
584 '\'' => Some(('\'', true)),
585 _ => None,
586 } {
587 index += 1;
588 let mut value = String::new();
589 while index < bytes.len() {
590 if bytes[index] == closing {
591 if closing != ']' && bytes.get(index + 1) == Some(&closing) {
593 value.push(closing);
594 index += 2;
595 continue;
596 }
597 index += 1;
598 break;
599 }
600 value.push(bytes[index]);
601 index += 1;
602 }
603 tokens.push(if is_string {
604 Tok::Str(value)
605 } else {
606 Tok::Quoted(value)
607 });
608 continue;
609 }
610
611 if ch.is_alphanumeric() || ch == '_' || ch == '$' {
612 let mut value = String::new();
613 while index < bytes.len()
614 && (bytes[index].is_alphanumeric() || bytes[index] == '_' || bytes[index] == '$')
615 {
616 value.push(bytes[index]);
617 index += 1;
618 }
619 tokens.push(Tok::Word(value.to_lowercase()));
620 continue;
621 }
622
623 tokens.push(Tok::Punct(ch));
624 index += 1;
625 }
626
627 tokens
628}
629
630fn canonical_tokens(_dialect: Dialect, sql: &str) -> Vec<String> {
635 let tokens = tokenize(sql);
636 let mut out = Vec::with_capacity(tokens.len());
637 let mut index = 0;
638
639 while index < tokens.len() {
640 if tokens[index].is_word("if")
641 && tokens.get(index + 1).is_some_and(|t| t.is_word("not"))
642 && tokens.get(index + 2).is_some_and(|t| t.is_word("exists"))
643 {
644 index += 3;
645 continue;
646 }
647 if tokens[index].is_word("concurrently") {
648 index += 1;
649 continue;
650 }
651 if tokens[index].is_punct(';') && index + 1 == tokens.len() {
652 break;
653 }
654 out.push(tokens[index].canonical());
655 index += 1;
656 }
657
658 out
659}
660
661const CONSTRAINT_KEYWORDS: [&str; 8] = [
664 "constraint",
665 "primary",
666 "unique",
667 "foreign",
668 "check",
669 "exclude",
670 "like",
671 "period",
672];
673
674#[must_use]
681pub fn classify_statement(dialect: Dialect, sql: &str) -> StatementTarget {
682 let tokens = tokenize(sql);
683 let mut index = 0;
684
685 if !tokens.first().is_some_and(|t| t.is_word("create")) {
686 return StatementTarget::Unclassified;
687 }
688 index += 1;
689
690 let mut unique = false;
691 loop {
692 match tokens.get(index) {
693 Some(t) if t.is_word("or") => {
694 index += 2;
696 }
697 Some(t) if t.is_word("unique") => {
698 unique = true;
699 index += 1;
700 }
701 Some(t)
702 if t.is_word("temp")
703 || t.is_word("temporary")
704 || t.is_word("unlogged")
705 || t.is_word("global")
706 || t.is_word("local") =>
707 {
708 index += 1;
709 }
710 _ => break,
711 }
712 }
713
714 let Some(kind) = tokens.get(index).and_then(Tok::ident).map(str::to_string) else {
715 return StatementTarget::Unclassified;
716 };
717 index += 1;
718
719 let skip_noise = |tokens: &[Tok], mut index: usize| {
721 loop {
722 match tokens.get(index) {
723 Some(t) if t.is_word("concurrently") => index += 1,
724 Some(t)
725 if t.is_word("if")
726 && tokens.get(index + 1).is_some_and(|t| t.is_word("not"))
727 && tokens.get(index + 2).is_some_and(|t| t.is_word("exists")) =>
728 {
729 index += 3;
730 }
731 _ => return index,
732 }
733 }
734 };
735
736 match kind.as_str() {
737 "table" => {
738 index = skip_noise(&tokens, index);
739 let Some((schema, name, next)) = read_qualified_name(&tokens, index) else {
740 return StatementTarget::Unclassified;
741 };
742 let Some(columns) = read_table_columns(&tokens, next) else {
743 return StatementTarget::Unclassified;
744 };
745 StatementTarget::Table {
746 schema,
747 name,
748 columns,
749 }
750 }
751 "index" => {
752 index = skip_noise(&tokens, index);
753 if tokens.get(index).is_some_and(|t| t.is_word("on")) {
755 return StatementTarget::Unclassified;
756 }
757 let Some((schema, name, mut next)) = read_qualified_name(&tokens, index) else {
758 return StatementTarget::Unclassified;
759 };
760 if !tokens.get(next).is_some_and(|t| t.is_word("on")) {
761 return StatementTarget::Unclassified;
762 }
763 next += 1;
764 let Some((_, _, after_table)) = read_qualified_name(&tokens, next) else {
765 return StatementTarget::Unclassified;
766 };
767 let mut cursor = after_table;
769 while cursor < tokens.len() && !tokens[cursor].is_punct('(') {
770 cursor += 1;
771 }
772 let Some(columns) = read_paren_list_heads(&tokens, cursor, false) else {
773 return StatementTarget::Unclassified;
774 };
775 StatementTarget::Index {
776 schema,
777 name,
778 unique,
779 columns,
780 }
781 }
782 "view" | "materialized" => {
783 let mut cursor = index;
784 if kind == "materialized" {
785 if !tokens.get(cursor).is_some_and(|t| t.is_word("view")) {
786 return StatementTarget::Unclassified;
787 }
788 cursor += 1;
789 }
790 cursor = skip_noise(&tokens, cursor);
791 let Some((schema, name, _)) = read_qualified_name(&tokens, cursor) else {
792 return StatementTarget::Unclassified;
793 };
794 StatementTarget::View { schema, name }
795 }
796 "type" if dialect == Dialect::PostgreSQL => {
797 let Some((schema, name, mut next)) = read_qualified_name(&tokens, index) else {
798 return StatementTarget::Unclassified;
799 };
800 if !tokens.get(next).is_some_and(|t| t.is_word("as")) {
801 return StatementTarget::Unclassified;
802 }
803 next += 1;
804 if !tokens.get(next).is_some_and(|t| t.is_word("enum")) {
805 return StatementTarget::Unclassified;
806 }
807 next += 1;
808 let Some(values) = read_enum_values(&tokens, next) else {
809 return StatementTarget::Unclassified;
810 };
811 StatementTarget::Enum {
812 schema,
813 name,
814 values,
815 }
816 }
817 _ => StatementTarget::Unclassified,
818 }
819}
820
821fn read_qualified_name(tokens: &[Tok], index: usize) -> Option<(Option<String>, String, usize)> {
825 let first = tokens.get(index)?.ident()?.to_string();
826 if tokens.get(index + 1).is_some_and(|t| t.is_punct('.')) {
827 let second = tokens.get(index + 2)?.ident()?.to_string();
828 return Some((Some(first), second, index + 3));
829 }
830 Some((None, first, index + 1))
831}
832
833fn read_paren_list_heads(
836 tokens: &[Tok],
837 index: usize,
838 skip_constraints: bool,
839) -> Option<Vec<String>> {
840 if !tokens.get(index)?.is_punct('(') {
841 return None;
842 }
843
844 let mut heads = Vec::new();
845 let mut depth = 0usize;
846 let mut entry_start = true;
847 let mut cursor = index;
848
849 while cursor < tokens.len() {
850 let token = &tokens[cursor];
851 if token.is_punct('(') {
852 depth += 1;
853 if depth == 1 {
854 entry_start = true;
855 cursor += 1;
856 continue;
857 }
858 } else if token.is_punct(')') {
859 depth -= 1;
860 if depth == 0 {
861 return Some(heads);
862 }
863 } else if depth == 1 && token.is_punct(',') {
864 entry_start = true;
865 cursor += 1;
866 continue;
867 }
868
869 if depth == 1 && entry_start {
870 entry_start = false;
871 if let Some(ident) = token.ident() {
872 let is_constraint = skip_constraints
873 && matches!(token, Tok::Word(_))
874 && CONSTRAINT_KEYWORDS.contains(&ident);
875 if !is_constraint {
876 heads.push(ident.to_string());
877 }
878 }
879 }
880
881 cursor += 1;
882 }
883
884 None
885}
886
887fn read_table_columns(tokens: &[Tok], index: usize) -> Option<Vec<String>> {
888 read_paren_list_heads(tokens, index, true)
889}
890
891fn read_enum_values(tokens: &[Tok], index: usize) -> Option<Vec<String>> {
892 if !tokens.get(index)?.is_punct('(') {
893 return None;
894 }
895 let mut values = Vec::new();
896 let mut cursor = index + 1;
897 while cursor < tokens.len() {
898 match &tokens[cursor] {
899 Tok::Str(value) => values.push(value.clone()),
900 Tok::Punct(')') => return Some(values),
901 Tok::Punct(',') => {}
902 _ => return None,
903 }
904 cursor += 1;
905 }
906 None
907}
908
909pub mod sqlite {
915 use super::{Catalog, CatalogObject, ObjectKind, StatementTarget, classify_statement};
916 use drizzle_types::Dialect;
917
918 pub const OBJECTS_QUERY: &str = "SELECT type, name, sql FROM sqlite_master \
922 WHERE type IN ('table', 'index', 'view', 'trigger') AND name NOT LIKE 'sqlite_%'";
923
924 #[must_use]
933 pub fn catalog(rows: &[(String, String, Option<String>)]) -> Catalog {
934 let mut catalog = Catalog::new();
935 for (kind, name, sql) in rows {
936 let kind = match kind.as_str() {
937 "table" => ObjectKind::Table,
938 "index" => ObjectKind::Index,
939 "view" => ObjectKind::View,
940 "trigger" => ObjectKind::Trigger,
941 _ => continue,
942 };
943
944 let parsed = sql
945 .as_deref()
946 .map(|sql| classify_statement(Dialect::SQLite, sql));
947 let members = match &parsed {
948 Some(
949 StatementTarget::Table { columns, .. } | StatementTarget::Index { columns, .. },
950 ) => columns.clone(),
951 _ => Vec::new(),
952 };
953 let unique = matches!(parsed, Some(StatementTarget::Index { unique: true, .. }));
954
955 catalog.push(CatalogObject {
956 kind,
957 schema: None,
958 name: name.clone(),
959 sql: sql.clone(),
960 members,
961 unique,
962 });
963 }
964 catalog
965 }
966}
967
968pub mod postgres {
985 use super::{Catalog, CatalogObject, ObjectKind};
986
987 pub const TABLES_QUERY: &str = "SELECT n.nspname::text, c.relname::text \
989 FROM pg_catalog.pg_class c \
990 JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \
991 WHERE c.relkind IN ('r', 'p') \
992 AND n.nspname NOT IN ('pg_catalog', 'information_schema') \
993 AND n.nspname NOT LIKE 'pg\\_%'";
994
995 pub const COLUMNS_QUERY: &str = "SELECT n.nspname::text, c.relname::text, a.attname::text \
998 FROM pg_catalog.pg_attribute a \
999 JOIN pg_catalog.pg_class c ON c.oid = a.attrelid \
1000 JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \
1001 WHERE c.relkind IN ('r', 'p') AND a.attnum > 0 AND NOT a.attisdropped \
1002 AND n.nspname NOT IN ('pg_catalog', 'information_schema') \
1003 AND n.nspname NOT LIKE 'pg\\_%' \
1004 ORDER BY n.nspname, c.relname, a.attnum";
1005
1006 pub const VIEWS_QUERY: &str = "SELECT n.nspname::text, c.relname::text \
1008 FROM pg_catalog.pg_class c \
1009 JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \
1010 WHERE c.relkind IN ('v', 'm') \
1011 AND n.nspname NOT IN ('pg_catalog', 'information_schema') \
1012 AND n.nspname NOT LIKE 'pg\\_%'";
1013
1014 pub const INDEXES_QUERY: &str = "SELECT n.nspname::text, c.relname::text, i.indisunique \
1016 FROM pg_catalog.pg_index i \
1017 JOIN pg_catalog.pg_class c ON c.oid = i.indexrelid \
1018 JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \
1019 WHERE n.nspname NOT IN ('pg_catalog', 'information_schema') \
1020 AND n.nspname NOT LIKE 'pg\\_%'";
1021
1022 pub const INDEX_COLUMNS_QUERY: &str = "SELECT n.nspname::text, c.relname::text, a.attname::text \
1030 FROM pg_catalog.pg_index i \
1031 JOIN pg_catalog.pg_class c ON c.oid = i.indexrelid \
1032 JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \
1033 JOIN LATERAL unnest(i.indkey) WITH ORDINALITY AS k(attnum, ord) ON true \
1034 JOIN pg_catalog.pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = k.attnum \
1035 WHERE n.nspname NOT IN ('pg_catalog', 'information_schema') \
1036 AND n.nspname NOT LIKE 'pg\\_%' \
1037 ORDER BY n.nspname, c.relname, k.ord";
1038
1039 pub const ENUMS_QUERY: &str = "SELECT n.nspname::text, t.typname::text, e.enumlabel::text \
1042 FROM pg_catalog.pg_type t \
1043 JOIN pg_catalog.pg_namespace n ON n.oid = t.typnamespace \
1044 JOIN pg_catalog.pg_enum e ON e.enumtypid = t.oid \
1045 WHERE n.nspname NOT IN ('pg_catalog', 'information_schema') \
1046 AND n.nspname NOT LIKE 'pg\\_%' \
1047 ORDER BY n.nspname, t.typname, e.enumsortorder";
1048
1049 #[must_use]
1064 pub fn catalog(
1065 tables: &[(String, String)],
1066 columns: &[(String, String, String)],
1067 views: &[(String, String)],
1068 indexes: &[(String, String, bool)],
1069 index_columns: &[(String, String, String)],
1070 enums: &[(String, String, String)],
1071 ) -> Catalog {
1072 fn members_of(rows: &[(String, String, String)], schema: &str, owner: &str) -> Vec<String> {
1074 rows.iter()
1075 .filter(|(row_schema, row_owner, _)| row_schema == schema && row_owner == owner)
1076 .map(|(_, _, member)| member.clone())
1077 .collect()
1078 }
1079
1080 let mut catalog = Catalog::new();
1081
1082 for (schema, name) in tables {
1083 catalog.push(CatalogObject {
1084 kind: ObjectKind::Table,
1085 schema: Some(schema.clone()),
1086 name: name.clone(),
1087 sql: None,
1088 members: members_of(columns, schema, name),
1089 unique: false,
1090 });
1091 }
1092
1093 for (schema, name) in views {
1094 catalog.push(CatalogObject {
1095 kind: ObjectKind::View,
1096 schema: Some(schema.clone()),
1097 name: name.clone(),
1098 sql: None,
1099 members: Vec::new(),
1100 unique: false,
1101 });
1102 }
1103
1104 for (schema, name, unique) in indexes {
1105 catalog.push(CatalogObject {
1106 kind: ObjectKind::Index,
1107 schema: Some(schema.clone()),
1108 name: name.clone(),
1109 sql: None,
1112 members: members_of(index_columns, schema, name),
1113 unique: *unique,
1114 });
1115 }
1116
1117 let mut enum_names: Vec<(String, String)> = Vec::new();
1118 for (schema, name, _) in enums {
1119 let key = (schema.clone(), name.clone());
1120 if !enum_names.contains(&key) {
1121 enum_names.push(key);
1122 }
1123 }
1124 for (schema, name) in enum_names {
1125 let members = members_of(enums, &schema, &name);
1126 catalog.push(CatalogObject {
1127 kind: ObjectKind::Enum,
1128 schema: Some(schema),
1129 name,
1130 sql: None,
1131 members,
1132 unique: false,
1133 });
1134 }
1135
1136 catalog
1137 }
1138}
1139
1140#[cfg(test)]
1141mod tests {
1142 use super::{
1143 Catalog, Disposition, ObjectKind, StatementTarget, classify_statement, plan, postgres,
1144 sqlite,
1145 };
1146 use crate::migrator::Migration;
1147 use drizzle_types::Dialect;
1148
1149 fn sqlite_table(name: &str, sql: &str) -> Catalog {
1151 sqlite::catalog(&[("table".to_string(), name.to_string(), Some(sql.to_string()))])
1152 }
1153
1154 fn migration(statements: &[&str]) -> Migration {
1155 Migration::with_hash(
1156 "20240101010101_test",
1157 "hash",
1158 1,
1159 statements.iter().map(|s| (*s).to_string()).collect(),
1160 )
1161 }
1162
1163 #[test]
1164 fn classifies_create_table_columns_without_constraints() {
1165 let target = classify_statement(
1166 Dialect::SQLite,
1167 "CREATE TABLE `users` (\n\t`id` integer PRIMARY KEY NOT NULL,\n\t`email` text NOT NULL,\n\tFOREIGN KEY (`id`) REFERENCES `other`(`id`),\n\tCONSTRAINT `uq` UNIQUE(`email`)\n)",
1168 );
1169 assert_eq!(
1170 target,
1171 StatementTarget::Table {
1172 schema: None,
1173 name: "users".to_string(),
1174 columns: vec!["id".to_string(), "email".to_string()],
1175 }
1176 );
1177 }
1178
1179 #[test]
1180 fn classifies_if_not_exists_and_schema_qualified_tables() {
1181 let target = classify_statement(
1182 Dialect::PostgreSQL,
1183 "CREATE TABLE IF NOT EXISTS \"app\".\"users\" (\"id\" serial PRIMARY KEY, \"name\" text)",
1184 );
1185 assert_eq!(
1186 target,
1187 StatementTarget::Table {
1188 schema: Some("app".to_string()),
1189 name: "users".to_string(),
1190 columns: vec!["id".to_string(), "name".to_string()],
1191 }
1192 );
1193 }
1194
1195 #[test]
1196 fn classifies_unique_and_concurrent_indexes() {
1197 assert_eq!(
1198 classify_statement(
1199 Dialect::SQLite,
1200 "CREATE UNIQUE INDEX `users_email_idx` ON `users` (`email`)"
1201 ),
1202 StatementTarget::Index {
1203 schema: None,
1204 name: "users_email_idx".to_string(),
1205 unique: true,
1206 columns: vec!["email".to_string()],
1207 }
1208 );
1209
1210 assert_eq!(
1211 classify_statement(
1212 Dialect::PostgreSQL,
1213 "CREATE INDEX CONCURRENTLY IF NOT EXISTS \"users_email_idx\" ON \"users\" USING btree (\"email\", \"name\")"
1214 ),
1215 StatementTarget::Index {
1216 schema: None,
1217 name: "users_email_idx".to_string(),
1218 unique: false,
1219 columns: vec!["email".to_string(), "name".to_string()],
1220 }
1221 );
1222 }
1223
1224 #[test]
1225 fn classifies_postgres_enum_types() {
1226 assert_eq!(
1227 classify_statement(
1228 Dialect::PostgreSQL,
1229 "CREATE TYPE \"public\".\"status\" AS ENUM('active', 'archived')"
1230 ),
1231 StatementTarget::Enum {
1232 schema: Some("public".to_string()),
1233 name: "status".to_string(),
1234 values: vec!["active".to_string(), "archived".to_string()],
1235 }
1236 );
1237 assert_eq!(
1239 classify_statement(Dialect::SQLite, "CREATE TYPE status AS ENUM('a')"),
1240 StatementTarget::Unclassified
1241 );
1242 }
1243
1244 #[test]
1245 fn leaves_non_create_statements_unclassified() {
1246 for sql in [
1247 "ALTER TABLE `users` ADD COLUMN `age` integer",
1248 "DROP TABLE `users`",
1249 "INSERT INTO `users` (`id`) VALUES (1)",
1250 "UPDATE `users` SET `id` = 2",
1251 "CREATE TRIGGER t AFTER INSERT ON users BEGIN SELECT 1; END",
1252 "PRAGMA foreign_keys=OFF",
1253 ] {
1254 assert_eq!(
1255 classify_statement(Dialect::SQLite, sql),
1256 StatementTarget::Unclassified,
1257 "unexpectedly classified: {sql}"
1258 );
1259 }
1260 }
1261
1262 #[test]
1263 fn skips_the_applied_prefix_and_executes_the_rest() {
1264 let migration = migration(&[
1265 "CREATE TABLE `a` (`id` integer PRIMARY KEY)",
1266 "CREATE TABLE `b` (`id` integer PRIMARY KEY)",
1267 ]);
1268
1269 let catalog = sqlite_table("a", "CREATE TABLE `a` (`id` integer PRIMARY KEY)");
1270
1271 let plan = plan(Dialect::SQLite, &migration, &catalog);
1272 assert_eq!(plan.skipped_count(), 1);
1273 assert!(plan.is_resolvable());
1274 assert!(matches!(
1275 plan.steps[0].disposition,
1276 Disposition::Skip { .. }
1277 ));
1278 assert_eq!(plan.steps[1].disposition, Disposition::Execute);
1279
1280 let executable = plan
1281 .into_executable("\"__drizzle_migrations\"")
1282 .expect("resolvable");
1283 assert_eq!(
1284 executable,
1285 vec!["CREATE TABLE `b` (`id` integer PRIMARY KEY)"]
1286 );
1287 }
1288
1289 #[test]
1290 fn matches_tables_structurally_when_formatting_differs() {
1291 let migration =
1292 migration(&["CREATE TABLE `a` (\n `id` integer PRIMARY KEY,\n `name` text\n)"]);
1293 let catalog = sqlite_table(
1294 "a",
1295 "CREATE TABLE \"a\" (\"id\" INTEGER PRIMARY KEY, \"name\" TEXT)",
1296 );
1297
1298 let plan = plan(Dialect::SQLite, &migration, &catalog);
1299 assert!(plan.is_resolvable());
1300 assert_eq!(plan.skipped_count(), 1);
1301 }
1302
1303 #[test]
1304 fn refuses_when_an_existing_table_does_not_match() {
1305 let migration = migration(&["CREATE TABLE `a` (`id` integer, `email` text)"]);
1306 let catalog = sqlite_table("a", "CREATE TABLE `a` (`id` integer)");
1307
1308 let plan = plan(Dialect::SQLite, &migration, &catalog);
1309 assert!(!plan.is_resolvable());
1310 let error = plan
1311 .into_executable("\"__drizzle_migrations\"")
1312 .expect_err("must refuse");
1313 assert!(error.to_string().contains("do not match"));
1314 assert!(error.to_string().contains("UPDATE"));
1315 }
1316
1317 #[test]
1318 fn refuses_unprovable_statements_inside_the_applied_prefix() {
1319 let migration = migration(&[
1320 "ALTER TABLE `a` ADD COLUMN `age` integer",
1321 "CREATE TABLE `b` (`id` integer)",
1322 ]);
1323 let catalog = Catalog::new();
1324
1325 let plan = plan(Dialect::SQLite, &migration, &catalog);
1326 assert!(!plan.is_resolvable());
1327 assert!(matches!(
1328 plan.steps[0].disposition,
1329 Disposition::Unresolvable { .. }
1330 ));
1331 }
1332
1333 #[test]
1334 fn executes_unprovable_statements_after_a_proven_gap() {
1335 let migration = migration(&[
1336 "CREATE TABLE `a` (`id` integer)",
1337 "CREATE TABLE `b` (`id` integer)",
1338 "ALTER TABLE `b` ADD COLUMN `age` integer",
1339 ]);
1340 let catalog = sqlite_table("a", "CREATE TABLE `a` (`id` integer)");
1341
1342 let plan = plan(Dialect::SQLite, &migration, &catalog);
1343 assert!(plan.is_resolvable(), "{:?}", plan.steps);
1344 let executable = plan.into_executable("\"t\"").expect("resolvable");
1345 assert_eq!(executable.len(), 2);
1346 }
1347
1348 #[test]
1349 fn postgres_catalog_round_trips_through_the_planner() {
1350 let catalog = postgres::catalog(
1351 &[("public".to_string(), "users".to_string())],
1352 &[
1353 ("public".to_string(), "users".to_string(), "id".to_string()),
1354 (
1355 "public".to_string(),
1356 "users".to_string(),
1357 "email".to_string(),
1358 ),
1359 ],
1360 &[],
1361 &[("public".to_string(), "users_email_idx".to_string(), true)],
1362 &[(
1363 "public".to_string(),
1364 "users_email_idx".to_string(),
1365 "email".to_string(),
1366 )],
1367 &[(
1368 "public".to_string(),
1369 "status".to_string(),
1370 "active".to_string(),
1371 )],
1372 );
1373
1374 assert!(
1375 catalog.find(ObjectKind::Table, None, "users").is_some(),
1376 "unqualified lookup resolves to public"
1377 );
1378
1379 let migration = migration(&[
1380 "CREATE TABLE \"users\" (\"id\" serial PRIMARY KEY, \"email\" text)",
1381 "CREATE UNIQUE INDEX CONCURRENTLY \"users_email_idx\" ON \"users\" (\"email\")",
1382 "CREATE TYPE \"public\".\"status\" AS ENUM('active')",
1383 "CREATE TABLE \"posts\" (\"id\" serial PRIMARY KEY)",
1384 ]);
1385
1386 let plan = plan(Dialect::PostgreSQL, &migration, &catalog);
1387 assert!(plan.is_resolvable(), "{:?}", plan.steps);
1388 assert_eq!(plan.skipped_count(), 3);
1389 let executable = plan
1390 .into_executable("\"drizzle\".\"t\"")
1391 .expect("resolvable");
1392 assert_eq!(executable.len(), 1);
1393 assert!(executable[0].contains("posts"));
1394 }
1395
1396 #[test]
1397 fn postgres_enum_label_mismatch_is_unresolvable() {
1398 let catalog = postgres::catalog(
1399 &[],
1400 &[],
1401 &[],
1402 &[],
1403 &[],
1404 &[(
1405 "public".to_string(),
1406 "status".to_string(),
1407 "active".to_string(),
1408 )],
1409 );
1410 let migration = migration(&["CREATE TYPE \"status\" AS ENUM('active', 'archived')"]);
1411
1412 let plan = plan(Dialect::PostgreSQL, &migration, &catalog);
1413 assert!(!plan.is_resolvable());
1414 }
1415
1416 #[test]
1417 fn postgres_index_uniqueness_mismatch_is_unresolvable() {
1418 let catalog = postgres::catalog(
1421 &[],
1422 &[],
1423 &[],
1424 &[("public".to_string(), "users_email_idx".to_string(), false)],
1425 &[(
1426 "public".to_string(),
1427 "users_email_idx".to_string(),
1428 "email".to_string(),
1429 )],
1430 &[],
1431 );
1432 let migration =
1433 migration(&["CREATE UNIQUE INDEX \"users_email_idx\" ON \"users\" (\"email\")"]);
1434
1435 let plan = plan(Dialect::PostgreSQL, &migration, &catalog);
1436 assert!(!plan.is_resolvable(), "{:?}", plan.steps);
1437 let text = plan
1438 .into_executable("\"t\"")
1439 .expect_err("must refuse")
1440 .to_string();
1441 assert!(text.contains("uniqueness differs"), "{text}");
1442 }
1443}