Repository: ariga/atlas
Pull request #1261 (2022-11-04 06:58:55)

sql/postgres: support unmarshaling operator_class to key parts


Patch:
diff --git a/sql/postgres/sqlspec.go b/sql/postgres/sqlspec.go
index 4131a9c6..a485fab9 100644
--- a/sql/postgres/sqlspec.go
+++ b/sql/postgres/sqlspec.go
@@ -303,7 +303,7 @@ func fixDefaultQuotes(spec *sqlspec.Column) error {
 
 // convertIndex converts a sqlspec.Index into a schema.Index.
 func convertIndex(spec *sqlspec.Index, t *schema.Table) (*schema.Index, error) {
-	idx, err := specutil.Index(spec, t)
+	idx, err := specutil.Index(spec, t, convertPart)
 	if err != nil {
 		return nil, err
 	}
@@ -347,6 +347,38 @@ func convertIndex(spec *sqlspec.Index, t *schema.Table) (*schema.Index, error) {
 	return idx, nil
 }
 
+func convertPart(spec *sqlspec.IndexPart, part *schema.IndexPart) error {
+	switch opc, ok := spec.Attr("ops"); {
+	case !ok:
+	case opc.IsRawExpr():
+		expr, err := opc.RawExpr()
+		if err != nil {
+			return err
+		}
+		switch parts := strings.Split(expr.X, "("); {
+		case len(parts) == 1 && parts[0] != "":
+			part.Attrs = append(part.Attrs, &IndexOpClass{Name: parts[0]})
+		case len(parts) == 2 && strings.HasSuffix(parts[1], ")"):
+			return mayAppendOps(part, parts[0], fmt.Sprintf("{%s}", parts[1][:len(parts[1])-1]), false)
+		default:
+			return fmt.Errorf("unexpected index.on.ops expression %q", expr.X)
+		}
+	case opc.IsRef():
+		name, err := opc.Ref()
+		if err != nil {
+			return err
+		}
+		part.Attrs = append(part.Attrs, &IndexOpClass{Name: name})
+	default:
+		name, err := opc.String()
+		if err != nil {
+			return err
+		}
+		part.Attrs = append(part.Attrs, &IndexOpClass{Name: name})
+	}
+	return nil
+}
+
 const defaultTimePrecision = 6
 
 // convertColumnType converts a sqlspec.Column into a concrete Postgres schema.Type.
@@ -532,7 +564,7 @@ func partAttr(idx *schema.Index, part *schema.IndexPart, spec *sqlspec.IndexPart
 	if len(op.Params) > 0 {
 		var kv []string
 		for _, e := range op.Params {
-			kv = append(kv, fmt.Sprintf("%s='%s'", e.N, e.V))
+			kv = append(kv, fmt.Sprintf("%s=%s", e.N, e.V))
 		}
 		spec.Extra.Attrs = append(
 			spec.Extra.Attrs,


Test patch:
diff --git a/sql/postgres/sqlspec_test.go b/sql/postgres/sqlspec_test.go
index 19492f1d..71723bb0 100644
--- a/sql/postgres/sqlspec_test.go
+++ b/sql/postgres/sqlspec_test.go
@@ -316,6 +316,95 @@ table "t" {
 	require.EqualValues(t, 2, idx.Attrs[1].(*IndexStorageParams).PagesPerRange)
 }
 
+func TestUnmarshalSpec_IndexOpClass(t *testing.T) {
+	const f = `table "users" {
+  schema = schema.test
+  column "a" {
+    null = false
+    type = text
+  }
+  column "b" {
+    null = false
+    type = tsvector
+  }
+  index "idx0" {
+    unique  = true
+    columns = [column.a, column.b]
+  }
+  index "idx1" {
+    unique = true
+    on {
+      column = column.a
+      ops    = text_pattern_ops
+    }
+    on {
+      column = column.b
+    }
+  }
+  index "idx2" {
+    unique  = true
+    columns = [column.a]
+  }
+  index "idx3" {
+    unique = true
+    on {
+      column = column.a
+      ops    = text_pattern_ops
+    }
+  }
+  index "idx4" {
+    unique = true
+    type   = GIST
+    on {
+      column = column.b
+      ops    = sql("tsvector_ops(siglen=1)")
+    }
+  }
+}
+schema "test" {
+}
+`
+	var s schema.Schema
+	err := EvalHCLBytes([]byte(f), &s, nil)
+	require.NoError(t, err)
+	// idx0
+	idx := s.Tables[0].Indexes[0]
+	require.Len(t, idx.Parts, 2)
+	require.Equal(t, "a", idx.Parts[0].C.Name)
+	require.Empty(t, idx.Parts[0].Attrs)
+	require.Equal(t, "b", idx.Parts[1].C.Name)
+	require.Empty(t, idx.Parts[1].Attrs)
+	// idx1
+	idx = s.Tables[0].Indexes[1]
+	require.Len(t, idx.Parts, 2)
+	require.Equal(t, "a", idx.Parts[0].C.Name)
+	require.Len(t, idx.Parts[0].Attrs, 1)
+	require.Equal(t, "text_pattern_ops", idx.Parts[0].Attrs[0].(*IndexOpClass).Name)
+	require.Equal(t, "b", idx.Parts[1].C.Name)
+	require.Empty(t, idx.Parts[1].Attrs)
+	// idx2
+	idx = s.Tables[0].Indexes[2]
+	require.Len(t, idx.Parts, 1)
+	require.Equal(t, "a", idx.Parts[0].C.Name)
+	require.Empty(t, idx.Parts[0].Attrs)
+	// idx3
+	idx = s.Tables[0].Indexes[3]
+	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, "text_pattern_ops", idx.Parts[0].Attrs[0].(*IndexOpClass).Name)
+	// idx4
+	idx = s.Tables[0].Indexes[4]
+	require.Len(t, idx.Parts, 1)
+	require.Equal(t, "b", idx.Parts[0].C.Name)
+	require.Len(t, idx.Parts[0].Attrs, 1)
+	opc := idx.Parts[0].Attrs[0].(*IndexOpClass)
+	require.Equal(t, "tsvector_ops", opc.Name)
+	require.Len(t, opc.Params, 1)
+	require.Equal(t, "siglen", opc.Params[0].N)
+	require.Equal(t, "1", opc.Params[0].V)
+}
+
 func TestUnmarshalSpec_Partitioned(t *testing.T) {
 	t.Run("Columns", func(t *testing.T) {
 		var (
@@ -692,7 +781,7 @@ func TestMarshalSpec_IndexOpClass(t *testing.T) {
     type   = GIST
     on {
       column = column.b
-      ops    = sql("tsvector_ops(siglen='1')")
+      ops    = sql("tsvector_ops(siglen=1)")
     }
   }
 }