Repository: ariga/atlas
Pull request #313 (2021-12-15 12:10:09)

sql: add support for IfExists/IfNotExists clauses in migration


Patch:
diff --git a/sql/internal/sqlx/diff.go b/sql/internal/sqlx/diff.go
index 84b439b3..f3abf477 100644
--- a/sql/internal/sqlx/diff.go
+++ b/sql/internal/sqlx/diff.go
@@ -308,17 +308,27 @@ func CommentChange(from, to []schema.Attr) schema.ChangeKind {
 	return schema.NoChange
 }
 
-// Has finds the first attribute in the attribute list that
+var (
+	attrsType   = reflect.TypeOf(([]schema.Attr)(nil))
+	clausesType = reflect.TypeOf(([]schema.Clause)(nil))
+	exprsType   = reflect.TypeOf(([]schema.Expr)(nil))
+)
+
+// Has finds the first element in the elements list that
 // matches target, and if so, sets target to that attribute
 // value and returns true.
-func Has(attrs []schema.Attr, target interface{}) bool {
-	rv := reflect.ValueOf(target)
-	if rv.Kind() != reflect.Ptr || rv.IsNil() {
+func Has(elements, target interface{}) bool {
+	ev := reflect.ValueOf(elements)
+	if t := ev.Type(); t != attrsType && t != clausesType && t != exprsType {
+		panic(fmt.Sprintf("unexpected elements type: %T", t))
+	}
+	tv := reflect.ValueOf(target)
+	if tv.Kind() != reflect.Ptr || tv.IsNil() {
 		panic("target must be a non-nil pointer")
 	}
-	for _, attr := range attrs {
-		if reflect.TypeOf(attr).AssignableTo(rv.Type()) {
-			rv.Elem().Set(reflect.ValueOf(attr).Elem())
+	for i := 0; i < ev.Len(); i++ {
+		if e := ev.Index(i).Elem(); e.Type().AssignableTo(tv.Type()) {
+			tv.Elem().Set(e.Elem())
 			return true
 		}
 	}
diff --git a/sql/mysql/migrate.go b/sql/mysql/migrate.go
index 410164d3..b742746b 100644
--- a/sql/mysql/migrate.go
+++ b/sql/mysql/migrate.go
@@ -53,6 +53,9 @@ func (m *migrate) topLevel(ctx context.Context, changes []schema.Change) ([]sche
 		switch c := c.(type) {
 		case *schema.AddSchema:
 			b := Build("CREATE DATABASE").Ident(c.S.Name)
+			if sqlx.Has(c.Extra, &schema.IfNotExists{}) {
+				b.P("IF NOT EXISTS")
+			}
 			if a := (schema.Charset{}); sqlx.Has(c.S.Attrs, &a) {
 				b.P("CHARACTER SET", a.V)
 			}
@@ -76,6 +79,9 @@ func (m *migrate) topLevel(ctx context.Context, changes []schema.Change) ([]sche
 // 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)
+	if sqlx.Has(add.Extra, &schema.IfNotExists{}) {
+		b.P("IF NOT EXISTS")
+	}
 	b.Wrap(func(b *sqlx.Builder) {
 		b.MapComma(add.T.Columns, func(i int, b *sqlx.Builder) {
 			m.column(b, add.T, add.T.Columns[i])
@@ -112,6 +118,9 @@ func (m *migrate) addTable(ctx context.Context, add *schema.AddTable) error {
 // dropTable builds and executes the query for dropping a table from a schema.
 func (m *migrate) dropTable(ctx context.Context, drop *schema.DropTable) error {
 	b := Build("DROP TABLE").Table(drop.T)
+	if sqlx.Has(drop.Extra, &schema.IfExists{}) {
+		b.P("IF EXISTS")
+	}
 	if _, err := m.ExecContext(ctx, b.String()); err != nil {
 		return fmt.Errorf("drop table: %w", err)
 	}
diff --git a/sql/postgres/migrate.go b/sql/postgres/migrate.go
index e8cd3557..4e35de99 100644
--- a/sql/postgres/migrate.go
+++ b/sql/postgres/migrate.go
@@ -51,7 +51,11 @@ func (m *migrate) topLevel(ctx context.Context, changes []schema.Change) ([]sche
 	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 {
+			b := Build("CREATE SCHEMA").Ident(c.S.Name)
+			if sqlx.Has(c.Extra, &schema.IfNotExists{}) {
+				b.P("IF NOT EXISTS")
+			}
+			if _, err := m.ExecContext(ctx, b.String()); err != nil {
 				return nil, fmt.Errorf("add schema: %w", err)
 			}
 		case *schema.DropSchema:
@@ -72,6 +76,9 @@ func (m *migrate) addTable(ctx context.Context, add *schema.AddTable) error {
 		return err
 	}
 	b := Build("CREATE TABLE").Table(add.T)
+	if sqlx.Has(add.Extra, &schema.IfNotExists{}) {
+		b.P("IF NOT EXISTS")
+	}
 	b.Wrap(func(b *sqlx.Builder) {
 		b.MapComma(add.T.Columns, func(i int, b *sqlx.Builder) {
 			m.column(b, add.T.Columns[i])
@@ -97,6 +104,9 @@ func (m *migrate) addTable(ctx context.Context, add *schema.AddTable) error {
 // dropTable builds and executes the query for dropping a table from a schema.
 func (m *migrate) dropTable(ctx context.Context, drop *schema.DropTable) error {
 	b := Build("DROP TABLE").Table(drop.T)
+	if sqlx.Has(drop.Extra, &schema.IfExists{}) {
+		b.P("IF EXISTS")
+	}
 	if _, err := m.ExecContext(ctx, b.String()); err != nil {
 		return fmt.Errorf("drop table: %w", err)
 	}
diff --git a/sql/schema/migrate.go b/sql/schema/migrate.go
index 24a23f31..c919e574 100644
--- a/sql/schema/migrate.go
+++ b/sql/schema/migrate.go
@@ -24,18 +24,33 @@ type (
 		change()
 	}
 
+	// Clause carries additional information that can be added
+	// to schema changes. The Clause interface can be implemented
+	// outside this package as follows:
+	//
+	//	type Authorization struct {
+	//		schema.Clause
+	//		UserName string
+	//	}
+	//
+	//	var c schema.Clause = &Authorization{UserName: "a8m"}
+	//
+	Clause interface {
+		clause()
+	}
+
 	// AddSchema describes a schema (named database) creation change.
 	// Unlike table creation, schemas and their elements are described
 	// with separate changes. For example, "AddSchema" and "AddTable"
 	AddSchema struct {
 		S     *Schema
-		Extra []Expr // Extra clauses and options.
+		Extra []Clause // Extra clauses and options.
 	}
 
 	// DropSchema describes a schema (named database) removal change.
 	DropSchema struct {
 		S     *Schema
-		Extra []Expr // Extra clauses and options.
+		Extra []Clause // Extra clauses and options.
 	}
 
 	// ModifySchema describes a modification change for schema attributes.
@@ -47,13 +62,13 @@ type (
 	// AddTable describes a table creation change.
 	AddTable struct {
 		T     *Table
-		Extra []Expr // Extra clauses and options.
+		Extra []Clause // Extra clauses and options.
 	}
 
 	// DropTable describes a table removal change.
 	DropTable struct {
 		T     *Table
-		Extra []Expr // Extra clauses.
+		Extra []Clause // Extra clauses.
 	}
 
 	// ModifyTable describes a table modification change.
@@ -124,6 +139,14 @@ type (
 	ModifyAttr struct {
 		From, To Attr
 	}
+
+	// IfExists represents a clause in a schema change that is commonly
+	// supported by multiple statements (e.g. DROP TABLE or DROP SCHEMA).
+	IfExists struct{}
+
+	// IfNotExists represents a clause in a schema change that is commonly
+	// supported by multiple statements (e.g. CREATE TABLE or CREATE SCHEMA).
+	IfNotExists struct{}
 )
 
 // A ChangeKind describes a change kind that can be combined
@@ -233,3 +256,7 @@ func (*ModifyColumn) change()     {}
 func (*AddForeignKey) change()    {}
 func (*DropForeignKey) change()   {}
 func (*ModifyForeignKey) change() {}
+
+// clauses.
+func (*IfExists) clause()    {}
+func (*IfNotExists) clause() {}
diff --git a/sql/sqlite/migrate.go b/sql/sqlite/migrate.go
index 24431d62..46d73e10 100644
--- a/sql/sqlite/migrate.go
+++ b/sql/sqlite/migrate.go
@@ -40,6 +40,9 @@ func (m *migrate) Exec(ctx context.Context, changes []schema.Change) (err error)
 // 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").Ident(add.T.Name)
+	if sqlx.Has(add.Extra, &schema.IfNotExists{}) {
+		b.P("IF NOT EXISTS")
+	}
 	b.Wrap(func(b *sqlx.Builder) {
 		b.MapComma(add.T.Columns, func(i int, b *sqlx.Builder) {
 			m.column(b, add.T.Columns[i])
@@ -66,7 +69,11 @@ func (m *migrate) addTable(ctx context.Context, add *schema.AddTable) error {
 // dropTable builds and executes the query for dropping a table from a schema.
 func (m *migrate) dropTable(ctx context.Context, drop *schema.DropTable) error {
 	if err := m.skipConstraints(ctx, func(ctx context.Context) error {
-		_, err := m.ExecContext(ctx, Build("DROP TABLE").Ident(drop.T.Name).String())
+		b := Build("DROP TABLE").Ident(drop.T.Name)
+		if sqlx.Has(drop.Extra, &schema.IfExists{}) {
+			b.P("IF EXISTS")
+		}
+		_, err := m.ExecContext(ctx, b.String())
 		return err
 	}); err != nil {
 		return fmt.Errorf("drop table: %w", err)


Test patch:
diff --git a/sql/mysql/migrate_test.go b/sql/mysql/migrate_test.go
index 72df312b..8679ce97 100644
--- a/sql/mysql/migrate_test.go
+++ b/sql/mysql/migrate_test.go
@@ -24,9 +24,9 @@ func TestMigrate_Exec(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`")).
+	mk.ExpectExec(sqltest.Escape("DROP TABLE `public`.`pets` IF EXISTS")).
 		WillReturnResult(sqlmock.NewResult(0, 0))
-	mk.ExpectExec(sqltest.Escape("CREATE TABLE `pets` (`a` int NOT NULL, `b` bigint NOT NULL, `c` bigint NULL, PRIMARY KEY (`a`, `b`), UNIQUE INDEX `b_c_unique` (`b`, `c`) COMMENT \"comment\")")).
+	mk.ExpectExec(sqltest.Escape("CREATE TABLE `pets` IF NOT EXISTS (`a` int NOT NULL, `b` bigint NOT NULL, `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))
@@ -40,7 +40,7 @@ func TestMigrate_Exec(t *testing.T) {
 		&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.DropTable{T: &schema.Table{Name: "pets", Schema: &schema.Schema{Name: "public"}}, Extra: []schema.Clause{&schema.IfExists{}}},
 		&schema.AddTable{
 			T: func() *schema.Table {
 				t := &schema.Table{
@@ -59,6 +59,9 @@ func TestMigrate_Exec(t *testing.T) {
 				}
 				return t
 			}(),
+			Extra: []schema.Clause{
+				&schema.IfNotExists{},
+			},
 		},
 	})
 	require.NoError(t, err)