Repository: ariga/atlas
Pull request #298 (2021-12-11 19:50:41)

sql/internal/specutil: adding TypeRegistry for TypeSpecs
Stacked on #297 

Adding a new container type for working with TypeSpecs. This is added to make the needed changes to mysql/postgres for type-safe type definitions in Atlas DDL. 


Patch:
diff --git a/schema/schemaspec/spec.go b/schema/schemaspec/spec.go
index 3593b217..278e5eba 100644
--- a/schema/schemaspec/spec.go
+++ b/schema/schemaspec/spec.go
@@ -67,7 +67,9 @@ type (
 
 	// TypeSpec represents a specification for defining a Type.
 	TypeSpec struct {
+		// Name is the identifier for the type in an Atlas DDL document.
 		Name       string
+		 // T is the database identifier for the type.
 		T          string
 		Attributes []*TypeAttr
 	}
diff --git a/sql/internal/specutil/types.go b/sql/internal/specutil/types.go
index 36fd6e10..495c8036 100644
--- a/sql/internal/specutil/types.go
+++ b/sql/internal/specutil/types.go
@@ -2,6 +2,7 @@ package specutil
 
 import (
 	"errors"
+	"fmt"
 	"strings"
 
 	"ariga.io/atlas/schema/schemaspec"
@@ -40,3 +41,43 @@ func PrintType(typ *schemaspec.Type, spec *schemaspec.TypeSpec) (string, error)
 	}
 	return typ.T + mid + suffix, nil
 }
