Repository: ariga/atlas
Pull request #1262 (2022-11-04 18:39:54)

sql/postgres: support operator class writing

Liooo: How timely, I was just searching for this and ended up here. Really appreciate the great work of you guys.

Patch:
diff --git a/sql/postgres/inspect.go b/sql/postgres/inspect.go
index 2f9f927b..c4789178 100644
--- a/sql/postgres/inspect.go
+++ b/sql/postgres/inspect.go
@@ -11,8 +11,10 @@ import (
 	"regexp"
 	"strconv"
 	"strings"
+	"sync"
 
 	"ariga.io/atlas/sql/internal/sqlx"
+	"ariga.io/atlas/sql/postgres/internal/postgresop"
 	"ariga.io/atlas/sql/schema"
 )
 
@@ -895,6 +897,68 @@ func (s *SerialType) sequence(t *schema.Table, c *schema.Column) string {
 	return fmt.Sprintf("%s_%s_seq", t.Name, c.Name)
 }
 
+var (
+	opsOnce    sync.Once
+	defaultOps map[postgresop.Class]bool
+)
+
+// DefaultFor reports if the operator_class is the default for the index part.
+func (o *IndexOpClass) DefaultFor(idx *schema.Index, part *schema.IndexPart) (bool, error) {
+	if len(o.Params) > 0 {
+		return false, nil
+	}
+	var it IndexType
+	// The type of the key need to be known in order to check
+	// if it is the default operator class.
+	if part.X != nil || !sqlx.Has(idx.Attrs, &it) || o.Default {
+		return false, nil
+	}
+	opsOnce.Do(func() {
+		defaultOps = make(map[postgresop.Class]bool, len(postgresop.Classes))
+		for _, op := range postgresop.Classes {
+			if op.Default {
+				defaultOps[postgresop.Class{Name: op.Name, Method: op.Method, Type: op.Type}] = true
+			}
+		}
+	})
+	var (
+		t   string
+		err error
+	)
+	switch typ := part.C.Type.Type.(type) {
+	case *schema.EnumType:
+		t = "anyenum"
+	case *ArrayType:
+		t = "anyarray"
+	default:
+		t, err = FormatType(typ)
+		if err != nil {
+			return false, fmt.Errorf("postgres: format operator-class type %T: %w", typ, err)
+		}
+	}
+	return defaultOps[postgresop.Class{Name: o.Name, Method: strings.ToUpper(it.T), Type: t}], nil
+}
+
+// String returns the string representation of the operator class.
+func (o *IndexOpClass) String() string {
+	if len(o.Params) == 0 {
+		return o.Name
+	}
+	var b strings.Builder
+	b.WriteString(o.Name)
+	b.WriteString("(")
+	for i, p := range o.Params {
+		if i > 0 {
+			b.WriteString(", ")
+		}
+		b.WriteString(p.N)
+		b.WriteString("=")
+		b.WriteString(p.V)
+	}
+	b.WriteString(")")
+	return b.String()
+}
+
 // newIndexStorage parses and returns the index storage parameters.
 func newIndexStorage(opts string) (*IndexStorageParams, error) {
 	params := &IndexStorageParams{}
diff --git a/sql/postgres/migrate.go b/sql/postgres/migrate.go
index 17c2dfe8..15f0d881 100644
--- a/sql/postgres/migrate.go
+++ b/sql/postgres/migrate.go
@@ -155,7 +155,9 @@ func (s *state) addTable(ctx context.Context, add *schema.AddTable) error {
 		})
 		if pk := add.T.PrimaryKey; pk != nil {
 			b.Comma().P("PRIMARY KEY")
-			s.indexParts(b, pk.Parts)
+			if err := s.indexParts(b, pk); err != nil {
+				errs = append(errs, err.Error())
+			}
 		}
 		if len(add.T.ForeignKeys) > 0 {
 			b.Comma()
@@ -184,7 +186,9 @@ func (s *state) addTable(ctx context.Context, add *schema.AddTable) error {
 		Comment: fmt.Sprintf("create %q table", add.T.Name),
 		Reverse: s.Build("DROP TABLE").Table(add.T).String(),
 	})
-	s.addIndexes(add.T, add.T.Indexes...)
+	if err := s.addIndexes(add.T, add.T.Indexes...); err != nil {
+		return err
+	}
 	s.addComments(add.T)
 	return nil
 }
@@ -377,7 +381,9 @@ func (s *state) alterTable(t *schema.Table, changes []schema.Change) error {
 				}
 			case *schema.AddIndex:
 				b.P("ADD CONSTRAINT").Ident(change.I.Name).P("UNIQUE")
-				s.indexParts(b, change.I.Parts)
+				if err := s.indexParts(b, change.I); err != nil {
+					return err
+				}
 				// Skip reversing this operation as it is the inverse of
 				// the operation below and should not be used besides this.
 			case *schema.DropIndex:
@@ -780,7 +786,7 @@ func (s *state) mayDropEnum(alter *alterChange, ns *schema.Schema, e *schema.Enu
 	return nil
 }
 
