Repository: SeaQL/sea-query (language: rust)

Issue:
Support for Postgres BTree index `INCLUDE (<col>, ...)`
<!--

Welcome! Thanks for suggesting features!

Do you want to ask a question? Are you looking for support?
Please ask us on
- Discord: https://discord.com/invite/uCPdDXzbdv
- GitHub Discussions: https://github.com/SeaQL/sea-query/discussions/new

Make sure you have a clear feature specification before open an issue. Alternatively, please start an "Idea" thread on GitHub Discussions and let's formulate the solution together.

-->

## Motivation

Postgres BTree indexex allow for an `INCLUDE` statement where the included columns values are stored in the leaf nodes of the index. This is an optimization to store data with the index for columns that are fetched commonly but, not necessarily filtered on.

## Proposed Solutions

```rust
Index::create()
        .name(name)
        .index_type(IndexType::BTree)
        .table(table)
        .col(col)
        .include(other_cols) // postgres feature gate
        .is_unique_key();
```

```sql
CREATE UNIQUE INDEX an_index ON table USING BTREE (col) INCLUDE (other_cols)
```

## Additional Information

I am unsure if this is supported on other databases.

[CREATE INDEX](https://www.postgresql.org/docs/current/sql-createindex.html)





Discussion / hints:

Fix patch:
diff --git a/src/backend/index_builder.rs b/src/backend/index_builder.rs
index 27389131b..e64864640 100644
--- a/src/backend/index_builder.rs
+++ b/src/backend/index_builder.rs
@@ -80,4 +80,8 @@ pub trait IndexBuilder: QuotedBuilder + TableRefBuilder {
     #[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) {}
+
+    #[doc(hidden)]
+    // Write INCLUDE clause for index. This function is only available in PostgreSQL.
+    fn prepare_include_columns(&self, _columns: &[SeaRc<dyn Iden>], _sql: &mut dyn SqlWriter) {}
 }
diff --git a/src/backend/postgres/index.rs b/src/backend/postgres/index.rs
index 0ff5b0635..be082b49f 100644
--- a/src/backend/postgres/index.rs
+++ b/src/backend/postgres/index.rs
@@ -26,6 +26,11 @@ impl IndexBuilder for PostgresQueryBuilder {
         }
 
         self.prepare_index_columns(&create.index.columns, sql);
+
+        if !create.include_columns.is_empty() {
+            write!(sql, " ").unwrap();
+            self.prepare_include_columns(&create.include_columns, sql);
+        }
     }
 
     fn prepare_index_create_statement(
@@ -61,6 +66,11 @@ impl IndexBuilder for PostgresQueryBuilder {
         write!(sql, " ").unwrap();
         self.prepare_index_columns(&create.index.columns, sql);
 
+        if !create.include_columns.is_empty() {
+            write!(sql, " ").unwrap();
+            self.prepare_include_columns(&create.include_columns, sql);
+        }
+
         if create.nulls_not_distinct {
             write!(sql, " NULLS NOT DISTINCT").unwrap();
         }
@@ -133,4 +143,16 @@ impl IndexBuilder for PostgresQueryBuilder {
     fn prepare_filter(&self, condition: &ConditionHolder, sql: &mut dyn SqlWriter) {
         self.prepare_condition(condition, "WHERE", sql);
     }
+
+    fn prepare_include_columns(&self, columns: &[SeaRc<dyn Iden>], sql: &mut dyn SqlWriter) {
+        write!(sql, "INCLUDE (").unwrap();
+        columns.iter().fold(true, |first, col| {
+            if !first {
+                write!(sql, ", ").unwrap();
+            }
+            col.prepare(sql.as_writer(), self.quote());
+            false
+        });
+        write!(sql, ")").unwrap();
+    }
 }
diff --git a/src/index/create.rs b/src/index/create.rs
index 714325ffb..3f95fc130 100644
--- a/src/index/create.rs
+++ b/src/index/create.rs
@@ -170,6 +170,23 @@ use super::common::*;
 ///     r#"CREATE INDEX "idx-glyph-aspect" ON "glyph" ("aspect" ASC) WHERE "glyph"."aspect" IN (3, 4)"#
 /// );
 /// ```
+///
+/// Index include non-key columns
+/// ```
+/// use sea_query::{tests_cfg::*, *};
+///
+/// let index = Index::create()
+///     .name("idx-font-name-include-language")
+///     .table(Font::Table)
+///     .col(Font::Name)
+///     .include(Font::Language)
+///     .to_owned();
+///
+/// assert_eq!(
+///     index.to_string(PostgresQueryBuilder),
+///     r#"CREATE INDEX "idx-font-name-include-language" ON "font" ("name") INCLUDE ("language")"#
+/// )
+/// ```
 #[derive(Default, Debug, Clone)]
 pub struct IndexCreateStatement {
     pub(crate) table: Option<TableRef>,
@@ -180,6 +197,7 @@ pub struct IndexCreateStatement {
     pub(crate) index_type: Option<IndexType>,
     pub(crate) if_not_exists: bool,
     pub(crate) r#where: ConditionHolder,
+    pub(crate) include_columns: Vec<DynIden>,
 }
 
 /// Specification of a table index
@@ -203,6 +221,7 @@ impl IndexCreateStatement {
             index_type: None,
             if_not_exists: false,
             r#where: ConditionHolder::new(),
+            include_columns: vec![],
         }
     }
 
@@ -270,6 +289,14 @@ impl IndexCreateStatement {
         self
     }
 
+    pub fn include<C>(&mut self, col: C) -> &mut Self
+    where
+        C: IntoIden,
+    {
+        self.include_columns.push(col.into_iden());
+        self
+    }
+
     pub fn is_primary_key(&self) -> bool {
         self.primary
     }
@@ -296,6 +323,7 @@ impl IndexCreateStatement {
             index_type: self.index_type.take(),
             if_not_exists: self.if_not_exists,
             r#where: self.r#where.clone(),
+            include_columns: self.include_columns.clone(),
         }
     }
 }


