Repository: ariga/atlas
Pull request #495 (2022-01-26 12:38:11)

sql/sqlspec: support unmarshaling index-part


Patch:
diff --git a/sql/internal/specutil/convert.go b/sql/internal/specutil/convert.go
index 15016d2d..67b14206 100644
--- a/sql/internal/specutil/convert.go
+++ b/sql/internal/specutil/convert.go
@@ -153,20 +153,42 @@ func Column(spec *sqlspec.Column, conv ConvertTypeFunc) (*schema.Column, error)
 
 // Index converts a sqlspec.Index to a schema.Index.
 func Index(spec *sqlspec.Index, parent *schema.Table) (*schema.Index, error) {
-	parts := make([]*schema.IndexPart, 0, len(spec.Columns))
-	for seqno, c := range spec.Columns {
-		cn, err := columnName(c)
-		if err != nil {
-			return nil, fmt.Errorf("specutil: failed converting column to index: %w", err)
+	parts := make([]*schema.IndexPart, 0, len(spec.Columns)+len(spec.Parts))
+	switch n, m := len(spec.Columns), len(spec.Parts); {
+	case n == 0 && m == 0:
+		return nil, fmt.Errorf("missing definition for index %q", spec.Name)
+	case n > 0 && m > 0:
+		return nil, fmt.Errorf(`multiple definitions for index %q, use "columns" or "on"`, spec.Name)
+	case n > 0:
+		for i, c := range spec.Columns {
+			c, err := column(parent, c)
+			if err != nil {
+				return nil, err
+			}
+			parts = append(parts, &schema.IndexPart{
+				SeqNo: i,
+				C:     c,
+			})
 		}
-		col, ok := parent.Column(cn)
-		if !ok {
-			return nil, fmt.Errorf("specutil: unknown column %q in table %q", cn, parent.Name)
+	case m > 0:
+		for i, p := range spec.Parts {
+			part := &schema.IndexPart{SeqNo: i, Desc: p.Desc}
+			switch {
+			case p.Column == nil && p.Expr == "":
+				return nil, fmt.Errorf(`"column" or "expr" are required for index %q at position %d`, spec.Name, i)
+			case p.Column != nil && p.Expr != "":
+				return nil, fmt.Errorf(`cannot use both "column" and "expr" in index %q at position %d`, spec.Name, i)
+			case p.Expr != "":
+				part.X = &schema.RawExpr{X: p.Expr}
+			case p.Column != nil:
+				c, err := column(parent, p.Column)
+				if err != nil {
+					return nil, err
+				}
+				part.C = c
+			}
+			parts = append(parts, part)
 		}
-		parts = append(parts, &schema.IndexPart{
-			SeqNo: seqno,
-			C:     col,
-		})
 	}
 	i := &schema.Index{
 		Name:   spec.Name,
@@ -192,17 +214,13 @@ func Check(spec *sqlspec.Check) (*schema.Check, error) {
 func PrimaryKey(spec *sqlspec.PrimaryKey, parent *schema.Table) (*schema.Index, error) {
 	parts := make([]*schema.IndexPart, 0, len(spec.Columns))
 	for seqno, c := range spec.Columns {
-		n, err := columnName(c)
+		c, err := column(parent, c)
 		if err != nil {
-			return nil, fmt.Errorf("sqlspec: cannot get column name %q as primary key for table %q", c.V, parent.Name)
-		}
-		pkc, ok := parent.Column(n)
-		if !ok {
-			return nil, fmt.Errorf("sqlspec: cannot set column %q as primary key for table %q", n, parent.Name)
+			return nil, nil
 		}
 		parts = append(parts, &schema.IndexPart{
 			SeqNo: seqno,
-			C:     pkc,
+			C:     c,
 		})
 	}
 	return &schema.Index{
@@ -262,15 +280,7 @@ func resolveCol(ref *schemaspec.Ref, sch *schema.Schema) (*schema.Column, error)
 	if !ok {
 		return nil, fmt.Errorf("sqlspec: table %q not found", t)
 	}
-	c, err := columnName(ref)
-	if err != nil {
-		return nil, fmt.Errorf("sqlspec: column %q not found", ref.V)
-	}
-	col, ok := tbl.Column(c)
-	if !ok {
-		return nil, fmt.Errorf("sqlspec: column %q not found in table %q", c, t)
-	}
-	return col, nil
+	return column(tbl, ref)
 }
 
 // FromRealm converts a schema.Realm into []sqlspec.Schema and []sqlspec.Table.
@@ -476,12 +486,16 @@ func SchemaName(ref *schemaspec.Ref) (string, error) {
 	return parts[1], nil
 }
 
-func columnName(ref *schemaspec.Ref) (string, error) {
+func column(t *schema.Table, ref *schemaspec.Ref) (*schema.Column, error) {
 	s := strings.Split(ref.V, "$column.")
 	if len(s) != 2 {
-		return "", fmt.Errorf("sqlspec: failed to extract column name from %q", ref)
+		return nil, fmt.Errorf("specutil: failed to extract column name from %q", ref)
 	}
-	return s[1], nil
+	c, ok := t.Column(s[1])
+	if !ok {
+		return nil, fmt.Errorf("specutil: unknown column %q in table %q", s[1], t.Name)
+	}
+	return c, nil
 }
 
 func tableName(ref *schemaspec.Ref) (string, error) {
diff --git a/sql/schema/dsl.go b/sql/schema/dsl.go
index ba9acb7c..e89e4870 100644
--- a/sql/schema/dsl.go
+++ b/sql/schema/dsl.go
@@ -571,14 +571,25 @@ func (i *Index) AddParts(parts ...*IndexPart) *Index {
 		if p.C != nil && !p.C.hasIndex(i) {
 			p.C.Indexes = append(p.C.Indexes, i)
 		}
+		p.SeqNo = len(i.Parts)
 		i.Parts = append(i.Parts, p)
 	}
 	return i
 }
 
 // NewIndexPart creates a new index part.
-func NewIndexPart(seqno int) *IndexPart {
-	return &IndexPart{SeqNo: seqno}
+func NewIndexPart() *IndexPart { return &IndexPart{} }
+
+// NewColumnPart creates a new index part with the given column.
+func NewColumnPart(c *Column) *IndexPart { return &IndexPart{C: c} }
+
+// NewExprPart creates a new index part with the given expression.
+func NewExprPart(x Expr) *IndexPart { return &IndexPart{X: x} }
+
+// SetDesc configures the "DESC" attribute of the key part.
+func (p *IndexPart) SetDesc(b bool) *IndexPart {
+	p.Desc = b
+	return p
 }
 
 // AddAttrs adds and additional attributes to the index-part.
diff --git a/sql/sqlspec/sqlspec.go b/sql/sqlspec/sqlspec.go
index 73ac81d3..b004b00a 100644
--- a/sql/sqlspec/sqlspec.go
+++ b/sql/sqlspec/sqlspec.go
@@ -44,10 +44,19 @@ type (
 	Index struct {
 		Name    string            `spec:",name"`
 		Unique  bool              `spec:"unique"`
+		Parts   []*IndexPart      `spec:"on"`
 		Columns []*schemaspec.Ref `spec:"columns"`
 		schemaspec.DefaultExtension
 	}
 
+	// IndexPart holds a specification for the index key part.
+	IndexPart struct {
+		Desc   bool            `spec:"desc"`
+		Column *schemaspec.Ref `spec:"column"`
+		Expr   string          `spec:"expr"`
+		schemaspec.DefaultExtension
+	}
+
 	// Check holds a specification for a check constraint on a table.
 	Check struct {
 		Name string `spec:",name"`
@@ -69,26 +78,6 @@ type (
 	Type string
 )
 
-// Type<X> are Types that represent database agnostic column types to be used
-// in Atlas DDL documents.
-const (
-	TypeInt     Type = "int"
-	TypeInt8    Type = "int8"
-	TypeInt16   Type = "int16"
-	TypeInt64   Type = "int64"
-	TypeUint    Type = "uint"
-	TypeUint8   Type = "uint8"
-	TypeUint16  Type = "uint16"
-	TypeUint64  Type = "uint64"
-	TypeString  Type = "string"
-	TypeBinary  Type = "binary"
-	TypeEnum    Type = "enum"
-	TypeBoolean Type = "boolean"
-	TypeDecimal Type = "decimal"
-	TypeFloat   Type = "float"
-	TypeTime    Type = "time"
-)
-
 func init() {
 	schemaspec.Register("table", &Table{})
 	schemaspec.Register("schema", &Schema{})


Test patch:
diff --git a/sql/internal/specutil/convert_test.go b/sql/internal/specutil/convert_test.go
index 80c78602..ec29043a 100644
--- a/sql/internal/specutil/convert_test.go
+++ b/sql/internal/specutil/convert_test.go
@@ -10,13 +10,6 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func TestRef_ColumnName(t *testing.T) {
-	ref := &schemaspec.Ref{V: "$table.accounts.$column.user_active"}
-	c, err := columnName(ref)
-	require.NoError(t, err)
-	require.Equal(t, "user_active", c)
-}
-
 func TestRef_TableName(t *testing.T) {
 	ref := &schemaspec.Ref{V: "$table.accounts.$column.user_active"}
 	c, err := tableName(ref)
diff --git a/sql/mysql/migrate_test.go b/sql/mysql/migrate_test.go
index 12edf0cf..94dc2e72 100644
--- a/sql/mysql/migrate_test.go
+++ b/sql/mysql/migrate_test.go
@@ -188,9 +188,7 @@ func TestPlanChanges(t *testing.T) {
 								Type: &schema.ColumnType{Type: &schema.StringType{T: "varchar(255)"}},
 								Indexes: []*schema.Index{
 									schema.NewIndex("name_index").
-										AddParts(schema.NewIndexPart(0).SetColumn(
-											schema.NewColumn("name"),
-										)),
+										AddParts(schema.NewColumnPart(schema.NewColumn("name"))),
 								},
 							}},
 					}
@@ -199,9 +197,7 @@ func TestPlanChanges(t *testing.T) {
 						Changes: []schema.Change{
 							&schema.DropIndex{
 								I: schema.NewIndex("name_index").
-									AddParts(schema.NewIndexPart(0).SetColumn(
-										schema.NewColumn("name"),
-									)),
+									AddParts(schema.NewColumnPart(schema.NewColumn("name"))),
 							},
 						},
 					}
diff --git a/sql/mysql/sqlspec_test.go b/sql/mysql/sqlspec_test.go
index 7b70d2e2..2d57265d 100644
--- a/sql/mysql/sqlspec_test.go
+++ b/sql/mysql/sqlspec_test.go
@@ -471,6 +471,47 @@ schema "test" {
 	require.EqualValues(t, expected, string(buf))
 }
 
+func TestUnmarshalSpec_IndexParts(t *testing.T) {
+	var (
+		s schema.Schema
+		f = `
+schema "test" {}
+table "users" {
+	schema = schema.test
+	column "name" {
+		type = text
+	}
+	index "idx" {
+		on {
+			column = table.users.column.name
+			desc = true
+		}
+		on {
+			expr = "lower(name)"
+		}
+	}
+}
+`
+	)
+	err := UnmarshalHCL([]byte(f), &s)
+	require.NoError(t, err)
+	c := schema.NewStringColumn("name", "text")
+	exp := schema.New("test").
+		AddTables(
+			schema.NewTable("users").
+				AddColumns(c).
+				AddIndexes(
+					schema.NewIndex("idx").
+						AddParts(
+							schema.NewColumnPart(c).SetDesc(true),
+							schema.NewExprPart(&schema.RawExpr{X: "lower(name)"}),
+						),
+				),
+		)
+	exp.Tables[0].Columns[0].Indexes = nil
+	require.EqualValues(t, exp, &s)
+}
+
 func TestMarshalSpec_TimePrecision(t *testing.T) {
 	s := schema.New("test").
 		AddTables(