-func (s *state) addIndexes(t *schema.Table, indexes ...*schema.Index) {
+func (s *state) addIndexes(t *schema.Table, indexes ...*schema.Index) error {
 	for _, idx := range indexes {
 		b := s.Build("CREATE")
 		if idx.Unique {
@@ -794,7 +800,9 @@ func (s *state) addIndexes(t *schema.Table, indexes ...*schema.Index) {
 			b.Ident(idx.Name)
 		}
 		b.P("ON").Table(t)
-		s.index(b, idx)
+		if err := s.index(b, idx); err != nil {
+			return err
+		}
 		s.append(&migrate.Change{
 			Cmd:     b.String(),
 			Comment: fmt.Sprintf("create index %q to table: %q", idx.Name, t.Name),
@@ -815,6 +823,7 @@ func (s *state) addIndexes(t *schema.Table, indexes ...*schema.Index) {
 			}(),
 		})
 	}
+	return nil
 }
 
 func (s *state) column(b *sqlx.Builder, t *schema.Table, c *schema.Column) error {
@@ -877,28 +886,41 @@ func (s *state) columnDefault(b *sqlx.Builder, c *schema.Column) {
 		}
 		b.P("DEFAULT", v)
 	case *schema.RawExpr:
-		// Ignore identity functions added by the differ.
+		// Ignore identity functions added by the Differ.
 		if _, ok := c.Type.Type.(*SerialType); !ok {
 			b.P("DEFAULT", x.X)
 		}
 	}
 }
 
-func (s *state) indexParts(b *sqlx.Builder, parts []*schema.IndexPart) {
+func (s *state) indexParts(b *sqlx.Builder, idx *schema.Index) (err error) {
 	b.Wrap(func(b *sqlx.Builder) {
-		b.MapComma(parts, func(i int, b *sqlx.Builder) {
-			switch part := parts[i]; {
+		err = b.MapCommaErr(idx.Parts, func(i int, b *sqlx.Builder) error {
+			switch part := idx.Parts[i]; {
 			case part.C != nil:
 				b.Ident(part.C.Name)
 			case part.X != nil:
 				b.WriteString(sqlx.MayWrap(part.X.(*schema.RawExpr).X))
 			}
-			s.partAttrs(b, parts[i])
+			return s.partAttrs(b, idx, idx.Parts[i])
 		})
 	})
+	return
 }
 
-func (s *state) partAttrs(b *sqlx.Builder, p *schema.IndexPart) {
+func (s *state) partAttrs(b *sqlx.Builder, idx *schema.Index, p *schema.IndexPart) error {
+	if c := (schema.Collation{}); sqlx.Has(p.Attrs, &c) {
+		b.P("COLLATE").Ident(c.V)
+	}
+	if op := (IndexOpClass{}); sqlx.Has(p.Attrs, &op) {
+		d, err := op.DefaultFor(idx, p)
+		if err != nil {
+			return err
+		}
+		if !d {
+			b.P(op.String())
+		}
+	}
 	if p.Desc {
 		b.P("DESC")
 	}
@@ -915,20 +937,23 @@ func (s *state) partAttrs(b *sqlx.Builder, p *schema.IndexPart) {
 			case !p.Desc && attr.NullsFirst:
 				b.P("NULL FIRST")
 			}
-		case *schema.Collation:
-			b.P("COLLATE").Ident(attr.V)
+		// Handled above.
+		case *IndexOpClass, *schema.Collation:
 		default:
-			panic(fmt.Sprintf("unexpected index part attribute: %T", attr))
+			return fmt.Errorf("postgres: unexpected index part attribute: %T", attr)
 		}
 	}
+	return nil
 }
 
-func (s *state) index(b *sqlx.Builder, idx *schema.Index) {
+func (s *state) index(b *sqlx.Builder, idx *schema.Index) error {
 	// Avoid appending the default method.
 	if t := (IndexType{}); sqlx.Has(idx.Attrs, &t) && strings.ToUpper(t.T) != IndexTypeBTree {
 		b.P("USING", t.T)
 	}
-	s.indexParts(b, idx.Parts)
+	if err := s.indexParts(b, idx); err != nil {
+		return err
+	}
 	if c := (IndexInclude{}); sqlx.Has(idx.Attrs, &c) {
 		b.P("INCLUDE")
 		b.Wrap(func(b *sqlx.Builder) {
@@ -957,9 +982,10 @@ func (s *state) index(b *sqlx.Builder, idx *schema.Index) {
 		switch attr.(type) {
 		case *schema.Comment, *ConType, *IndexType, *IndexInclude, *Concurrently, *IndexPredicate, *IndexStorageParams:
 		default:
-			panic(fmt.Sprintf("unexpected index attribute: %T", attr))
+			return fmt.Errorf("postgres: unexpected index attribute: %T", attr)
 		}
 	}
+	return nil
 }
 
 func (s *state) fks(b *sqlx.Builder, fks ...*schema.ForeignKey) {
diff --git a/sql/postgres/sqlspec.go b/sql/postgres/sqlspec.go
index a485fab9..08a21847 100644
--- a/sql/postgres/sqlspec.go
+++ b/sql/postgres/sqlspec.go
@@ -9,7 +9,6 @@ import (
 	"reflect"
 	"strconv"
 	"strings"
-	"sync"
 
 	"ariga.io/atlas/schemahcl"
 	"ariga.io/atlas/sql/internal/specutil"
@@ -561,29 +560,17 @@ func partAttr(idx *schema.Index, part *schema.IndexPart, spec *sqlspec.IndexPart
 	if !sqlx.Has(part.Attrs, &op) {
 		return nil
 	}
-	if len(op.Params) > 0 {
-		var kv []string
-		for _, e := range op.Params {
-			kv = append(kv, fmt.Sprintf("%s=%s", e.N, e.V))
-		}
-		spec.Extra.Attrs = append(
-			spec.Extra.Attrs,
-			schemahcl.RawAttr("ops", fmt.Sprintf("%s(%s)", op.Name, strings.Join(kv, ", "))),
-		)
-		return nil
-	}
-	var it IndexType
-	// The type of the key need to be known in order to check
-	// if it is the default operator class.
-	if part.X != nil || !sqlx.Has(idx.Attrs, &it) || op.Default {
-		return nil
-	}
-	d, err := defaultOp(op.Name, it.T, part.C.Type.Type)
-	if err != nil {
-		return fmt.Errorf("postgres: checking if operator_class is default: %w", err)
+	if d, err := op.DefaultFor(idx, part); err != nil || d {
+		return err
 	}
-	if !d {
-		spec.Extra.Attrs = append(spec.Extra.Attrs, specutil.VarAttr("ops", op.Name))
+	switch d, err := op.DefaultFor(idx, part); {
+	case err != nil:
+		return err
+	case d:
+	case len(op.Params) > 0:
+		spec.Extra.Attrs = append(spec.Extra.Attrs, schemahcl.RawAttr("ops", op.String()))
+	default:
+		spec.Extra.Attrs = append(spec.Extra.Attrs, specutil.VarAttr("ops", op.String()))
 	}
 	return nil
 }
@@ -791,36 +778,3 @@ func formatTime() schemahcl.TypeSpecOption {
 
 // generatedType returns the default and only type for a generated column.
 func generatedType(string) string { return "STORED" }
-
-var (
-	opsOnce    sync.Once
-	defaultOps map[postgresop.Class]bool
-)
-
-// defaultOp reports if the given operator class name is the default for the index part.
-func defaultOp(name, method string, typ schema.Type) (bool, error) {
-	opsOnce.Do(func() {
-		defaultOps = make(map[postgresop.Class]bool, len(postgresop.Classes))
-		for _, op := range postgresop.Classes {
-			if op.Default {
-				defaultOps[postgresop.Class{Name: op.Name, Method: op.Method, Type: op.Type}] = true
-			}
-		}
-	})
-	var (
-		t   string
-		err error
-	)
-	switch typ := typ.(type) {
-	case *schema.EnumType:
-		t = "anyenum"
-	case *ArrayType:
-		t = "anyarray"
-	default:
-		t, err = FormatType(typ)
-		if err != nil {
-			return false, err
-		}
-	}
-	return defaultOps[postgresop.Class{Name: name, Method: strings.ToUpper(method), Type: t}], nil
-}


Test patch:
diff --git a/sql/postgres/migrate_test.go b/sql/postgres/migrate_test.go
index e2789dc0..d3b73d18 100644
--- a/sql/postgres/migrate_test.go
+++ b/sql/postgres/migrate_test.go
@@ -406,6 +406,19 @@ func TestPlanChanges(t *testing.T) {
 									},
 								},
 							},
+							&schema.AddIndex{
+								I: &schema.Index{
+									Name: "operator_class",
+									Parts: []*schema.IndexPart{
+										{C: users.Columns[0], Attrs: []schema.Attr{&IndexOpClass{Name: "int8_bloom_ops"}}},
+										{C: users.Columns[0], Attrs: []schema.Attr{&IndexOpClass{Name: "int8_minmax_ops"}}},
+										{C: users.Columns[0], Attrs: []schema.Attr{&IndexOpClass{Name: "int8_minmax_multi_ops", Params: []struct{ N, V string }{{"values_per_range", "8"}}}}},
+									},
+									Attrs: []schema.Attr{
+										&IndexType{T: IndexTypeBRIN},
+									},
+								},
+							},
 						},
 					}
 				}(),
@@ -434,6 +447,10 @@ func TestPlanChanges(t *testing.T) {
 						Cmd:     `CREATE INDEX CONCURRENTLY "concurrently" ON "users" ("id")`,
 						Reverse: `DROP INDEX CONCURRENTLY "concurrently"`,
 					},
+					{
+						Cmd:     `CREATE INDEX "operator_class" ON "users" USING BRIN ("id" int8_bloom_ops, "id", "id" int8_minmax_multi_ops(values_per_range=8))`,
+						Reverse: `DROP INDEX "operator_class"`,
+					},
 					{
 						Cmd:     `COMMENT ON COLUMN "users" ."name" IS 'foo'`,
 						Reverse: `COMMENT ON COLUMN "users" ."name" IS ''`,