Repository: ariga/atlas
Pull request #1021 (2022-08-07 12:50:02)

sql/migrate: allow running 'migrate diff' against a schema connection


Patch:
diff --git a/cmd/atlas/internal/cmdapi/migrate.go b/cmd/atlas/internal/cmdapi/migrate.go
index 9490a982..9d6ce67a 100644
--- a/cmd/atlas/internal/cmdapi/migrate.go
+++ b/cmd/atlas/internal/cmdapi/migrate.go
@@ -419,26 +419,31 @@ func CmdMigrateDiffRun(cmd *cobra.Command, args []string) error {
 	}
 	// Get a state reader for the desired state.
 	desired, err := to(cmd.Context(), dev)
-	if src, ok := desired.(io.Closer); ok {
-		defer src.Close()
-	}
 	if err != nil {
 		return err
 	}
+	defer desired.Close()
 	f, err := formatter()
 	if err != nil {
 		return err
 	}
+	opts := []migrate.PlannerOption{migrate.PlanFormat(f)}
+	if dev.URL.Schema != "" {
+		opts = append(opts, migrate.PlanSchema(dev.URL.Schema))
+	}
 	// Plan the changes and create a new migration file.
-	pl := migrate.NewPlanner(dev.Driver, dir, migrate.WithFormatter(f))
+	pl := migrate.NewPlanner(dev.Driver, dir, opts...)
 	var name string
 	if len(args) > 0 {
 		name = args[0]
 	}