Test patch:
diff --git a/tests/postgres/index.rs b/tests/postgres/index.rs
index d51eb1905..f62fe7c1b 100644
--- a/tests/postgres/index.rs
+++ b/tests/postgres/index.rs
@@ -98,6 +98,22 @@ fn create_7() {
     );
 }
 
+#[test]
+fn create_8() {
+    assert_eq!(
+        Index::create()
+            .name("idx-font-name-include-id-language")
+            .table(Font::Table)
+            .col(Font::Name)
+            .include(Font::Id)
+            .include(Font::Language)
+            .unique()
+            .nulls_not_distinct()
+            .to_string(PostgresQueryBuilder),
+        r#"CREATE UNIQUE INDEX "idx-font-name-include-id-language" ON "font" ("name") INCLUDE ("id", "language") NULLS NOT DISTINCT"#
+    );
+}
+
 #[test]
 fn drop_1() {
     assert_eq!(
diff --git a/tests/postgres/table.rs b/tests/postgres/table.rs
index d26fa2c3f..8eb2591b8 100644
--- a/tests/postgres/table.rs
+++ b/tests/postgres/table.rs
@@ -611,3 +611,40 @@ fn create_16() {
         .join(" ")
     );
 }
+
+#[test]
+fn create_17() {
+    assert_eq!(
+        Table::create()
+        .table(Font::Table)
+        .col(
+            ColumnDef::new(Font::Id)
+                .integer()
+                .not_null()
+                .primary_key()
+                .auto_increment(),
+        )
+        .col(ColumnDef::new(Font::Name).string())
+        .col(ColumnDef::new(Font::Variant).string_len(255).not_null())
+        .col(ColumnDef::new(Font::Language).string_len(255).not_null())
+        .index(
+            Index::create()
+                .name("idx-font-name-include-language")
+                .unique()
+                .nulls_not_distinct()
+                .col(Font::Name)
+                .include(Font::Language),
+        )
+        .to_string(PostgresQueryBuilder),
+        [
+            r#"CREATE TABLE "font" ("#,
+            r#""id" serial NOT NULL PRIMARY KEY,"#,
+            r#""name" varchar,"#,
+            r#""variant" varchar(255) NOT NULL,"#,
+            r#""language" varchar(255) NOT NULL,"#,
+            r#"CONSTRAINT "idx-font-name-include-language" UNIQUE NULLS NOT DISTINCT ("name") INCLUDE ("language")"#,
+            r#")"#,
+        ]
+        .join(" ")
+    );
+}