Skip to content

Commit

Permalink
sql/postgres: restore-realm function should be called multiple times (#…
Browse files Browse the repository at this point in the history
a8m authored Jan 22, 2025
1 parent 9102f98 commit 399dfdc
Showing 3 changed files with 103 additions and 27 deletions.
14 changes: 0 additions & 14 deletions sql/internal/sqlx/sqlx.go
Original file line number Diff line number Diff line change
@@ -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 {
45 changes: 32 additions & 13 deletions sql/postgres/driver.go
Original file line number Diff line number Diff line change
@@ -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
71 changes: 71 additions & 0 deletions sql/postgres/driver_test.go
Original file line number Diff line number Diff line change
@@ -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
}

0 comments on commit 399dfdc

Please sign in to comment.