Repository: ariga/atlas
Pull request #727 (2022-04-21 13:37:32)

sql/sqlite: finish generated columns implementation 


Patch:
diff --git a/sql/mysql/diff.go b/sql/mysql/diff.go
index d75be297..c503bb13 100644
--- a/sql/mysql/diff.go
+++ b/sql/mysql/diff.go
@@ -381,7 +381,7 @@ func (*diff) generatedChanged(from, to *schema.Column) (bool, error) {
 		fromX, toX     schema.GeneratedExpr
 		fromHas, toHas = sqlx.Has(from.Attrs, &fromX), sqlx.Has(to.Attrs, &toX)
 	)
-	if !fromHas && !toHas || fromHas && toHas && fromX.Expr == toX.Expr && storedOrVirtual(fromX.Type) == storedOrVirtual(toX.Type) {
+	if !fromHas && !toHas || fromHas && toHas && sqlx.MayWrap(fromX.Expr) == sqlx.MayWrap(toX.Expr) && storedOrVirtual(fromX.Type) == storedOrVirtual(toX.Type) {
 		return false, nil
 	}
 	return true, checkChangeGenerated(from, to)
diff --git a/sql/postgres/diff.go b/sql/postgres/diff.go
index bb793db7..18fc96d6 100644
--- a/sql/postgres/diff.go
+++ b/sql/postgres/diff.go
@@ -89,7 +89,7 @@ func (d *diff) defaultChanged(from, to *schema.Column) (bool, error) {
 func (*diff) generatedChanged(from, to *schema.Column) (bool, error) {
 	var fromX, toX schema.GeneratedExpr
 	switch fromHas, toHas := sqlx.Has(from.Attrs, &fromX), sqlx.Has(to.Attrs, &toX); {
-	case fromHas && toHas && fromX.Expr != toX.Expr:
+	case fromHas && toHas && sqlx.MayWrap(fromX.Expr) != sqlx.MayWrap(toX.Expr):
 		return false, fmt.Errorf("changing the generation expression for a column %q is not supported", from.Name)
 	case !fromHas && toHas:
 		return false, fmt.Errorf("changing column %q to generated column is not supported (drop and add is required)", from.Name)
diff --git a/sql/sqlite/diff.go b/sql/sqlite/diff.go
index 8001abab..bfbaf7c9 100644
--- a/sql/sqlite/diff.go
+++ b/sql/sqlite/diff.go
@@ -55,10 +55,13 @@ func (d *diff) ColumnChange(from, to *schema.Column) (schema.ChangeKind, error)
 	if changed := d.defaultChanged(from, to); changed {
 		change |= schema.ChangeDefault
 	}
+	if d.generatedChanged(from, to) {
+		change |= schema.ChangeGenerated
+	}
 	return change, nil
 }
 
-// typeChanged reports if the a column type was changed.
+// typeChanged reports if the column type was changed.
 func (d *diff) typeChanged(from, to *schema.Column) (bool, error) {
 	fromT, toT := from.Type.Type, to.Type.Type
 	if fromT == nil || toT == nil {
@@ -68,8 +71,7 @@ func (d *diff) typeChanged(from, to *schema.Column) (bool, error) {
 	return reflect.TypeOf(fromT) != reflect.TypeOf(toT), nil
 }
 
-// defaultChanged reports if the a default value of a column
-// type was changed.
+// defaultChanged reports if the default value of a column was changed.
 func (d *diff) defaultChanged(from, to *schema.Column) bool {
 	d1, ok1 := sqlx.DefaultValue(from)
 	d2, ok2 := sqlx.DefaultValue(to)
@@ -84,6 +86,15 @@ func (d *diff) defaultChanged(from, to *schema.Column) bool {
 	return err1 != nil || err2 != nil || x1 != x2
 }
 
+// generatedChanged reports if the generated expression of a column was changed.
+func (*diff) generatedChanged(from, to *schema.Column) bool {
+	var (
+		fromX, toX     schema.GeneratedExpr
+		fromHas, toHas = sqlx.Has(from.Attrs, &fromX), sqlx.Has(to.Attrs, &toX)
+	)
+	return fromHas != toHas || fromHas && (sqlx.MayWrap(fromX.Expr) != sqlx.MayWrap(toX.Expr) || storedOrVirtual(fromX.Type) != storedOrVirtual(toX.Type))
+}
+
 // IsGeneratedIndexName reports if the index name was generated by the database.
 // See: https://github.com/sqlite/sqlite/blob/e937df8/src/build.c#L3583.
 func (d *diff) IsGeneratedIndexName(t *schema.Table, idx *schema.Index) bool {
diff --git a/sql/sqlite/migrate.go b/sql/sqlite/migrate.go
index 5436d7d4..0c529f9e 100644
--- a/sql/sqlite/migrate.go
+++ b/sql/sqlite/migrate.go
@@ -182,7 +182,11 @@ func (s *state) modifyTable(ctx context.Context, modify *schema.ModifyTable) err
 }
 
 func (s *state) column(b *sqlx.Builder, c *schema.Column) error {
-	b.Ident(c.Name).P(mustFormat(c.Type.Type))
+	t, err := FormatType(c.Type.Type)
+	if err != nil {
+		return err
+	}
+	b.Ident(c.Name).P(t)
 	if !c.Type.Null {
 		b.P("NOT")
 	}
@@ -194,8 +198,15 @@ func (s *state) column(b *sqlx.Builder, c *schema.Column) error {
 		}
 		b.P("DEFAULT", x)
 	}
-	if sqlx.Has(c.Attrs, &AutoIncrement{}) {
+	switch hasA, hasX := sqlx.Has(c.Attrs, &AutoIncrement{}), sqlx.Has(c.Attrs, &schema.GeneratedExpr{}); {
+	case hasA && hasX:
+		return fmt.Errorf("both autoincrement and generation expression specified for column %q", c.Name)
+	case hasA:
 		b.P("PRIMARY KEY AUTOINCREMENT")
+	case hasX:
+		x := &schema.GeneratedExpr{}
+		sqlx.Has(c.Attrs, x)
+		b.P("AS", sqlx.MayWrap(x.Expr), x.Type)
 	}
 	return nil
 }
@@ -429,6 +440,10 @@ func alterable(modify *schema.ModifyTable) bool {
 			if len(change.C.Indexes) > 0 || len(change.C.ForeignKeys) > 0 || change.C.Default != nil {
 				return false
 			}
+			// Only VIRTUAL generated columns can be added using ALTER TABLE.
+			if x := (schema.GeneratedExpr{}); sqlx.Has(change.C.Attrs, &x) && storedOrVirtual(x.Type) == stored {
+				return false
+			}
 		default:
 			return false
 		}
diff --git a/sql/sqlite/sqlspec.go b/sql/sqlite/sqlspec.go
index 42c5be56..73029e25 100644
--- a/sql/sqlite/sqlspec.go
+++ b/sql/sqlite/sqlspec.go
@@ -3,6 +3,7 @@ package sqlite
 import (
 	"reflect"
 	"strconv"
+	"strings"
 
 	"ariga.io/atlas/schema/schemaspec"
 	"ariga.io/atlas/schema/schemaspec/schemahcl"
@@ -60,6 +61,9 @@ func convertColumn(spec *sqlspec.Column, _ *schema.Table) (*schema.Column, error
 			c.AddAttrs(&AutoIncrement{})
 		}
 	}
+	if err := specutil.ConvertGenExpr(spec.Remain(), c, storedOrVirtual); err != nil {
+		return nil, err
+	}
 	return c, nil
 }
 
@@ -105,6 +109,9 @@ func columnSpec(c *schema.Column, _ *schema.Table) (*sqlspec.Column, error) {
 	if sqlx.Has(c.Attrs, &AutoIncrement{}) {
 		s.Extra.Attrs = append(s.Extra.Attrs, specutil.BoolAttr("auto_increment", true))
 	}
+	if x := (schema.GeneratedExpr{}); sqlx.Has(c.Attrs, &x) {
+		s.Extra.Children = append(s.Extra.Children, specutil.FromGenExpr(x, storedOrVirtual))
+	}
 	return s, nil
 }
 
@@ -169,3 +176,12 @@ var (
 		return MarshalSpec(v, hclState)
 	})
 )
+
+// storedOrVirtual returns a STORED or VIRTUAL
+// generated type option based on the given string.
+func storedOrVirtual(s string) string {
+	if s = strings.ToUpper(s); s == "" {
+		return virtual
+	}
+	return s
+}


Test patch:
diff --git a/sql/mysql/diff_test.go b/sql/mysql/diff_test.go
index 65d3e491..bba06bfb 100644
--- a/sql/mysql/diff_test.go
+++ b/sql/mysql/diff_test.go
@@ -244,6 +244,8 @@ func TestDiff_TableDiff(t *testing.T) {
 							SetGeneratedExpr(&schema.GeneratedExpr{Expr: "1", Type: "VIRTUAL"}),
 						schema.NewIntColumn("c5", "int").
 							SetGeneratedExpr(&schema.GeneratedExpr{Expr: "1", Type: "PERSISTENT"}),
+						schema.NewIntColumn("c6", "int").
+							SetGeneratedExpr(&schema.GeneratedExpr{Expr: "1", Type: "PERSISTENT"}),
 					)
 				to = schema.NewTable("t1").
 					SetSchema(s).
@@ -261,6 +263,8 @@ func TestDiff_TableDiff(t *testing.T) {
 							SetGeneratedExpr(&schema.GeneratedExpr{Expr: "1"}),
 						schema.NewIntColumn("c5", "int").
 							SetGeneratedExpr(&schema.GeneratedExpr{Expr: "1", Type: "STORED"}),
+						schema.NewIntColumn("c6", "int").
+							SetGeneratedExpr(&schema.GeneratedExpr{Expr: "(1)", Type: "PERSISTENT"}),
 					)
 			)
 			return testcase{
diff --git a/sql/sqlite/diff_test.go b/sql/sqlite/diff_test.go
index c44ecd64..aa326fab 100644
--- a/sql/sqlite/diff_test.go
+++ b/sql/sqlite/diff_test.go
@@ -135,6 +135,47 @@ func TestDiff_TableDiff(t *testing.T) {
 				},
 			}
 		}(),
+		func() testcase {
+			var (
+				s    = schema.New("public")
+				from = schema.NewTable("t1").
+					SetSchema(s).
+					AddColumns(
+						schema.NewIntColumn("c1", "int"),
+						schema.NewIntColumn("c2", "int").
+							SetGeneratedExpr(&schema.GeneratedExpr{Expr: "1", Type: "STORED"}),
+						schema.NewIntColumn("c3", "int").
+							SetGeneratedExpr(&schema.GeneratedExpr{Expr: "1"}),
+						schema.NewIntColumn("c4", "int").
+							SetGeneratedExpr(&schema.GeneratedExpr{Expr: "1", Type: "VIRTUAL"}),
+					)
+				to = schema.NewTable("t1").
+					SetSchema(s).
+					AddColumns(
+						// Add generated expression.
+						schema.NewIntColumn("c1", "int").
+							SetGeneratedExpr(&schema.GeneratedExpr{Expr: "1", Type: "STORED"}),
+						// Drop generated expression.
+						schema.NewIntColumn("c2", "int"),
+						// Modify generated expression.
+						schema.NewIntColumn("c3", "int").
+							SetGeneratedExpr(&schema.GeneratedExpr{Expr: "2"}),
+						// No change.
+						schema.NewIntColumn("c4", "int").
+							SetGeneratedExpr(&schema.GeneratedExpr{Expr: "1"}),
+					)
+			)
+			return testcase{
+				name: "modify column generated",
+				from: from,
+				to:   to,
+				wantChanges: []schema.Change{
+					&schema.ModifyColumn{From: from.Columns[0], To: to.Columns[0], Change: schema.ChangeGenerated},
+					&schema.ModifyColumn{From: from.Columns[1], To: to.Columns[1], Change: schema.ChangeGenerated},
+					&schema.ModifyColumn{From: from.Columns[2], To: to.Columns[2], Change: schema.ChangeGenerated},
+				},
+			}
+		}(),
 		func() testcase {
 			var (
 				from = &schema.Table{
diff --git a/sql/sqlite/migrate_test.go b/sql/sqlite/migrate_test.go
index 2a659c54..f244504f 100644
--- a/sql/sqlite/migrate_test.go
+++ b/sql/sqlite/migrate_test.go
@@ -7,6 +7,7 @@ package sqlite
 import (
 	"context"
 	"database/sql"
+	"strconv"
 	"testing"
 
 	"ariga.io/atlas/sql/internal/sqltest"
@@ -127,6 +128,62 @@ func TestPlanChanges(t *testing.T) {
 				},
 			},
 		},
+		// Add VIRTUAL column.
+		{
+			changes: []schema.Change{
+				func() schema.Change {
+					users := schema.NewTable("users").
+						AddColumns(schema.NewIntColumn("id", "bigint"))
+					return &schema.ModifyTable{
+						T: users,
+						Changes: []schema.Change{
+							&schema.AddColumn{
+								C: schema.NewIntColumn("nid", "bigint").
+									SetGeneratedExpr(&schema.GeneratedExpr{Expr: "1"}),
+							},
+						},
+					}
+				}(),
+			},
+			plan: &migrate.Plan{
+				Transactional: true,
+				Changes: []*migrate.Change{
+					{Cmd: "ALTER TABLE `users` ADD COLUMN `nid` bigint NOT NULL AS (1)"},
+				},
+			},
+		},
+		// Add STORED column.
+		{
+			changes: []schema.Change{
+				func() schema.Change {
+					users := schema.NewTable("users").
+						AddColumns(
+							schema.NewIntColumn("id", "bigint"),
+							schema.NewIntColumn("nid", "bigint").
+								SetGeneratedExpr(&schema.GeneratedExpr{Expr: "1", Type: "STORED"}),
+						)
+					return &schema.ModifyTable{
+						T: users,
+						Changes: []schema.Change{
+							&schema.AddColumn{
+								C: users.Columns[1],
+							},
+						},
+					}
+				}(),
+			},
+			plan: &migrate.Plan{
+				Transactional: true,
+				Changes: []*migrate.Change{
+					{Cmd: "PRAGMA foreign_keys = off"},
+					{Cmd: "CREATE TABLE `new_users` (`id` bigint NOT NULL, `nid` bigint NOT NULL AS (1) STORED)", Reverse: "DROP TABLE `new_users`"},
+					{Cmd: "INSERT INTO `new_users` (`id`) SELECT `id` FROM `users`"},
+					{Cmd: "DROP TABLE `users`"},
+					{Cmd: "ALTER TABLE `new_users` RENAME TO `users`"},
+					{Cmd: "PRAGMA foreign_keys = on"},
+				},
+			},
+		},
 		{
 			changes: []schema.Change{
 				func() schema.Change {
@@ -165,23 +222,25 @@ func TestPlanChanges(t *testing.T) {
 			},
 		},
 	}
-	for _, tt := range tests {
-		db, mk, err := sqlmock.New()
-		require.NoError(t, err)
-		m := mock{mk}
-		m.systemVars("3.36.0")
-		if tt.mock != nil {
-			tt.mock(m)
-		}
-		drv, err := Open(db)
-		require.NoError(t, err)
-		plan, err := drv.PlanChanges(context.Background(), "plan", tt.changes)
-		require.NoError(t, err)
-		require.Equal(t, tt.plan.Reversible, plan.Reversible)
-		require.Equal(t, tt.plan.Transactional, plan.Transactional)
-		for i, c := range plan.Changes {
-			require.Equal(t, tt.plan.Changes[i].Cmd, c.Cmd)
-			require.Equal(t, tt.plan.Changes[i].Reverse, c.Reverse)
-		}
+	for i, tt := range tests {
+		t.Run(strconv.Itoa(i), func(t *testing.T) {
+			db, mk, err := sqlmock.New()
+			require.NoError(t, err)
+			m := mock{mk}
+			m.systemVars("3.36.0")
+			if tt.mock != nil {
+				tt.mock(m)
+			}
+			drv, err := Open(db)
+			require.NoError(t, err)
+			plan, err := drv.PlanChanges(context.Background(), "plan", tt.changes)
+			require.NoError(t, err)
+			require.Equal(t, tt.plan.Reversible, plan.Reversible)
+			require.Equal(t, tt.plan.Transactional, plan.Transactional)
+			for i, c := range plan.Changes {
+				require.Equal(t, tt.plan.Changes[i].Cmd, c.Cmd)
+				require.Equal(t, tt.plan.Changes[i].Reverse, c.Reverse)
+			}
+		})
 	}
 }