-	switch plan, err := pl.Plan(cmd.Context(), name, desired); {
+	var cerr migrate.NotCleanError
+	switch plan, err := pl.Plan(cmd.Context(), name, desired.StateReader); {
 	case errors.Is(err, migrate.ErrNoPlan):
 		cmd.Println("The migration directory is synced with the desired state, no changes to be made")
 		return nil
+	case errors.As(err, &cerr) && dev.URL.Schema == "" && desired.Schema != "":
+		return fmt.Errorf("dev database is not clean (%s). Add a schema to the URL to limit the scope of the connection", cerr.Reason)
 	case err != nil:
 		return err
 	default:
@@ -474,7 +479,7 @@ func CmdMigrateNewRun(_ *cobra.Command, args []string) error {
 	if len(args) > 0 {
 		name = args[0]
 	}
-	return migrate.NewPlanner(nil, dir, migrate.WithFormatter(f)).WritePlan(&migrate.Plan{Name: name})
+	return migrate.NewPlanner(nil, dir, migrate.PlanFormat(f)).WritePlan(&migrate.Plan{Name: name})
 }
 
 // CmdMigrateStatusRun is the command executed when running the CLI with 'migrate status' args.
@@ -721,8 +726,14 @@ func dir(create bool) (migrate.Dir, error) {
 	return d, err
 }
 
+type target struct {
+	migrate.StateReader        // desired state.
+	io.Closer                  // optional close function.
+	Schema              string // in case we work on a single schema.
+}
+
 // to returns a migrate.StateReader for the given to flag.
-func to(ctx context.Context, client *sqlclient.Client) (migrate.StateReader, error) {
+func to(ctx context.Context, dev *sqlclient.Client) (*target, error) {
 	scheme, err := selectScheme(MigrateFlags.ToURLs)
 	if err != nil {
 		return nil, err
@@ -739,7 +750,7 @@ func to(ctx context.Context, client *sqlclient.Client) (migrate.StateReader, err
 		if err != nil {
 			return nil, err
 		}
-		if err := client.Eval(parsed, realm, nil); err != nil {
+		if err := dev.Eval(parsed, realm, nil); err != nil {
 			return nil, err
 		}
 		if len(schemas) > 0 {
@@ -754,28 +765,51 @@ func to(ctx context.Context, client *sqlclient.Client) (migrate.StateReader, err
 				}
 			}
 		}
-		if norm, ok := client.Driver.(schema.Normalizer); ok {
+		// In case the dev connection is bound to a specific schema, we require the
+		// desired schema to contain only one schema. Thus, executing diff will be
+		// done on the content of these two schema and not the whole realm.
+		if dev.URL.Schema != "" && len(realm.Schemas) > 1 {
+			return nil, fmt.Errorf("cannot use HCL with more than 1 schema when dev-url is limited to schema %q", dev.URL.Schema)
+		}
+		if norm, ok := dev.Driver.(schema.Normalizer); ok && len(realm.Schemas) > 0 {
 			realm, err = norm.NormalizeRealm(ctx, realm)
 			if err != nil {
 				return nil, err
 			}
 		}
-		return migrate.Realm(realm), nil
+		t := &target{StateReader: migrate.Realm(realm), Closer: io.NopCloser(nil)}
+		if len(realm.Schemas) == 1 {
+			t.Schema = realm.Schemas[0].Name
+		}
+		return t, nil
 	default: // database connection
 		client, err := sqlclient.Open(ctx, MigrateFlags.ToURLs[0])
 		if err != nil {
 			return nil, err
 		}
-		if client.URL.Schema != "" {
-			schemas = append(schemas, client.URL.Schema)
+		t := &target{Closer: client}
+		switch s := client.URL.Schema; {
+		// Connection to a specific schema.
+		case s != "":
+			if len(schemas) > 1 || len(schemas) == 1 && schemas[0] != s {
+				return nil, fmt.Errorf("cannot specify schemas with a schema connection to %q", s)
+			}
+			t.Schema = s
+			t.StateReader = migrate.SchemaConn(client, s, &schema.InspectOptions{})
+		// A single schema is selected.
+		case len(schemas) == 1:
+			t.Schema = schemas[0]
+			t.StateReader = migrate.SchemaConn(client, schemas[0], &schema.InspectOptions{})
+		// Multiple or all schemas.
+		default:
+			// In case the dev connection is limited to a single schema,
+			// but we compare it to entire database.
+			if dev.URL.Schema != "" {
+				return nil, fmt.Errorf("cannot use database-url without a schema when dev-url is limited to %q", dev.URL.Schema)
+			}
+			t.StateReader = migrate.RealmConn(client, &schema.InspectRealmOption{Schemas: schemas})
 		}
-		return struct {
-			migrate.StateReader
-			io.Closer
-		}{
-			Closer:      client,
-			StateReader: migrate.Conn(client, &schema.InspectRealmOption{Schemas: schemas}),
-		}, nil
+		return t, nil
 	}
 }
 
diff --git a/sql/migrate/migrate.go b/sql/migrate/migrate.go
index 9a13c689..e508f804 100644
--- a/sql/migrate/migrate.go
+++ b/sql/migrate/migrate.go
@@ -120,21 +120,33 @@ func Schema(s *schema.Schema) StateReader {
 	})
 }
 
-// Conn returns a StateReader for a Driver.
-func Conn(drv Driver, opts *schema.InspectRealmOption) StateReader {
+// RealmConn returns a StateReader for a Driver connected to a database.
+func RealmConn(drv Driver, opts *schema.InspectRealmOption) StateReader {
 	return StateReaderFunc(func(ctx context.Context) (*schema.Realm, error) {
 		return drv.InspectRealm(ctx, opts)
 	})
 }
 
+// SchemaConn returns a StateReader for a Driver connected to a schema.
+func SchemaConn(drv Driver, name string, opts *schema.InspectOptions) StateReader {
+	return StateReaderFunc(func(ctx context.Context) (*schema.Realm, error) {
+		s, err := drv.InspectSchema(ctx, name, opts)
+		if err != nil {
+			return nil, err
+		}
+		return Schema(s).ReadState(ctx)
+	})
+}
+
 type (
 	// Planner can plan the steps to take to migrate from one state to another. It uses the enclosed Dir to write
 	// those changes to versioned migration files.
 	Planner struct {
-		drv Driver    // driver to use
-		dir Dir       // where migration files are stored and read from
-		fmt Formatter // how to format a plan to migration files
-		sum bool      // whether to create a sum file for the migration directory
+		drv    Driver    // driver to use
+		dir    Dir       // where migration files are stored and read from
+		fmt    Formatter // how to format a plan to migration files
+		schema string    // optional schema to use for inspection
+		sum    bool      // whether to create a sum file for the migration directory
 	}
 
 	// PlannerOption allows managing a Planner using functional arguments.
@@ -236,20 +248,37 @@ func NewPlanner(drv Driver, dir Dir, opts ...PlannerOption) *Planner {
 	return p
 }
 
-// WithFormatter sets the Formatter of a Planner.
-func WithFormatter(fmt Formatter) PlannerOption {
+// PlanSchema allows setting a schema name to inspect
+// after replaying the migration directory.
+func PlanSchema(name string) PlannerOption {
+	return func(p *Planner) {
+		p.schema = name
+	}
+}
+
+// PlanFormat sets the Formatter of a Planner.
+func PlanFormat(fmt Formatter) PlannerOption {
 	return func(p *Planner) {
 		p.fmt = fmt
 	}
 }
 
-// DisableChecksum disables the hash-sum functionality for the migration directory.
-func DisableChecksum() PlannerOption {
+// PlanWithoutChecksum disables the hash-sum functionality for the migration directory.
+func PlanWithoutChecksum() PlannerOption {
 	return func(p *Planner) {
 		p.sum = false
 	}
 }
 
+var (
+	// WithFormatter calls PlanFormat.
+	// Deprecated: use PlanFormat instead.
+	WithFormatter = PlanFormat
+	// DisableChecksum calls PlanWithoutChecksum.
+	// Deprecated: use PlanWithoutChecksum instead.
+	DisableChecksum = PlanWithoutChecksum
+)
+
 // Plan calculates the migration Plan required for moving the current state (from) state to
 // the next state (to). A StateReader can be a directory, static schema elements or a Driver connection.
 func (p *Planner) Plan(ctx context.Context, name string, to StateReader) (*Plan, error) {
@@ -272,7 +301,22 @@ func (p *Planner) Plan(ctx context.Context, name string, to StateReader) (*Plan,
 	if err != nil {
 		return nil, err
 	}
-	changes, err := p.drv.RealmDiff(current, desired)
+	var changes []schema.Change
+	switch {
+	case p.schema != "":
+		switch s, ok := current.Schema(p.schema); {
+		case !ok:
+			return nil, fmt.Errorf("missing schema %q in realm after replaying migration directory", p.schema)
+		case len(desired.Schemas) == 0:
+			return nil, errors.New("missing schema definition in desired state")
+		case len(desired.Schemas) > 1:
+			return nil, fmt.Errorf("found %d schema definitions in desired state on schema mode", len(desired.Schemas))
+		default:
+			changes, err = p.drv.SchemaDiff(s, desired.Schemas[0])
+		}
+	default:
+		changes, err = p.drv.RealmDiff(current, desired)
+	}
 	if err != nil {
 		return nil, err
 	}


Test patch:
diff --git a/cmd/atlas/internal/cmdapi/migrate_test.go b/cmd/atlas/internal/cmdapi/migrate_test.go
index 46cdf474..1ecd8f65 100644
--- a/cmd/atlas/internal/cmdapi/migrate_test.go
+++ b/cmd/atlas/internal/cmdapi/migrate_test.go
@@ -270,6 +270,7 @@ func TestMigrate_ApplyBaseline(t *testing.T) {
 
 func TestMigrate_Diff(t *testing.T) {
 	p := t.TempDir()
+	to := hclURL(t)
 
 	// Will create migration directory if not existing.
 	MigrateFlags.Force = false
@@ -278,7 +279,8 @@ func TestMigrate_Diff(t *testing.T) {
 		"name",
 		"--dir", "file://"+filepath.Join(p, "migrations"),
 		"--dev-url", openSQLite(t, ""),
-		"--to", hclURL(t))
+		"--to", to,
+	)
 	require.NoError(t, err)
 	require.FileExists(t, filepath.Join(p, "migrations", fmt.Sprintf("%s_name.sql", time.Now().UTC().Format("20060102150405"))))
 
@@ -289,7 +291,7 @@ func TestMigrate_Diff(t *testing.T) {
 		"name",
 		"--dir", "file://"+p,
 		"--dev-url", openSQLite(t, "create table t (c int);"),
-		"--to", hclURL(t),
+		"--to", to,
 	)
 	require.ErrorAs(t, err, &migrate.NotCleanError{})
 	require.ErrorContains(t, err, "found table \"t\"")
@@ -300,7 +302,7 @@ func TestMigrate_Diff(t *testing.T) {
 		"name",
 		"--dir", "file://"+p,
 		"--dev-url", openSQLite(t, ""),
-		"--to", hclURL(t),
+		"--to", to,
 	)
 	require.NoError(t, err)
 	require.Zero(t, s)
@@ -324,7 +326,7 @@ func TestMigrate_Diff(t *testing.T) {
 		"name",
 		"--dir", "file://"+t.TempDir(),
 		"--dev-url", fmt.Sprintf("sqlitelockdiff://file:%s?cache=shared&_fk=1", filepath.Join(p, "test.db")),
-		"--to", hclURL(t),
+		"--to", to,
 	)
 	require.True(t, strings.HasPrefix(s, "Error: "+errLock.Error()))
 	require.ErrorIs(t, err, errLock)
diff --git a/sql/migrate/migrate_test.go b/sql/migrate/migrate_test.go
index 11828f01..d43be935 100644
--- a/sql/migrate/migrate_test.go
+++ b/sql/migrate/migrate_test.go
@@ -87,6 +87,39 @@ func TestPlanner_Plan(t *testing.T) {
 	require.Equal(t, drv.plan, plan)
 }
 
+func TestPlanner_PlanSchemaMode(t *testing.T) {
+	var (
+		drv = &lockMockDriver{&mockDriver{}}
+		ctx = context.Background()
+	)
+	d, err := migrate.NewLocalDir(t.TempDir())
+	require.NoError(t, err)
+
+	// Schema is missing in dev connection.
+	pl := migrate.NewPlanner(drv, d, migrate.PlanSchema("test"))
+	plan, err := pl.Plan(ctx, "empty", migrate.Realm(nil))
+	require.EqualError(t, err, `missing schema "test" in realm after replaying migration directory`)
+	require.Nil(t, plan)
+
+	drv.realm = *schema.NewRealm(schema.New("test"))
+	pl = migrate.NewPlanner(drv, d, migrate.PlanSchema("test"))
+	plan, err = pl.Plan(ctx, "empty", migrate.Realm(schema.NewRealm()))
+	require.EqualError(t, err, `missing schema definition in desired state`)
+	require.Nil(t, plan)
+
+	drv.realm = *schema.NewRealm(schema.New("test"))
+	pl = migrate.NewPlanner(drv, d, migrate.PlanSchema("test"))
+	plan, err = pl.Plan(ctx, "multi", migrate.Realm(schema.NewRealm(schema.New("test"), schema.New("dev"))))
+	require.EqualError(t, err, `found 2 schema definitions in desired state on schema mode`)
+	require.Nil(t, plan)
+
+	drv.realm = *schema.NewRealm(schema.New("test"))
+	pl = migrate.NewPlanner(drv, d, migrate.PlanSchema("test"))
+	plan, err = pl.Plan(ctx, "multi", migrate.Realm(schema.NewRealm(schema.New("test"))))
+	require.ErrorIs(t, err, migrate.ErrNoPlan)
+	require.Nil(t, plan)
+}
+
 func TestExecutor_ReadState(t *testing.T) {
 	ctx := context.Background()
 	d, err := migrate.NewLocalDir("testdata/migrate")
@@ -117,6 +150,7 @@ func TestExecutor_ReadState(t *testing.T) {
 	drv.locks = make(map[string]struct{})
 
 	// Does not work if database is not clean.
+	drv.dirty = true
 	drv.realm = schema.Realm{Schemas: []*schema.Schema{{Name: "schema"}}}
 	_, err = ex.ReadState(ctx)
 	require.ErrorAs(t, err, &migrate.NotCleanError{})
@@ -457,6 +491,10 @@ func (m *mockDriver) InspectRealm(context.Context, *schema.InspectRealmOption) (
 	return &m.realm, nil
 }
 
+func (m *mockDriver) SchemaDiff(_, _ *schema.Schema) ([]schema.Change, error) {
+	return m.changes, nil
+}
+
 func (m *mockDriver) RealmDiff(_, _ *schema.Realm) ([]schema.Change, error) {
 	return m.changes, nil
 }
@@ -471,7 +509,7 @@ func (m *mockDriver) ApplyChanges(_ context.Context, changes []schema.Change) er
 }
 
 func (m *mockDriver) Snapshot(context.Context) (migrate.RestoreFunc, error) {
-	if len(m.realm.Schemas) > 0 {
+	if m.dirty {
 		return nil, migrate.NotCleanError{}
 	}
 	realm := m.realm