Repository: ariga/atlas
Pull request #1273 (2022-11-07 20:44:30)

sql/postgres: expose UnmarshalText for parsing operator classes
Will be used by Ent index annnotations


Patch:
diff --git a/sql/postgres/inspect.go b/sql/postgres/inspect.go
index 8ece5326..59536d53 100644
--- a/sql/postgres/inspect.go
+++ b/sql/postgres/inspect.go
@@ -5,6 +5,7 @@
 package postgres
 
 import (
+	"bytes"
 	"context"
 	"database/sql"
 	"fmt"
@@ -413,16 +414,8 @@ func mayAppendOps(part *schema.IndexPart, name string, params string, defaults b
 		return nil
 	}
 	op := &IndexOpClass{Name: name, Default: defaults}
-	switch {
-	case len(params) == 0:
-	case len(params) < 2, params[0] != '{', params[len(params)-1] != '}':
-		return fmt.Errorf("postgres: unexpected operator class parameters format: %q", params)
-	default:
-		for _, e := range strings.Split(params[1:len(params)-1], ",") {
-			if kv := strings.Split(strings.TrimSpace(e), "="); len(kv) == 2 {
-				op.Params = append(op.Params, struct{ N, V string }{N: kv[0], V: kv[1]})
-			}
-		}
+	if err := op.parseParams(params); err != nil {
+		return err
 	}
 	part.Attrs = append(part.Attrs, op)
 	return nil
@@ -975,6 +968,34 @@ func (o *IndexOpClass) String() string {
 	return b.String()
 }
 
+// UnmarshalText parses the operator class from its string representation.
+func (o *IndexOpClass) UnmarshalText(text []byte) error {
+	i := bytes.IndexByte(text, '(')
+	if i == -1 {
+		o.Name = string(text)
+		return nil
+	}
+	o.Name = string(text[:i])
+	return o.parseParams(string(text[i:]))
+}
+
+// parseParams parses index class parameters defined in HCL or returned
+// from the database. For example: '{k=v}', '(k1=v1,k2=v2)'.
+func (o *IndexOpClass) parseParams(kv string) error {
+	switch {
+	case kv == "":
+	case strings.HasPrefix(kv, "(") && strings.HasSuffix(kv, ")"), strings.HasPrefix(kv, "{") && strings.HasSuffix(kv, "}"):
+		for _, e := range strings.Split(kv[1:len(kv)-1], ",") {
+			if kv := strings.Split(strings.TrimSpace(e), "="); len(kv) == 2 {
+				o.Params = append(o.Params, struct{ N, V string }{N: kv[0], V: kv[1]})
+			}
+		}
+	default:
+		return fmt.Errorf("postgres: unexpected operator class parameters format: %q", kv)
+	}
+	return nil
+}
+
 // newIndexStorage parses and returns the index storage parameters.
 func newIndexStorage(opts string) (*IndexStorageParams, error) {
 	params := &IndexStorageParams{}
diff --git a/sql/postgres/sqlspec.go b/sql/postgres/sqlspec.go
index 0b27f0e6..9ad89ac0 100644
--- a/sql/postgres/sqlspec.go
+++ b/sql/postgres/sqlspec.go
@@ -354,13 +354,12 @@ func convertPart(spec *sqlspec.IndexPart, part *schema.IndexPart) error {
 		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)
+		var op IndexOpClass
+		if err := op.UnmarshalText([]byte(expr.X)); err != nil {
+			return fmt.Errorf("unexpected index.on.ops expression %q: %w", expr.X, err)
+		}
+		if op.Name != "" {
+			part.Attrs = append(part.Attrs, &op)
 		}
 	case opc.IsRef():
 		name, err := opc.Ref()


Test patch:
diff --git a/sql/postgres/inspect_test.go b/sql/postgres/inspect_test.go
index 0bedd59f..94da2c39 100644
--- a/sql/postgres/inspect_test.go
+++ b/sql/postgres/inspect_test.go
@@ -674,6 +674,19 @@ public
 	}(), realm)
 }
 
+func TestIndexOpClass_UnmarshalText(t *testing.T) {
+	var op IndexOpClass
+	require.NoError(t, op.UnmarshalText([]byte("int4_ops")))
+	require.EqualValues(t, IndexOpClass{Name: "int4_ops"}, op)
+
+	op = IndexOpClass{}
+	require.NoError(t, op.UnmarshalText([]byte("")))
+	require.EqualValues(t, IndexOpClass{}, op)
+
+	require.NoError(t, op.UnmarshalText([]byte("int4_ops(siglen=1)")))
+	require.EqualValues(t, IndexOpClass{Name: "int4_ops", Params: []struct{ N, V string }{{"siglen", "1"}}}, op)
+}
+
 type mock struct {
 	sqlmock.Sqlmock
 }