From 8f8d8b23a05c357fb85fb905221c03f3f75e8f23 Mon Sep 17 00:00:00 2001
From: taozhi8833998 <taozhi8833998@163.com>
Date: Fri, 22 Nov 2024 14:25:33 +0800
Subject: [PATCH] feat: support include columns in create index for pg

---
 pegjs/postgresql.pegjs | 13 +++++++++++++
 src/create.js          |  2 +-
 test/postgres.spec.js  |  7 +++++++
 3 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/pegjs/postgresql.pegjs b/pegjs/postgresql.pegjs
index 3e2c16b3..c3543f37 100644
--- a/pegjs/postgresql.pegjs
+++ b/pegjs/postgresql.pegjs
@@ -1057,6 +1057,16 @@ create_sequence_definition_list
     return createList(head, tail, 1)
 }
 
+include_column
+  = k:'INCLUDE'i __ LPAREN __ c:column_list __ RPAREN {
+    // => { type: 'include', keyword: 'include', columns: column_list }
+    return {
+      type: k.toLowerCase(),
+      keyword: k.toLowerCase(),
+      columns:c,
+    }
+  }
+  
 create_index_stmt
   = a:KW_CREATE __
   kw:KW_UNIQUE? __
@@ -1068,6 +1078,7 @@ create_index_stmt
   ta:table_name __
   um:index_type? __
   LPAREN __ cols:column_order_list __ RPAREN __
+  include:include_column? __
   wr:(KW_WITH __ LPAREN __ index_options_list __ RPAREN)? __
   ts:(KW_TABLESPACE __ ident_name)? __
   w:where_clause? __ {
@@ -1083,6 +1094,7 @@ create_index_stmt
       table: table_name;
       index_using?: index_type;
       index_columns: column_order[];
+      include?: column_list_items;
       with?: index_option[];
       with_before_where: true;
       tablespace?: {type: 'origin'; value: string; }
@@ -1104,6 +1116,7 @@ create_index_stmt
           table: ta,
           index_using: um,
           index_columns: cols,
+          include,
           with: wr && wr[4],
           with_before_where: true,
           tablespace: ts && { type: 'origin', value: ts[2] },
diff --git a/src/create.js b/src/create.js
index c3665a31..bed93ec1 100644
--- a/src/create.js
+++ b/src/create.js
@@ -182,7 +182,7 @@ function createIndexToSQL(stmt) {
     with: withExpr, with_before_where: withBeforeWhere,
   } = stmt
   const withIndexOpt = withExpr && `WITH (${indexOptionListToSQL(withExpr).join(', ')})`
-  const includeColumns = include && `${toUpper(include.keyword)} (${include.columns.map(col => identifierToSql(col)).join(', ')})`
+  const includeColumns = include && `${toUpper(include.keyword)} (${include.columns.map(col => (typeof col === 'string' ? identifierToSql(col) : exprToSQL(col))).join(', ')})`
   let indexName = index
   if (index) {
     indexName = typeof index === 'string' ? identifierToSql(index) : [identifierToSql(index.schema), identifierToSql(index.name)].filter(hasVal).join('.')
diff --git a/test/postgres.spec.js b/test/postgres.spec.js
index cced53ed..35f67ec5 100644
--- a/test/postgres.spec.js
+++ b/test/postgres.spec.js
@@ -1693,6 +1693,13 @@ describe('Postgres', () => {
         'CREATE UNIQUE INDEX IF NOT EXISTS "public_i_locations_pkey" ON "public"."i_locations" (id)'
       ]
     },
+    {
+      title: 'create index with include clause',
+      sql: [
+        'CREATE INDEX ON tableName (supplier, amount) INCLUDE(id);',
+        'CREATE INDEX ON "tableName" (supplier, amount) INCLUDE (id)'
+      ]
+    },
   ]
   function neatlyNestTestedSQL(sqlList){
     sqlList.forEach(sqlInfo => {