+
+// TypeRegistry is a collection of *schemaspec.TypeSpec.
+type TypeRegistry struct {
+	r []*schemaspec.TypeSpec
+}
+
+// Register adds one or more TypeSpec to the registry.
+func (r *TypeRegistry) Register(specs ...*schemaspec.TypeSpec) error {
+	for _, s := range specs {
+		if _, exists := r.Find(s.T); exists {
+			return fmt.Errorf("specutil: type with T of %q already registered", s.T)
+		}
+		if _, exists := r.FindByName(s.Name); exists {
+			return fmt.Errorf("specutil: type with name of %q already registered", s.T)
+		}
+	}
+	r.r = append(r.r, specs...)
+	return nil
+}
+
+// FindByName searches the registry for types that have the provided name.
+func (r *TypeRegistry) FindByName(name string) (*schemaspec.TypeSpec, bool) {
+	for _, current := range r.r {
+		if current.Name == name {
+			return current, true
+		}
+	}
+	return nil, false
+}
+
+// Find searches the registry for types that have the provided T.
+func (r *TypeRegistry) Find(t string) (*schemaspec.TypeSpec, bool) {
+	for _, current := range r.r {
+		if current.T == t {
+			return current, true
+		}
+	}
+	return nil, false
+}
+
diff --git a/sql/mysql/migrate.go b/sql/mysql/migrate.go
index 5f014411..91c7c8ed 100644
--- a/sql/mysql/migrate.go
+++ b/sql/mysql/migrate.go
@@ -25,7 +25,11 @@ func (d *Driver) Migrate() schema.Execer {
 // Exec executes the changes on the database. An error is returned
 // if one of the operations fail, or a change is not supported.
 func (m *migrate) Exec(ctx context.Context, changes []schema.Change) error {
-	planned, err := sqlx.DetachCycles(changes)
+	planned, err := m.topLevel(ctx, changes)
+	if err != nil {
+		return err
+	}
+	planned, err = sqlx.DetachCycles(planned)
 	if err != nil {
 		return err
 	}
@@ -47,6 +51,33 @@ func (m *migrate) Exec(ctx context.Context, changes []schema.Change) error {
 	return nil
 }
 
+// topLevel executes first the changes for creating or dropping schemas (top-level schema elements).
+func (m *migrate) topLevel(ctx context.Context, changes []schema.Change) ([]schema.Change, error) {
+	planned := make([]schema.Change, 0, len(changes))
+	for _, c := range changes {
+		switch c := c.(type) {
+		case *schema.AddSchema:
+			b := Build("CREATE DATABASE").Ident(c.S.Name)
+			if a := (schema.Charset{}); sqlx.Has(c.S.Attrs, &a) {
+				b.P("CHARACTER SET", a.V)
+			}
+			if a := (schema.Collation{}); sqlx.Has(c.S.Attrs, &a) {
+				b.P("COLLATE", a.V)
+			}
+			if _, err := m.ExecContext(ctx, b.String()); err != nil {
+				return nil, fmt.Errorf("add schema: %w", err)
+			}
+		case *schema.DropSchema:
+			if _, err := m.ExecContext(ctx, Build("DROP DATABASE").Ident(c.S.Name).String()); err != nil {
+				return nil, fmt.Errorf("drop schema: %w", err)
+			}
+		default:
+			planned = append(planned, c)
+		}
+	}
+	return planned, nil
+}
+
 // addTable builds and executes the query for creating a table in a schema.
 func (m *migrate) addTable(ctx context.Context, add *schema.AddTable) error {
 	b := Build("CREATE TABLE").Table(add.T)
diff --git a/sql/postgres/migrate.go b/sql/postgres/migrate.go
index 82932ae1..8ade837d 100644
--- a/sql/postgres/migrate.go
+++ b/sql/postgres/migrate.go
@@ -24,7 +24,11 @@ func (d *Driver) Migrate() schema.Execer {
 // Exec executes the changes on the database. An error is returned
 // if one of the operations fail, or a change is not supported.
 func (m *migrate) Exec(ctx context.Context, changes []schema.Change) error {
-	planned, err := sqlx.DetachCycles(changes)
+	planned, err := m.topLevel(ctx, changes)
+	if err != nil {
+		return err
+	}
+	planned, err = sqlx.DetachCycles(planned)
 	if err != nil {
 		return err
 	}
@@ -46,6 +50,26 @@ func (m *migrate) Exec(ctx context.Context, changes []schema.Change) error {
 	return nil
 }
 
+// topLevel executes first the changes for creating or dropping schemas (top-level schema elements).
+func (m *migrate) topLevel(ctx context.Context, changes []schema.Change) ([]schema.Change, error) {
+	planned := make([]schema.Change, 0, len(changes))
+	for _, c := range changes {
+		switch c := c.(type) {
+		case *schema.AddSchema:
+			if _, err := m.ExecContext(ctx, Build("CREATE SCHEMA").Ident(c.S.Name).String()); err != nil {
+				return nil, fmt.Errorf("add schema: %w", err)
+			}
+		case *schema.DropSchema:
+			if _, err := m.ExecContext(ctx, Build("DROP SCHEMA").Ident(c.S.Name).String()); err != nil {
+				return nil, fmt.Errorf("drop schema: %w", err)
+			}
+		default:
+			planned = append(planned, c)
+		}
+	}
+	return planned, nil
+}
+
 // addTable builds and executes the query for creating a table in a schema.
 func (m *migrate) addTable(ctx context.Context, add *schema.AddTable) error {
 	// Create enum types before using them in the `CREATE TABLE` statement.


Test patch:
diff --git a/sql/internal/specutil/types_test.go b/sql/internal/specutil/types_test.go
index 408d942d..36e81cff 100644
--- a/sql/internal/specutil/types_test.go
+++ b/sql/internal/specutil/types_test.go
@@ -50,3 +50,15 @@ func TestTypePrint(t *testing.T) {
 		})
 	}
 }
+
+func TestRegistry(t *testing.T) {
+	r := &TypeRegistry{}
+	text := &schemaspec.TypeSpec{Name: "text", T: "text"}
+	err := r.Register(text)
+	require.NoError(t, err)
+	err = r.Register(text)
+	require.EqualError(t, err, `specutil: type with T of "text" already registered`)
+	spec, ok := r.FindByName("text")
+	require.True(t, ok)
+	require.EqualValues(t, spec, text)
+}
diff --git a/sql/mysql/migrate_test.go b/sql/mysql/migrate_test.go
index 0d80e9f0..7ac58226 100644
--- a/sql/mysql/migrate_test.go
+++ b/sql/mysql/migrate_test.go
@@ -18,6 +18,10 @@ import (
 func TestMigrate_Exec(t *testing.T) {
 	migrate, mk, err := newMigrate("8.0.13")
 	require.NoError(t, err)
+	mk.ExpectExec(sqltest.Escape("CREATE DATABASE `test` CHARACTER SET latin")).
+		WillReturnResult(sqlmock.NewResult(0, 1))
+	mk.ExpectExec(sqltest.Escape("DROP DATABASE `atlas`")).
+		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`")).
@@ -33,6 +37,8 @@ func TestMigrate_Exec(t *testing.T) {
 	mk.ExpectExec(sqltest.Escape("CREATE TABLE `comments` (`id` bigint NOT NULL, `post_id` bigint NULL, CONSTRAINT `comment` FOREIGN KEY (`post_id`) REFERENCES `posts` (`id`))")).
 		WillReturnResult(sqlmock.NewResult(0, 0))
 	err = migrate.Exec(context.Background(), []schema.Change{
+		&schema.AddSchema{S: &schema.Schema{Name: "test", Attrs: []schema.Attr{&schema.Charset{V: "latin"}}}},
+		&schema.DropSchema{S: &schema.Schema{Name: "atlas", Attrs: []schema.Attr{&schema.Charset{V: "latin"}}}},
 		&schema.DropTable{T: &schema.Table{Name: "users"}},
 		&schema.DropTable{T: &schema.Table{Name: "pets", Schema: &schema.Schema{Name: "public"}}},
 		&schema.AddTable{