From 4f2502df540c07499a0c5ceef2662fb073a34268 Mon Sep 17 00:00:00 2001 From: tianzhou Date: Tue, 7 Apr 2026 02:48:58 -0700 Subject: [PATCH] fix: preserve INCLUDE columns on CREATE INDEX (#385) The index inspector was using indnatts (total columns) for both key and include columns, causing INCLUDE columns to be merged into regular indexed columns. Now uses indnkeyatts to extract only key columns, and separately extracts include columns from positions indnkeyatts+1 through indnatts. Co-Authored-By: Claude Opus 4.6 (1M context) --- internal/diff/index.go | 12 +++++++++++ internal/diff/table.go | 10 +++++++++ ir/inspector.go | 21 ++++++++++--------- ir/ir.go | 11 +++++----- ir/queries/queries.sql | 17 ++++++++++----- ir/queries/queries.sql.go | 21 ++++++++++++++----- testdata/diff/create_index/add_index/diff.sql | 2 ++ testdata/diff/create_index/add_index/new.sql | 2 ++ .../diff/create_index/add_index/plan.json | 6 ++++++ testdata/diff/create_index/add_index/plan.sql | 2 ++ testdata/diff/create_index/add_index/plan.txt | 3 +++ 11 files changed, 82 insertions(+), 25 deletions(-) diff --git a/internal/diff/index.go b/internal/diff/index.go index 330875b3..f50f28b6 100644 --- a/internal/diff/index.go +++ b/internal/diff/index.go @@ -121,6 +121,18 @@ func generateIndexSQLWithName(index *ir.Index, indexName string, targetSchema st } builder.WriteString(")") + // INCLUDE columns (non-key columns stored in the index) + if len(index.IncludeColumns) > 0 { + builder.WriteString(" INCLUDE (") + for i, col := range index.IncludeColumns { + if i > 0 { + builder.WriteString(", ") + } + builder.WriteString(col) + } + builder.WriteString(")") + } + // NULLS NOT DISTINCT for unique indexes (PostgreSQL 15+) if index.NullsNotDistinct && index.Type == ir.IndexTypeUnique { builder.WriteString(" NULLS NOT DISTINCT") diff --git a/internal/diff/table.go b/internal/diff/table.go index 76a1a182..89dfff01 100644 --- a/internal/diff/table.go +++ b/internal/diff/table.go @@ -1511,6 +1511,16 @@ func indexesStructurallyEqual(oldIndex, newIndex *ir.Index) bool { } } + // Compare INCLUDE columns + if len(oldIndex.IncludeColumns) != len(newIndex.IncludeColumns) { + return false + } + for i, oldCol := range oldIndex.IncludeColumns { + if oldCol != newIndex.IncludeColumns[i] { + return false + } + } + return true } diff --git a/ir/inspector.go b/ir/inspector.go index ed39b26d..4542ca3e 100644 --- a/ir/inspector.go +++ b/ir/inspector.go @@ -746,16 +746,17 @@ func (i *Inspector) buildIndexes(ctx context.Context, schema *IR, targetSchema s } index := &Index{ - Schema: schemaName, - Table: tableName, - Name: indexName, - Type: indexType, - Method: method, - IsPartial: isPartial, - IsExpression: hasExpressions, - Where: "", - Comment: comment, - Columns: []*IndexColumn{}, + Schema: schemaName, + Table: tableName, + Name: indexName, + Type: indexType, + Method: method, + IncludeColumns: indexRow.IncludeColumns, + IsPartial: isPartial, + IsExpression: hasExpressions, + Where: "", + Comment: comment, + Columns: []*IndexColumn{}, } // Check for NULLS NOT DISTINCT (PostgreSQL 15+) diff --git a/ir/ir.go b/ir/ir.go index 9022bb0b..d6e0149c 100644 --- a/ir/ir.go +++ b/ir/ir.go @@ -251,11 +251,12 @@ type Index struct { Type IndexType `json:"type"` Method string `json:"method"` // btree, hash, gin, gist, etc. Columns []*IndexColumn `json:"columns"` - IsPartial bool `json:"is_partial"` // has a WHERE clause - IsExpression bool `json:"is_expression"` // functional/expression index - Where string `json:"where,omitempty"` // partial index condition - NullsNotDistinct bool `json:"nulls_not_distinct,omitempty"` // NULLS NOT DISTINCT (PG15+) - Comment string `json:"comment,omitempty"` + IncludeColumns []string `json:"include_columns,omitempty"` // INCLUDE columns (non-key) + IsPartial bool `json:"is_partial"` // has a WHERE clause + IsExpression bool `json:"is_expression"` // functional/expression index + Where string `json:"where,omitempty"` // partial index condition + NullsNotDistinct bool `json:"nulls_not_distinct,omitempty"` // NULLS NOT DISTINCT (PG15+) + Comment string `json:"comment,omitempty"` } // IndexColumn represents a column within an index diff --git a/ir/queries/queries.sql b/ir/queries/queries.sql index 78bcf0ce..3a699fa3 100644 --- a/ir/queries/queries.sql +++ b/ir/queries/queries.sql @@ -436,10 +436,11 @@ WITH index_base AS ( ELSE false END as has_expressions, COALESCE(d.description, '') AS index_comment, + idx.indnkeyatts as num_key_columns, idx.indnatts as num_columns, ARRAY( SELECT pg_get_indexdef(idx.indexrelid, k::int, true) - FROM generate_series(1, idx.indnatts) k + FROM generate_series(1, idx.indnkeyatts) k ) as column_definitions, ARRAY( SELECT @@ -447,16 +448,20 @@ WITH index_base AS ( WHEN (idx.indoption[k-1] & 1) = 1 THEN 'DESC' ELSE 'ASC' END - FROM generate_series(1, idx.indnatts) k + FROM generate_series(1, idx.indnkeyatts) k ) as column_directions, ARRAY( SELECT CASE WHEN opc.opcdefault THEN '' -- Omit default operator classes ELSE COALESCE(opc.opcname, '') END - FROM generate_series(1, idx.indnatts) k + FROM generate_series(1, idx.indnkeyatts) k LEFT JOIN pg_opclass opc ON opc.oid = idx.indclass[k-1] - ) as column_opclasses + ) as column_opclasses, + ARRAY( + SELECT pg_get_indexdef(idx.indexrelid, k::int, true) + FROM generate_series(idx.indnkeyatts + 1, idx.indnatts) k + ) as include_columns FROM pg_index idx JOIN pg_class i ON i.oid = idx.indexrelid JOIN pg_class t ON t.oid = idx.indrelid @@ -488,10 +493,12 @@ SELECT sp.partial_predicate, ib.has_expressions, ib.index_comment, + ib.num_key_columns, ib.num_columns, ib.column_definitions, ib.column_directions, - ib.column_opclasses + ib.column_opclasses, + ib.include_columns FROM index_base ib CROSS JOIN LATERAL ( SELECT diff --git a/ir/queries/queries.sql.go b/ir/queries/queries.sql.go index 28efd3e9..00617e03 100644 --- a/ir/queries/queries.sql.go +++ b/ir/queries/queries.sql.go @@ -1692,10 +1692,11 @@ WITH index_base AS ( ELSE false END as has_expressions, COALESCE(d.description, '') AS index_comment, + idx.indnkeyatts as num_key_columns, idx.indnatts as num_columns, ARRAY( SELECT pg_get_indexdef(idx.indexrelid, k::int, true) - FROM generate_series(1, idx.indnatts) k + FROM generate_series(1, idx.indnkeyatts) k ) as column_definitions, ARRAY( SELECT @@ -1703,16 +1704,20 @@ WITH index_base AS ( WHEN (idx.indoption[k-1] & 1) = 1 THEN 'DESC' ELSE 'ASC' END - FROM generate_series(1, idx.indnatts) k + FROM generate_series(1, idx.indnkeyatts) k ) as column_directions, ARRAY( SELECT CASE WHEN opc.opcdefault THEN '' -- Omit default operator classes ELSE COALESCE(opc.opcname, '') END - FROM generate_series(1, idx.indnatts) k + FROM generate_series(1, idx.indnkeyatts) k LEFT JOIN pg_opclass opc ON opc.oid = idx.indclass[k-1] - ) as column_opclasses + ) as column_opclasses, + ARRAY( + SELECT pg_get_indexdef(idx.indexrelid, k::int, true) + FROM generate_series(idx.indnkeyatts + 1, idx.indnatts) k + ) as include_columns FROM pg_index idx JOIN pg_class i ON i.oid = idx.indexrelid JOIN pg_class t ON t.oid = idx.indrelid @@ -1744,10 +1749,12 @@ SELECT sp.partial_predicate, ib.has_expressions, ib.index_comment, + ib.num_key_columns, ib.num_columns, ib.column_definitions, ib.column_directions, - ib.column_opclasses + ib.column_opclasses, + ib.include_columns FROM index_base ib CROSS JOIN LATERAL ( SELECT @@ -1772,10 +1779,12 @@ type GetIndexesForSchemaRow struct { PartialPredicate sql.NullString `db:"partial_predicate" json:"partial_predicate"` HasExpressions sql.NullBool `db:"has_expressions" json:"has_expressions"` IndexComment sql.NullString `db:"index_comment" json:"index_comment"` + NumKeyColumns int16 `db:"num_key_columns" json:"num_key_columns"` NumColumns int16 `db:"num_columns" json:"num_columns"` ColumnDefinitions []string `db:"column_definitions" json:"column_definitions"` ColumnDirections []string `db:"column_directions" json:"column_directions"` ColumnOpclasses []string `db:"column_opclasses" json:"column_opclasses"` + IncludeColumns []string `db:"include_columns" json:"include_columns"` } // GetIndexesForSchema retrieves all indexes for a specific schema @@ -1803,10 +1812,12 @@ func (q *Queries) GetIndexesForSchema(ctx context.Context, dollar_1 sql.NullStri &i.PartialPredicate, &i.HasExpressions, &i.IndexComment, + &i.NumKeyColumns, &i.NumColumns, pq.Array(&i.ColumnDefinitions), pq.Array(&i.ColumnDirections), pq.Array(&i.ColumnOpclasses), + pq.Array(&i.IncludeColumns), ); err != nil { return nil, err } diff --git a/testdata/diff/create_index/add_index/diff.sql b/testdata/diff/create_index/add_index/diff.sql index fbfb4851..a8e8a8c7 100644 --- a/testdata/diff/create_index/add_index/diff.sql +++ b/testdata/diff/create_index/add_index/diff.sql @@ -7,6 +7,8 @@ CREATE TABLE IF NOT EXISTS users ( CREATE INDEX IF NOT EXISTS idx_users_email ON users (email varchar_pattern_ops); +CREATE INDEX IF NOT EXISTS idx_users_email_include ON users (email) INCLUDE (name); + CREATE UNIQUE INDEX IF NOT EXISTS idx_users_email_unique ON users (email) NULLS NOT DISTINCT; CREATE INDEX IF NOT EXISTS idx_users_id ON users (id); diff --git a/testdata/diff/create_index/add_index/new.sql b/testdata/diff/create_index/add_index/new.sql index 466264db..4fd0e1bd 100644 --- a/testdata/diff/create_index/add_index/new.sql +++ b/testdata/diff/create_index/add_index/new.sql @@ -12,3 +12,5 @@ CREATE INDEX idx_users_id ON public.users (id); CREATE INDEX "public.idx_users" ON public.users (email, name); -- Test NULLS NOT DISTINCT (issue #355) CREATE UNIQUE INDEX idx_users_email_unique ON public.users (email) NULLS NOT DISTINCT; +-- Test INCLUDE columns (issue #385) +CREATE INDEX idx_users_email_include ON public.users (email) INCLUDE (name); diff --git a/testdata/diff/create_index/add_index/plan.json b/testdata/diff/create_index/add_index/plan.json index 05ead951..dd7946a9 100644 --- a/testdata/diff/create_index/add_index/plan.json +++ b/testdata/diff/create_index/add_index/plan.json @@ -20,6 +20,12 @@ "operation": "create", "path": "public.users.idx_users_email" }, + { + "sql": "CREATE INDEX IF NOT EXISTS idx_users_email_include ON users (email) INCLUDE (name);", + "type": "table.index", + "operation": "create", + "path": "public.users.idx_users_email_include" + }, { "sql": "CREATE UNIQUE INDEX IF NOT EXISTS idx_users_email_unique ON users (email) NULLS NOT DISTINCT;", "type": "table.index", diff --git a/testdata/diff/create_index/add_index/plan.sql b/testdata/diff/create_index/add_index/plan.sql index fbfb4851..a8e8a8c7 100644 --- a/testdata/diff/create_index/add_index/plan.sql +++ b/testdata/diff/create_index/add_index/plan.sql @@ -7,6 +7,8 @@ CREATE TABLE IF NOT EXISTS users ( CREATE INDEX IF NOT EXISTS idx_users_email ON users (email varchar_pattern_ops); +CREATE INDEX IF NOT EXISTS idx_users_email_include ON users (email) INCLUDE (name); + CREATE UNIQUE INDEX IF NOT EXISTS idx_users_email_unique ON users (email) NULLS NOT DISTINCT; CREATE INDEX IF NOT EXISTS idx_users_id ON users (id); diff --git a/testdata/diff/create_index/add_index/plan.txt b/testdata/diff/create_index/add_index/plan.txt index 39873c50..6fd5ef57 100644 --- a/testdata/diff/create_index/add_index/plan.txt +++ b/testdata/diff/create_index/add_index/plan.txt @@ -6,6 +6,7 @@ Summary by type: Tables: + users + idx_users_email (index) + + idx_users_email_include (index) + idx_users_email_unique (index) + idx_users_id (index) + idx_users_name (index) @@ -23,6 +24,8 @@ CREATE TABLE IF NOT EXISTS users ( CREATE INDEX IF NOT EXISTS idx_users_email ON users (email varchar_pattern_ops); +CREATE INDEX IF NOT EXISTS idx_users_email_include ON users (email) INCLUDE (name); + CREATE UNIQUE INDEX IF NOT EXISTS idx_users_email_unique ON users (email) NULLS NOT DISTINCT; CREATE INDEX IF NOT EXISTS idx_users_id ON users (id);