From eae7d390a9e054fe98cb2538c70e14e1708f8574 Mon Sep 17 00:00:00 2001 From: Ariel Mashraki <7413593+a8m@users.noreply.github.com> Date: Tue, 8 Nov 2022 14:08:10 +0200 Subject: [PATCH] dialect/sql: add support for index operator-class using atlasgo.io (#3073) --- dialect/entsql/annotation.go | 73 +++++++++++++++++++ dialect/sql/schema/mysql.go | 2 +- dialect/sql/schema/postgres.go | 28 ++++++- entc/gen/template/migrate/schema.tmpl | 15 ++++ .../migrate/entv2/migrate/schema.go | 10 +++ entc/integration/migrate/entv2/schema/user.go | 5 ++ entc/integration/migrate/migrate_test.go | 10 +++ go.mod | 2 +- go.sum | 2 + 9 files changed, 144 insertions(+), 3 deletions(-) diff --git a/dialect/entsql/annotation.go b/dialect/entsql/annotation.go index f3e3a58d5..dbf816820 100644 --- a/dialect/entsql/annotation.go +++ b/dialect/entsql/annotation.go @@ -261,6 +261,33 @@ type IndexAnnotation struct { // Types map[string]string + // OpClass defines the operator class for a single string column index. + // In PostgreSQL, the following annotation maps to: + // + // index.Fields("column"). + // Annotation( + // entsql.IndexType("BRIN"), + // entsql.OpClass("int8_bloom_ops"), + // ) + // + // CREATE INDEX "table_column" ON "table" USING BRIN ("column" int8_bloom_ops) + // + OpClass string + + // OpClassColumns defines operator-classes for a multi-column index. + // In PostgreSQL, the following annotation maps to: + // + // index.Fields("c1", "c2", "c3"). + // Annotation( + // entsql.IndexType("BRIN"), + // entsql.OpClassColumn("c1", "int8_bloom_ops"), + // entsql.OpClassColumn("c2", "int8_minmax_multi_ops(values_per_range=8)"), + // ) + // + // CREATE INDEX "table_column" ON "table" USING BRIN ("c1" int8_bloom_ops, "c2" int8_minmax_multi_ops(values_per_range=8), "c3") + // + OpClassColumns map[string]string + // IndexWhere allows configuring partial indexes in SQLite and PostgreSQL. // Read more: https://postgresql.org/docs/current/indexes-partial.html. // @@ -307,6 +334,41 @@ func PrefixColumn(name string, prefix uint) *IndexAnnotation { } } +// OpClass defines the operator class for a single string column index. +// In PostgreSQL, the following annotation maps to: +// +// index.Fields("column"). +// Annotation( +// entsql.IndexType("BRIN"), +// entsql.OpClass("int8_bloom_ops"), +// ) +// +// CREATE INDEX "table_column" ON "table" USING BRIN ("column" int8_bloom_ops) +func OpClass(op string) *IndexAnnotation { + return &IndexAnnotation{ + OpClass: op, + } +} + +// OpClassColumn returns a new index annotation with column operator +// class for multi-column indexes. In PostgreSQL, the following annotation maps to: +// +// index.Fields("c1", "c2", "c3"). +// Annotation( +// entsql.IndexType("BRIN"), +// entsql.OpClassColumn("c1", "int8_bloom_ops"), +// entsql.OpClassColumn("c2", "int8_minmax_multi_ops(values_per_range=8)"), +// ) +// +// CREATE INDEX "table_column" ON "table" USING BRIN ("c1" int8_bloom_ops, "c2" int8_minmax_multi_ops(values_per_range=8), "c3") +func OpClassColumn(name, op string) *IndexAnnotation { + return &IndexAnnotation{ + OpClassColumns: map[string]string{ + name: op, + }, + } +} + // Desc returns a new index annotation with the DESC clause for a // single column index. In MySQL, the following annotation maps to: // @@ -423,6 +485,17 @@ func (a IndexAnnotation) Merge(other schema.Annotation) schema.Annotation { a.PrefixColumns[column] = prefix } } + if ant.OpClass != "" { + a.OpClass = ant.OpClass + } + if ant.OpClassColumns != nil { + if a.OpClassColumns == nil { + a.OpClassColumns = make(map[string]string) + } + for column, op := range ant.OpClassColumns { + a.OpClassColumns[column] = op + } + } if ant.Desc { a.Desc = ant.Desc } diff --git a/dialect/sql/schema/mysql.go b/dialect/sql/schema/mysql.go index 5aef0e729..6d246b813 100644 --- a/dialect/sql/schema/mysql.go +++ b/dialect/sql/schema/mysql.go @@ -764,7 +764,7 @@ func (d *MySQL) indexModified(old, new *Index) bool { return false } -// indexParts returns a map holding the sub_part mapping if exist. +// indexParts returns a map holding the sub_part mapping if exists. func indexParts(idx *Index) map[string]uint { parts := make(map[string]uint) if idx.Annotation == nil { diff --git a/dialect/sql/schema/postgres.go b/dialect/sql/schema/postgres.go index 530dcce07..56b8bca7e 100644 --- a/dialect/sql/schema/postgres.go +++ b/dialect/sql/schema/postgres.go @@ -777,13 +777,39 @@ func (d *Postgres) atIncrementT(t *schema.Table, v int64) { t.AddAttrs(&postgres.Identity{Sequence: &postgres.Sequence{Start: v}}) } +// indexOpClass returns a map holding the operator-class mapping if exists. +func indexOpClass(idx *Index) map[string]string { + opc := make(map[string]string) + if idx.Annotation == nil { + return opc + } + // If operator-class (without a name) was defined on + // the annotation, map it to the single column index. + if idx.Annotation.OpClass != "" && len(idx.Columns) == 1 { + opc[idx.Columns[0].Name] = idx.Annotation.OpClass + } + for column, op := range idx.Annotation.OpClassColumns { + opc[column] = op + } + return opc +} + func (d *Postgres) atIndex(idx1 *Index, t2 *schema.Table, idx2 *schema.Index) error { + opc := indexOpClass(idx1) for _, c1 := range idx1.Columns { c2, ok := t2.Column(c1.Name) if !ok { return fmt.Errorf("unexpected index %q column: %q", idx1.Name, c1.Name) } - idx2.AddParts(&schema.IndexPart{C: c2}) + part := &schema.IndexPart{C: c2} + if v, ok := opc[c1.Name]; ok { + var op postgres.IndexOpClass + if err := op.UnmarshalText([]byte(v)); err != nil { + return fmt.Errorf("unmarshaling operator-class %q for column %q: %v", v, c1.Name, err) + } + part.Attrs = append(part.Attrs, &op) + } + idx2.AddParts(part) } if t, ok := indexType(idx1, dialect.Postgres); ok { idx2.AddAttrs(&postgres.IndexType{T: t}) diff --git a/entc/gen/template/migrate/schema.tmpl b/entc/gen/template/migrate/schema.tmpl index d14b1e5e0..a6c40a50f 100644 --- a/entc/gen/template/migrate/schema.tmpl +++ b/entc/gen/template/migrate/schema.tmpl @@ -115,6 +115,21 @@ var ( {{- end }} }, {{- end }} + {{- with $ant.OpClass }} + OpClass: "{{ . }}", + {{- end }} + {{- with $keys := keys $ant.OpClassColumns }} + OpClassColumns: map[string]string{ + {{- range $k := $keys }} + {{- /* Use the column reference instead of using raw string. */}} + {{- range $i, $c := $t.Columns }} + {{- if eq $k $c.Name }} + {{ $columns }}[{{ $i }}].Name: "{{ index $ant.OpClassColumns $k }}", + {{ end }} + {{- end }} + {{- end }} + }, + {{- end }} {{- with $ant.Desc }} Desc: {{ . }}, {{- end }} diff --git a/entc/integration/migrate/entv2/migrate/schema.go b/entc/integration/migrate/entv2/migrate/schema.go index 1d7ac0786..0f76a971f 100644 --- a/entc/integration/migrate/entv2/migrate/schema.go +++ b/entc/integration/migrate/entv2/migrate/schema.go @@ -224,6 +224,16 @@ var ( Where: "active", }, }, + { + Name: "user_age_phone", + Unique: false, + Columns: []*schema.Column{UsersColumns[4], UsersColumns[8]}, + Annotation: &entsql.IndexAnnotation{ + OpClassColumns: map[string]string{ + UsersColumns[8].Name: "bpchar_pattern_ops", + }, + }, + }, }, } // FriendsColumns holds the columns for the "friends" table. diff --git a/entc/integration/migrate/entv2/schema/user.go b/entc/integration/migrate/entv2/schema/user.go index 9ac3f6ddb..97e6cc26e 100644 --- a/entc/integration/migrate/entv2/schema/user.go +++ b/entc/integration/migrate/entv2/schema/user.go @@ -161,6 +161,11 @@ func (User) Indexes() []ent.Index { Annotations( entsql.IndexWhere("active"), ), + // For PostgreSQL, operator classes can be configured for each field. + index.Fields("age", "phone"). + Annotations( + entsql.OpClassColumn("phone", "bpchar_pattern_ops"), + ), } } diff --git a/entc/integration/migrate/migrate_test.go b/entc/integration/migrate/migrate_test.go index 44be1ced0..efd305f46 100644 --- a/entc/integration/migrate/migrate_test.go +++ b/entc/integration/migrate/migrate_test.go @@ -133,6 +133,7 @@ func TestPostgres(t *testing.T) { CheckConstraint(t, clientv2) TimePrecision(t, drv, "SELECT datetime_precision FROM information_schema.columns WHERE table_name = $1 AND column_name = $2") PartialIndexes(t, drv, "select indexdef from pg_indexes where indexname=$1", "CREATE INDEX user_phone ON public.users USING btree (phone) WHERE active") + IndexOpClass(t, drv) if version != "10" { IncludeColumns(t, drv) } @@ -669,6 +670,15 @@ func IncludeColumns(t *testing.T, drv *sql.Driver) { require.Equal(t, d, "CREATE INDEX user_workplace ON public.users USING btree (workplace) INCLUDE (nickname)") } +func IndexOpClass(t *testing.T, drv *sql.Driver) { + rows, err := drv.QueryContext(context.Background(), "select indexdef from pg_indexes where indexname='user_age_phone'") + require.NoError(t, err) + d, err := sql.ScanString(rows) + require.NoError(t, err) + require.NoError(t, rows.Close()) + require.Equal(t, d, "CREATE INDEX user_age_phone ON public.users USING btree (age, phone bpchar_pattern_ops)") +} + func PartialIndexes(t *testing.T, drv *sql.Driver, query, def string) { rows, err := drv.QueryContext(context.Background(), query, "user_phone") require.NoError(t, err) diff --git a/go.mod b/go.mod index ffb8f3088..2f4a2f911 100644 --- a/go.mod +++ b/go.mod @@ -3,7 +3,7 @@ module entgo.io/ent go 1.19 require ( - ariga.io/atlas v0.7.3-0.20221011160332-3ca609863edd + ariga.io/atlas v0.8.2-0.20221108073928-ba5d4f596240 github.com/DATA-DOG/go-sqlmock v1.5.0 github.com/go-openapi/inflect v0.19.0 github.com/go-sql-driver/mysql v1.6.0 diff --git a/go.sum b/go.sum index 4dae8acf0..adc4b35e1 100644 --- a/go.sum +++ b/go.sum @@ -1,5 +1,7 @@ ariga.io/atlas v0.7.3-0.20221011160332-3ca609863edd h1:c3F2jvvEZzsoH/KUpDNhTsCVeUPnpXaF8kADZvUSiU0= ariga.io/atlas v0.7.3-0.20221011160332-3ca609863edd/go.mod h1:ft47uSh5hWGDCmQC9DsztZg6Xk+KagM5Ts/mZYKb9JE= +ariga.io/atlas v0.8.2-0.20221108073928-ba5d4f596240 h1:Skxqk163AiuhtDEmAfF2/dvaDJTB7rNKEiGP+m1fOHw= +ariga.io/atlas v0.8.2-0.20221108073928-ba5d4f596240/go.mod h1:ft47uSh5hWGDCmQC9DsztZg6Xk+KagM5Ts/mZYKb9JE= cloud.google.com/go v0.26.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw= github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU= github.com/DATA-DOG/go-sqlmock v1.5.0 h1:Shsta01QNfFxHCfpW6YH2STWB0MudeXXEWMr20OEh60=