dialect/sql: add support for index operator-class using atlasgo.io (#3073)

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
 	}

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 {

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})