Repository: supabase/cli
Pull request #396 (2022-08-22 03:49:42)

fix: implement db reset using pgx
## What kind of change does this PR introduce?

enhancement

## What is the current behavior?

- `db reset` runs psql inside a container to apply migration scripts.
- reset failure will wipe out existing data in local postgres

## What is the new behavior?

- opens a pgx connection to apply migrations: more efficient than copy + exec seed file in container
- supports line specific debug output with `--debug`
- applies migrations to a separate clean database before switching over to `postgres`
- use psql command to disconnect database: more robust than disconnecting docker network
- update terminate sql to check active replication slots of specified database only

## Additional context

Add any other context or screenshots.



Patch:
diff --git a/cmd/db.go b/cmd/db.go
index cd7d002b..6b907d47 100644
--- a/cmd/db.go
+++ b/cmd/db.go
@@ -132,7 +132,8 @@ var (
 		Use:   "reset",
 		Short: "Resets the local database to current migrations",
 		RunE: func(cmd *cobra.Command, args []string) error {
-			return reset.Run()
+			ctx, _ := signal.NotifyContext(cmd.Context(), os.Interrupt)
+			return reset.Run(ctx, afero.NewOsFs())
 		},
 	}
 )
diff --git a/internal/db/reset/reset.go b/internal/db/reset/reset.go
index 021b476e..401e4cb9 100644
--- a/internal/db/reset/reset.go
+++ b/internal/db/reset/reset.go
@@ -1,320 +1,139 @@
 package reset
 
 import (
-	"bytes"
 	"context"
 	"errors"
 	"fmt"
-	"io"
 	"os"
-	"regexp"
-	"strconv"
-	"strings"
 
-	"github.com/charmbracelet/bubbles/progress"
-	"github.com/charmbracelet/bubbles/spinner"
-	tea "github.com/charmbracelet/bubbletea"
-	"github.com/charmbracelet/lipgloss"
-	"github.com/docker/docker/api/types/network"
-	"github.com/docker/docker/pkg/stdcopy"
-	"github.com/muesli/reflow/wrap"
+	"github.com/jackc/pgconn"
+	"github.com/jackc/pgerrcode"
+	"github.com/jackc/pgx/v4"
+	"github.com/spf13/afero"
+	"github.com/spf13/viper"
+	"github.com/supabase/cli/internal/db/diff"
+	"github.com/supabase/cli/internal/debug"
 	"github.com/supabase/cli/internal/utils"
+	"github.com/supabase/cli/internal/utils/parser"
 )
 
