Repository: ariga/atlas
Pull request #602 (2022-02-27 11:14:16)

sql: add twin-driver and support postgres normalize


Patch:
diff --git a/sql/mysql/normalize.go b/sql/internal/sqlx/twin.go
similarity index 51%
rename from sql/mysql/normalize.go
rename to sql/internal/sqlx/twin.go
index 9b70e748..2a785d22 100644
--- a/sql/mysql/normalize.go
+++ b/sql/internal/sqlx/twin.go
@@ -2,7 +2,7 @@
 // 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
+package sqlx
 
 import (
 	"context"
@@ -10,17 +10,32 @@ import (
 	"fmt"
 	"time"
 
+	"ariga.io/atlas/sql/migrate"
 	"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)
-		}
-	}
+// TwinDriver is a driver that provides additional functionality
+// to interact with the twin/dev database.
+type TwinDriver struct {
+	// A Driver connected to the twin database.
+	migrate.Driver
+
+	// MaxNameLen configures the max length of object names in
+	// the connected database (e.g. 64 in MySQL). Longer names
+	// are trimmed and suffixed with their hash.
+	MaxNameLen int
+
+	// DropClause holds optional clauses that
+	// can be added to the DropSchema change.
+	DropClause []schema.Clause
+}
+
+// NormalizeRealm implements the schema.Normalizer interface.
+//
+// The implementation converts schema objects in "natural form" (e.g. HCL or DSL)
+// to their "normal presentation" in the database, by creating them temporarily in
+// a "twin database", and then inspects them from there.
+func (t *TwinDriver) NormalizeRealm(ctx context.Context, r *schema.Realm) (nr *schema.Realm, err error) {
 	var (
 		twins   = make(map[string]string)
 		changes = make([]schema.Change, 0, len(r.Schemas))
@@ -30,7 +45,7 @@ func (d *Driver) NormalizeRealm(ctx context.Context, r *schema.Realm) (nr *schem
 		}
 	)
 	for _, s := range r.Schemas {
-		twin := twinName(s.Name)
+		twin := t.formatName(s.Name)
 		twins[twin] = s.Name
 		s.Name = twin
 		opts.Schemas = append(opts.Schemas, s.Name)
@@ -38,7 +53,7 @@ func (d *Driver) NormalizeRealm(ctx context.Context, r *schema.Realm) (nr *schem
 		// 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{}}})
+		reverse = append(reverse, &schema.DropSchema{S: st, Extra: append(t.DropClause, &schema.IfExists{})})
 		for _, t := range s.Tables {
 			// If objects are not strongly connected.
 			if t.Schema != s {
@@ -56,46 +71,44 @@ func (d *Driver) NormalizeRealm(ctx context.Context, r *schema.Realm) (nr *schem
 	// the source realm to its initial state.
 	defer func() {
 		patch(r)
-		uerr := d.ApplyChanges(ctx, reverse)
+		uerr := t.ApplyChanges(ctx, reverse)
 		if err != nil {
 			err = fmt.Errorf("%w: %v", err, uerr)
 		}
 		err = uerr
 	}()
-	if err := d.ApplyChanges(ctx, changes); err != nil {
+	if err := t.ApplyChanges(ctx, changes); err != nil {
 		return nil, err
 	}
-	if nr, err = d.InspectRealm(ctx, opts); err != nil {
+	if nr, err = t.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) {
+// NormalizeSchema returns the normal representation of the given database. See NormalizeRealm for more info.
+func (t *TwinDriver) 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)
+	nr, err := t.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 nil, fmt.Errorf("missing normalized schema %q", s.Name)
 	}
 	return ns, nil
 }
 
-const maxLen = 64
-
-func twinName(name string) string {
+func (t *TwinDriver) formatName(name string) string {
 	twin := fmt.Sprintf("atlas_twin_%s_%d", name, time.Now().Unix())
-	if len(twin) <= maxLen {
+	if t.MaxNameLen == 0 || len(twin) <= t.MaxNameLen {
 		return twin
 	}
-	return fmt.Sprintf("%s_%x", twin[:maxLen-33], md5.Sum([]byte(twin)))
+	return fmt.Sprintf("%s_%x", twin[:t.MaxNameLen-1-md5.Size*2], md5.Sum([]byte(twin)))
 }
diff --git a/sql/mysql/driver.go b/sql/mysql/driver.go
index 4665fbc0..140351cf 100644
--- a/sql/mysql/driver.go
+++ b/sql/mysql/driver.go
@@ -62,6 +62,20 @@ func Open(db schema.ExecQuerier) (*Driver, error) {
 	}, nil
 }
 
+func (d *Driver) twin() *sqlx.TwinDriver {
+	return &sqlx.TwinDriver{Driver: d, MaxNameLen: 64}
+}
+
+// NormalizeRealm returns the normal representation of the given database.
+func (d *Driver) NormalizeRealm(ctx context.Context, r *schema.Realm) (*schema.Realm, error) {
+	return d.twin().NormalizeRealm(ctx, r)
+}
+
+// NormalizeSchema returns the normal representation of the given database.
+func (d *Driver) NormalizeSchema(ctx context.Context, s *schema.Schema) (*schema.Schema, error) {
+	return d.twin().NormalizeSchema(ctx, s)
+}
+
 // supportsCheck reports if the connected database supports
 // the CHECK clause, and return the querying for getting them.
 func (d *conn) supportsCheck() (string, bool) {
diff --git a/sql/postgres/driver.go b/sql/postgres/driver.go
index 2a03aa32..35363279 100644
--- a/sql/postgres/driver.go
+++ b/sql/postgres/driver.go
@@ -65,6 +65,20 @@ func Open(db schema.ExecQuerier) (*Driver, error) {
 	}, nil
 }
 
+func (d *Driver) twin() *sqlx.TwinDriver {
+	return &sqlx.TwinDriver{Driver: d, MaxNameLen: 63, DropClause: []schema.Clause{&Cascade{}}}
+}
+
+// NormalizeRealm returns the normal representation of the given database.
+func (d *Driver) NormalizeRealm(ctx context.Context, r *schema.Realm) (*schema.Realm, error) {
+	return d.twin().NormalizeRealm(ctx, r)
+}
+
+// NormalizeSchema returns the normal representation of the given database.
+func (d *Driver) NormalizeSchema(ctx context.Context, s *schema.Schema) (*schema.Schema, error) {
+	return d.twin().NormalizeSchema(ctx, s)
+}
+
 // Standard column types (and their aliases) as defined in
 // PostgreSQL codebase/website.
 const (


Test patch:
diff --git a/sql/internal/sqlx/twin_test.go b/sql/internal/sqlx/twin_test.go
new file mode 100644
index 00000000..058d5902
--- /dev/null
+++ b/sql/internal/sqlx/twin_test.go
@@ -0,0 +1,59 @@
+// 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 sqlx
+
+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 (
+		drv = &mockDriver{
+			realm: schema.NewRealm(schema.New("test").SetCharset("utf8mb4")),
+		}
+		twin = &TwinDriver{
+			Driver:     drv,
+			MaxNameLen: 64,
+		}
+	)
+	normal, err := twin.NormalizeRealm(context.Background(), schema.NewRealm(schema.New("test")))
+	require.NoError(t, err)
+	require.Equal(t, normal, drv.realm)
+
+	require.Len(t, drv.schemas, 1)
+	require.True(t, strings.HasPrefix(drv.schemas[0], "atlas_twin_test_"))
+
+	require.Len(t, drv.changes, 2, "expect 2 calls (create and drop)")
+	require.Len(t, drv.changes[0], 1)
+	require.Equal(t, &schema.AddSchema{S: schema.New(drv.schemas[0])}, drv.changes[0][0])
+	require.Len(t, drv.changes[1], 1)
+	require.Equal(t, &schema.DropSchema{S: schema.New(drv.schemas[0]), Extra: []schema.Clause{&schema.IfExists{}}}, drv.changes[1][0])
+}
+
+type mockDriver struct {
+	migrate.Driver
+	// Inspect.
+	schemas []string
+	realm   *schema.Realm
+	// Apply.
+	changes [][]schema.Change
+}
+
+func (m *mockDriver) InspectRealm(_ context.Context, opts *schema.InspectRealmOption) (*schema.Realm, error) {
+	m.schemas = append(m.schemas, opts.Schemas...)
+	return m.realm, nil
+}
+
+func (m *mockDriver) ApplyChanges(_ context.Context, changes []schema.Change) error {
+	m.changes = append(m.changes, changes)
+	return nil
+}
diff --git a/sql/mysql/normalize_test.go b/sql/mysql/normalize_test.go
deleted file mode 100644
index 24aae6a4..00000000
--- a/sql/mysql/normalize_test.go
+++ /dev/null
@@ -1,62 +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 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
-}