Repository: ariga/atlas
Pull request #980 (2022-07-23 19:07:45)

schemahcl: fixed invalid identity token

giautm: > Well done! Can you add an integration test using testscript to verify this works e2e on postgres/mysql?

yes, please me try. :D 

Patch:
diff --git a/schemahcl/hcl.go b/schemahcl/hcl.go
index 4f2fb77c..347b6bc7 100644
--- a/schemahcl/hcl.go
+++ b/schemahcl/hcl.go
@@ -425,12 +425,10 @@ func (s *State) writeAttr(attr *Attr, body *hclwrite.Body) error {
 	attr = normalizeLiterals(attr)
 	switch v := attr.V.(type) {
 	case *Ref:
-		expr := strings.ReplaceAll(v.V, "$", "")
-		body.SetAttributeRaw(attr.K, hclRawTokens(expr))
+		body.SetAttributeRaw(attr.K, hclRefTokens(v.V))
 	case *Type:
 		if v.IsRef {
-			expr := strings.ReplaceAll(v.T, "$", "")
-			body.SetAttributeRaw(attr.K, hclRawTokens(expr))
+			body.SetAttributeRaw(attr.K, hclRefTokens(v.T))
 			break
 		}
 		spec, ok := s.findTypeSpec(v.T)
@@ -456,19 +454,18 @@ func (s *State) writeAttr(attr *Attr, body *hclwrite.Body) error {
 		if v.V == nil {
 			return nil
 		}
-		lst := make([]string, 0, len(v.V))
+		lst := make([]hclwrite.Tokens, 0, len(v.V))
 		for _, item := range v.V {
 			switch v := item.(type) {
 			case *Ref:
-				expr := strings.ReplaceAll(v.V, "$", "")
-				lst = append(lst, expr)
+				lst = append(lst, hclRefTokens(v.V))
 			case *LiteralValue:
-				lst = append(lst, v.V)
+				lst = append(lst, hclRawTokens(v.V))
 			default:
 				return fmt.Errorf("cannot write elem type %T of attr %q to HCL list", v, attr)
 			}
 		}
-		body.SetAttributeRaw(attr.K, hclRawList(lst))
+		body.SetAttributeRaw(attr.K, hclList(lst))
 	default:
 		return fmt.Errorf("schemacl: unknown literal type %T", v)
 	}
@@ -541,6 +538,39 @@ func findAttr(attrs []*Attr, k string) (*Attr, bool) {
 	return nil, false
 }
 
+func hclRefTokens(ref string) hclwrite.Tokens {
+	var t []*hclwrite.Token
+	for i, s := range strings.Split(ref, ".") {
+		// Ignore the first $ as token for reference.
+		if len(s) > 1 && s[0] == '$' {
+			s = s[1:]
+		}
+		switch {
+		case i == 0:
+			t = append(t, hclRawTokens(s)...)
+		case hclsyntax.ValidIdentifier(s):
+			t = append(t, &hclwrite.Token{
+				Type:  hclsyntax.TokenDot,
+				Bytes: []byte{'.'},
+			}, &hclwrite.Token{
+				Type:  hclsyntax.TokenIdent,
+				Bytes: []byte(s),
+			})
+		default:
+			t = append(t, &hclwrite.Token{
+				Type:  hclsyntax.TokenOBrack,
+				Bytes: []byte{'['},
+			})
+			t = append(t, hclwrite.TokensForValue(cty.StringVal(s))...)
+			t = append(t, &hclwrite.Token{
+				Type:  hclsyntax.TokenCBrack,
+				Bytes: []byte{']'},
+			})
+		}
+	}
+	return t
+}
+
 func hclRawTokens(s string) hclwrite.Tokens {
 	return hclwrite.Tokens{
 		&hclwrite.Token{
@@ -550,7 +580,7 @@ func hclRawTokens(s string) hclwrite.Tokens {
 	}
 }
 
-func hclRawList(items []string) hclwrite.Tokens {
+func hclList(items []hclwrite.Tokens) hclwrite.Tokens {
 	t := hclwrite.Tokens{&hclwrite.Token{
 		Type:  hclsyntax.TokenOBrack,
 		Bytes: []byte("["),
@@ -559,7 +589,7 @@ func hclRawList(items []string) hclwrite.Tokens {
 		if i > 0 {
 			t = append(t, &hclwrite.Token{Type: hclsyntax.TokenComma, Bytes: []byte(",")})
 		}
-		t = append(t, &hclwrite.Token{Type: hclsyntax.TokenIdent, Bytes: []byte(item)})
+		t = append(t, item...)
 	}
 	t = append(t, &hclwrite.Token{
 		Type:  hclsyntax.TokenCBrack,
diff --git a/sql/sqlspec/migratespec.go b/sql/sqlspec/migratespec.go
deleted file mode 100644
index 9b81e321..00000000
--- a/sql/sqlspec/migratespec.go
+++ /dev/null
@@ -1,35 +0,0 @@
-// Copyright 2021-present The Atlas Authors. All rights reserved.
-// This source code is licensed under the Apache 2.0 license found
-// in the LICENSE file in the root directory of this source tree.
-
-package sqlspec
-
-import (
-	"ariga.io/atlas/schemahcl"
-)
-
-type (
-	// Change is the interface implemented by change specifications. Change instances are supposed
-	// to be mappable to schema.Change instances.
-	Change interface {
-		change()
-	}
-
-	// ModifyTable is a specification for a schema.ModifyTable.
-	ModifyTable struct {
-		Change
-		Table   string   `spec:"table"`
-		Changes []Change `spec:""`
-	}
-
-	// AddColumn is a specification for a schema.AddColumn.
-	AddColumn struct {
-		Change
-		Column *Column `spec:"column"`
-	}
-)
-
-func init() {
-	schemahcl.Register("modify_table", &ModifyTable{})
-	schemahcl.Register("add_column", &AddColumn{})
-}


Test patch:
diff --git a/internal/integration/hclsqlspec/hclsqlspec_test.go b/internal/integration/hclsqlspec/hclsqlspec_test.go
index 3dac0075..fc31558f 100644
--- a/internal/integration/hclsqlspec/hclsqlspec_test.go
+++ b/internal/integration/hclsqlspec/hclsqlspec_test.go
@@ -344,6 +344,8 @@ table "accounts" {
 	require.NoError(t, err)
 }
 
+var hcl = schemahcl.New(schemahcl.WithTypes(postgres.TypeRegistry.Specs()))
+
 func TestMarshalTopLevel(t *testing.T) {
 	c := &sqlspec.Schema{
 		Name: "schema",
diff --git a/internal/integration/hclsqlspec/migrate_test.go b/internal/integration/hclsqlspec/migrate_test.go
deleted file mode 100644
index cdd55502..00000000
--- a/internal/integration/hclsqlspec/migrate_test.go
+++ /dev/null
@@ -1,47 +0,0 @@
-// Copyright 2021-present The Atlas Authors. All rights reserved.
-// This source code is licensed under the Apache 2.0 license found
-// in the LICENSE file in the root directory of this source tree.
-
-package hclsqlspec
-
-import (
-	"testing"
-
-	"ariga.io/atlas/schemahcl"
-	"ariga.io/atlas/sql/postgres"
-	"ariga.io/atlas/sql/sqlspec"
-
-	"github.com/stretchr/testify/require"
-)
-
-var hcl = schemahcl.New(schemahcl.WithTypes(postgres.TypeRegistry.Specs()))
-
-func TestMigrate(t *testing.T) {
-	f := `
-modify_table {
-	table = "users"
-	add_column {
-		column "id" {
-			type = int
-		}
-	}
-}
-`
-	var test struct {
-		Changes []sqlspec.Change `spec:""`
-	}
-	err := hcl.EvalBytes([]byte(f), &test, nil)
-	require.NoError(t, err)
-	require.EqualValues(t, &sqlspec.ModifyTable{
-		Table: "users",
-		Changes: []sqlspec.Change{
-			&sqlspec.AddColumn{
-				Column: &sqlspec.Column{
-					Name: "id",
-					Null: false,
-					Type: &schemahcl.Type{T: "int"},
-				},
-			},
-		},
-	}, test.Changes[0])
-}
diff --git a/internal/integration/testdata/mysql/column-generated-inspect.txt b/internal/integration/testdata/mysql/column-generated-inspect.txt
index f830743a..2dc389e2 100644
--- a/internal/integration/testdata/mysql/column-generated-inspect.txt
+++ b/internal/integration/testdata/mysql/column-generated-inspect.txt
@@ -29,6 +29,20 @@ table "users" {
             expr = "b"
         }
     }
+    column "d e" {
+        null = false
+        type = varchar(255)
+    }
+    column "d $" {
+        null = false
+        type = varchar(255)
+    }
+    index "idx_1" {
+        columns = [column["d e"]]
+    }
+    index "idx_2" {
+        columns = [column["d $"]]
+    }
 }
 
 -- 1.inspect.hcl --
@@ -54,6 +68,20 @@ table "users" {
       type = STORED
     }
   }
+  column "d e" {
+    null = false
+    type = varchar(255)
+  }
+  column "d $" {
+    null = false
+    type = varchar(255)
+  }
+  index "idx_1" {
+    columns = [column["d e"]]
+  }
+  index "idx_2" {
+    columns = [column["d $"]]
+  }
 }
 schema "$db" {
   charset = "$charset"
diff --git a/sql/mysql/sqlspec_test.go b/sql/mysql/sqlspec_test.go
index 5e26c466..1420b13d 100644
--- a/sql/mysql/sqlspec_test.go
+++ b/sql/mysql/sqlspec_test.go
@@ -585,16 +585,23 @@ table "users" {
 
 func TestMarshalSpec_IndexParts(t *testing.T) {
 	c := schema.NewStringColumn("name", "text")
+	c2 := schema.NewStringColumn("Full Name", "text")
 	s := schema.New("test").
 		AddTables(
 			schema.NewTable("users").
-				AddColumns(c).
+				AddColumns(c, c2).
 				AddIndexes(
 					schema.NewIndex("idx").
 						AddParts(
 							schema.NewColumnPart(c).SetDesc(true).AddAttrs(&SubPart{Len: 10}),
 							schema.NewExprPart(&schema.RawExpr{X: "lower(name)"}),
 						),
+					schema.NewIndex("idx2").
+						AddParts(
+							schema.NewColumnPart(c2).SetDesc(true).AddAttrs(&SubPart{Len: 10}),
+						),
+					schema.NewIndex("idx3").
+						AddParts(schema.NewColumnPart(c2)),
 				),
 		)
 	buf, err := MarshalHCL(s)
@@ -605,6 +612,10 @@ func TestMarshalSpec_IndexParts(t *testing.T) {
     null = false
     type = text
   }
+  column "Full Name" {
+    null = false
+    type = text
+  }
   index "idx" {
     on {
       desc   = true
@@ -615,11 +626,21 @@ func TestMarshalSpec_IndexParts(t *testing.T) {
       expr = "lower(name)"
     }
   }
+  index "idx2" {
+    on {
+      desc   = true
+      column = column["Full Name"]
+      prefix = 10
+    }
+  }
+  index "idx3" {
+    columns = [column["Full Name"]]
+  }
 }
 schema "test" {
 }
 `
-	require.EqualValues(t, exp, buf)
+	require.EqualValues(t, exp, string(buf))
 }
 
 func TestMarshalSpec_TimePrecision(t *testing.T) {