Support for Partial Indexes (Filtered Indexes)

Cargo.toml

@@ -115,7 +115,7 @@ smallvec = "1"
 strip-ansi-escapes = "0.1"
 structopt = "0.3"
 sqlformat = "0.2"
-sqlparser = "0.60"
+sqlparser = { version = "0.61", features = ["visitor"] }
 sqlx-core = "0.8"
 sqlx-sqlite = "0.8"
 tempfile = "3"

libs/sql-ddl/src/postgres.rs

@@ -347,6 +347,7 @@ pub struct CreateIndex<'a> {
     pub table_reference: &'a dyn Display,
     pub columns: Vec<IndexColumn<'a>>,
     pub using: Option<IndexAlgorithm>,
+    pub where_clause: Option<Cow<'a, str>>,
 }
 
 impl Display for CreateIndex<'_> {
@@ -388,7 +389,13 @@ impl Display for CreateIndex<'_> {
             })
             .join(", ", f)?;
 
-        f.write_str(")")
+        f.write_str(")")?;
+
+        if let Some(predicate) = &self.where_clause {
+            write!(f, " WHERE {}", predicate)?;
+        }
+
+        Ok(())
     }
 }
 
@@ -432,6 +439,7 @@ mod tests {
             table_reference: &PostgresIdentifier::Simple(Cow::Borrowed("Cat")),
             columns,
             using: None,
+            where_clause: None,
         };
 
         assert_eq!(
@@ -450,6 +458,7 @@ mod tests {
             table_reference: &PostgresIdentifier::Simple(Cow::Borrowed("Cat")),
             columns,
             using: Some(IndexAlgorithm::Hash),
+            where_clause: None,
         };
 
         assert_eq!(
@@ -479,6 +488,7 @@ mod tests {
             table_reference: &PostgresIdentifier::Simple("Cat".into()),
             columns,
             using: None,
+            where_clause: None,
         };
 
         assert_eq!(
@@ -487,6 +497,25 @@ mod tests {
         )
     }
 
+    #[test]
+    fn create_partial_unique_index() {
+        let columns = vec![IndexColumn::new("name")];
+
+        let create_index = CreateIndex {
+            is_unique: true,
+            index_name: "meow_idx".into(),
+            table_reference: &PostgresIdentifier::Simple(Cow::Borrowed("Cat")),
+            columns,
+            using: None,
+            where_clause: Some("status = 'active'".into()),
+        };
+
+        assert_eq!(
+            create_index.to_string(),
+            "CREATE UNIQUE INDEX \"meow_idx\" ON \"Cat\"(\"name\") WHERE status = 'active'"
+        )
+    }
+
     #[test]
     fn full_alter_table_add_foreign_key() {
         let alter_table = AlterTable {

libs/sql-ddl/src/sqlite.rs

@@ -148,6 +148,53 @@ impl Display for Column<'_> {
     }
 }
 
+/// A column in an index definition.
+#[derive(Debug, Default)]
+pub struct IndexColumn<'a> {
+    pub name: Cow<'a, str>,
+    pub sort_order: Option<crate::SortOrder>,
+}
+
+/// Create an index statement.
+#[derive(Debug)]
+pub struct CreateIndex<'a> {
+    pub index_name: Cow<'a, str>,
+    pub table_name: Cow<'a, str>,
+    pub columns: Vec<IndexColumn<'a>>,
+    pub is_unique: bool,
+    pub where_clause: Option<Cow<'a, str>>,
+}
+
+impl Display for CreateIndex<'_> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(
+            f,
+            "CREATE {uniqueness}INDEX \"{index_name}\" ON \"{table_name}\"(",
+            uniqueness = if self.is_unique { "UNIQUE " } else { "" },
+            index_name = self.index_name,
+            table_name = self.table_name,
+        )?;
+
+        for (i, c) in self.columns.iter().enumerate() {
+            if i > 0 {
+                f.write_str(", ")?;
+            }
+            write!(f, "\"{}\"", c.name)?;
+            if let Some(sort_order) = c.sort_order {
+                write!(f, " {}", sort_order.as_ref())?;
+            }
+        }
+
+        f.write_str(")")?;
+
+        if let Some(predicate) = &self.where_clause {
+            write!(f, " WHERE {}", predicate)?;
+        }
+
+        Ok(())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -270,4 +317,42 @@ mod tests {
 
         assert_eq!(create_table.to_string(), expected.trim_matches('\n'))
     }
+
+    #[test]
+    fn create_unique_index() {
+        let create_index = CreateIndex {
+            index_name: "idx_name".into(),
+            table_name: "Cat".into(),
+            columns: vec![IndexColumn {
+                name: "name".into(),
+                sort_order: None,
+            }],
+            is_unique: true,
+            where_clause: None,
+        };
+
+        assert_eq!(
+            create_index.to_string(),
+            r#"CREATE UNIQUE INDEX "idx_name" ON "Cat"("name")"#
+        )
+    }
+
+    #[test]
+    fn create_partial_unique_index() {
+        let create_index = CreateIndex {
+            index_name: "idx_name".into(),
+            table_name: "Cat".into(),
+            columns: vec![IndexColumn {
+                name: "name".into(),
+                sort_order: None,
+            }],
+            is_unique: true,
+            where_clause: Some("status = 'active'".into()),
+        };
+
+        assert_eq!(
+            create_index.to_string(),
+            r#"CREATE UNIQUE INDEX "idx_name" ON "Cat"("name") WHERE status = 'active'"#
+        )
+    }
 }