-// TODO: Handle cleanup on SIGINT/SIGTERM.
-func Run() error {
+func Run(ctx context.Context, fsys afero.Fs) error {
 	// Sanity checks.
 	{
-		if err := utils.AssertSupabaseStartIsRunning(); err != nil {
+		if err := utils.LoadConfigFS(fsys); err != nil {
 			return err
 		}
-
-		branch, err := utils.GetCurrentBranch()
-		if err != nil {
+		if err := utils.AssertSupabaseDbIsRunning(); err != nil {
 			return err
 		}
-		currBranch = branch
 	}
 
-	s := spinner.NewModel()
-	s.Spinner = spinner.Dot
-	s.Style = lipgloss.NewStyle().Foreground(lipgloss.Color("205"))
-	p := utils.NewProgram(model{spinner: s})
+	branch, err := utils.GetCurrentBranchFS(fsys)
+	if err != nil {
+		// Assume we are on main branch
+		branch = "main"
+	}
 
-	errCh := make(chan error, 1)
-	go func() {
-		errCh <- run(p)
-		p.Send(tea.Quit())
-	}()
+	var opts []func(*pgx.ConnConfig)
+	if viper.GetBool("DEBUG") {
+		opts = append(opts, debug.SetupPGX)
+	}
 
-	if err := p.Start(); err != nil {
+	fmt.Fprintln(os.Stderr, "Resetting database...")
+	if err := diff.ResetDatabase(ctx, utils.DbId, branch); err != nil {
 		return err
 	}
-	if errors.Is(ctx.Err(), context.Canceled) {
-		return errors.New("Aborted " + utils.Aqua("supabase db reset") + ".")
-	}
-	if err := <-errCh; err != nil {
+
+	fmt.Fprintln(os.Stderr, "Initialising schema...")
+	url := fmt.Sprintf("postgresql://postgres:postgres@localhost:%d/%s", utils.Config.Db.Port, branch)
+	if err := diff.ApplyMigrations(ctx, url, fsys, opts...); err != nil {
 		return err
 	}
 
-	fmt.Println("Finished " + utils.Aqua("supabase db reset") + " on branch " + utils.Aqua(currBranch) + ".")
-	return nil
-}
-
-var (
-	ctx, cancelCtx = context.WithCancel(context.Background())
-
-	currBranch string
-)
-
-func run(p utils.Program) error {
-	defer cleanup()
-
-	// 1. Prevent new db connections to be established while db is recreated.
-	if err := utils.Docker.NetworkDisconnect(ctx, utils.NetId, utils.DbId, false); err != nil {
+	if err := SeedDatabase(ctx, url, fsys); err != nil && !errors.Is(err, os.ErrNotExist) {
 		return err
 	}
 
-	p.Send(utils.StatusMsg("Resetting database..."))
-
-	if err := func() error {
-		// 2. Recreate db.
-		{
-			out, err := utils.DockerExec(ctx, utils.DbId, []string{
-				"sh", "-c", `psql --set ON_ERROR_STOP=on postgresql://postgres:postgres@localhost/template1 <<'EOSQL'
-BEGIN;
-` + fmt.Sprintf(utils.TerminateDbSqlFmt, "postgres") + `
-COMMIT;
-DROP DATABASE postgres;
-CREATE DATABASE postgres;
-EOSQL
-`,
-			})
-			if err != nil {
-				return err
-			}
-			var errBuf bytes.Buffer
-			if _, err := stdcopy.StdCopy(io.Discard, &errBuf, out); err != nil {
-				return err
-			}
-			if errBuf.Len() > 0 {
-				return errors.New("Error resetting database: " + errBuf.String())
-			}
-		}
-
-		// 3. Apply initial schema + migrations + seed.
-
-		p.Send(utils.StatusMsg("Setting up initial schema..."))
-		{
-			out, err := utils.DockerExec(ctx, utils.DbId, []string{
-				"sh", "-c", `PGOPTIONS='--client-min-messages=error' psql postgresql://postgres:postgres@localhost/postgres <<'EOSQL'
-BEGIN;
-` + utils.InitialSchemaSql + `
-COMMIT;
-EOSQL
-`,
-			})
-			if err != nil {
-				return err
-			}
-			var errBuf bytes.Buffer
-			if _, err := stdcopy.StdCopy(io.Discard, &errBuf, out); err != nil {
-				return err
-			}
-			if errBuf.Len() > 0 {
-				return errors.New("Error resetting database: " + errBuf.String())
-			}
-		}
-
-		if err := utils.MkdirIfNotExist("supabase/migrations"); err != nil {
-			return err
-		}
-		migrations, err := os.ReadDir("supabase/migrations")
-		if err != nil {
-			return err
-		}
-
-		for i, migration := range migrations {
-			// NOTE: To handle backward-compatibility. `<timestamp>_init.sql` as
-			// the first migration (prev versions of the CLI) is deprecated.
-			if i == 0 {
-				matches := regexp.MustCompile(`([0-9]{14})_init\.sql`).FindStringSubmatch(migration.Name())
-				if len(matches) == 2 {
-					if timestamp, err := strconv.ParseUint(matches[1], 10, 64); err != nil {
-						return err
-					} else if timestamp < 20211209000000 {
-						continue
-					}
-				}
-			}
-
-			p.Send(utils.StatusMsg("Applying migration " + utils.Bold(migration.Name()) + "..."))
-
-			content, err := os.ReadFile("supabase/migrations/" + migration.Name())
-			if err != nil {
-				return err
-			}
-
-			out, err := utils.DockerExec(ctx, utils.DbId, []string{
-				"sh", "-c", `PGOPTIONS='--client-min-messages=error' psql postgresql://postgres:postgres@localhost/postgres <<'EOSQL'
-BEGIN;
-` + string(content) + `
-COMMIT;
-EOSQL
-`,
-			})
-			if err != nil {
-				return err
-			}
-			var errBuf bytes.Buffer
-			if _, err := stdcopy.StdCopy(io.Discard, &errBuf, out); err != nil {
-				return err
-			}
-			if errBuf.Len() > 0 {
-				return errors.New("Error resetting database: " + errBuf.String())
-			}
-		}
-
-		p.Send(utils.StatusMsg("Applying " + utils.Bold("supabase/seed.sql") + "..."))
-		{
-			content, err := os.ReadFile("supabase/seed.sql")
-			if errors.Is(err, os.ErrNotExist) {
-				// skip
-			} else if err != nil {
-				return err
-			}
-
-			err = utils.DockerAddFile(ctx, utils.DbId, "seed.sql", content)
-
-			if err != nil {
-				return err
-			} else {
-				out, err := utils.DockerExec(ctx, utils.DbId, []string{
-					"psql", "postgresql://postgres:postgres@localhost/postgres", "-f", "/tmp/seed.sql",
-				})
-				if err != nil {
-					return err
-				}
-				var errBuf bytes.Buffer
-				if _, err := stdcopy.StdCopy(io.Discard, &errBuf, out); err != nil {
-					return err
-				}
-				if errBuf.Len() > 0 {
-					return errors.New("Error resetting database: " + errBuf.String())
-				}
-			}
-		}
-
-		// Reload PostgREST schema cache.
-		{
-			// Need to connect for PostgREST to connect.
-			if err := utils.Docker.NetworkConnect(ctx, utils.NetId, utils.DbId, &network.EndpointSettings{}); err != nil {
-				fmt.Fprintf(os.Stderr, "Error reconnecting database: %v", err)
-				return nil
-			}
-
-			if err := utils.Docker.ContainerKill(ctx, utils.RestId, "SIGUSR1"); err != nil {
-				fmt.Fprintf(os.Stderr, "Error reloading PostgREST schema cache: %v", err)
-				return nil
-			}
-		}
-
-		return nil
-	}(); err != nil {
-		_ = os.RemoveAll("supabase/.branches/" + currBranch)
+	fmt.Fprintln(os.Stderr, "Activating branch...")
+	if err := ActivateDatabase(ctx, branch); err != nil {
 		return err
 	}
 
-	return nil
-}
-
-func cleanup() {
-	_ = utils.Docker.NetworkConnect(context.Background(), utils.NetId, utils.DbId, &network.EndpointSettings{})
-}
-
-type model struct {
-	spinner     spinner.Model
-	status      string
-	progress    *progress.Model
-	psqlOutputs []string
+	// Reload PostgREST schema cache.
+	if err := utils.Docker.ContainerKill(ctx, utils.RestId, "SIGUSR1"); err != nil {
+		fmt.Fprintf(os.Stderr, "Error reloading PostgREST schema cache: %v", err)
+	}
 
-	width int
+	fmt.Fprintln(os.Stderr, "Finished "+utils.Aqua("supabase db reset")+" on branch "+utils.Aqua(branch)+".")
+	return nil
 }
 
-func (m model) Init() tea.Cmd {
-	return spinner.Tick
+func SeedDatabase(ctx context.Context, url string, fsys afero.Fs, options ...func(*pgx.ConnConfig)) error {
+	sql, err := fsys.Open(utils.SeedDataPath)
+	if err != nil {
+		return err
+	}
+	defer sql.Close()
+	fmt.Fprintln(os.Stderr, "Seeding data...")
+	// Parse connection url
+	config, err := pgx.ParseConfig(url)
+	if err != nil {
+		return err
+	}
+	// Apply config overrides
+	for _, op := range options {
+		op(config)
+	}
+	// Connect to database
+	conn, err := pgx.ConnectConfig(ctx, config)
+	if err != nil {
+		return err
+	}
+	defer conn.Close(ctx)
+	// Batch seed commands
+	batch := diff.ToBatchQuery(parser.Split(sql))
+	if err := conn.SendBatch(ctx, &batch).Close(); err != nil {
+		return err
+	}
+	return nil
 }
 
-func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
-	switch msg := msg.(type) {
-	case tea.KeyMsg:
-		switch msg.Type {
-		case tea.KeyCtrlC:
-			// Stop future runs
-			cancelCtx()
-			// Stop current runs
-			cleanup()
-			return m, tea.Quit
-		default:
-			return m, nil
-		}
-	case tea.WindowSizeMsg:
-		m.width = msg.Width
-		return m, nil
-	case spinner.TickMsg:
-		spinnerModel, cmd := m.spinner.Update(msg)
-		m.spinner = spinnerModel
-		return m, cmd
-	case progress.FrameMsg:
-		if m.progress == nil {
-			return m, nil
-		}
-
-		tmp, cmd := m.progress.Update(msg)
-		progressModel := tmp.(progress.Model)
-		m.progress = &progressModel
-		return m, cmd
-	case utils.StatusMsg:
-		m.status = string(msg)
-		return m, nil
-	case utils.ProgressMsg:
-		if msg == nil {
-			m.progress = nil
-			return m, nil
-		}
-
-		if m.progress == nil {
-			progressModel := progress.NewModel(progress.WithGradient("#1c1c1c", "#34b27b"))
-			m.progress = &progressModel
-		}
-
-		return m, m.progress.SetPercent(*msg)
-	case utils.PsqlMsg:
-		if msg == nil {
-			m.psqlOutputs = []string{}
-			return m, nil
-		}
-
-		m.psqlOutputs = append(m.psqlOutputs, *msg)
-		if len(m.psqlOutputs) > 5 {
-			m.psqlOutputs = m.psqlOutputs[1:]
+func ActivateDatabase(ctx context.Context, branch string, options ...func(*pgx.ConnConfig)) error {
+	// Parse connection url
+	url := fmt.Sprintf("postgresql://postgres:postgres@localhost:%d/template1", utils.Config.Db.Port)
+	config, err := pgx.ParseConfig(url)
+	if err != nil {
+		return err
+	}
+	// Apply config overrides
+	for _, op := range options {
+		op(config)
+	}
+	// Connect to database
+	conn, err := pgx.ConnectConfig(ctx, config)
+	if err != nil {
+		return err
+	}
+	defer conn.Close(ctx)
+	// Must be executed separately because running in transaction is unsupported
+	disconn := "ALTER DATABASE postgres ALLOW_CONNECTIONS false;"
+	if _, err := conn.Exec(ctx, disconn); err != nil {
+		var pgErr *pgconn.PgError
+		if errors.As(err, &pgErr) && pgErr.Code != pgerrcode.InvalidCatalogName {
+			return err
 		}
-		return m, nil
-	default:
-		return m, nil
 	}
-}
-
-func (m model) View() string {
-	var progress string
-	if m.progress != nil {
-		progress = "\n\n" + m.progress.View()
+	term := fmt.Sprintf(utils.TerminateDbSqlFmt, "postgres")
+	if _, err := conn.Exec(ctx, term); err != nil {
+		return err
 	}
-
-	var psqlOutputs string
-	if len(m.psqlOutputs) > 0 {
-		psqlOutputs = "\n\n" + strings.Join(m.psqlOutputs, "\n")
+	drop := "DROP DATABASE IF EXISTS postgres WITH (FORCE);"
+	if _, err := conn.Exec(ctx, drop); err != nil {
+		return err
 	}
-
-	return wrap.String(m.spinner.View()+m.status+progress+psqlOutputs, m.width)
+	swap := "ALTER DATABASE " + branch + " RENAME TO postgres;"
+	if _, err := conn.Exec(ctx, swap); err != nil {
+		return err
+	}
+	return nil
 }
diff --git a/internal/utils/misc.go b/internal/utils/misc.go
index f47782db..5bb0e172 100644
--- a/internal/utils/misc.go
+++ b/internal/utils/misc.go
@@ -48,8 +48,9 @@ const (
 	TerminateDbSqlFmt = `
 SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = '%[1]s';
 -- Wait for WAL sender to drop replication slot.
-DO 'BEGIN WHILE (SELECT COUNT(*) FROM pg_replication_slots) > 0 LOOP END LOOP; END';
-`
+DO 'BEGIN WHILE (
+	SELECT COUNT(*) FROM pg_replication_slots WHERE database = ''%[1]s''
+) > 0 LOOP END LOOP; END';`
 	AnonKey        = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJzdXBhYmFzZS1kZW1vIiwicm9sZSI6ImFub24ifQ.625_WdcF3KHqz5amU0x2X5WWHP-OEs_4qj0ssLNHzTs"
 	ServiceRoleKey = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJzdXBhYmFzZS1kZW1vIiwicm9sZSI6InNlcnZpY2Vfcm9sZSJ9.vI9obAHOGyVVKa3pD--kJlyxp-Z2zV9UUMAhKpNLAcU"
 
@@ -58,6 +59,7 @@ DO 'BEGIN WHILE (SELECT COUNT(*) FROM pg_replication_slots) > 0 LOOP END LOOP; E
 	RemoteDbPath   = "supabase/.temp/remote-db-url"
 	CurrBranchPath = "supabase/.branches/_current_branch"
 	MigrationsDir  = "supabase/migrations"
+	SeedDataPath   = "supabase/seed.sql"
 )
 
 var (


Test patch:
diff --git a/internal/db/diff/migra_test.go b/internal/db/diff/migra_test.go
index 3278b115..82c17b2a 100644
--- a/internal/db/diff/migra_test.go
+++ b/internal/db/diff/migra_test.go
@@ -83,7 +83,7 @@ func TestApplyMigrations(t *testing.T) {
 		conn := pgtest.NewConn()
 		defer conn.Close(t)
 		conn.Query(query).
-			ReplyError(pgerrcode.DuplicateTable, "table \"test\" already exists")
+			ReplyError(pgerrcode.DuplicateTable, `relation "test" already exists`)
 		// Run test
 		assert.Error(t, ApplyMigrations(context.Background(), postgresUrl, fsys, conn.Intercept))
 	})
diff --git a/internal/db/reset/reset_test.go b/internal/db/reset/reset_test.go
new file mode 100644
index 00000000..a5e05ee2
--- /dev/null
+++ b/internal/db/reset/reset_test.go
@@ -0,0 +1,214 @@
+package reset
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net/http"
+	"os"
+	"testing"
+
+	"github.com/docker/docker/api/types"
+	"github.com/docker/docker/client"
+	"github.com/jackc/pgerrcode"
+	"github.com/spf13/afero"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/supabase/cli/internal/testing/apitest"
+	"github.com/supabase/cli/internal/testing/pgtest"
+	"github.com/supabase/cli/internal/utils"
+	"gopkg.in/h2non/gock.v1"
+)
+
+func TestResetCommand(t *testing.T) {
+	const version = "1.41"
+
+	t.Run("throws error on missing config", func(t *testing.T) {
+		assert.Error(t, Run(context.Background(), afero.NewMemMapFs()))
+	})
+
+	t.Run("throws error on db is not started", func(t *testing.T) {
+		// Setup in-memory fs
+		fsys := afero.NewMemMapFs()
+		require.NoError(t, utils.WriteConfig(fsys, false))
+		// Setup mock docker
+		require.NoError(t, client.WithHTTPClient(http.DefaultClient)(utils.Docker))
+		defer gock.OffAll()
+		gock.New("http:///var/run/docker.sock").
+			Head("/_ping").
+			Reply(http.StatusOK).
+			SetHeader("API-Version", version).
+			SetHeader("OSType", "linux")
+		gock.New("http:///var/run/docker.sock").
+			Get("/v" + version + "/containers").
+			Reply(http.StatusServiceUnavailable)
+		// Run test
+		assert.Error(t, Run(context.Background(), fsys))
+		// Validate api
+		assert.Empty(t, apitest.ListUnmatchedRequests())
+	})
+
+	t.Run("throws error on failure to exec", func(t *testing.T) {
+		// Setup in-memory fs
+		fsys := afero.NewMemMapFs()
+		require.NoError(t, utils.WriteConfig(fsys, false))
+		// Setup mock docker
+		require.NoError(t, client.WithHTTPClient(http.DefaultClient)(utils.Docker))
+		defer gock.OffAll()
+		gock.New("http:///var/run/docker.sock").
+			Head("/_ping").
+			Reply(http.StatusOK).
+			SetHeader("API-Version", version).
+			SetHeader("OSType", "linux")
+		gock.New("http:///var/run/docker.sock").
+			Get("/v" + version + "/containers").
+			Reply(200).
+			JSON(types.ContainerJSON{})
+		gock.New("http:///var/run/docker.sock").
+			Post("/v" + version + "/containers").
+			Reply(http.StatusServiceUnavailable)
+		// Run test
+		assert.Error(t, Run(context.Background(), fsys))
+		// Validate api
+		assert.Empty(t, apitest.ListUnmatchedRequests())
+	})
+}
+
+func TestSeedDatabase(t *testing.T) {
+	const postgresUrl = "postgresql://postgres:password@localhost:5432/postgres"
+
+	t.Run("seeds from file", func(t *testing.T) {
+		// Setup in-memory fs
+		fsys := afero.NewMemMapFs()
+		// Setup seed file
+		sql := "INSERT INTO employees(name) VALUES ('Alice')"
+		require.NoError(t, afero.WriteFile(fsys, utils.SeedDataPath, []byte(sql), 0644))
+		// Setup mock postgres
+		conn := pgtest.NewConn()
+		defer conn.Close(t)
+		conn.Query(sql).
+			Reply("INSERT 0 1")
+		// Run test
+		assert.NoError(t, SeedDatabase(context.Background(), postgresUrl, fsys, conn.Intercept))
+	})
+
+	t.Run("throws error on missing seed", func(t *testing.T) {
+		// Run test
+		err := SeedDatabase(context.Background(), postgresUrl, afero.NewMemMapFs())
+		// Check error
+		if assert.Error(t, err) {
+			assert.True(t, errors.Is(err, os.ErrNotExist))
+		}
+	})
+
+	t.Run("throws error on malformed url", func(t *testing.T) {
+		// Setup in-memory fs
+		fsys := afero.NewMemMapFs()
+		// Setup seed file
+		_, err := fsys.Create(utils.SeedDataPath)
+		require.NoError(t, err)
+		// Run test
+		assert.Error(t, SeedDatabase(context.Background(), "malformed", fsys))
+	})
+
+	t.Run("throws error on postgres", func(t *testing.T) {
+		// Setup in-memory fs
+		fsys := afero.NewMemMapFs()
+		// Setup seed file
+		sql := "INSERT INTO employees(name) VALUES ('Alice')"
+		require.NoError(t, afero.WriteFile(fsys, utils.SeedDataPath, []byte(sql), 0644))
+		// Setup mock postgres
+		conn := pgtest.NewConn()
+		defer conn.Close(t)
+		conn.Query(sql).
+			ReplyError(pgerrcode.NotNullViolation, `null value in column "age" of relation "employees"`)
+		// Run test
+		assert.Error(t, SeedDatabase(context.Background(), postgresUrl, fsys, conn.Intercept))
+	})
+}
+
+func TestActivateDatabase(t *testing.T) {
+	const branch = "main"
+
+	t.Run("activates main branch", func(t *testing.T) {
+		utils.Config.Db.Port = 54322
+		// Setup mock postgres
+		conn := pgtest.NewConn()
+		defer conn.Close(t)
+		conn.Query("ALTER DATABASE postgres ALLOW_CONNECTIONS false;").
+			Reply("ALTER DATABASE").
+			Query(fmt.Sprintf(utils.TerminateDbSqlFmt, "postgres")).
+			Reply("DO").
+			Query("DROP DATABASE IF EXISTS postgres WITH (FORCE);").
+			Reply("DROP DATABASE").
+			Query("ALTER DATABASE " + branch + " RENAME TO postgres;").
+			Reply("ALTER DATABASE")
+		// Run test
+		assert.NoError(t, ActivateDatabase(context.Background(), branch, conn.Intercept))
+	})
+
+	t.Run("continues on missing database", func(t *testing.T) {
+		utils.Config.Db.Port = 54322
+		// Setup mock postgres
+		conn := pgtest.NewConn()
+		defer conn.Close(t)
+		conn.Query("ALTER DATABASE postgres ALLOW_CONNECTIONS false;").
+			ReplyError(pgerrcode.InvalidCatalogName, `database "postgres" does not exist`).
+			Query(fmt.Sprintf(utils.TerminateDbSqlFmt, "postgres")).
+			ReplyError(pgerrcode.UndefinedTable, `relation "pg_stat_activity" does not exist`)
+		// Run test
+		assert.Error(t, ActivateDatabase(context.Background(), branch, conn.Intercept))
+	})
+
+	t.Run("throws error on invalid port", func(t *testing.T) {
+		// Default port 0
+		assert.Error(t, ActivateDatabase(context.Background(), branch))
+		// Setup invalid port
+		utils.Config.Db.Port = 65536
+		// Run test
+		assert.Error(t, ActivateDatabase(context.Background(), branch))
+	})
+
+	t.Run("throws error on failure to disconnect", func(t *testing.T) {
+		utils.Config.Db.Port = 54322
+		// Setup mock postgres
+		conn := pgtest.NewConn()
+		defer conn.Close(t)
+		conn.Query("ALTER DATABASE postgres ALLOW_CONNECTIONS false;").
+			ReplyError(pgerrcode.InvalidParameterValue, `cannot disallow connections for current database`)
+		// Run test
+		assert.Error(t, ActivateDatabase(context.Background(), branch, conn.Intercept))
+	})
+
+	t.Run("throws error on failure to drop", func(t *testing.T) {
+		utils.Config.Db.Port = 54322
+		// Setup mock postgres
+		conn := pgtest.NewConn()
+		defer conn.Close(t)
+		conn.Query("ALTER DATABASE postgres ALLOW_CONNECTIONS false;").
+			Reply("ALTER DATABASE").
+			Query(fmt.Sprintf(utils.TerminateDbSqlFmt, "postgres")).
+			Reply("DO").
+			Query("DROP DATABASE IF EXISTS postgres WITH (FORCE);").
+			ReplyError(pgerrcode.ObjectInUse, `database "postgres" is used by an active logical replication slot`)
+		// Run test
+		assert.Error(t, ActivateDatabase(context.Background(), branch, conn.Intercept))
+	})
+
+	t.Run("throws error on failure to swap", func(t *testing.T) {
+		utils.Config.Db.Port = 54322
+		// Setup mock postgres
+		conn := pgtest.NewConn()
+		defer conn.Close(t)
+		conn.Query("ALTER DATABASE postgres ALLOW_CONNECTIONS false;").
+			Reply("ALTER DATABASE").
+			Query(fmt.Sprintf(utils.TerminateDbSqlFmt, "postgres")).
+			Reply("DO").
+			Query("DROP DATABASE IF EXISTS postgres WITH (FORCE);").
+			Reply("DROP DATABASE").
+			Query("ALTER DATABASE "+branch+" RENAME TO postgres;").
+			ReplyError(pgerrcode.DuplicateDatabase, `database "postgres" already exists`)
+		// Run test
+		assert.Error(t, ActivateDatabase(context.Background(), branch, conn.Intercept))
+	})
+}