Repository: ariga/atlas
Pull request #3337 (2025-01-22 10:06:04)

sql/postgres: restore-realm function can be called multiple times
Fixed https://github.com/ariga/atlas/issues/3336


Patch:
diff --git a/sql/internal/sqlx/sqlx.go b/sql/internal/sqlx/sqlx.go
index ff381ce8..8e9b0571 100644
--- a/sql/internal/sqlx/sqlx.go
+++ b/sql/internal/sqlx/sqlx.go
@@ -78,20 +78,6 @@ func ScanNullBool(rows *sql.Rows) (sql.NullBool, error) {
 	return b, ScanOne(rows, &b)
 }
 
-// ScanStrings scans sql.Rows into a slice of strings and closes it at the end.
-func ScanStrings(rows *sql.Rows) ([]string, error) {
-	defer rows.Close()
-	var vs []string
-	for rows.Next() {
-		var v sql.NullString
-		if err := rows.Scan(&v); err != nil {
-			return nil, err
-		}
-		vs = append(vs, v.String)
-	}
-	return vs, nil
-}
-
 type (
 	// ScanStringer groups the fmt.Stringer and sql.Scanner interfaces.
 	ScanStringer interface {
diff --git a/sql/postgres/driver.go b/sql/postgres/driver.go
index 56511200..1994b1a9 100644
--- a/sql/postgres/driver.go
+++ b/sql/postgres/driver.go
@@ -7,7 +7,6 @@ package postgres
 import (
 	"context"
 	"database/sql"
-	"errors"
 	"fmt"
 	"hash/fnv"
 	"math/rand"
@@ -237,20 +236,40 @@ func (d *Driver) SchemaRestoreFunc(desired *schema.Schema) migrate.RestoreFunc {
 
 // RealmRestoreFunc returns a function that restores the given realm to its desired state.
 func (d *Driver) RealmRestoreFunc(desired *schema.Realm) migrate.RestoreFunc {
-	return func(ctx context.Context) (err error) {
-		// Default behavior for Postgres is to have a single "public" schema.
-		// In that case, all other schemas are dropped, but this one is cleared
-		// object by object. To keep process faster, we drop the schema and recreate it.
-		if !d.crdb && len(desired.Schemas) == 1 && desired.Schemas[0].Name == "public" {
-			if pb := desired.Schemas[0]; len(pb.Tables)+len(pb.Views)+len(pb.Funcs)+len(pb.Procs)+len(pb.Objects) == 0 {
-				desired = &schema.Realm{Attrs: desired.Attrs, Objects: desired.Objects}
-				defer func() {
-					err = errors.Join(err, d.ApplyChanges(ctx, []schema.Change{
-						&schema.AddSchema{S: pb, Extra: []schema.Clause{&schema.IfExists{}}},
-					}))
-				}()
+	// Default behavior for Postgres is to have a single "public" schema.
+	// In that case, all other schemas are dropped, but this one is cleared
+	// object by object. To keep process faster, we drop the schema and recreate it.
+	if !d.crdb && len(desired.Schemas) == 1 && desired.Schemas[0].Name == "public" {
+		if pb := desired.Schemas[0]; len(pb.Tables)+len(pb.Views)+len(pb.Funcs)+len(pb.Procs)+len(pb.Objects) == 0 {
+			return func(ctx context.Context) error {
+				current, err := d.InspectRealm(ctx, nil)
+				if err != nil {
+					return err
+				}
+				changes, err := d.RealmDiff(current, desired)
+				if err != nil {
+					return err
+				}
+				// If there is no diff, do nothing.
+				if len(changes) == 0 {
+					return nil
+				}
+				// Else, prefer to drop the public schema and apply
+				// database changes instead of executing changes one by one.
+				if changes, err = d.RealmDiff(current, &schema.Realm{Attrs: desired.Attrs, Objects: desired.Objects}); err != nil {
+					return err
+				}
+				if err := d.ApplyChanges(ctx, withCascade(changes)); err != nil {
+					return err
+				}
+				// Recreate the public schema.
+				return d.ApplyChanges(ctx, []schema.Change{
+					&schema.AddSchema{S: pb, Extra: []schema.Clause{&schema.IfExists{}}},
+				})
 			}
 		}
+	}
+	return func(ctx context.Context) (err error) {
 		current, err := d.InspectRealm(ctx, nil)
 		if err != nil {
 			return err


Test patch:
diff --git a/sql/postgres/driver_test.go b/sql/postgres/driver_test.go
index 4a0025bd..6882a9f7 100644
--- a/sql/postgres/driver_test.go
+++ b/sql/postgres/driver_test.go
@@ -167,6 +167,62 @@ func TestDriver_Version(t *testing.T) {
 	require.Equal(t, "130000", drv.(vr).Version())
 }
 
+func TestDriver_RealmRestoreFunc(t *testing.T) {
+	var (
+		apply   = &mockPlanApplier{}
+		inspect = &mockInspector{}
+		drv     = &Driver{
+			Inspector:   inspect,
+			Differ:      DefaultDiff,
+			conn:        &conn{schema: "test"},
+			PlanApplier: apply,
+		}
+	)
+	f := drv.RealmRestoreFunc(schema.NewRealm().AddSchemas(schema.New("public")))
+
+	// No changes.
+	inspect.realm = schema.NewRealm().AddSchemas(schema.New("public"))
+	err := f(context.Background())
+	require.NoError(t, err)
+	require.Empty(t, apply.applied)
+
+	// Schema changes.
+	inspect.realm = schema.NewRealm().AddSchemas(schema.New("public").AddTables(schema.NewTable("t1")))
+	err = f(context.Background())
+	require.NoError(t, err)
+	require.Len(t, apply.applied, 2)
+	drop, ok := apply.applied[0].(*schema.DropSchema)
+	require.True(t, ok)
+	require.Equal(t, "public", drop.S.Name)
+	create, ok := apply.applied[1].(*schema.AddSchema)
+	require.True(t, ok)
+	require.Equal(t, "public", create.S.Name)
+
+	// Recreate the public schema.
+	apply.applied = nil
+	inspect.realm = schema.NewRealm().AddSchemas(schema.New("test").AddTables(schema.NewTable("t1")))
+	err = f(context.Background())
+	require.NoError(t, err)
+	require.Len(t, apply.applied, 2)
+	drop, ok = apply.applied[0].(*schema.DropSchema)
+	require.True(t, ok)
+	require.Equal(t, "test", drop.S.Name)
+	create, ok = apply.applied[1].(*schema.AddSchema)
+	require.True(t, ok)
+	require.Equal(t, "public", create.S.Name)
+
+	// Non-public changes.
+	apply.applied = nil
+	f = drv.RealmRestoreFunc(schema.NewRealm().AddSchemas(schema.New("test")))
+	inspect.realm = schema.NewRealm().AddSchemas(schema.New("test").AddTables(schema.NewTable("t1")))
+	err = f(context.Background())
+	require.NoError(t, err)
+	require.Len(t, apply.applied, 1)
+	dropT, ok := apply.applied[0].(*schema.DropTable)
+	require.True(t, ok)
+	require.Equal(t, "t1", dropT.T.Name)
+}
+
 type mockInspector struct {
 	schema.Inspector
 	realm  *schema.Realm
@@ -183,3 +239,18 @@ func (m *mockInspector) InspectSchema(context.Context, string, *schema.InspectOp
 func (m *mockInspector) InspectRealm(context.Context, *schema.InspectRealmOption) (*schema.Realm, error) {
 	return m.realm, nil
 }
+
+type mockPlanApplier struct {
+	planned []schema.Change
+	applied []schema.Change
+}
+
+func (m *mockPlanApplier) PlanChanges(_ context.Context, _ string, planned []schema.Change, _ ...migrate.PlanOption) (*migrate.Plan, error) {
+	m.planned = append(m.planned, planned...)
+	return nil, nil
+}
+
+func (m *mockPlanApplier) ApplyChanges(_ context.Context, applied []schema.Change, _ ...migrate.PlanOption) error {
+	m.applied = append(m.applied, applied...)
+	return nil
+}