Repository: ariga/atlas
Pull request #2660 (2024-04-02 19:38:05)

sql/postgres: support marshal/eval custom index op_class and types


Patch:
diff --git a/sql/postgres/internal/postgresop/postgresop.go b/sql/postgres/internal/postgresop/postgresop.go
index 203ee142..af177c78 100644
--- a/sql/postgres/internal/postgresop/postgresop.go
+++ b/sql/postgres/internal/postgresop/postgresop.go
@@ -4,6 +4,8 @@
 
 package postgresop
 
+import "sync"
+
 // Class describes an index operator class.
 type Class struct {
 	Name    string // operator class name
@@ -209,3 +211,20 @@ var Classes = []*Class{
 	{Name: "spgist_geometry_ops_3d", Method: "SPGIST", Type: "geometry", Default: false},
 	{Name: "spgist_geometry_ops_nd", Method: "SPGIST", Type: "geometry", Default: false},
 }
+
+var (
+	mapOnce sync.Once
+	byName map[string]struct{}
+)
+
+// HasClass reports if the given operator class name exists.
+func HasClass(name string) bool {
+	mapOnce.Do(func() {
+		byName = make(map[string]struct{}, len(Classes))
+		for _, c := range Classes {
+			byName[c.Name] = struct{}{}
+		}
+	})
+	_, ok := byName[name]
+	return ok
+}
diff --git a/sql/postgres/sqlspec.go b/sql/postgres/sqlspec.go
index a2fdebc8..6fcb61b6 100644
--- a/sql/postgres/sqlspec.go
+++ b/sql/postgres/sqlspec.go
@@ -662,7 +662,14 @@ func indexSpec(idx *schema.Index) (*sqlspec.Index, error) {
 	}
 	// Avoid printing the index type if it is the default.
 	if i := (IndexType{}); sqlx.Has(idx.Attrs, &i) && strings.ToUpper(i.T) != IndexTypeBTree {
-		spec.Extra.Attrs = append(spec.Extra.Attrs, specutil.VarAttr("type", strings.ToUpper(i.T)))
+		var attr *schemahcl.Attr
+		switch strings.ToUpper(i.T) {
+		case IndexTypeBRIN, IndexTypeHash, IndexTypeGIN, IndexTypeGiST, IndexTypeSPGiST:
+			attr = specutil.VarAttr("type", strings.ToUpper(i.T))
+		default:
+			attr = schemahcl.StringAttr("type", i.T)
+		}
+		spec.Extra.Attrs = append(spec.Extra.Attrs, attr)
 	}
 	if i := (IndexPredicate{}); sqlx.Has(idx.Attrs, &i) && i.P != "" {
 		spec.Extra.Attrs = append(spec.Extra.Attrs, specutil.VarAttr("where", strconv.Quote(i.P)))
@@ -699,8 +706,10 @@ func partAttr(idx *schema.Index, part *schema.IndexPart, spec *sqlspec.IndexPart
 	case d:
 	case len(op.Params) > 0:
 		spec.Extra.Attrs = append(spec.Extra.Attrs, schemahcl.RawAttr("ops", op.String()))
-	default:
+	case postgresop.HasClass(op.String()):
 		spec.Extra.Attrs = append(spec.Extra.Attrs, specutil.VarAttr("ops", op.String()))
+	default:
+		spec.Extra.Attrs = append(spec.Extra.Attrs, schemahcl.StringAttr("ops", op.String()))
 	}
 	return nil
 }


Test patch:
diff --git a/sql/postgres/sqlspec_test.go b/sql/postgres/sqlspec_test.go
index 99ff6929..eedc3e14 100644
--- a/sql/postgres/sqlspec_test.go
+++ b/sql/postgres/sqlspec_test.go
@@ -65,6 +65,13 @@ table "table" {
 		where = "active"
 		comment = "index comment"
 	}
+	index "index_hnsw" {
+		type = "HNSW"
+		columns = [
+			table.table.column.col,
+			table.table.column.age,
+		]
+	}
 	foreign_key "accounts" {
 		columns = [
 			table.table.column.account_name,
@@ -240,6 +247,17 @@ enum "account_type" {
 				&IndexPredicate{P: "active"},
 			},
 		},
+		{
+			Name:  "index_hnsw",
+			Table: exp.Tables[0],
+			Parts: []*schema.IndexPart{
+				{SeqNo: 0, C: exp.Tables[0].Columns[0]},
+				{SeqNo: 1, C: exp.Tables[0].Columns[1]},
+			},
+			Attrs: []schema.Attr{
+				&IndexType{T: "HNSW"},
+			},
+		},
 	}
 	exp.Tables[0].ForeignKeys = []*schema.ForeignKey{
 		{
@@ -540,9 +558,51 @@ table "t" {
 		require.NoError(t, EvalHCLBytes([]byte(r), &s, nil))
 		idx = s.Tables[0].Indexes[0]
 		require.Equal(t, IndexTypeGiST, idx.Attrs[0].(*IndexType).T)
+
+		s = schema.Schema{}
+		r = fmt.Sprintf(f, `"HNSW"`)
+		require.NoError(t, EvalHCLBytes([]byte(r), &s, nil))
+		idx = s.Tables[0].Indexes[0]
+		require.Equal(t, "HNSW", idx.Attrs[0].(*IndexType).T)
 	})
 }
 
+func TestMarshalSpec_IndexType(t *testing.T) {
+	s := schema.New("public").
+		AddTables(
+			schema.NewTable("t1").
+				AddColumns(
+					schema.NewIntColumn("id", "int"),
+				).
+				AddIndexes(
+					schema.NewIndex("i1").AddAttrs(&IndexType{T: IndexTypeHash}),
+					schema.NewIndex("i2").AddAttrs(&IndexType{T: "HNSW"}),
+					schema.NewIndex("i2").AddAttrs(&IndexType{T: IndexTypeBTree}), // Default.
+				),
+		)
+	buf, err := MarshalHCL(s)
+	require.NoError(t, err)
+	require.Equal(t, `table "t1" {
+  schema = schema.public
+  column "id" {
+    null = false
+    type = int
+  }
+  index "i1" {
+    type = HASH
+  }
+  index "i2" {
+    type = "HNSW"
+  }
+  index "i2" {
+  }
+}
+schema "public" {
+}
+`,
+		string(buf))
+}
+
 func TestUnmarshalSpec_BRINIndex(t *testing.T) {
 	f := `
 schema "s" {}
@@ -610,6 +670,13 @@ func TestUnmarshalSpec_IndexOpClass(t *testing.T) {
       ops    = sql("tsvector_ops(siglen=1)")
     }
   }
+  index "idx5" {
+    unique = true
+    on {
+      column = column.a
+      ops    = "vector_cosine_ops"
+    }
+  }
 }
 schema "test" {
 }
@@ -653,6 +720,12 @@ schema "test" {
 	require.Len(t, opc.Params, 1)
 	require.Equal(t, "siglen", opc.Params[0].N)
 	require.Equal(t, "1", opc.Params[0].V)
+	// idx5
+	idx = s.Tables[0].Indexes[5]
+	require.Len(t, idx.Parts, 1)
+	require.Equal(t, "a", idx.Parts[0].C.Name)
+	require.Len(t, idx.Parts[0].Attrs, 1)
+	require.Equal(t, "vector_cosine_ops", idx.Parts[0].Attrs[0].(*IndexOpClass).Name)
 }
 
 func TestUnmarshalSpec_Partitioned(t *testing.T) {
@@ -1036,6 +1109,14 @@ func TestMarshalSpec_IndexOpClass(t *testing.T) {
 				&IndexType{T: IndexTypeGiST},
 			},
 		},
+		{
+			Name:   "idx5",
+			Table:  s.Tables[0],
+			Unique: true,
+			Parts: []*schema.IndexPart{
+				{SeqNo: 0, C: s.Tables[0].Columns[0], Attrs: []schema.Attr{&IndexOpClass{Name: "vector_cosine_ops"}}},
+			},
+		},
 	}
 	buf, err := MarshalSpec(s, hclState)
 	require.NoError(t, err)
@@ -1082,6 +1163,13 @@ func TestMarshalSpec_IndexOpClass(t *testing.T) {
       ops    = sql("tsvector_ops(siglen=1)")
     }
   }
+  index "idx5" {
+    unique = true
+    on {
+      column = column.a
+      ops    = "vector_cosine_ops"
+    }
+  }
 }
 schema "test" {
 }