Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit cfb0388

Browse files
a8ma8m
authored
dialect/sql/entsql: support including non-key columns in postgres indexes (#2912)
1 parent 2cdb627 commit cfb0388

File tree

9 files changed

+90
-11
lines changed

9 files changed

+90
-11
lines changed

‎dialect/entsql/annotation.go‎

Lines changed: 31 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -213,7 +213,7 @@ type IndexAnnotation struct {
213213
//
214214
Desc bool
215215

216-
// DescColumns defines the DESC clause for columns in a multi column index.
216+
// DescColumns defines the DESC clause for columns in multi-column index.
217217
// In MySQL, the following annotation maps to:
218218
//
219219
// index.Fields("c1", "c2", "c3").
@@ -225,6 +225,18 @@ type IndexAnnotation struct {
225225
//
226226
DescColumns map[string]bool
227227

228+
// IncludeColumns defines the INCLUDE clause for the index.
229+
// Works only in Postgres and its definition is as follows:
230+
//
231+
// index.Fields("c1").
232+
// Annotation(
233+
// entsql.IncludeColumns("c2"),
234+
// )
235+
//
236+
// CREATE INDEX "table_column" ON "table"("c1") INCLUDE ("c2")
237+
//
238+
IncludeColumns []string
239+
228240
// Type defines the type of the index.
229241
// In MySQL, the following annotation maps to:
230242
//
@@ -257,7 +269,6 @@ type IndexAnnotation struct {
257269
// Annotation(entsql.Prefix(100))
258270
//
259271
// CREATE INDEX `table_column` ON `table`(`column`(100))
260-
//
261272
func Prefix(prefix uint) *IndexAnnotation {
262273
return &IndexAnnotation{
263274
Prefix: prefix,
@@ -274,7 +285,6 @@ func Prefix(prefix uint) *IndexAnnotation {
274285
// )
275286
//
276287
// CREATE INDEX `table_c1_c2_c3` ON `table`(`c1`(100), `c2`(200), `c3`)
277-
//
278288
func PrefixColumn(name string, prefix uint) *IndexAnnotation {
279289
return &IndexAnnotation{
280290
PrefixColumns: map[string]uint{
@@ -290,7 +300,6 @@ func PrefixColumn(name string, prefix uint) *IndexAnnotation {
290300
// Annotation(entsql.Desc())
291301
//
292302
// CREATE INDEX `table_column` ON `table`(`column` DESC)
293-
//
294303
func Desc() *IndexAnnotation {
295304
return &IndexAnnotation{
296305
Desc: true,
@@ -306,7 +315,6 @@ func Desc() *IndexAnnotation {
306315
// )
307316
//
308317
// CREATE INDEX `table_c1_c2_c3` ON `table`(`c1` DESC, `c2` DESC, `c3`)
309-
//
310318
func DescColumns(names ...string) *IndexAnnotation {
311319
ant := &IndexAnnotation{
312320
DescColumns: make(map[string]bool, len(names)),
@@ -317,7 +325,20 @@ func DescColumns(names ...string) *IndexAnnotation {
317325
return ant
318326
}
319327

320-
// Type defines the type of the index.
328+
// IncludeColumns defines the INCLUDE clause for the index.
329+
// Works only in Postgres and its definition is as follows:
330+
//
331+
// index.Fields("c1").
332+
// Annotation(
333+
// entsql.IncludeColumns("c2"),
334+
// )
335+
//
336+
// CREATE INDEX "table_column" ON "table"("c1") INCLUDE ("c2")
337+
func IncludeColumns(names ...string) *IndexAnnotation {
338+
return &IndexAnnotation{IncludeColumns: names}
339+
}
340+
341+
// IndexType defines the type of the index.
321342
// In MySQL, the following annotation maps to:
322343
//
323344
// index.Fields("c1").
@@ -326,12 +347,11 @@ func DescColumns(names ...string) *IndexAnnotation {
326347
// )
327348
//
328349
// CREATE FULLTEXT INDEX `table_c1` ON `table`(`c1`)
329-
//
330350
func IndexType(t string) *IndexAnnotation {
331351
return &IndexAnnotation{Type: t}
332352
}
333353

334-
// Types is like the Type option but allows mapping an index-type per dialect.
354+
// IndexTypes is like the Type option but allows mapping an index-type per dialect.
335355
//
336356
// index.Fields("c1").
337357
// Annotations(
@@ -340,7 +360,6 @@ func IndexType(t string) *IndexAnnotation {
340360
// dialect.Postgres: "GIN",
341361
// }),
342362
// )
343-
//
344363
func IndexTypes(types map[string]string) *IndexAnnotation {
345364
return &IndexAnnotation{Types: types}
346365
}
@@ -385,6 +404,9 @@ func (a IndexAnnotation) Merge(other schema.Annotation) schema.Annotation {
385404
a.DescColumns[column] = desc
386405
}
387406
}
407+
if ant.IncludeColumns != nil {
408+
a.IncludeColumns = append(a.IncludeColumns, ant.IncludeColumns...)
409+
}
388410
if ant.Type != "" {
389411
a.Type = ant.Type
390412
}

‎dialect/sql/schema/postgres.go‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -785,6 +785,17 @@ func (d *Postgres) atIndex(idx1 *Index, t2 *schema.Table, idx2 *schema.Index) er
785785
if t, ok := indexType(idx1, dialect.Postgres); ok {
786786
idx2.AddAttrs(&postgres.IndexType{T: t})
787787
}
788+
if ant, supportsInclude := idx1.Annotation, compareVersions(d.version, "11.0.0") >= 0; ant != nil && len(ant.IncludeColumns) > 0 && supportsInclude {
789+
columns := make([]*schema.Column, len(ant.IncludeColumns))
790+
for i, ic := range ant.IncludeColumns {
791+
c, ok := t2.Column(ic)
792+
if !ok {
793+
return fmt.Errorf("include column %q was not found for index %q", ic, idx1.Name)
794+
}
795+
columns[i] = c
796+
}
797+
idx2.AddAttrs(&postgres.IndexInclude{Columns: columns})
798+
}
788799
return nil
789800
}
790801

‎doc/md/schema-indexes.md‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -214,7 +214,7 @@ CREATE INDEX `users_c1_c2_c3` ON `users`(`c1`(100), `c2`(200), `c3`)
214214

215215
## Atlas Support
216216

217-
Starting with v0.10, Ent supports running migration with [Atlas](migrate.md#atlas-integration). This option provides
217+
Starting with v0.10, Ent running migration with [Atlas](https://github.com/ariga/atlas). This option provides
218218
more control on indexes such as, configuring their types or define indexes in a reverse order.
219219
```go
220220
func (User) Indexes() []ent.Index {
@@ -234,6 +234,12 @@ func (User) Indexes() []ent.Index {
234234
dialect.Postgres: "GIN",
235235
}),
236236
),
237+
// For PostgreSQL, we can include in the index
238+
// non-key columns.
239+
index.Fields("workplace").
240+
Annotations(
241+
entsql.IncludeColumns("address"),
242+
),
237243
}
238244
}
239245
```

‎entc/gen/template/migrate/schema.tmpl‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -130,6 +130,17 @@ var (
130130
{{- end }}
131131
},
132132
{{- end }}
133+
{{- with $ant.IncludeColumns }}
134+
IncludeColumns: []string{
135+
{{- range $ic := . }}
136+
{{- range $i, $c := $t.Columns }}
137+
{{- if eq $ic $c.Name }}
138+
{{ $columns }}[{{ $i }}].Name,
139+
{{ end }}
140+
{{- end }}
141+
{{- end }}
142+
},
143+
{{- end }}
133144
{{- with $ant.Type }}
134145
Type: "{{ . }}",
135146
{{- end }}

‎entc/integration/migrate/entv2/migrate/schema.go‎

Lines changed: 10 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎entc/integration/migrate/entv2/schema/user.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -148,6 +148,11 @@ func (User) Indexes() []ent.Index {
148148
dialect.MySQL: "FULLTEXT",
149149
}),
150150
),
151+
// For PostgreSQL, we can include in the index non-key columns.
152+
index.Fields("workplace").
153+
Annotations(
154+
entsql.IncludeColumns("nickname"),
155+
),
151156
}
152157
}
153158

‎entc/integration/migrate/migrate_test.go‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,9 @@ func TestPostgres(t *testing.T) {
109109
V1ToV2(t, drv.Dialect(), clientv1, clientv2)
110110
CheckConstraint(t, clientv2)
111111
TimePrecision(t, drv, "SELECT datetime_precision FROM information_schema.columns WHERE table_name = $1 AND column_name = $2")
112+
if version != "10" {
113+
IncludeColumns(t, drv)
114+
}
112115

113116
vdrv, err := sql.Open(dialect.Postgres, dsn+" dbname=versioned_migrate")
114117
require.NoError(t, err, "connecting to versioned migrate database")
@@ -562,6 +565,15 @@ func TimePrecision(t *testing.T, drv *sql.Driver, query string) {
562565
require.NoError(t, rows.Close())
563566
}
564567

568+
func IncludeColumns(t *testing.T, drv *sql.Driver) {
569+
rows, err := drv.QueryContext(context.Background(), "select indexdef from pg_indexes where indexname='user_workplace'")
570+
require.NoError(t, err)
571+
d, err := sql.ScanString(rows)
572+
require.NoError(t, err)
573+
require.NoError(t, rows.Close())
574+
require.Equal(t, d, "CREATE INDEX user_workplace ON public.users USING btree (workplace) INCLUDE (nickname)")
575+
}
576+
565577
func idRange(t *testing.T, id, l, h int) {
566578
require.Truef(t, id > l && id < h, "id %d should be between %d to %d", id, l, h)
567579
}

‎go.mod‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ module entgo.io/ent
33
go 1.19
44

55
require (
6-
ariga.io/atlas v0.6.2-0.20220819114704-2060066abac7
6+
ariga.io/atlas v0.6.5-0.20220907173155-3332f3c1b8c9
77
github.com/DATA-DOG/go-sqlmock v1.5.0
88
github.com/go-openapi/inflect v0.19.0
99
github.com/go-sql-driver/mysql v1.6.0

‎go.sum‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
ariga.io/atlas v0.6.2-0.20220819114704-2060066abac7 h1:qhVEfrV5Z9XZyQJxgogBq6c2pjWUxGLU7bvFeEDY0DA=
22
ariga.io/atlas v0.6.2-0.20220819114704-2060066abac7/go.mod h1:ft47uSh5hWGDCmQC9DsztZg6Xk+KagM5Ts/mZYKb9JE=
3+
ariga.io/atlas v0.6.5-0.20220907173155-3332f3c1b8c9 h1:hb7cCS3+idkvWRxKIiH0pBiyO9tJ9gRiecY+ohA//VY=
4+
ariga.io/atlas v0.6.5-0.20220907173155-3332f3c1b8c9/go.mod h1:ft47uSh5hWGDCmQC9DsztZg6Xk+KagM5Ts/mZYKb9JE=
35
cloud.google.com/go v0.26.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw=
46
github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
57
github.com/DATA-DOG/go-sqlmock v1.5.0 h1:Shsta01QNfFxHCfpW6YH2STWB0MudeXXEWMr20OEh60=

0 commit comments

Comments
 (0)