implement (#826)

SeaQL · Feb 16, 2025 · 8c40438 · 8c40438
1 parent e03e987
commit 8c40438
Show file tree

Hide file tree

Showing 5 changed files with 107 additions and 0 deletions.
diff --git 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
@@ -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
@@ -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(),
         }
     }
 }

diff --git 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
@@ -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(" ")
+    );
+}