diff --git a/internal/pgdb/pgdb_index.go b/internal/pgdb/pgdb_index.go index f87da4e..9485fd2 100644 --- a/internal/pgdb/pgdb_index.go +++ b/internal/pgdb/pgdb_index.go @@ -57,6 +57,50 @@ func (module *Module) getMessageIndexes(ctx pgsgo.Context, m pgs.Message, ix *im return rv } +// resolveFieldPath resolves a dot-separated field path (e.g., "secret_trait.status") +// to the fully qualified column name through nested message structures. +func resolveFieldPath(m pgs.Message, fieldPath string) string { + path := strings.Split(fieldPath, ".") + message := m + resolution := "" + for i, p := range path { + lastP := i == len(path)-1 + + if !lastP { + f := fieldByName(message, p) + resolution += getNestedName(f) + message = f.Type().Embed() + continue + } + + name := "" + if f, ok := tryFieldByName(message, p); ok { + var err error + name, err = getColumnName(f) + if err != nil { + panic(err) + } + } else { + for _, oo := range message.RealOneOfs() { + if oo.Name().String() == p { + var err error + name, err = getColumnOneOfName(oo) + if err != nil { + panic(err) + } + break + } + } + } + if name == "" { + panic(fmt.Errorf("could not find field: %s in %s", fieldPath, m.FullyQualifiedName())) + } + + resolution += name + } + return resolution +} + func (module *Module) extraIndexes(ctx pgsgo.Context, m pgs.Message, ix *importTracker, idx *pgdb_v1.MessageOptions_Index) *indexContext { indexName, err := getIndexName(m, idx.GetName()) if err != nil { @@ -75,51 +119,25 @@ func (module *Module) extraIndexes(ctx pgsgo.Context, m pgs.Message, ix *importT rv.DB.Method = idx.GetMethod() for _, fieldName := range idx.GetColumns() { - path := strings.Split(fieldName, ".") - message := m - resolution := "" - for i, p := range path { - lastP := i == len(path)-1 - - if !lastP { - f := fieldByName(message, p) - resolution += getNestedName(f) - message = f.Type().Embed() - continue - } + rv.DB.Columns = append(rv.DB.Columns, resolveFieldPath(m, fieldName)) + } - name := "" - // could be a real field! - if f, ok := tryFieldByName(message, p); ok { - name, err = getColumnName(f) - if err != nil { - panic(err) - } - } else { - // look in oneofs! - for _, oo := range message.RealOneOfs() { - if oo.Name().String() == p { - name, err = getColumnOneOfName(oo) - if err != nil { - panic(err) - } - break - } - } - } - if name == "" { - panic(fmt.Errorf("could not find field for index: %s", path)) - } + for _, fieldName := range idx.GetIncludeColumns() { + rv.DB.IncludeColumns = append(rv.DB.IncludeColumns, resolveFieldPath(m, fieldName)) + } - resolution += name - rv.DB.Columns = append(rv.DB.Columns, resolution) - } + if len(rv.DB.IncludeColumns) > 0 && rv.DB.Method != pgdb_v1.MessageOptions_Index_INDEX_METHOD_BTREE { + panic(fmt.Errorf("include_columns is only valid for BTREE indexes, got %s on index %s", rv.DB.Method.String(), idx.GetName())) } if idx.GetBitHammingOps() && idx.GetMethod() == pgdb_v1.MessageOptions_Index_INDEX_METHOD_HNSW_COSINE { rv.DB.OverrideExpression = fmt.Sprintf("pb$%s bit_hamming_ops", rv.DB.Columns[0]) } + if idx.GetPartialDeletedAtIsNull() && len(idx.GetWhere()) > 0 { + panic(fmt.Errorf("index %s: cannot use both partial_deleted_at_is_null and where", idx.GetName())) + } + if idx.GetPartialDeletedAtIsNull() { if f, ok := tryFieldByName(m, "deleted_at"); ok { name, err := getColumnName(f) @@ -132,9 +150,27 @@ func (module *Module) extraIndexes(ctx pgsgo.Context, m pgs.Message, ix *importT name, ) } else { - panic(fmt.Sprintf("%s ould not find field for partial index: deleted_at", m.FullyQualifiedName())) + panic(fmt.Sprintf("%s could not find field for partial index: deleted_at", m.FullyQualifiedName())) + } + } + + if len(idx.GetWhere()) > 0 { + parts := make([]string, 0, len(idx.GetWhere())) + for _, pred := range idx.GetWhere() { + if pred.GetColumn() == "" { + panic(fmt.Errorf("index %s: where predicate missing column", idx.GetName())) + } + if pred.GetOp() == "" { + panic(fmt.Errorf("index %s: where predicate missing op for column %s", idx.GetName(), pred.GetColumn())) + } + col := resolveFieldPath(m, pred.GetColumn()) + parts = append(parts, fmt.Sprintf( + `" + io.ColumnName("%s") + " %s`, + col, pred.GetOp())) } + rv.DB.WherePredicate = strings.Join(parts, ` AND "`) } + return rv } diff --git a/internal/pgdb/templates/descriptor.tmpl b/internal/pgdb/templates/descriptor.tmpl index 83a0535..88f6b68 100644 --- a/internal/pgdb/templates/descriptor.tmpl +++ b/internal/pgdb/templates/descriptor.tmpl @@ -236,6 +236,7 @@ func (d *{{.ReceiverType}}) Indexes(opts ...pgdb_v1.IndexOptionsFunc) []*pgdb_v1 IsUnique: {{ .DB.IsUnique }}, IsDropped: {{ .DB.IsDropped }}, Columns: []string{ {{- range .DB.Columns -}} io.ColumnName("{{- . -}}"), {{- end -}} }, + IncludeColumns: []string{ {{- range .DB.IncludeColumns -}} io.ColumnName("{{- . -}}"), {{- end -}} }, OverrideExpression: "{{ .DB.OverrideExpression }}", WherePredicate: "{{ .DB.WherePredicate }}", }) diff --git a/pgdb/v1/descriptor.go b/pgdb/v1/descriptor.go index abdcc21..db5c77e 100644 --- a/pgdb/v1/descriptor.go +++ b/pgdb/v1/descriptor.go @@ -288,6 +288,8 @@ type Index struct { // OverrideExpression if set, this string is used to render indexes contents, instead of the Columns list. OverrideExpression string WherePredicate string + // IncludeColumns are non-key columns stored in index leaf pages for covering indexes (Index Only Scans). + IncludeColumns []string } type Statistic struct { diff --git a/pgdb/v1/pgdb.pb.go b/pgdb/v1/pgdb.pb.go index fcfe024..b5579bf 100644 --- a/pgdb/v1/pgdb.pb.go +++ b/pgdb/v1/pgdb.pb.go @@ -1,6 +1,6 @@ // Code generated by protoc-gen-go. DO NOT EDIT. // versions: -// protoc-gen-go v1.36.11 +// protoc-gen-go v1.36.6 // protoc (unknown) // source: pgdb/v1/pgdb.proto @@ -1204,13 +1204,15 @@ func (b0 MessageOptions_StorageParameters_builder) Build() *MessageOptions_Stora } type MessageOptions_Index struct { - state protoimpl.MessageState `protogen:"opaque.v1"` - xxx_hidden_Name string `protobuf:"bytes,1,opt,name=name,proto3"` - xxx_hidden_Method MessageOptions_Index_IndexMethod `protobuf:"varint,2,opt,name=method,proto3,enum=pgdb.v1.MessageOptions_Index_IndexMethod"` - xxx_hidden_Columns []string `protobuf:"bytes,3,rep,name=columns,proto3"` - xxx_hidden_Dropped bool `protobuf:"varint,4,opt,name=dropped,proto3"` - xxx_hidden_PartialDeletedAtIsNull bool `protobuf:"varint,5,opt,name=partial_deleted_at_is_null,json=partialDeletedAtIsNull,proto3"` - xxx_hidden_BitHammingOps bool `protobuf:"varint,6,opt,name=bit_hamming_ops,json=bitHammingOps,proto3"` + state protoimpl.MessageState `protogen:"opaque.v1"` + xxx_hidden_Name string `protobuf:"bytes,1,opt,name=name,proto3"` + xxx_hidden_Method MessageOptions_Index_IndexMethod `protobuf:"varint,2,opt,name=method,proto3,enum=pgdb.v1.MessageOptions_Index_IndexMethod"` + xxx_hidden_Columns []string `protobuf:"bytes,3,rep,name=columns,proto3"` + xxx_hidden_Dropped bool `protobuf:"varint,4,opt,name=dropped,proto3"` + xxx_hidden_PartialDeletedAtIsNull bool `protobuf:"varint,5,opt,name=partial_deleted_at_is_null,json=partialDeletedAtIsNull,proto3"` + xxx_hidden_BitHammingOps bool `protobuf:"varint,6,opt,name=bit_hamming_ops,json=bitHammingOps,proto3"` + xxx_hidden_IncludeColumns []string `protobuf:"bytes,7,rep,name=include_columns,json=includeColumns,proto3"` + xxx_hidden_Where *[]*MessageOptions_IndexPredicate `protobuf:"bytes,8,rep,name=where,proto3"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -1282,6 +1284,22 @@ func (x *MessageOptions_Index) GetBitHammingOps() bool { return false } +func (x *MessageOptions_Index) GetIncludeColumns() []string { + if x != nil { + return x.xxx_hidden_IncludeColumns + } + return nil +} + +func (x *MessageOptions_Index) GetWhere() []*MessageOptions_IndexPredicate { + if x != nil { + if x.xxx_hidden_Where != nil { + return *x.xxx_hidden_Where + } + } + return nil +} + func (x *MessageOptions_Index) SetName(v string) { x.xxx_hidden_Name = v } @@ -1306,6 +1324,14 @@ func (x *MessageOptions_Index) SetBitHammingOps(v bool) { x.xxx_hidden_BitHammingOps = v } +func (x *MessageOptions_Index) SetIncludeColumns(v []string) { + x.xxx_hidden_IncludeColumns = v +} + +func (x *MessageOptions_Index) SetWhere(v []*MessageOptions_IndexPredicate) { + x.xxx_hidden_Where = &v +} + type MessageOptions_Index_builder struct { _ [0]func() // Prevents comparability and use of unkeyed literals for the builder. @@ -1318,6 +1344,10 @@ type MessageOptions_Index_builder struct { PartialDeletedAtIsNull bool // adds bit_hamming_ops to the index for HNSW_COSINE indexes BitHammingOps bool + // non-key columns stored in index leaf pages for covering index (Index Only Scan) + IncludeColumns []string + // Partial index predicates. ANDed together to form the WHERE clause. + Where []*MessageOptions_IndexPredicate } func (b0 MessageOptions_Index_builder) Build() *MessageOptions_Index { @@ -1330,6 +1360,81 @@ func (b0 MessageOptions_Index_builder) Build() *MessageOptions_Index { x.xxx_hidden_Dropped = b.Dropped x.xxx_hidden_PartialDeletedAtIsNull = b.PartialDeletedAtIsNull x.xxx_hidden_BitHammingOps = b.BitHammingOps + x.xxx_hidden_IncludeColumns = b.IncludeColumns + x.xxx_hidden_Where = &b.Where + return m0 +} + +// IndexPredicate represents a single partial index predicate. +// Column is a field path (resolved at codegen time), Op is a raw SQL operator fragment. +type MessageOptions_IndexPredicate struct { + state protoimpl.MessageState `protogen:"opaque.v1"` + xxx_hidden_Column string `protobuf:"bytes,1,opt,name=column,proto3"` + xxx_hidden_Op string `protobuf:"bytes,2,opt,name=op,proto3"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *MessageOptions_IndexPredicate) Reset() { + *x = MessageOptions_IndexPredicate{} + mi := &file_pgdb_v1_pgdb_proto_msgTypes[5] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *MessageOptions_IndexPredicate) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*MessageOptions_IndexPredicate) ProtoMessage() {} + +func (x *MessageOptions_IndexPredicate) ProtoReflect() protoreflect.Message { + mi := &file_pgdb_v1_pgdb_proto_msgTypes[5] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +func (x *MessageOptions_IndexPredicate) GetColumn() string { + if x != nil { + return x.xxx_hidden_Column + } + return "" +} + +func (x *MessageOptions_IndexPredicate) GetOp() string { + if x != nil { + return x.xxx_hidden_Op + } + return "" +} + +func (x *MessageOptions_IndexPredicate) SetColumn(v string) { + x.xxx_hidden_Column = v +} + +func (x *MessageOptions_IndexPredicate) SetOp(v string) { + x.xxx_hidden_Op = v +} + +type MessageOptions_IndexPredicate_builder struct { + _ [0]func() // Prevents comparability and use of unkeyed literals for the builder. + + Column string + Op string +} + +func (b0 MessageOptions_IndexPredicate_builder) Build() *MessageOptions_IndexPredicate { + m0 := &MessageOptions_IndexPredicate{} + b, x := &b0, m0 + _, _ = b, x + x.xxx_hidden_Column = b.Column + x.xxx_hidden_Op = b.Op return m0 } @@ -1345,7 +1450,7 @@ type MessageOptions_Stat struct { func (x *MessageOptions_Stat) Reset() { *x = MessageOptions_Stat{} - mi := &file_pgdb_v1_pgdb_proto_msgTypes[5] + mi := &file_pgdb_v1_pgdb_proto_msgTypes[6] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1357,7 +1462,7 @@ func (x *MessageOptions_Stat) String() string { func (*MessageOptions_Stat) ProtoMessage() {} func (x *MessageOptions_Stat) ProtoReflect() protoreflect.Message { - mi := &file_pgdb_v1_pgdb_proto_msgTypes[5] + mi := &file_pgdb_v1_pgdb_proto_msgTypes[6] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1482,7 +1587,7 @@ var File_pgdb_v1_pgdb_proto protoreflect.FileDescriptor const file_pgdb_v1_pgdb_proto_rawDesc = "" + "\n" + - "\x12pgdb/v1/pgdb.proto\x12\apgdb.v1\x1a google/protobuf/descriptor.proto\"\xb5\x14\n" + + "\x12pgdb/v1/pgdb.proto\x12\apgdb.v1\x1a google/protobuf/descriptor.proto\"\xd6\x15\n" + "\x0eMessageOptions\x12\x1a\n" + "\bdisabled\x18\x01 \x01(\bR\bdisabled\x12X\n" + "\x12storage_parameters\x18\v \x01(\v2).pgdb.v1.MessageOptions.StorageParametersR\x11storageParameters\x127\n" + @@ -1525,20 +1630,25 @@ const file_pgdb_v1_pgdb_proto_rawDesc = "" + "\x1c_autovacuum_freeze_table_ageB\r\n" + "\v_fillfactorB\x15\n" + "\x13_toast_tuple_targetB\x15\n" + - "\x13_autovacuum_enabled\x1a\x8c\x03\n" + + "\x13_autovacuum_enabled\x1a\xf3\x03\n" + "\x05Index\x12\x12\n" + "\x04name\x18\x01 \x01(\tR\x04name\x12A\n" + "\x06method\x18\x02 \x01(\x0e2).pgdb.v1.MessageOptions.Index.IndexMethodR\x06method\x12\x18\n" + "\acolumns\x18\x03 \x03(\tR\acolumns\x12\x18\n" + "\adropped\x18\x04 \x01(\bR\adropped\x12:\n" + "\x1apartial_deleted_at_is_null\x18\x05 \x01(\bR\x16partialDeletedAtIsNull\x12&\n" + - "\x0fbit_hamming_ops\x18\x06 \x01(\bR\rbitHammingOps\"\x93\x01\n" + + "\x0fbit_hamming_ops\x18\x06 \x01(\bR\rbitHammingOps\x12'\n" + + "\x0finclude_columns\x18\a \x03(\tR\x0eincludeColumns\x12<\n" + + "\x05where\x18\b \x03(\v2&.pgdb.v1.MessageOptions.IndexPredicateR\x05where\"\x93\x01\n" + "\vIndexMethod\x12\x1c\n" + "\x18INDEX_METHOD_UNSPECIFIED\x10\x00\x12\x16\n" + "\x12INDEX_METHOD_BTREE\x10\x01\x12\x14\n" + "\x10INDEX_METHOD_GIN\x10\x02\x12\x1a\n" + "\x16INDEX_METHOD_BTREE_GIN\x10\x03\x12\x1c\n" + - "\x18INDEX_METHOD_HNSW_COSINE\x10\x04\x1a\x80\x02\n" + + "\x18INDEX_METHOD_HNSW_COSINE\x10\x04\x1a8\n" + + "\x0eIndexPredicate\x12\x16\n" + + "\x06column\x18\x01 \x01(\tR\x06column\x12\x0e\n" + + "\x02op\x18\x02 \x01(\tR\x02op\x1a\x80\x02\n" + "\x04Stat\x12\x12\n" + "\x04name\x18\x01 \x01(\tR\x04name\x12<\n" + "\x05kinds\x18\x02 \x03(\x0e2&.pgdb.v1.MessageOptions.Stat.StatsKindR\x05kinds\x12\x18\n" + @@ -1593,7 +1703,7 @@ const file_pgdb_v1_pgdb_proto_rawDesc = "" + "\vcom.pgdb.v1B\tPgdbProtoP\x01Z*github.com/ductone/protoc-gen-pgdb/pgdb/v1\xa2\x02\x03PXX\xaa\x02\aPgdb.V1\xca\x02\aPgdb\\V1\xe2\x02\x13Pgdb\\V1\\GPBMetadata\xea\x02\bPgdb::V1b\x06proto3" var file_pgdb_v1_pgdb_proto_enumTypes = make([]protoimpl.EnumInfo, 7) -var file_pgdb_v1_pgdb_proto_msgTypes = make([]protoimpl.MessageInfo, 6) +var file_pgdb_v1_pgdb_proto_msgTypes = make([]protoimpl.MessageInfo, 7) var file_pgdb_v1_pgdb_proto_goTypes = []any{ (VectorElementType)(0), // 0: pgdb.v1.VectorElementType (MessageOptions_PartitionedByDateRange)(0), // 1: pgdb.v1.MessageOptions.PartitionedByDateRange @@ -1607,33 +1717,35 @@ var file_pgdb_v1_pgdb_proto_goTypes = []any{ (*EnumValueOptions)(nil), // 9: pgdb.v1.EnumValueOptions (*MessageOptions_StorageParameters)(nil), // 10: pgdb.v1.MessageOptions.StorageParameters (*MessageOptions_Index)(nil), // 11: pgdb.v1.MessageOptions.Index - (*MessageOptions_Stat)(nil), // 12: pgdb.v1.MessageOptions.Stat - (*descriptorpb.MessageOptions)(nil), // 13: google.protobuf.MessageOptions - (*descriptorpb.FieldOptions)(nil), // 14: google.protobuf.FieldOptions - (*descriptorpb.EnumValueOptions)(nil), // 15: google.protobuf.EnumValueOptions + (*MessageOptions_IndexPredicate)(nil), // 12: pgdb.v1.MessageOptions.IndexPredicate + (*MessageOptions_Stat)(nil), // 13: pgdb.v1.MessageOptions.Stat + (*descriptorpb.MessageOptions)(nil), // 14: google.protobuf.MessageOptions + (*descriptorpb.FieldOptions)(nil), // 15: google.protobuf.FieldOptions + (*descriptorpb.EnumValueOptions)(nil), // 16: google.protobuf.EnumValueOptions } var file_pgdb_v1_pgdb_proto_depIdxs = []int32{ 10, // 0: pgdb.v1.MessageOptions.storage_parameters:type_name -> pgdb.v1.MessageOptions.StorageParameters 11, // 1: pgdb.v1.MessageOptions.indexes:type_name -> pgdb.v1.MessageOptions.Index 1, // 2: pgdb.v1.MessageOptions.partitioned_by_date_range:type_name -> pgdb.v1.MessageOptions.PartitionedByDateRange - 12, // 3: pgdb.v1.MessageOptions.stats:type_name -> pgdb.v1.MessageOptions.Stat + 13, // 3: pgdb.v1.MessageOptions.stats:type_name -> pgdb.v1.MessageOptions.Stat 4, // 4: pgdb.v1.FieldOptions.full_text_type:type_name -> pgdb.v1.FieldOptions.FullTextType 5, // 5: pgdb.v1.FieldOptions.full_text_weight:type_name -> pgdb.v1.FieldOptions.FullTextWeight 6, // 6: pgdb.v1.FieldOptions.message_behavior:type_name -> pgdb.v1.FieldOptions.MessageBehavior 0, // 7: pgdb.v1.EnumValueOptions.vector_element_type:type_name -> pgdb.v1.VectorElementType 2, // 8: pgdb.v1.MessageOptions.Index.method:type_name -> pgdb.v1.MessageOptions.Index.IndexMethod - 3, // 9: pgdb.v1.MessageOptions.Stat.kinds:type_name -> pgdb.v1.MessageOptions.Stat.StatsKind - 13, // 10: pgdb.v1.msg:extendee -> google.protobuf.MessageOptions - 14, // 11: pgdb.v1.options:extendee -> google.protobuf.FieldOptions - 15, // 12: pgdb.v1.enum:extendee -> google.protobuf.EnumValueOptions - 7, // 13: pgdb.v1.msg:type_name -> pgdb.v1.MessageOptions - 8, // 14: pgdb.v1.options:type_name -> pgdb.v1.FieldOptions - 9, // 15: pgdb.v1.enum:type_name -> pgdb.v1.EnumValueOptions - 16, // [16:16] is the sub-list for method output_type - 16, // [16:16] is the sub-list for method input_type - 13, // [13:16] is the sub-list for extension type_name - 10, // [10:13] is the sub-list for extension extendee - 0, // [0:10] is the sub-list for field type_name + 12, // 9: pgdb.v1.MessageOptions.Index.where:type_name -> pgdb.v1.MessageOptions.IndexPredicate + 3, // 10: pgdb.v1.MessageOptions.Stat.kinds:type_name -> pgdb.v1.MessageOptions.Stat.StatsKind + 14, // 11: pgdb.v1.msg:extendee -> google.protobuf.MessageOptions + 15, // 12: pgdb.v1.options:extendee -> google.protobuf.FieldOptions + 16, // 13: pgdb.v1.enum:extendee -> google.protobuf.EnumValueOptions + 7, // 14: pgdb.v1.msg:type_name -> pgdb.v1.MessageOptions + 8, // 15: pgdb.v1.options:type_name -> pgdb.v1.FieldOptions + 9, // 16: pgdb.v1.enum:type_name -> pgdb.v1.EnumValueOptions + 17, // [17:17] is the sub-list for method output_type + 17, // [17:17] is the sub-list for method input_type + 14, // [14:17] is the sub-list for extension type_name + 11, // [11:14] is the sub-list for extension extendee + 0, // [0:11] is the sub-list for field type_name } func init() { file_pgdb_v1_pgdb_proto_init() } @@ -1648,7 +1760,7 @@ func file_pgdb_v1_pgdb_proto_init() { GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: unsafe.Slice(unsafe.StringData(file_pgdb_v1_pgdb_proto_rawDesc), len(file_pgdb_v1_pgdb_proto_rawDesc)), NumEnums: 7, - NumMessages: 6, + NumMessages: 7, NumExtensions: 3, NumServices: 0, }, diff --git a/pgdb/v1/schema_sql.go b/pgdb/v1/schema_sql.go index 3370abb..3f2d974 100644 --- a/pgdb/v1/schema_sql.go +++ b/pgdb/v1/schema_sql.go @@ -62,6 +62,13 @@ func index2sql(desc Descriptor, idx *Index) string { }), ", \n")) } _, _ = buf.WriteString("\n)\n") + if len(idx.IncludeColumns) > 0 { + _, _ = buf.WriteString("INCLUDE (\n") + _, _ = buf.WriteString(strings.Join(slice.Convert(idx.IncludeColumns, func(in string) string { + return ` "` + in + `"` + }), ", \n")) + _, _ = buf.WriteString("\n)\n") + } if idx.WherePredicate != "" { _, _ = buf.WriteString("WHERE ") _, _ = buf.WriteString(idx.WherePredicate) diff --git a/pgdb/v1/schema_sql_test.go b/pgdb/v1/schema_sql_test.go index 067ab09..f185d6a 100644 --- a/pgdb/v1/schema_sql_test.go +++ b/pgdb/v1/schema_sql_test.go @@ -2,6 +2,7 @@ package v1 import ( "bytes" + "strings" "testing" ) @@ -552,3 +553,329 @@ SET ( }) } } + +func TestIndex2SQL_IncludeColumns(t *testing.T) { + desc := &mockDescriptor{tableName: "pb_app_resource"} + + t.Run("btree with include columns", func(t *testing.T) { + got := index2sql(desc, &Index{ + Name: "idx_covering", + Method: MessageOptions_Index_INDEX_METHOD_BTREE, + Columns: []string{"pb$tenant_id", "pb$app_id"}, + IncludeColumns: []string{"pb$access_config_id", "pb$id"}, + }) + + assertExact(t, got, + "CREATE INDEX CONCURRENTLY IF NOT EXISTS\n"+ + " \"idx_covering\"\n"+ + "ON\n"+ + " \"pb_app_resource\"\n"+ + "USING\n"+ + " BTREE\n"+ + "(\n"+ + " \"pb$tenant_id\", \n"+ + " \"pb$app_id\"\n"+ + ")\n"+ + "INCLUDE (\n"+ + " \"pb$access_config_id\", \n"+ + " \"pb$id\"\n"+ + ")\n") + }) + + t.Run("single include column", func(t *testing.T) { + got := index2sql(desc, &Index{ + Name: "idx_single", + Method: MessageOptions_Index_INDEX_METHOD_BTREE, + Columns: []string{"pb$tenant_id"}, + IncludeColumns: []string{"pb$id"}, + }) + + assertContains(t, got, "INCLUDE (\n \"pb$id\"\n)") + }) + + t.Run("no include columns omits INCLUDE", func(t *testing.T) { + got := index2sql(desc, &Index{ + Name: "idx_plain", + Method: MessageOptions_Index_INDEX_METHOD_BTREE, + Columns: []string{"pb$tenant_id"}, + }) + + assertNotContains(t, got, "INCLUDE") + }) + + t.Run("empty include columns slice omits INCLUDE", func(t *testing.T) { + got := index2sql(desc, &Index{ + Name: "idx_empty", + Method: MessageOptions_Index_INDEX_METHOD_BTREE, + Columns: []string{"pb$tenant_id"}, + IncludeColumns: []string{}, + }) + + assertNotContains(t, got, "INCLUDE") + }) + + t.Run("dropped index ignores include columns", func(t *testing.T) { + got := index2sql(desc, &Index{ + Name: "idx_dropped", + IsDropped: true, + Method: MessageOptions_Index_INDEX_METHOD_BTREE, + Columns: []string{"pb$tenant_id"}, + IncludeColumns: []string{"pb$id"}, + }) + + assertContains(t, got, "DROP INDEX") + assertNotContains(t, got, "INCLUDE") + }) + + t.Run("partitioned table omits CONCURRENTLY", func(t *testing.T) { + partDesc := &mockDescriptor{tableName: "pb_partitioned", isPartitioned: true} + got := index2sql(partDesc, &Index{ + Name: "idx_part", + Method: MessageOptions_Index_INDEX_METHOD_BTREE, + Columns: []string{"pb$tenant_id"}, + IncludeColumns: []string{"pb$id"}, + }) + + assertNotContains(t, got, "CONCURRENTLY") + assertContains(t, got, "INCLUDE") + }) +} + +func TestIndex2SQL_WherePredicate(t *testing.T) { + desc := &mockDescriptor{tableName: "pb_app_resource"} + + t.Run("IS NULL predicate", func(t *testing.T) { + got := index2sql(desc, &Index{ + Name: "idx_alive", + Method: MessageOptions_Index_INDEX_METHOD_BTREE, + Columns: []string{"pb$tenant_id", "pb$app_id"}, + WherePredicate: `"pb$deleted_at" IS NULL`, + }) + + assertContains(t, got, `WHERE "pb$deleted_at" IS NULL`) + assertNotContains(t, got, "INCLUDE") + }) + + t.Run("IS NOT NULL predicate", func(t *testing.T) { + got := index2sql(desc, &Index{ + Name: "idx_deleted", + Method: MessageOptions_Index_INDEX_METHOD_BTREE, + Columns: []string{"pb$tenant_id"}, + WherePredicate: `"pb$deleted_at" IS NOT NULL`, + }) + + assertContains(t, got, `WHERE "pb$deleted_at" IS NOT NULL`) + }) + + t.Run("EQUALS predicate", func(t *testing.T) { + got := index2sql(desc, &Index{ + Name: "idx_active", + Method: MessageOptions_Index_INDEX_METHOD_BTREE, + Columns: []string{"pb$tenant_id"}, + WherePredicate: `"pb$is_active" = true`, + }) + + assertContains(t, got, `WHERE "pb$is_active" = true`) + }) + + t.Run("multiple predicates ANDed", func(t *testing.T) { + got := index2sql(desc, &Index{ + Name: "idx_compound", + Method: MessageOptions_Index_INDEX_METHOD_BTREE, + Columns: []string{"pb$tenant_id"}, + WherePredicate: `"pb$deleted_at" IS NULL AND "pb$is_active" = true`, + }) + + assertContains(t, got, `WHERE "pb$deleted_at" IS NULL AND "pb$is_active" = true`) + }) + + t.Run("no predicate omits WHERE", func(t *testing.T) { + got := index2sql(desc, &Index{ + Name: "idx_no_where", + Method: MessageOptions_Index_INDEX_METHOD_BTREE, + Columns: []string{"pb$tenant_id"}, + }) + + assertNotContains(t, got, "WHERE") + }) +} + +func TestIndex2SQL_CombinedFeatures(t *testing.T) { + desc := &mockDescriptor{tableName: "pb_app_resource"} + + t.Run("columns + INCLUDE + WHERE exact output", func(t *testing.T) { + got := index2sql(desc, &Index{ + Name: "idx_full", + Method: MessageOptions_Index_INDEX_METHOD_BTREE, + Columns: []string{"pb$tenant_id", "pb$app_id", "pb$app_resource_type_id"}, + IncludeColumns: []string{"pb$access_config_id", "pb$id"}, + WherePredicate: `pb$deleted_at IS NULL`, + }) + + assertExact(t, got, + "CREATE INDEX CONCURRENTLY IF NOT EXISTS\n"+ + " \"idx_full\"\n"+ + "ON\n"+ + " \"pb_app_resource\"\n"+ + "USING\n"+ + " BTREE\n"+ + "(\n"+ + " \"pb$tenant_id\", \n"+ + " \"pb$app_id\", \n"+ + " \"pb$app_resource_type_id\"\n"+ + ")\n"+ + "INCLUDE (\n"+ + " \"pb$access_config_id\", \n"+ + " \"pb$id\"\n"+ + ")\n"+ + "WHERE pb$deleted_at IS NULL\n") + }) + + t.Run("INCLUDE appears before WHERE", func(t *testing.T) { + got := index2sql(desc, &Index{ + Name: "idx_order", + Method: MessageOptions_Index_INDEX_METHOD_BTREE, + Columns: []string{"pb$tenant_id"}, + IncludeColumns: []string{"pb$id"}, + WherePredicate: `pb$deleted_at IS NULL`, + }) + + includePos := strings.Index(got, "INCLUDE") + wherePos := strings.Index(got, "WHERE") + if includePos == -1 || wherePos == -1 { + t.Fatalf("expected both INCLUDE and WHERE, got:\n%s", got) + } + if includePos >= wherePos { + t.Errorf("INCLUDE (pos %d) should appear before WHERE (pos %d)", includePos, wherePos) + } + }) + + t.Run("unique index with INCLUDE and WHERE", func(t *testing.T) { + got := index2sql(desc, &Index{ + Name: "idx_unique_covering", + IsUnique: true, + Method: MessageOptions_Index_INDEX_METHOD_BTREE, + Columns: []string{"pb$tenant_id", "pb$email"}, + IncludeColumns: []string{"pb$id"}, + WherePredicate: `pb$deleted_at IS NULL`, + }) + + assertContains(t, got, "CREATE UNIQUE INDEX") + assertContains(t, got, "INCLUDE") + assertContains(t, got, "WHERE") + }) + + t.Run("override expression bypasses columns but INCLUDE still renders", func(t *testing.T) { + got := index2sql(desc, &Index{ + Name: "idx_override", + Method: MessageOptions_Index_INDEX_METHOD_BTREE, + Columns: []string{"pb$data"}, + OverrideExpression: "pb$data jsonb_path_ops", + IncludeColumns: []string{"pb$id"}, + }) + + assertContains(t, got, "pb$data jsonb_path_ops") + assertContains(t, got, "INCLUDE") + assertNotContains(t, got, "\"pb$data\"") + }) +} + +func TestIndex2SQL_DeletedExclusionEquivalence(t *testing.T) { + desc := &mockDescriptor{tableName: "pb_app_resource"} + io := NewIndexOptions(nil) + + // The old system (partial_deleted_at_is_null: true) and the new system + // (where: [{column: "deleted_at", op: "IS NULL"}]) both emit identical + // generated Go code: + // WherePredicate: "" + io.ColumnName("deleted_at") + " IS NULL" + // + // Build the predicate just like the generated code does so the test + // stays correct even if the column-name prefix changes. + deletedPredicate := io.ColumnName("deleted_at") + " IS NULL" + + t.Run("basic equivalence", func(t *testing.T) { + // Old system: partial_deleted_at_is_null: true + oldSQL := index2sql(desc, &Index{ + Name: "idx_alive", + Method: MessageOptions_Index_INDEX_METHOD_BTREE, + Columns: []string{io.ColumnName("tenant_id"), io.ColumnName("app_id")}, + WherePredicate: deletedPredicate, + }) + + // New system: where: [{column: "deleted_at", op: "IS NULL"}] + newSQL := index2sql(desc, &Index{ + Name: "idx_alive", + Method: MessageOptions_Index_INDEX_METHOD_BTREE, + Columns: []string{io.ColumnName("tenant_id"), io.ColumnName("app_id")}, + WherePredicate: deletedPredicate, + }) + + assertExact(t, oldSQL, newSQL) + assertContains(t, oldSQL, "WHERE "+deletedPredicate) + }) + + t.Run("with include columns", func(t *testing.T) { + oldSQL := index2sql(desc, &Index{ + Name: "idx_alive_covering", + Method: MessageOptions_Index_INDEX_METHOD_BTREE, + Columns: []string{io.ColumnName("tenant_id"), io.ColumnName("app_id")}, + IncludeColumns: []string{io.ColumnName("id")}, + WherePredicate: deletedPredicate, + }) + + newSQL := index2sql(desc, &Index{ + Name: "idx_alive_covering", + Method: MessageOptions_Index_INDEX_METHOD_BTREE, + Columns: []string{io.ColumnName("tenant_id"), io.ColumnName("app_id")}, + IncludeColumns: []string{io.ColumnName("id")}, + WherePredicate: deletedPredicate, + }) + + assertExact(t, oldSQL, newSQL) + assertContains(t, oldSQL, "INCLUDE") + assertContains(t, oldSQL, "WHERE "+deletedPredicate) + }) + + t.Run("exact SQL output", func(t *testing.T) { + got := index2sql(desc, &Index{ + Name: "idx_alive", + Method: MessageOptions_Index_INDEX_METHOD_BTREE, + Columns: []string{io.ColumnName("tenant_id"), io.ColumnName("app_id")}, + WherePredicate: deletedPredicate, + }) + + assertExact(t, got, + "CREATE INDEX CONCURRENTLY IF NOT EXISTS\n"+ + " \"idx_alive\"\n"+ + "ON\n"+ + " \"pb_app_resource\"\n"+ + "USING\n"+ + " BTREE\n"+ + "(\n"+ + " \""+io.ColumnName("tenant_id")+"\", \n"+ + " \""+io.ColumnName("app_id")+"\"\n"+ + ")\n"+ + "WHERE "+deletedPredicate+"\n") + }) +} + +func assertExact(t *testing.T, got, expected string) { + t.Helper() + if got != expected { + t.Errorf("index2sql mismatch.\ngot:\n%s\nexpected:\n%s", got, expected) + } +} + +func assertContains(t *testing.T, got, substr string) { + t.Helper() + if !strings.Contains(got, substr) { + t.Errorf("expected output to contain %q, got:\n%s", substr, got) + } +} + +func assertNotContains(t *testing.T, got, substr string) { + t.Helper() + if strings.Contains(got, substr) { + t.Errorf("expected output to NOT contain %q, got:\n%s", substr, got) + } +} diff --git a/proto/pgdb/v1/pgdb.proto b/proto/pgdb/v1/pgdb.proto index 7b4b84d..891b8a1 100644 --- a/proto/pgdb/v1/pgdb.proto +++ b/proto/pgdb/v1/pgdb.proto @@ -53,6 +53,17 @@ message MessageOptions { bool partial_deleted_at_is_null = 5; // adds bit_hamming_ops to the index for HNSW_COSINE indexes bool bit_hamming_ops = 6; + // non-key columns stored in index leaf pages for covering index (Index Only Scan) + repeated string include_columns = 7; + // Partial index predicates. ANDed together to form the WHERE clause. + repeated IndexPredicate where = 8; + } + + // IndexPredicate represents a single partial index predicate. + // Column is a field path (resolved at codegen time), Op is a raw SQL operator fragment. + message IndexPredicate { + string column = 1; + string op = 2; } repeated Index indexes = 2; // defaults to `tenant_id`. Must be set if an object does not have a