feat: basic CREATE/DROP index support

Author: jvatic

src/diff.rs

@@ -2,7 +2,8 @@ use std::{cmp::Ordering, collections::HashSet};
 
 use sqlparser::ast::{
     AlterTableOperation, AlterType, AlterTypeAddValue, AlterTypeAddValuePosition,
-    AlterTypeOperation, CreateTable, Ident, ObjectName, Statement, UserDefinedTypeRepresentation,
+    AlterTypeOperation, CreateIndex, CreateTable, Ident, ObjectName, ObjectType, Statement,
+    UserDefinedTypeRepresentation,
 };
 
 pub trait Diff: Sized {
@@ -21,6 +22,7 @@ impl Diff for Vec<Statement> {
                 // CreateTable: compare against another CreateTable with the same name
                 // TODO: handle renames (e.g. use comments to tag a previous name for a table in a schema)
                 Statement::CreateTable(a) => find_and_compare_create_table(sa, a, other),
+                Statement::CreateIndex(a) => find_and_compare_create_index(sa, a, other),
                 Statement::CreateType { name, .. } => find_and_compare_create_type(sa, name, other),
                 Statement::CreateExtension {
                     name,
@@ -37,6 +39,10 @@ impl Diff for Vec<Statement> {
                         Statement::CreateTable(a) => a.name == b.name,
                         _ => false,
                     }),
+                    Statement::CreateIndex(b) => self.iter().find(|sa| match sa {
+                        Statement::CreateIndex(a) => a.name == b.name,
+                        _ => false,
+                    }),
                     Statement::CreateType { name: b_name, .. } => self.iter().find(|sa| match sa {
                         Statement::CreateType { name: a_name, .. } => a_name == b_name,
                         _ => false,
@@ -107,6 +113,36 @@ fn find_and_compare_create_table(
     )
 }
 
+fn find_and_compare_create_index(
+    sa: &Statement,
+    a: &CreateIndex,
+    other: &[Statement],
+) -> Option<Vec<Statement>> {
+    find_and_compare(
+        sa,
+        other,
+        |sb| match sb {
+            Statement::CreateIndex(b) => a.name == b.name,
+            _ => false,
+        },
+        || {
+            let name = a
+                .name
+                .clone()
+                .expect("can't drop an unnamed index, please ensure all indexes are named");
+            Some(vec![Statement::Drop {
+                object_type: sqlparser::ast::ObjectType::Index,
+                if_exists: a.if_not_exists,
+                names: vec![name],
+                cascade: false,
+                restrict: false,
+                purge: false,
+                temporary: false,
+            }])
+        },
+    )
+}
+
 fn find_and_compare_create_type(
     sa: &Statement,
     a_name: &ObjectName,
@@ -170,6 +206,10 @@ impl Diff for Statement {
                 Self::CreateTable(b) => compare_create_table(a, b),
                 _ => None,
             },
+            Self::CreateIndex(a) => match other {
+                Self::CreateIndex(b) => compare_create_index(a, b),
+                _ => None,
+            },
             Self::CreateType {
                 name: a_name,
                 representation: a_rep,
@@ -233,6 +273,30 @@ fn compare_create_table(a: &CreateTable, b: &CreateTable) -> Option<Vec<Statemen
     }])
 }
 
+fn compare_create_index(a: &CreateIndex, b: &CreateIndex) -> Option<Vec<Statement>> {
+    if a == b {
+        return None;
+    }
+
+    if a.name.is_none() || b.name.is_none() {
+        panic!("can't diff unnamed indexes!");
+    }
+    let name = a.name.clone().unwrap();
+
+    Some(vec![
+        Statement::Drop {
+            object_type: ObjectType::Index,
+            if_exists: a.if_not_exists,
+            names: vec![name],
+            cascade: false,
+            restrict: false,
+            purge: false,
+            temporary: false,
+        },
+        Statement::CreateIndex(b.clone()),
+    ])
+}
+
 fn compare_create_type(
     a_name: &ObjectName,
     a_rep: &UserDefinedTypeRepresentation,

src/lib.rs

@@ -244,6 +244,27 @@ mod tests {
         );
     }
 
+    #[test]
+    fn diff_create_index() {
+        run_test_cases(
+            vec![
+                TestCase {
+                    dialect: Dialect::Generic,
+                    sql_a: "CREATE UNIQUE INDEX title_idx ON films (title);",
+                    sql_b: "CREATE UNIQUE INDEX title_idx ON films ((lower(title)));",
+                    expect: "DROP INDEX title_idx;\n\nCREATE UNIQUE INDEX title_idx ON films((lower(title)));",
+                },
+                TestCase {
+                    dialect: Dialect::Generic,
+                    sql_a: "CREATE UNIQUE INDEX IF NOT EXISTS title_idx ON films (title);",
+                    sql_b: "CREATE UNIQUE INDEX IF NOT EXISTS title_idx ON films ((lower(title)));",
+                    expect: "DROP INDEX IF EXISTS title_idx;\n\nCREATE UNIQUE INDEX IF NOT EXISTS title_idx ON films((lower(title)));",
+                },
+            ],
+            |ast_a, ast_b| ast_a.diff(&ast_b).unwrap(),
+        );
+    }
+
     #[test]
     fn diff_create_type() {
         run_test_cases(
@@ -409,6 +430,31 @@ mod tests {
         ], |ast_a, ast_b| ast_a.migrate(&ast_b).unwrap());
     }
 
+    #[test]
+    fn apply_create_index() {
+        run_test_cases(vec![
+            TestCase {
+                dialect: Dialect::Generic,
+                sql_a: "CREATE UNIQUE INDEX title_idx ON films (title);",
+                sql_b: "CREATE INDEX code_idx ON films (code);",
+                expect: "CREATE UNIQUE INDEX title_idx ON films(title);\n\nCREATE INDEX code_idx ON films(code);",
+            },
+            TestCase {
+                dialect: Dialect::Generic,
+                sql_a: "CREATE UNIQUE INDEX title_idx ON films (title);",
+                sql_b: "DROP INDEX title_idx;",
+                expect: "",
+            },
+            TestCase {
+                dialect: Dialect::Generic,
+                sql_a: "CREATE UNIQUE INDEX title_idx ON films (title);",
+                sql_b: "DROP INDEX title_idx;CREATE INDEX code_idx ON films (code);",
+                expect: "CREATE INDEX code_idx ON films(code);",
+            },
+        ],
+        |ast_a, ast_b| ast_a.migrate(&ast_b).unwrap());
+    }
+
     #[test]
     fn apply_alter_create_type() {
         run_test_cases(

src/migration.rs

@@ -30,6 +30,19 @@ impl Migrate for Vec<Statement> {
                             _ => false,
                         })
                         .map_or(Some(orig), |sb| sa.migrate(sb)),
+                    Statement::CreateIndex(a) => other
+                        .iter()
+                        .find(|sb| match sb {
+                            Statement::Drop {
+                                object_type, names, ..
+                            } => {
+                                *object_type == ObjectType::Index
+                                    && names.len() == 1
+                                    && Some(&names[0]) == a.name.as_ref()
+                            }
+                            _ => false,
+                        })
+                        .map_or(Some(orig), |sb| sa.migrate(sb)),
                     Statement::CreateType { name, .. } => other
                         .iter()
                         .find(|sb| match sb {
@@ -56,9 +69,10 @@ impl Migrate for Vec<Statement> {
             })
             // CREATE table etc.
             .chain(other.iter().filter_map(|sb| match sb {
-                Statement::CreateTable(_) => Some(sb.clone()),
-                Statement::CreateType { .. } => Some(sb.clone()),
-                Statement::CreateExtension { .. } => Some(sb.clone()),
+                Statement::CreateTable(_)
+                | Statement::CreateIndex { .. }
+                | Statement::CreateType { .. }
+                | Statement::CreateExtension { .. } => Some(sb.clone()),
                 _ => None,
             }))
             .collect();
@@ -92,6 +106,20 @@ impl Migrate for Statement {
                 }
                 _ => todo!("handle migrating statement: {:?}", other),
             },
+            Self::CreateIndex(a) => match other {
+                Self::Drop {
+                    object_type, names, ..
+                } => {
+                    let name = a.name.clone().expect("can't migrate unnamed index: {a}");
+                    if *object_type == ObjectType::Index && names.contains(&name) {
+                        None
+                    } else {
+                        // DROP statement is for another index
+                        Some(Self::CreateIndex(a))
+                    }
+                }
+                _ => todo!("handle migrating statement: {:?}", other),
+            },
             Self::CreateType {
                 name,
                 representation,
@@ -127,7 +155,7 @@ impl Migrate for Statement {
                 }
                 _ => todo!("handle migrating statement: {:?}", other),
             },
-            _ => todo!("handle migrating statement: {:?}", other),
+            _ => todo!("handle migrating statement: {self:?} -> {other:?}"),
         }
     }
 }