support partial index (#478)

kyoto7250 · tyt2y3 · web-flow · commit ee234f143dec · 2024-09-28T09:48:00.000+01:00
* support partial index

* add tests for partial index

* update examples for partial index

* import once in src/index/create.rs

* fix an output in examples/postgres/Readme.md

* cargo fmt for  examples

* use bigdecimal:0.4.5

---------

Co-authored-by: Chris Tsang &lt;chris.2y3@outlook.com&gt;
diff --git a/examples/diesel_postgres/Readme.md b/examples/diesel_postgres/Readme.md
@@ -1,7 +1,7 @@
 # SeaQuery Diesel Postgres example
 
 > WARN: If you enable `with-bigdecimal`, you HAVE to update the version used by default by `diesel`
-> otherwise it will fail to build. Use `cargo update -p bigdecimal:0.4.2 --precise 0.3.1`.
+> otherwise it will fail to build. Use `cargo update -p bigdecimal:0.4.5 --precise 0.3.1`.
 
 Running:
 
diff --git a/examples/diesel_sqlite/src/main.rs b/examples/diesel_sqlite/src/main.rs
@@ -5,7 +5,10 @@ use diesel::deserialize::{self, FromSql};
 use diesel::sql_types::BigInt;
 use diesel::sql_types::{Blob, Text};
 use diesel::{Connection, QueryableByName, RunQueryDsl, SqliteConnection};
-use sea_query::{Alias, ColumnDef, Expr, Func, Iden, Order, Query, SqliteQueryBuilder, Table};
+use sea_query::{
+    Alias, ColumnDef, ConditionalStatement, Expr, Func, Iden, Index, Order, Query,
+    SqliteQueryBuilder, Table,
+};
 use sea_query_diesel::DieselBinder;
 use serde_json::json;
 use time::macros::{date, time};
@@ -40,6 +43,13 @@ fn main() {
             .col(ColumnDef::new(Character::Meta).json().not_null())
             .col(ColumnDef::new(Character::Created).date_time())
             .build(SqliteQueryBuilder),
+        Index::create()
+            .name("partial_index_small_font")
+            .if_not_exists()
+            .table(Character::Table)
+            .col(Character::FontSize)
+            .and_where(Expr::col(Character::FontSize).lt(11).not())
+            .build(SqliteQueryBuilder),
     ]
     .join("; ");
 
diff --git a/examples/postgres/Readme.md b/examples/postgres/Readme.md
@@ -7,8 +7,8 @@ cargo run
 
 Example output:
 ```
-DROP TABLE IF EXISTS "document"; CREATE TABLE IF NOT EXISTS "document" ( "id" serial NOT NULL PRIMARY KEY, "json_field" jsonb, "timestamp" timestamp )
-Create table document: ()
+DROP TABLE IF EXISTS "document"; CREATE TABLE IF NOT EXISTS "document" ( "id" serial NOT NULL PRIMARY KEY, "uuid" uuid, "json_field" jsonb, "timestamp" timestamp, "timestamp_with_time_zone" timestamp with time zone, "decimal" decimal, "array" integer[] ); CREATE INDEX "partial_index_small_decimal" ON "document" ("decimal") WHERE NOT "decimal" < 11
+Create table document: Ok(())
 
 Insert into document: Ok(1)
 
diff --git a/examples/postgres/src/main.rs b/examples/postgres/src/main.rs
@@ -2,7 +2,10 @@ use chrono::{DateTime, FixedOffset, NaiveDate, NaiveDateTime};
 
 use postgres::{Client, NoTls, Row};
 use rust_decimal::Decimal;
-use sea_query::{ColumnDef, ColumnType, Iden, Order, PostgresQueryBuilder, Query, Table};
+use sea_query::{
+    ColumnDef, ColumnType, ConditionalStatement, Expr, Iden, Index, Order, PostgresQueryBuilder,
+    Query, Table,
+};
 use sea_query_postgres::PostgresBinder;
 use time::{
     macros::{date, offset, time},
@@ -37,6 +40,12 @@ fn main() {
             .col(ColumnDef::new(Document::Decimal).decimal())
             .col(ColumnDef::new(Document::Array).array(ColumnType::Integer))
             .build(PostgresQueryBuilder),
+        Index::create()
+            .name("partial_index_small_decimal")
+            .table(Document::Table)
+            .col(Document::Decimal)
+            .and_where(Expr::col(Document::Decimal).lt(11).not())
+            .build(PostgresQueryBuilder),
     ]
     .join("; ");
 
diff --git a/examples/rusqlite/src/main.rs b/examples/rusqlite/src/main.rs
@@ -1,7 +1,9 @@
 use chrono::{NaiveDate, NaiveDateTime};
 use rusqlite::{Connection, Result, Row};
-use sea_query::{ColumnDef, Expr, Func, Iden, Order, Query, SqliteQueryBuilder, Table};
-
+use sea_query::{
+    ColumnDef, ConditionalStatement, Expr, Func, Iden, Index, Order, Query, SqliteQueryBuilder,
+    Table,
+};
 use sea_query_rusqlite::RusqliteBinder;
 use serde_json::{json, Value as Json};
 use time::{
@@ -37,6 +39,12 @@ fn main() -> Result<()> {
             .col(ColumnDef::new(Character::Meta).json())
             .col(ColumnDef::new(Character::Created).date_time())
             .build(SqliteQueryBuilder),
+        Index::create()
+            .name("partial_index_small_font")
+            .table(Character::Table)
+            .col(Character::FontSize)
+            .and_where(Expr::col(Character::FontSize).lt(11).not())
+            .build(SqliteQueryBuilder),
     ]
     .join("; ");
 
diff --git a/examples/sqlx_postgres/Readme.md b/examples/sqlx_postgres/Readme.md
@@ -9,6 +9,8 @@ Example output:
 ```
 Create table character: Ok(PgQueryResult { rows_affected: 0 })
 
+Create partial index: Ok(PgQueryResult { rows_affected: 0 })
+
 Insert into character: Ok(PgQueryResult { rows_affected: 1 })
 
 Select one from character:
diff --git a/examples/sqlx_postgres/src/main.rs b/examples/sqlx_postgres/src/main.rs
@@ -2,7 +2,8 @@ use bigdecimal::{BigDecimal, FromPrimitive};
 use chrono::{NaiveDate, NaiveDateTime};
 use rust_decimal::Decimal;
 use sea_query::{
-    ColumnDef, Expr, Func, Iden, OnConflict, Order, PostgresQueryBuilder, Query, Table,
+    ColumnDef, ConditionalStatement, Expr, Func, Iden, Index, OnConflict, Order,
+    PostgresQueryBuilder, Query, Table,
 };
 use sea_query_binder::SqlxBinder;
 use sqlx::{PgPool, Row};
@@ -50,6 +51,18 @@ async fn main() {
     let result = sqlx::query(&sql).execute(&mut *pool).await;
     println!("Create table character: {result:?}\n");
 
+    // Partial Index
+    let partial_index = Index::create()
+        .if_not_exists()
+        .name("partial_index_small_font")
+        .table(Character::Table)
+        .col(Character::FontSize)
+        .and_where(Expr::col(Character::FontSize).lt(11).not())
+        .build(PostgresQueryBuilder);
+
+    let index = sqlx::query(&partial_index).execute(&mut *pool).await;
+    println!("Create partial index: {index:?}\n");
+
     // Create
 
     let (sql, values) = Query::insert()
diff --git a/examples/sqlx_sqlite/Readme.md b/examples/sqlx_sqlite/Readme.md
@@ -9,6 +9,8 @@ Example output:
 ```
 Create table character: Ok(SqliteQueryResult { changes: 0, last_insert_rowid: 0 })
 
+Create partial index: Ok(SqliteQueryResult { changes: 0, last_insert_rowid: 0 })
+
 Insert into character: last_insert_id = 1
 
 Select one from character:
diff --git a/examples/sqlx_sqlite/src/main.rs b/examples/sqlx_sqlite/src/main.rs
@@ -1,5 +1,8 @@
 use chrono::{NaiveDate, NaiveDateTime};
-use sea_query::{ColumnDef, Expr, Func, Iden, OnConflict, Order, Query, SqliteQueryBuilder, Table};
+use sea_query::{
+    ColumnDef, ConditionalStatement, Expr, Func, Iden, Index, OnConflict, Order, Query,
+    SqliteQueryBuilder, Table,
+};
 use sea_query_binder::SqlxBinder;
 use serde_json::{json, Value as Json};
 use sqlx::{Row, SqlitePool};
@@ -34,6 +37,17 @@ async fn main() {
     let result = sqlx::query(&sql).execute(&pool).await;
     println!("Create table character: {result:?}\n");
 
+    // Partial Index
+    let partial_index = Index::create()
+        .name("partial_index_small_font")
+        .table(Character::Table)
+        .col(Character::FontSize)
+        .and_where(Expr::col(Character::FontSize).lt(11).not())
+        .build(SqliteQueryBuilder);
+
+    let index = sqlx::query(&partial_index).execute(&pool).await;
+    println!("Create partial index: {index:?}\n");
+
     // Create
     let (sql, values) = Query::insert()
         .into_table(Character::Table)
diff --git a/src/backend/index_builder.rs b/src/backend/index_builder.rs
@@ -23,6 +23,8 @@ pub trait IndexBuilder: QuotedBuilder + TableRefBuilder {
         self.prepare_index_prefix(create, sql);
 
         self.prepare_index_columns(&create.index.columns, sql);
+
+        self.prepare_filter(&create.r#where, sql);
     }
 
     /// Translate [`IndexCreateStatement`] into SQL statement.
@@ -74,4 +76,8 @@ pub trait IndexBuilder: QuotedBuilder + TableRefBuilder {
         });
         write!(sql, ")").unwrap();
     }
+
+    #[doc(hidden)]
+    // Write WHERE clause for partial index. This function is not available in MySQL.
+    fn prepare_filter(&self, _condition: &ConditionHolder, _sql: &mut dyn SqlWriter) {}
 }
diff --git a/src/backend/postgres/index.rs b/src/backend/postgres/index.rs
@@ -64,6 +64,7 @@ impl IndexBuilder for PostgresQueryBuilder {
         if create.nulls_not_distinct {
             write!(sql, " NULLS NOT DISTINCT").unwrap();
         }
+        self.prepare_filter(&create.r#where, sql);
     }
 
     fn prepare_table_ref_index_stmt(&self, table_ref: &TableRef, sql: &mut dyn SqlWriter) {
@@ -128,4 +129,8 @@ impl IndexBuilder for PostgresQueryBuilder {
             write!(sql, "UNIQUE ").unwrap();
         }
     }
+
+    fn prepare_filter(&self, condition: &ConditionHolder, sql: &mut dyn SqlWriter) {
+        self.prepare_condition(condition, "WHERE", sql);
+    }
 }
diff --git a/src/backend/sqlite/index.rs b/src/backend/sqlite/index.rs
@@ -32,6 +32,7 @@ impl IndexBuilder for SqliteQueryBuilder {
 
         write!(sql, " ").unwrap();
         self.prepare_index_columns(&create.index.columns, sql);
+        self.prepare_filter(&create.r#where, sql);
     }
 
     fn prepare_table_ref_index_stmt(&self, table_ref: &TableRef, sql: &mut dyn SqlWriter) {
@@ -69,4 +70,8 @@ impl IndexBuilder for SqliteQueryBuilder {
     }
 
     fn write_column_index_prefix(&self, _col_prefix: &Option<u32>, _sql: &mut dyn SqlWriter) {}
+
+    fn prepare_filter(&self, condition: &ConditionHolder, sql: &mut dyn SqlWriter) {
+        self.prepare_condition(condition, "WHERE", sql);
+    }
 }
diff --git a/src/index/create.rs b/src/index/create.rs
@@ -1,6 +1,7 @@
 use inherent::inherent;
 
 use crate::{backend::SchemaBuilder, types::*, SchemaStatementBuilder};
+use crate::{ConditionHolder, ConditionalStatement, IntoCondition};
 
 use super::common::*;
 
@@ -148,6 +149,27 @@ use super::common::*;
 ///     r#"CREATE INDEX "idx-glyph-aspect" ON "glyph" ("aspect" ASC)"#
 /// );
 /// ```
+///
+/// Partial Index with prefix and order
+/// ```
+/// use sea_query::{tests_cfg::*, *};
+///
+/// let index = Index::create()
+///     .name("idx-glyph-aspect")
+///     .table(Glyph::Table)
+///     .col((Glyph::Aspect, 64, IndexOrder::Asc))
+///     .and_where(Expr::col((Glyph::Table, Glyph::Aspect)).is_in(vec![3, 4]))
+///     .to_owned();
+///
+/// assert_eq!(
+///     index.to_string(PostgresQueryBuilder),
+///     r#"CREATE INDEX "idx-glyph-aspect" ON "glyph" ("aspect" (64) ASC) WHERE "glyph"."aspect" IN (3, 4)"#
+/// );
+/// assert_eq!(
+///     index.to_string(SqliteQueryBuilder),
+///     r#"CREATE INDEX "idx-glyph-aspect" ON "glyph" ("aspect" ASC) WHERE "glyph"."aspect" IN (3, 4)"#
+/// );
+/// ```
 #[derive(Default, Debug, Clone)]
 pub struct IndexCreateStatement {
     pub(crate) table: Option<TableRef>,
@@ -157,6 +179,7 @@ pub struct IndexCreateStatement {
     pub(crate) nulls_not_distinct: bool,
     pub(crate) index_type: Option<IndexType>,
     pub(crate) if_not_exists: bool,
+    pub(crate) r#where: ConditionHolder,
 }
 
 /// Specification of a table index
@@ -171,7 +194,16 @@ pub enum IndexType {
 impl IndexCreateStatement {
     /// Construct a new [`IndexCreateStatement`]
     pub fn new() -> Self {
-        Self::default()
+        Self {
+            table: None,
+            index: Default::default(),
+            primary: false,
+            unique: false,
+            nulls_not_distinct: false,
+            index_type: None,
+            if_not_exists: false,
+            r#where: ConditionHolder::new(),
+        }
     }
 
     /// Create index if index not exists
@@ -263,6 +295,7 @@ impl IndexCreateStatement {
             nulls_not_distinct: self.nulls_not_distinct,
             index_type: self.index_type.take(),
             if_not_exists: self.if_not_exists,
+            r#where: self.r#where.clone(),
         }
     }
 }
@@ -283,3 +316,18 @@ impl SchemaStatementBuilder for IndexCreateStatement {
 
     pub fn to_string<T: SchemaBuilder>(&self, schema_builder: T) -> String;
 }
+
+impl ConditionalStatement for IndexCreateStatement {
+    fn and_or_where(&mut self, condition: LogicalChainOper) -> &mut Self {
+        self.r#where.add_and_or(condition);
+        self
+    }
+
+    fn cond_where<C>(&mut self, condition: C) -> &mut Self
+    where
+        C: IntoCondition,
+    {
+        self.r#where.add_condition(condition.into_condition());
+        self
+    }
+}
diff --git a/tests/postgres/index.rs b/tests/postgres/index.rs
@@ -83,6 +83,21 @@ fn create_6() {
     );
 }
 
+#[test]
+fn create_7() {
+    assert_eq!(
+        Index::create()
+            .unique()
+            .nulls_not_distinct()
+            .name("partial-index-glyph-image-not-null")
+            .table(Glyph::Table)
+            .col(Glyph::Image)
+            .and_where(Expr::col(Glyph::Image).is_not_null())
+            .to_string(PostgresQueryBuilder),
+        r#"CREATE UNIQUE INDEX "partial-index-glyph-image-not-null" ON "glyph" ("image") NULLS NOT DISTINCT WHERE "image" IS NOT NULL"#
+    );
+}
+
 #[test]
 fn drop_1() {
     assert_eq!(
diff --git a/tests/sqlite/index.rs b/tests/sqlite/index.rs
@@ -42,6 +42,21 @@ fn create_3() {
     );
 }
 
+#[test]
+fn create_4() {
+    assert_eq!(
+        Index::create()
+            .if_not_exists()
+            .unique()
+            .name("partial-index-glyph-image-not-null")
+            .table(Glyph::Table)
+            .col(Glyph::Image)
+            .and_where(Expr::col(Glyph::Image).is_not_null())
+            .to_string(SqliteQueryBuilder),
+        r#"CREATE UNIQUE INDEX IF NOT EXISTS "partial-index-glyph-image-not-null" ON "glyph" ("image") WHERE "image" IS NOT NULL"#
+    );
+}
+
 #[test]
 fn drop_1() {
     assert_eq!(

Original file line number	Diff line number	Diff line change
`@@ -23,6 +23,8 @@ pub trait IndexBuilder: QuotedBuilder + TableRefBuilder {`
`23`	`23`	`self.prepare_index_prefix(create, sql);`
`24`	`24`
`25`	`25`	`self.prepare_index_columns(&create.index.columns, sql);`
	`26`	`+`
	`27`	`+ self.prepare_filter(&create.r#where, sql);`
`26`	`28`	`}`
`27`	`29`
`28`	`30`	/// Translate [`IndexCreateStatement`] into SQL statement.
`@@ -74,4 +76,8 @@ pub trait IndexBuilder: QuotedBuilder + TableRefBuilder {`
`74`	`76`	`});`
`75`	`77`	`write!(sql, ")").unwrap();`
`76`	`78`	`}`
	`79`	`+`
	`80`	`+ #[doc(hidden)]`
	`81`	`+ // Write WHERE clause for partial index. This function is not available in MySQL.`
	`82`	`+ fn prepare_filter(&self, _condition: &ConditionHolder, _sql: &mut dyn SqlWriter) {}`
`77`	`83`	`}`