Repository: ariga/atlas
Pull request #600 (2022-02-23 21:45:16)

sql/postgres: support capital tables
…cause "to_regclass" will always return null if a table has capital letters


Patch:
diff --git a/sql/mysql/migrate.go b/sql/mysql/migrate.go
index 4377edff..589b38c6 100644
--- a/sql/mysql/migrate.go
+++ b/sql/mysql/migrate.go
@@ -95,10 +95,11 @@ func (s *state) topLevel(changes []schema.Change) ([]schema.Change, error) {
 	for _, c := range changes {
 		switch c := c.(type) {
 		case *schema.AddSchema:
-			b := Build("CREATE DATABASE").Ident(c.S.Name)
+			b := Build("CREATE DATABASE")
 			if sqlx.Has(c.Extra, &schema.IfNotExists{}) {
 				b.P("IF NOT EXISTS")
 			}
+			b.Ident(c.S.Name)
 			// Schema was created with CHARSET and it is not the default database character set.
 			if a := (schema.Charset{}); sqlx.Has(c.S.Attrs, &a) && a.V != "" && a.V != s.charset {
 				b.P("CHARSET", a.V)
@@ -114,10 +115,11 @@ func (s *state) topLevel(changes []schema.Change) ([]schema.Change, error) {
 				Comment: fmt.Sprintf("add new schema named %q", c.S.Name),
 			})
 		case *schema.DropSchema:
-			b := Build("DROP DATABASE").Ident(c.S.Name)
+			b := Build("DROP DATABASE")
 			if sqlx.Has(c.Extra, &schema.IfExists{}) {
 				b.P("IF EXISTS")
 			}
+			b.Ident(c.S.Name)
 			s.append(&migrate.Change{
 				Cmd:     b.String(),
 				Source:  c,
@@ -200,11 +202,12 @@ func (s *state) modifySchema(modify *schema.ModifySchema) error {
 func (s *state) addTable(add *schema.AddTable) error {
 	var (
 		errors []string
-		b      = Build("CREATE TABLE").Table(add.T)
+		b      = Build("CREATE TABLE")
 	)
 	if sqlx.Has(add.Extra, &schema.IfNotExists{}) {
 		b.P("IF NOT EXISTS")
 	}
+	b.Table(add.T)
 	b.Wrap(func(b *sqlx.Builder) {
 		b.MapComma(add.T.Columns, func(i int, b *sqlx.Builder) {
 			if err := s.column(b, add.T, add.T.Columns[i]); err != nil {
@@ -251,10 +254,11 @@ func (s *state) addTable(add *schema.AddTable) error {
 // dropTable builds and appends the migrate.Change
 // for dropping a table from a schema.
 func (s *state) dropTable(drop *schema.DropTable) {
-	b := Build("DROP TABLE").Table(drop.T)
+	b := Build("DROP TABLE")
 	if sqlx.Has(drop.Extra, &schema.IfExists{}) {
 		b.P("IF EXISTS")
 	}
+	b.Table(drop.T)
 	s.append(&migrate.Change{
 		Cmd:     b.String(),
 		Source:  drop,
diff --git a/sql/mysql/normalize.go b/sql/mysql/normalize.go
new file mode 100644
index 00000000..9b70e748
--- /dev/null
+++ b/sql/mysql/normalize.go
@@ -0,0 +1,101 @@
+// 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 mysql
+
+import (
+	"context"
+	"crypto/md5"
+	"fmt"
+	"time"
+
+	"ariga.io/atlas/sql/schema"
+)
+
+// NormalizeRealm returns the normal representation of the given database.
+func (d *Driver) NormalizeRealm(ctx context.Context, r *schema.Realm) (nr *schema.Realm, err error) {
+	for _, s := range r.Schemas {
+		switch s.Name {
+		case "mysql", "information_schema", "performance_schema", "sys":
+			return nil, fmt.Errorf("sql/mysql: normalizing internal schema %q is not supported", s.Name)
+		}
+	}
+	var (
+		twins   = make(map[string]string)
+		changes = make([]schema.Change, 0, len(r.Schemas))
+		reverse = make([]schema.Change, 0, len(r.Schemas))
+		opts    = &schema.InspectRealmOption{
+			Schemas: make([]string, 0, len(r.Schemas)),
+		}
+	)
+	for _, s := range r.Schemas {
+		twin := twinName(s.Name)
+		twins[twin] = s.Name
+		s.Name = twin
+		opts.Schemas = append(opts.Schemas, s.Name)
+		// Skip adding the schema.IfNotExists clause
+		// to fail if the schema exists.
+		st := schema.New(twin).AddAttrs(s.Attrs...)
+		changes = append(changes, &schema.AddSchema{S: st})
+		reverse = append(reverse, &schema.DropSchema{S: st, Extra: []schema.Clause{&schema.IfExists{}}})
+		for _, t := range s.Tables {
+			// If objects are not strongly connected.
+			if t.Schema != s {
+				t.Schema = s
+			}
+			changes = append(changes, &schema.AddTable{T: t})
+		}
+	}
+	patch := func(r *schema.Realm) {
+		for _, s := range r.Schemas {
+			s.Name = twins[s.Name]
+		}
+	}
+	// Delete the twin resources, and return
+	// the source realm to its initial state.
+	defer func() {
+		patch(r)
+		uerr := d.ApplyChanges(ctx, reverse)
+		if err != nil {
+			err = fmt.Errorf("%w: %v", err, uerr)
+		}
+		err = uerr
+	}()
+	if err := d.ApplyChanges(ctx, changes); err != nil {
+		return nil, err
+	}
+	if nr, err = d.InspectRealm(ctx, opts); err != nil {
+		return nil, err
+	}
+	patch(nr)
+	return nr, nil
+}
+
+// NormalizeSchema returns the normal representation of the given database.
+func (d *Driver) NormalizeSchema(ctx context.Context, s *schema.Schema) (*schema.Schema, error) {
+	r := &schema.Realm{}
+	if s.Realm != nil {
+		r.Attrs = s.Realm.Attrs
+	}
+	r.Schemas = append(r.Schemas, s)
+	nr, err := d.NormalizeRealm(ctx, r)
+	if err != nil {
+		return nil, err
+	}
+	ns, ok := nr.Schema(s.Name)
+	if !ok {
+		return nil, fmt.Errorf("sql/mysql: missing normalized schema %q", s.Name)
+	}
+	return ns, nil
+}
+
+const maxLen = 64
+
+func twinName(name string) string {
+	twin := fmt.Sprintf("atlas_twin_%s_%d", name, time.Now().Unix())
+	if len(twin) <= maxLen {
+		return twin
+	}
+	return fmt.Sprintf("%s_%x", twin[:maxLen-33], md5.Sum([]byte(twin)))
+}
diff --git a/sql/postgres/inspect.go b/sql/postgres/inspect.go
index 489ece46..7297c8ae 100644
--- a/sql/postgres/inspect.go
+++ b/sql/postgres/inspect.go
@@ -681,6 +681,16 @@ type (
 		schema.Attr
 		Columns []string
 	}
+
+	// Cascade clause for appending to the different DROP changes.
+	// For example, A schema.DropSchema with this clause will be
+	// applied as follows:
+	//
+	//	DROP SCHEMA <name> CASCADE
+	//
+	Cascade struct {
+		schema.Clause
+	}
 )
 
 const (
@@ -701,7 +711,7 @@ SELECT
 	pg_catalog.obj_description(t2.oid, 'pg_class') AS COMMENT
 FROM
 	INFORMATION_SCHEMA.TABLES AS t1
-	JOIN pg_catalog.pg_class AS t2 ON t2.oid = to_regclass(t1.table_schema || '.' || t1.table_name)::oid
+	JOIN pg_catalog.pg_class AS t2 ON t2.oid = to_regclass(t1.table_schema || '.' || quote_ident(t1.table_name))::oid
 WHERE
 	t1.table_type = 'BASE TABLE'
 	AND t1.table_schema IN (%s)
@@ -715,7 +725,7 @@ SELECT
 	pg_catalog.obj_description(t2.oid, 'pg_class') AS COMMENT
 FROM
 	INFORMATION_SCHEMA.TABLES AS t1
-	JOIN pg_catalog.pg_class AS t2 ON t2.oid = to_regclass(t1.table_schema || '.' || t1.table_name)::oid
+	JOIN pg_catalog.pg_class AS t2 ON t2.oid = to_regclass(t1.table_schema || '.' || quote_ident(t1.table_name))::oid
 WHERE
 	t1.table_type = 'BASE TABLE'
 	AND t1.table_schema IN (%s)
diff --git a/sql/postgres/migrate.go b/sql/postgres/migrate.go
index ee9119ca..7a7f258b 100644
--- a/sql/postgres/migrate.go
+++ b/sql/postgres/migrate.go
@@ -86,10 +86,11 @@ func (s *state) topLevel(changes []schema.Change) []schema.Change {
 	for _, c := range changes {
 		switch c := c.(type) {
 		case *schema.AddSchema:
-			b := Build("CREATE SCHEMA").Ident(c.S.Name)
+			b := Build("CREATE SCHEMA")
 			if sqlx.Has(c.Extra, &schema.IfNotExists{}) {
 				b.P("IF NOT EXISTS")
 			}
+			b.Ident(c.S.Name)
 			s.append(&migrate.Change{
 				Cmd:     b.String(),
 				Source:  c,
@@ -97,9 +98,13 @@ func (s *state) topLevel(changes []schema.Change) []schema.Change {
 				Comment: fmt.Sprintf("Add new schema named %q", c.S.Name),
 			})
 		case *schema.DropSchema:
-			b := Build("DROP SCHEMA").Ident(c.S.Name)
+			b := Build("DROP SCHEMA")
 			if sqlx.Has(c.Extra, &schema.IfExists{}) {
-				b.P("IF NOT EXISTS")
+				b.P("IF EXISTS")
+			}
+			b.Ident(c.S.Name)
+			if sqlx.Has(c.Extra, &Cascade{}) {
+				b.P("CASCADE")
 			}
 			s.append(&migrate.Change{
 				Cmd:     b.String(),
@@ -119,10 +124,11 @@ func (s *state) addTable(ctx context.Context, add *schema.AddTable) error {
 	if err := s.addTypes(ctx, add.T.Columns...); err != nil {
 		return err
 	}
-	b := Build("CREATE TABLE").Table(add.T)
+	b := Build("CREATE TABLE")
 	if sqlx.Has(add.Extra, &schema.IfNotExists{}) {
 		b.P("IF NOT EXISTS")
 	}
+	b.Table(add.T)
 	b.Wrap(func(b *sqlx.Builder) {
 		b.MapComma(add.T.Columns, func(i int, b *sqlx.Builder) {
 			s.column(b, add.T.Columns[i])
@@ -155,10 +161,11 @@ func (s *state) addTable(ctx context.Context, add *schema.AddTable) error {
 
 // dropTable builds and executes the query for dropping a table from a schema.
 func (s *state) dropTable(drop *schema.DropTable) {
-	b := Build("DROP TABLE").Table(drop.T)
+	b := Build("DROP TABLE")
 	if sqlx.Has(drop.Extra, &schema.IfExists{}) {
 		b.P("IF EXISTS")
 	}
+	b.Table(drop.T)
 	s.append(&migrate.Change{
 		Cmd:     b.String(),
 		Source:  drop,
diff --git a/sql/schema/inspect.go b/sql/schema/inspect.go
index 2ad0deef..930cb6b7 100644
--- a/sql/schema/inspect.go
+++ b/sql/schema/inspect.go
@@ -64,3 +64,15 @@ type (
 		InspectRealm(ctx context.Context, opts *InspectRealmOption) (*Realm, error)
 	}
 )
+
+// Normalizer is the interface implemented by the different database drivers for
+// "normalizing" schema objects. i.e. converting schema objects defined in natural
+// form to their representation in the database. Thus, two schema objects are equal
+// if their normal forms are equal.
+type Normalizer interface {
+	// NormalizeSchema returns the normal representation of a schema.
+	NormalizeSchema(context.Context, *Schema) (*Schema, error)
+
+	// NormalizeRealm returns the normal representation of a database.
+	NormalizeRealm(context.Context, *Realm) (*Realm, error)
+}


Test patch:
diff --git a/sql/mysql/migrate_test.go b/sql/mysql/migrate_test.go
index b24db365..75b844db 100644
--- a/sql/mysql/migrate_test.go
+++ b/sql/mysql/migrate_test.go
@@ -25,9 +25,9 @@ func TestMigrate_ApplyChanges(t *testing.T) {
 		WillReturnResult(sqlmock.NewResult(0, 1))
 	mk.ExpectExec(sqltest.Escape("DROP TABLE `users`")).
 		WillReturnResult(sqlmock.NewResult(0, 0))
-	mk.ExpectExec(sqltest.Escape("DROP TABLE `public`.`pets` IF EXISTS")).
+	mk.ExpectExec(sqltest.Escape("DROP TABLE IF EXISTS `public`.`pets`")).
 		WillReturnResult(sqlmock.NewResult(0, 0))
-	mk.ExpectExec(sqltest.Escape("CREATE TABLE `pets` IF NOT EXISTS (`a` int NOT NULL DEFAULT (int(rand())), `b` bigint NOT NULL DEFAULT 1, `c` bigint NULL, PRIMARY KEY (`a`, `b`), UNIQUE INDEX `b_c_unique` (`b`, `c`) COMMENT \"comment\")")).
+	mk.ExpectExec(sqltest.Escape("CREATE TABLE IF NOT EXISTS `pets` (`a` int NOT NULL DEFAULT (int(rand())), `b` bigint NOT NULL DEFAULT 1, `c` bigint NULL, PRIMARY KEY (`a`, `b`), UNIQUE INDEX `b_c_unique` (`b`, `c`) COMMENT \"comment\")")).
 		WillReturnResult(sqlmock.NewResult(0, 0))
 	mk.ExpectExec(sqltest.Escape("ALTER TABLE `users` DROP INDEX `id_spouse_id`")).
 		WillReturnResult(sqlmock.NewResult(0, 0))
@@ -176,6 +176,24 @@ func TestPlanChanges(t *testing.T) {
 		wantPlan *migrate.Plan
 		wantErr  bool
 	}{
+		{
+			input: []schema.Change{
+				&schema.AddSchema{S: schema.New("test").SetCharset("utf8mb4"), Extra: []schema.Clause{&schema.IfNotExists{}}},
+				&schema.DropSchema{S: schema.New("test").SetCharset("utf8mb4"), Extra: []schema.Clause{&schema.IfExists{}}},
+			},
+			wantPlan: &migrate.Plan{
+				Reversible: false,
+				Changes: []*migrate.Change{
+					{
+						Cmd:     "CREATE DATABASE IF NOT EXISTS `test` CHARSET utf8mb4",
+						Reverse: "DROP DATABASE `test`",
+					},
+					{
+						Cmd: "DROP DATABASE IF EXISTS `test`",
+					},
+				},
+			},
+		},
 		{
 			input: []schema.Change{
 				func() schema.Change {
diff --git a/sql/mysql/normalize_test.go b/sql/mysql/normalize_test.go
new file mode 100644
index 00000000..24aae6a4
--- /dev/null
+++ b/sql/mysql/normalize_test.go
@@ -0,0 +1,62 @@
+// 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 mysql
+
+import (
+	"context"
+	"strings"
+	"testing"
+
+	"ariga.io/atlas/sql/migrate"
+	"ariga.io/atlas/sql/schema"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestDriver_NormalizeRealm(t *testing.T) {
+	var (
+		apply   = &mockApply{}
+		inspect = &mockInspect{
+			realm: schema.NewRealm(schema.New("test").SetCharset("utf8mb4")),
+		}
+		drv = &Driver{
+			Inspector:   inspect,
+			PlanApplier: apply,
+		}
+	)
+	normal, err := drv.NormalizeRealm(context.Background(), schema.NewRealm(schema.New("test")))
+	require.NoError(t, err)
+	require.Equal(t, normal, inspect.realm)
+
+	require.Len(t, inspect.schemas, 1)
+	require.True(t, strings.HasPrefix(inspect.schemas[0], "atlas_twin_test_"))
+
+	require.Len(t, apply.changes, 2, "expect 2 calls (create and drop)")
+	require.Len(t, apply.changes[0], 1)
+	require.Equal(t, &schema.AddSchema{S: schema.New(inspect.schemas[0])}, apply.changes[0][0])
+	require.Len(t, apply.changes[1], 1)
+	require.Equal(t, &schema.DropSchema{S: schema.New(inspect.schemas[0]), Extra: []schema.Clause{&schema.IfExists{}}}, apply.changes[1][0])
+}
+
+type mockInspect struct {
+	schema.Inspector
+	schemas []string
+	realm   *schema.Realm
+}
+
+func (m *mockInspect) InspectRealm(_ context.Context, opts *schema.InspectRealmOption) (*schema.Realm, error) {
+	m.schemas = append(m.schemas, opts.Schemas...)
+	return m.realm, nil
+}
+
+type mockApply struct {
+	migrate.PlanApplier
+	changes [][]schema.Change
+}
+
+func (m *mockApply) ApplyChanges(_ context.Context, changes []schema.Change) error {
+	m.changes = append(m.changes, changes)
+	return nil
+}
diff --git a/sql/postgres/migrate_test.go b/sql/postgres/migrate_test.go
index 9a7a4d8f..346ef00e 100644
--- a/sql/postgres/migrate_test.go
+++ b/sql/postgres/migrate_test.go
@@ -22,6 +22,29 @@ func TestPlanChanges(t *testing.T) {
 		mock    func(mock)
 		plan    *migrate.Plan
 	}{
+		{
+			changes: []schema.Change{
+				&schema.AddSchema{S: schema.New("test"), Extra: []schema.Clause{&schema.IfNotExists{}}},
+				&schema.DropSchema{S: schema.New("test"), Extra: []schema.Clause{&schema.IfExists{}}},
+				&schema.DropSchema{S: schema.New("test"), Extra: []schema.Clause{&Cascade{}}},
+			},
+			plan: &migrate.Plan{
+				Reversible:    false,
+				Transactional: true,
+				Changes: []*migrate.Change{
+					{
+						Cmd:     `CREATE SCHEMA IF NOT EXISTS "test"`,
+						Reverse: `DROP SCHEMA "test"`,
+					},
+					{
+						Cmd: `DROP SCHEMA IF EXISTS "test"`,
+					},
+					{
+						Cmd: `DROP SCHEMA "test" CASCADE`,
+					},
+				},
+			},
+		},
 		{
 			changes: []schema.Change{
 				func() schema.Change {