Repository: supabase/cli
Pull request #465 (2022-09-17 08:17:47)

fix: restart database after issuing pg_terminate_backend
## What kind of change does this PR introduce?

Bug fix

## What is the current behavior?

pg_terminate_backend induces pg_cron scheduler to shutdown

## What is the new behavior?

- restart postgres as recommended https://github.com/citusdata/pg_cron/issues/99 
- refactors branch switch to use pgx
- add unit tests
- adds health probe 

## Additional context

Add any other context or screenshots.


coveralls: ## Pull Request Test Coverage Report for [Build 3072675396](https://coveralls.io/builds/52551868)

* **8** of **49**   **(16.33%)**  changed or added relevant lines in **5** files are covered.
* **3** unchanged lines in **2** files lost coverage.
* Overall coverage increased (+**0.2%**) to **36.561%**

---

|  Changes Missing Coverage | Covered Lines | Changed/Added Lines | % |
| :-----|--------------|--------|---: |
| [cmd/db.go](https://coveralls.io/builds/52551868/source?filename=cmd%2Fdb.go#L66) | 0 | 2 | 0.0%
| [internal/start/start.go](https://coveralls.io/builds/52551868/source?filename=internal%2Fstart%2Fstart.go#L273) | 0 | 2 | 0.0%
| [internal/db/branch/switch_/switch_.go](https://coveralls.io/builds/52551868/source?filename=internal%2Fdb%2Fbranch%2Fswitch_%2Fswitch_.go#L18) | 0 | 37 | 0.0%
<!-- | **Total:** | **8** | **49** | **16.33%** | -->

|  Files with Coverage Reduction | New Missed Lines | % |
| :-----|--------------|--: |
| [cmd/db.go](https://coveralls.io/builds/52551868/source?filename=cmd%2Fdb.go#L68) | 1 | 46.34% |
| [internal/db/branch/switch_/switch_.go](https://coveralls.io/builds/52551868/source?filename=internal%2Fdb%2Fbranch%2Fswitch_%2Fswitch_.go#L45) | 2 | 0% |
<!-- | **Total:** | **3** |  | -->

|  Totals | [![Coverage Status](https://coveralls.io/builds/52551868/badge)](https://coveralls.io/builds/52551868) |
| :-- | --: |
| Change from base [Build 3068522683](https://coveralls.io/builds/52536625): |  0.2% |
| Covered Lines: | 1805 |
| Relevant Lines: | 4937 |

---
##### 💛  - [Coveralls](https://coveralls.io)

soedirgo: We should also do a db healthcheck after restarting, otherwise some services might have issues connecting:
```
2022-09-19T06:48:28: PM2 log: Launching in no daemon mode
2022-09-19T06:48:28: PM2 log: App [server:0] starting in -fork mode-
2022-09-19T06:48:28: PM2 log: App [server:0] online
running migrations
Error: connect ECONNREFUSED 172.31.0.2:5432
    at TCPConnectWrap.afterConnect [as oncomplete] (node:net:1300:16) {
  errno: -111,
  code: 'ECONNREFUSED',
  syscall: 'connect',
  address: '172.31.0.2',
  port: 5432
}
```
It'd be nice if we can configure PM2 to restart on error though
sweatybridge: added restart for storage-api and a todo to improve it.

Patch:
diff --git a/cmd/db.go b/cmd/db.go
index aca4d610..85aeb5f7 100644
--- a/cmd/db.go
+++ b/cmd/db.go
@@ -63,7 +63,8 @@ var (
 		Short: "Switch the active branch",
 		Args:  cobra.ExactArgs(1),
 		RunE: func(cmd *cobra.Command, args []string) error {
-			return switch_.Run(args[0])
+			ctx, _ := signal.NotifyContext(cmd.Context(), os.Interrupt)
+			return switch_.Run(ctx, args[0], afero.NewOsFs())
 		},
 	}
 
diff --git a/internal/db/branch/list/list.go b/internal/db/branch/list/list.go
index b54a8faa..72a8e50b 100644
--- a/internal/db/branch/list/list.go
+++ b/internal/db/branch/list/list.go
@@ -21,7 +21,7 @@ func Run(fsys afero.Fs, out io.Writer) error {
 
 	currBranch, _ := utils.GetCurrentBranchFS(fsys)
 	for _, branch := range branches {
-		if branch.Name() == "_current_branch" {
+		if branch.Name() == filepath.Base(utils.CurrBranchPath) {
 			continue
 		}
 
diff --git a/internal/db/branch/switch_/switch_.go b/internal/db/branch/switch_/switch_.go
index f52b66c8..4e800765 100644
--- a/internal/db/branch/switch_/switch_.go
+++ b/internal/db/branch/switch_/switch_.go
@@ -1,84 +1,86 @@
 package switch_
 
 import (
-	"bytes"
 	"context"
 	"errors"
 	"fmt"
-	"io"
 	"os"
+	"path/filepath"
 
-	"github.com/docker/docker/api/types/network"
-	"github.com/docker/docker/pkg/stdcopy"
+	"github.com/jackc/pgx/v4"
+	"github.com/spf13/afero"
+	"github.com/supabase/cli/internal/db/lint"
+	"github.com/supabase/cli/internal/db/reset"
 	"github.com/supabase/cli/internal/utils"
 )
 
-var ctx = context.Background()
-
-func Run(target string) error {
-	if err := utils.AssertSupabaseStartIsRunning(); err != nil {
-		return err
+func Run(ctx context.Context, target string, fsys afero.Fs, options ...func(*pgx.ConnConfig)) error {
+	// 1. Sanity checks
+	{
+		if err := utils.AssertSupabaseCliIsSetUpFS(fsys); err != nil {
+			return err
+		}
+		if err := utils.LoadConfigFS(fsys); err != nil {
+			return err
+		}
+		if err := utils.AssertSupabaseDbIsRunning(); err != nil {
+			return err
+		}
+		if target != "main" && utils.IsBranchNameReserved(target) {
+			return errors.New("Cannot switch branch " + utils.Aqua(target) + ": branch name is reserved.")
+		}
+		branchPath := filepath.Join(filepath.Dir(utils.CurrBranchPath), target)
+		if _, err := fsys.Stat(branchPath); errors.Is(err, os.ErrNotExist) {
+			return errors.New("Branch " + utils.Aqua(target) + " does not exist.")
+		} else if err != nil {
+			return err
+		}
 	}
 
-	errBranchNotExist := errors.New("Branch " + utils.Aqua(target) + " does not exist.")
-
-	branches, err := os.ReadDir("supabase/.branches")
-	if errors.Is(err, os.ErrNotExist) {
-		return errBranchNotExist
+	// 2. Check current branch
+	currBranch, err := utils.GetCurrentBranchFS(fsys)
+	if err != nil {
+		// Assume we are on main branch
+		currBranch = "main"
 	}
 
-	for _, branch := range branches {
-		if branch.Name() == "_current_branch" {
-			continue
-		}
-
-		if branch.Name() == target {
-			currBranch, err := utils.GetCurrentBranch()
-			if err != nil {
-				currBranch = "main"
-			}
-
-			if err := os.WriteFile(utils.CurrBranchPath, []byte(target), 0644); err != nil {
-				return err
-			}
-
-			// Prevent new db connections to be established while db is recreated.
-			if err := utils.Docker.NetworkDisconnect(ctx, utils.NetId, utils.DbId, false); err != nil {
-				return err
-			}
-
-			// Recreate current branch.
-			{
-				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;
-ALTER DATABASE postgres RENAME TO "` + currBranch + `";
-ALTER DATABASE "` + target + `" RENAME TO 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 switching to branch " + utils.Aqua(target) + ": " + errBuf.String())
-				}
-			}
+	// 3. Switch Postgres database
+	if currBranch == target {
+		fmt.Println("Already on branch " + utils.Aqua(target) + ".")
+	} else if err := switchDatabase(ctx, currBranch, target, options...); err != nil {
+		return errors.New("Error switching to branch " + utils.Aqua(target) + ": " + err.Error())
+	} else {
+		fmt.Println("Switched to branch " + utils.Aqua(target) + ".")
+	}
 
-			if err := utils.Docker.NetworkConnect(ctx, utils.NetId, utils.DbId, &network.EndpointSettings{}); err != nil {
-				return err
-			}
+	// 4. Update current branch
+	if err := afero.WriteFile(fsys, utils.CurrBranchPath, []byte(target), 0644); err != nil {
+		return errors.New("Unable to update local branch file. Fix by running: echo '" + target + "' > " + utils.CurrBranchPath)
+	}
+	return nil
+}
 
-			fmt.Println("Switched to branch " + utils.Aqua(target) + ".")
-			return nil
+func switchDatabase(ctx context.Context, source, target string, options ...func(*pgx.ConnConfig)) error {
+	conn, err := lint.ConnectLocalPostgres(ctx, "localhost", utils.Config.Db.Port, "template1", options...)
+	if err != nil {
+		return err
+	}
+	defer conn.Close(ctx)
+	if err := reset.DisconnectClients(ctx, conn); err != nil {
+		return err
+	}
+	defer reset.RestartDatabase(context.Background())
+	backup := "ALTER DATABASE postgres RENAME TO " + source + ";"
+	if _, err := conn.Exec(ctx, backup); err != nil {
+		return err
+	}
+	rename := "ALTER DATABASE " + target + " RENAME TO postgres;"
+	if _, err := conn.Exec(ctx, rename); err != nil {
+		rollback := "ALTER DATABASE " + source + " RENAME TO postgres;"
+		if _, err := conn.Exec(ctx, rollback); err != nil {
+			fmt.Fprintln(os.Stderr, "Failed to rollback database:", err)
 		}
+		return err
 	}
-
-	return errBranchNotExist
+	return nil
 }
diff --git a/internal/db/reset/reset.go b/internal/db/reset/reset.go
index e0dfe868..c7875d8b 100644
--- a/internal/db/reset/reset.go
+++ b/internal/db/reset/reset.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"os"
 	"strings"
+	"time"
 
 	"github.com/jackc/pgconn"
 	"github.com/jackc/pgerrcode"
@@ -18,6 +19,10 @@ import (
 	"github.com/supabase/cli/internal/utils/parser"
 )
 
+var (
+	healthTimeout = 5 * time.Second
+)
+
 func Run(ctx context.Context, fsys afero.Fs) error {
 	// Sanity checks.
 	{
@@ -122,6 +127,22 @@ func ActivateDatabase(ctx context.Context, branch string, options ...func(*pgx.C
 		return err
 	}
 	defer conn.Close(ctx)
+	if err := DisconnectClients(ctx, conn); err != nil {
+		return err
+	}
+	defer RestartDatabase(context.Background())
+	drop := "DROP DATABASE IF EXISTS postgres WITH (FORCE);"
+	if _, err := conn.Exec(ctx, drop); err != nil {
+		return err
+	}
+	swap := "ALTER DATABASE " + branch + " RENAME TO postgres;"
+	if _, err := conn.Exec(ctx, swap); err != nil {
+		return err
+	}
+	return nil
+}
+
+func DisconnectClients(ctx context.Context, conn *pgx.Conn) error {
 	// 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 {
@@ -134,13 +155,29 @@ func ActivateDatabase(ctx context.Context, branch string, options ...func(*pgx.C
 	if _, err := conn.Exec(ctx, term); err != nil {
 		return err
 	}
-	drop := "DROP DATABASE IF EXISTS postgres WITH (FORCE);"
-	if _, err := conn.Exec(ctx, drop); err != nil {
-		return err
-	}
-	swap := "ALTER DATABASE " + branch + " RENAME TO postgres;"
-	if _, err := conn.Exec(ctx, swap); err != nil {
-		return err
-	}
 	return nil
 }
+
+func RestartDatabase(ctx context.Context) {
+	// Some extensions must be manually restarted after pg_terminate_backend
+	// Ref: https://github.com/citusdata/pg_cron/issues/99
+	if err := utils.Docker.ContainerRestart(ctx, utils.DbId, nil); err != nil {
+		fmt.Fprintln(os.Stderr, "Failed to restart database:", err)
+		return
+	}
+	// Poll for healthy database
+	now := time.Now()
+	expiry := now.Add(healthTimeout)
+	ticker := time.NewTicker(time.Second)
+	defer ticker.Stop()
+	for t := now; t.Before(expiry); t = <-ticker.C {
+		if resp, err := utils.Docker.ContainerInspect(ctx, utils.DbId); err == nil && resp.State.Health.Status == "healthy" {
+			// TODO: update storage-api to handle postgres restarts
+			if err := utils.Docker.ContainerRestart(ctx, utils.StorageId, nil); err != nil {
+				fmt.Fprintln(os.Stderr, "Failed to restart storage-api:", err)
+			}
+			return
+		}
+	}
+	fmt.Fprintln(os.Stderr, "Database is not healthy.")
+}
diff --git a/internal/start/start.go b/internal/start/start.go
index cbf56968..7dd5fd77 100644
--- a/internal/start/start.go
+++ b/internal/start/start.go
@@ -12,6 +12,7 @@ import (
 	"strconv"
 	"strings"
 	"text/template"
+	"time"
 
 	"github.com/charmbracelet/bubbles/progress"
 	"github.com/charmbracelet/bubbles/spinner"
@@ -22,6 +23,7 @@ import (
 	"github.com/docker/docker/pkg/stdcopy"
 	"github.com/docker/go-connections/nat"
 	"github.com/muesli/reflow/wrap"
+	"github.com/supabase/cli/internal/db/reset"
 	"github.com/supabase/cli/internal/utils"
 )
 
@@ -254,6 +256,12 @@ func run(p utils.Program, ctx context.Context) error {
 				Image: utils.GetRegistryImageUrl(utils.DbImage),
 				Env:   []string{"POSTGRES_PASSWORD=postgres"},
 				Cmd:   cmd,
+				Healthcheck: &container.HealthConfig{
+					Test:     []string{"CMD", "pg_isready", "-u", "postgres", "-h", "localhost", "-p", "5432"},
+					Interval: 2 * time.Second,
+					Timeout:  2 * time.Second,
+					Retries:  10,
+				},
 				Labels: map[string]string{
 					"com.supabase.cli.project":   utils.Config.ProjectId,
 					"com.docker.compose.project": utils.Config.ProjectId,
@@ -269,7 +277,7 @@ func run(p utils.Program, ctx context.Context) error {
 		}
 
 		out, err := utils.DockerExec(ctx, utils.DbId, []string{
-			"sh", "-c", "until pg_isready --host $(hostname --ip-address); do sleep 0.1; done " +
+			"sh", "-c", "until pg_isready --username postgres --host $(hostname --ip-address); do sleep 0.1; done " +
 				`&& psql --username postgres --host 127.0.0.1 <<'EOSQL'
 BEGIN;
 ` + utils.GlobalsSql + `
@@ -491,6 +499,7 @@ EOSQL
 			if err != nil {
 				return err
 			}
+			defer reset.RestartDatabase(context.Background())
 			var errBuf bytes.Buffer
 			if _, err := stdcopy.StdCopy(io.Discard, &errBuf, out); err != nil {
 				return err


Test patch:
diff --git a/internal/db/branch/switch_/switch__test.go b/internal/db/branch/switch_/switch__test.go
new file mode 100644
index 00000000..bac85716
--- /dev/null
+++ b/internal/db/branch/switch_/switch__test.go
@@ -0,0 +1,335 @@
+package switch_
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"path/filepath"
+	"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 TestSwitchCommand(t *testing.T) {
+	const version = "1.41"
+
+	t.Run("switches local branch", func(t *testing.T) {
+		// Setup in-memory fs
+		fsys := afero.NewMemMapFs()
+		require.NoError(t, utils.WriteConfig(fsys, false))
+		// Setup target branch
+		branch := "target"
+		branchPath := filepath.Join(filepath.Dir(utils.CurrBranchPath), branch)
+		require.NoError(t, fsys.Mkdir(branchPath, 0755))
+		require.NoError(t, afero.WriteFile(fsys, utils.CurrBranchPath, []byte("main"), 0644))
+		// 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.StatusOK).
+			JSON(types.ContainerJSON{})
+		gock.New("http:///var/run/docker.sock").
+			Post("/v" + version + "/containers").
+			Reply(http.StatusServiceUnavailable)
+		// 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("ALTER DATABASE postgres RENAME TO main;").
+			Reply("ALTER DATABASE").
+			Query("ALTER DATABASE " + branch + " RENAME TO postgres;").
+			Reply("ALTER DATABASE")
+		// Run test
+		assert.NoError(t, Run(context.Background(), branch, fsys, conn.Intercept))
+		// Validate output
+		assert.Empty(t, apitest.ListUnmatchedRequests())
+		contents, err := afero.ReadFile(fsys, utils.CurrBranchPath)
+		assert.NoError(t, err)
+		assert.Equal(t, []byte(branch), contents)
+	})
+
+	t.Run("throws error on missing config", func(t *testing.T) {
+		// Setup in-memory fs
+		fsys := afero.NewMemMapFs()
+		// Run test
+		err := Run(context.Background(), "target", fsys)
+		// Check error
+		assert.ErrorContains(t, err, "Have you set up the project with supabase init?")
+	})
+
+	t.Run("throws error on malformed config", func(t *testing.T) {
+		// Setup in-memory fs
+		fsys := afero.NewMemMapFs()
+		require.NoError(t, afero.WriteFile(fsys, utils.ConfigPath, []byte("malformed"), 0644))
+		// Run test
+		err := Run(context.Background(), "target", fsys)
+		// Check error
+		assert.ErrorContains(t, err, "Failed to read config: toml")
+	})
+
+	t.Run("throws error on missing database", 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
+		err := Run(context.Background(), "target", fsys)
+		// Check error
+		assert.ErrorContains(t, err, "supabase start is not running.")
+		assert.Empty(t, apitest.ListUnmatchedRequests())
+	})
+
+	t.Run("throws error on reserved branch", 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.StatusOK).
+			JSON(types.ContainerJSON{})
+		// Run test
+		err := Run(context.Background(), utils.ShadowDbName, fsys)
+		// Check error
+		assert.ErrorContains(t, err, "branch name is reserved.")
+		assert.Empty(t, apitest.ListUnmatchedRequests())
+	})
+
+	t.Run("throws error on missing branch", 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.StatusOK).
+			JSON(types.ContainerJSON{})
+		// Run test
+		err := Run(context.Background(), "main", fsys)
+		// Check error
+		assert.ErrorContains(t, err, "Branch main does not exist.")
+		assert.Empty(t, apitest.ListUnmatchedRequests())
+	})
+
+	t.Run("noop on current branch", 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.StatusOK).
+			JSON(types.ContainerJSON{})
+		// Setup target branch
+		branch := "main"
+		branchPath := filepath.Join(filepath.Dir(utils.CurrBranchPath), branch)
+		require.NoError(t, fsys.Mkdir(branchPath, 0755))
+		// Run test
+		assert.NoError(t, Run(context.Background(), branch, fsys))
+		// Check error
+		assert.Empty(t, apitest.ListUnmatchedRequests())
+		contents, err := afero.ReadFile(fsys, utils.CurrBranchPath)
+		assert.NoError(t, err)
+		assert.Equal(t, []byte(branch), contents)
+	})
+
+	t.Run("throws error on failure to switch", 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.StatusOK).
+			JSON(types.ContainerJSON{})
+		// Setup target branch
+		branch := "target"
+		branchPath := filepath.Join(filepath.Dir(utils.CurrBranchPath), branch)
+		require.NoError(t, fsys.Mkdir(branchPath, 0755))
+		// Setup mock postgres
+		conn := pgtest.NewConn()
+		// Run test
+		err := Run(context.Background(), branch, fsys, conn.Intercept)
+		// Check error
+		assert.ErrorContains(t, err, "Error switching to branch target: conn closed")
+		assert.Empty(t, apitest.ListUnmatchedRequests())
+	})
+
+	t.Run("throws error on failure to write", 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.StatusOK).
+			JSON(types.ContainerJSON{})
+		// Setup target branch
+		branch := "main"
+		branchPath := filepath.Join(filepath.Dir(utils.CurrBranchPath), branch)
+		require.NoError(t, fsys.Mkdir(branchPath, 0755))
+		// Run test
+		err := Run(context.Background(), branch, afero.NewReadOnlyFs(fsys))
+		// Check error
+		assert.ErrorContains(t, err, "Unable to update local branch file.")
+		assert.Empty(t, apitest.ListUnmatchedRequests())
+	})
+}
+
+func TestSwitchDatabase(t *testing.T) {
+	const version = "1.41"
+
+	t.Run("throws error on failure to connect", func(t *testing.T) {
+		// Setup invalid port
+		utils.Config.Db.Port = 0
+		// Run test
+		err := switchDatabase(context.Background(), "main", "target")
+		// Check error
+		assert.ErrorContains(t, err, "invalid port")
+	})
+
+	t.Run("throws error on failure to disconnect", func(t *testing.T) {
+		// Setup valid config
+		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
+		err := switchDatabase(context.Background(), "main", "target", conn.Intercept)
+		// Check error
+		assert.ErrorContains(t, err, pgerrcode.InvalidParameterValue)
+		assert.Empty(t, apitest.ListUnmatchedRequests())
+	})
+
+	t.Run("throws error on failure to backup", func(t *testing.T) {
+		// Setup valid config
+		utils.DbId = "test-switch"
+		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("ALTER DATABASE postgres RENAME TO main;").
+			ReplyError(pgerrcode.DuplicateDatabase, `database "main" already exists`)
+		// 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").
+			Post("/v" + version + "/containers/" + utils.DbId + "/restart").
+			Reply(http.StatusServiceUnavailable)
+		// Run test
+		err := switchDatabase(context.Background(), "main", "target", conn.Intercept)
+		// Check error
+		assert.ErrorContains(t, err, pgerrcode.DuplicateDatabase)
+		assert.Empty(t, apitest.ListUnmatchedRequests())
+	})
+
+	t.Run("throws error on failure to rename", func(t *testing.T) {
+		// Setup valid config
+		utils.DbId = "test-switch"
+		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("ALTER DATABASE postgres RENAME TO main;").
+			Reply("ALTER DATABASE").
+			Query("ALTER DATABASE target RENAME TO postgres;").
+			ReplyError(pgerrcode.InvalidCatalogName, `database "target" does not exist`).
+			// Attempt to rollback
+			Query("ALTER DATABASE main RENAME TO postgres;").
+			ReplyError(pgerrcode.DuplicateDatabase, `database "postgres" already exists`)
+		// 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").
+			Post("/v" + version + "/containers/" + utils.DbId + "/restart").
+			Reply(http.StatusServiceUnavailable)
+		// Run test
+		err := switchDatabase(context.Background(), "main", "target", conn.Intercept)
+		// Check error
+		assert.ErrorContains(t, err, pgerrcode.InvalidCatalogName)
+		assert.Empty(t, apitest.ListUnmatchedRequests())
+	})
+}
diff --git a/internal/db/reset/reset_test.go b/internal/db/reset/reset_test.go
index a5e05ee2..f8e4965b 100644
--- a/internal/db/reset/reset_test.go
+++ b/internal/db/reset/reset_test.go
@@ -7,6 +7,7 @@ import (
 	"net/http"
 	"os"
 	"testing"
+	"time"
 
 	"github.com/docker/docker/api/types"
 	"github.com/docker/docker/client"
@@ -212,3 +213,60 @@ func TestActivateDatabase(t *testing.T) {
 		assert.Error(t, ActivateDatabase(context.Background(), branch, conn.Intercept))
 	})
 }
+
+func TestRestartDatabase(t *testing.T) {
+	const version = "1.41"
+
+	t.Run("restarts storage api", func(t *testing.T) {
+		utils.DbId = "test-reset"
+		// 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").
+			Post("/v" + version + "/containers/" + utils.DbId + "/restart").
+			Reply(http.StatusOK)
+		gock.New("http:///var/run/docker.sock").
+			Get("/v" + version + "/containers/" + utils.DbId + "/json").
+			Reply(http.StatusOK).
+			JSON(types.ContainerJSON{ContainerJSONBase: &types.ContainerJSONBase{
+				State: &types.ContainerState{
+					Health: &types.Health{
+						Status: "healthy",
+					},
+				},
+			}})
+		utils.StorageId = "test-storage"
+		gock.New("http:///var/run/docker.sock").
+			Post("/v" + version + "/containers/" + utils.StorageId + "/restart").
+			Reply(http.StatusServiceUnavailable)
+		// Run test
+		RestartDatabase(context.Background())
+		// Check error
+		assert.Empty(t, apitest.ListUnmatchedRequests())
+	})
+
+	t.Run("timeout health check", func(t *testing.T) {
+		utils.DbId = "test-reset"
+		healthTimeout = 0 * time.Second
+		// 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").
+			Post("/v" + version + "/containers/" + utils.DbId + "/restart").
+			Reply(http.StatusOK)
+		// Run test
+		RestartDatabase(context.Background())
+		// Check error
+		assert.Empty(t, apitest.ListUnmatchedRequests())
+	})
+}