Repository: ariga/atlas
Pull request #1040 (2022-08-10 11:31:36)

sql/internal/sqlx: add check to prevent misuse of schema qualifiers


Patch:
diff --git a/sql/internal/sqlx/plan.go b/sql/internal/sqlx/plan.go
index 6be2ec43..a6a96b4b 100644
--- a/sql/internal/sqlx/plan.go
+++ b/sql/internal/sqlx/plan.go
@@ -248,3 +248,37 @@ func isDropped(changes []schema.Change, t *schema.Table) bool {
 	}
 	return false
 }
+
+// CheckChangesScope checks that changes can be applied
+// on a schema scope (connection).
+func CheckChangesScope(changes []schema.Change) error {
+	names := make(map[string]struct{})
+	for _, c := range changes {
+		var t *schema.Table
+		switch c := c.(type) {
+		case *schema.AddSchema, *schema.ModifySchema, *schema.DropSchema:
+			return fmt.Errorf("%T is not allowed when migration plan is scoped to one schema", c)
+		case *schema.AddTable:
+			t = c.T
+		case *schema.ModifyTable:
+			t = c.T
+		case *schema.DropTable:
+			t = c.T
+		default:
+			continue
+		}
+		if t.Schema != nil && t.Schema.Name != "" {
+			names[t.Schema.Name] = struct{}{}
+		}
+		for _, c := range t.Columns {
+			e, ok := c.Type.Type.(*schema.EnumType)
+			if ok && e.Schema != nil && e.Schema.Name != "" {
+				names[t.Schema.Name] = struct{}{}
+			}
+		}
+	}
+	if len(names) > 1 {
+		return fmt.Errorf("found %d schemas when migration plan is scoped to one", len(names))
+	}
+	return nil
+}
diff --git a/sql/mysql/migrate.go b/sql/mysql/migrate.go
index c2d4bae1..80ad48ca 100644
--- a/sql/mysql/migrate.go
+++ b/sql/mysql/migrate.go
@@ -67,6 +67,11 @@ type state struct {
 // plan builds the migration plan for applying the
 // given changes on the attached connection.
 func (s *state) plan(changes []schema.Change) error {
+	if s.SchemaQualifier != nil {
+		if err := sqlx.CheckChangesScope(changes); err != nil {
+			return err
+		}
+	}
 	planned, err := s.topLevel(changes)
 	if err != nil {
 		return err
diff --git a/sql/postgres/migrate.go b/sql/postgres/migrate.go
index 561ba016..4056fe1c 100644
--- a/sql/postgres/migrate.go
+++ b/sql/postgres/migrate.go
@@ -68,6 +68,11 @@ type state struct {
 // Exec executes the changes on the database. An error is returned
 // if one of the operations fail, or a change is not supported.
 func (s *state) plan(ctx context.Context, changes []schema.Change) error {
+	if s.SchemaQualifier != nil {
+		if err := sqlx.CheckChangesScope(changes); err != nil {
+			return err
+		}
+	}
 	planned := s.topLevel(changes)
 	planned, err := sqlx.DetachCycles(planned)
 	if err != nil {


Test patch:
diff --git a/sql/internal/sqlx/plan_test.go b/sql/internal/sqlx/plan_test.go
index 64fb5085..7976b97f 100644
--- a/sql/internal/sqlx/plan_test.go
+++ b/sql/internal/sqlx/plan_test.go
@@ -92,3 +92,23 @@ func TestDetachCycles(t *testing.T) {
 	workplaces.ForeignKeys = nil
 	require.Equal(t, deletion, planned[2:])
 }
+
+func TestCheckChangesScope(t *testing.T) {
+	err := CheckChangesScope([]schema.Change{
+		&schema.AddSchema{},
+	})
+	require.EqualError(t, err, "*schema.AddSchema is not allowed when migration plan is scoped to one schema")
+	err = CheckChangesScope([]schema.Change{
+		&schema.ModifySchema{},
+	})
+	require.EqualError(t, err, "*schema.ModifySchema is not allowed when migration plan is scoped to one schema")
+	err = CheckChangesScope([]schema.Change{
+		&schema.DropSchema{},
+	})
+	require.EqualError(t, err, "*schema.DropSchema is not allowed when migration plan is scoped to one schema")
+	err = CheckChangesScope([]schema.Change{
+		&schema.AddTable{T: schema.NewTable("users").SetSchema(schema.New("s1"))},
+		&schema.AddTable{T: schema.NewTable("users").SetSchema(schema.New("s2"))},
+	})
+	require.EqualError(t, err, "found 2 schemas when migration plan is scoped to one")
+}
diff --git a/sql/mysql/migrate_test.go b/sql/mysql/migrate_test.go
index 73dba6b0..55b834c3 100644
--- a/sql/mysql/migrate_test.go
+++ b/sql/mysql/migrate_test.go
@@ -922,6 +922,17 @@ func TestPlanChanges(t *testing.T) {
 				},
 			},
 		},
+		// Empty qualifier in multi-schema mode should fail.
+		{
+			changes: []schema.Change{
+				&schema.AddTable{T: schema.NewTable("t1").SetSchema(schema.New("s1")).AddColumns(schema.NewIntColumn("a", "int"))},
+				&schema.AddTable{T: schema.NewTable("t2").SetSchema(schema.New("s2")).AddColumns(schema.NewIntColumn("a", "int"))},
+			},
+			options: []migrate.PlanOption{
+				func(o *migrate.PlanOptions) { o.SchemaQualifier = new(string) },
+			},
+			wantErr: true,
+		},
 	}
 	for i, tt := range tests {
 		t.Run(strconv.Itoa(i), func(t *testing.T) {
diff --git a/sql/postgres/migrate_test.go b/sql/postgres/migrate_test.go
index f85bb68e..5c385f76 100644
--- a/sql/postgres/migrate_test.go
+++ b/sql/postgres/migrate_test.go
@@ -1097,6 +1097,17 @@ func TestPlanChanges(t *testing.T) {
 				},
 			},
 		},
+		// Empty qualifier in multi-schema mode should fail.
+		{
+			changes: []schema.Change{
+				&schema.AddTable{T: schema.NewTable("t1").SetSchema(schema.New("s1")).AddColumns(schema.NewIntColumn("a", "int"))},
+				&schema.AddTable{T: schema.NewTable("t2").SetSchema(schema.New("s2")).AddColumns(schema.NewIntColumn("a", "int"))},
+			},
+			options: []migrate.PlanOption{
+				func(o *migrate.PlanOptions) { o.SchemaQualifier = new(string) },
+			},
+			wantErr: true,
+		},
 	}
 	for i, tt := range tests {
 		t.Run(strconv.Itoa(i), func(t *testing.T) {