Repository: supabase/cli
Pull request #317 (2022-07-27 09:58:10)

test: db remote set command
## What kind of change does this PR introduce?

unit test, mocks were setup using actual pg requests / responses captured via a [MITM proxy](https://github.com/jackc/pgmock/blob/master/pgmockproxy/main.go)

## What is the current behavior?

Please link any relevant issues here.

## What is the new behavior?

- mocks postgres using `pgmock`
- refactors sql statements to const
- switch pgx to simple text protocol: [reasons](https://github.com/supabase/cli/pull/317/files#diff-6b4053a9760d66e9434239790cc06f9c93317040cf101c963d572df25c257a7bR50)
- eliminates [version check query](https://github.com/supabase/cli/pull/317/files#diff-6b4053a9760d66e9434239790cc06f9c93317040cf101c963d572df25c257a7bL39) by reading from pg startup message
-  fixes an [index out of bounds](https://github.com/supabase/cli/pull/317/files#diff-6b4053a9760d66e9434239790cc06f9c93317040cf101c963d572df25c257a7bL97) issue with regex substring match
- coverage: 89.7%

## Additional context

Add any other context or screenshots.


soedirgo: Thoughts about using [embedded Postgres](https://github.com/fergusstrange/embedded-postgres) instead of mocks? I haven't looked into the runtime costs though
sweatybridge: It can be a good solution for integration test which we currently lacks. For unit tests, I'd prefer mocks because it's easier to setup initial state and reproduce edge cases.
soedirgo: In absence of the concern over perf hit, my concern with using mocks for this is 1) coupling of test to implementation, and 2) fidelity.

E.g. if we were to change how we query the db in `set.go`, the unit test would fail, even if it would produce a correct result on a real db. We're testing the interaction here, when what we're really interested in is the behavior. Using embedded Postgres would be more like how we're using `afero` for faking FS.

Embedded Postgres is also a real db, which might uncover subtle things which might not occur on a stub (e.g. some query might work only on a specific Postgres version).

> forces me to understand deeper about postgres driver / protocol, rather than treating it as a blackbox

I think that's great! But maybe this is more suitable as a profiling approach rather than testing? I don't mind having some profiling-specific code (like logging stuff through `pgmockproxy`), we can even enable these once we implement `DEBUG=1`, but if we use it in tests I worry if we would "ossify" the codebase since it adds friction to refactoring
sweatybridge: > E.g. if we were to change how we query the db in set.go, the unit test would fail, even if it would produce a correct result on a real db.

I think we want to be explicit about invariants and post-conditions of postgres queries because they affect how our cli handles responses from database. There could be numerous sql statements that return the same result, but as long as the post-conditions are maintained, ie. `Reply(...)`, we don't need to make more changes to the test or application logic.

We need to update the test if we are changing the response handling logic completely. For eg. we may decide to use batched query instead. In that case, it's actually better for the test to fail because the post-conditions that our application is built upon are no longer valid. We should then find out what are the new post-conditions guaranteed by postgres for batched query. This will guide us in covering most edge cases in the application code.

> We're testing the interaction here, when what we're really interested in is the behavior.

We should find out the behaviour of postgres before writing any code (using psql or an equivalent tool) and make them explicit in tests. These are the post-conditions that are least likely to break because every client relies on them.

> Embedded Postgres is also a real db, which might uncover subtle things which might not occur on a stub (e.g. some query might work only on a specific Postgres version).

If we are interested in regression between multiple postgres versions, integration test is the way to go. We can run cli command end-to-end to validate the happy paths. These integration tests won't be exhaustive though. If it does uncover a regression, we can write a unit test for that specific version by mocking the reply.

> 1) coupling of test to implementation, and 2) fidelity

So overall, I'm aiming to cover most, if not all, code paths at the unit test level. I think the current level of fidelity is needed to ensure functional correctness of the cli program. Since we are interacting with an external system, it's ok for the unit test to support a new version of the external program.

---
On a separate note,

> Using embedded Postgres would be more like how we're using afero for faking FS.

While `afero` solves some of our FS testing needs, its `MemMapFs` implementation is not (and will never be) completely consistent with OS FS. For eg, we found out that file permissions are ignored and `WriteFile` doesn't error on missing directories. This makes it impossible to test those error paths. Ideally, I'd want to mock the `fs.FS` interface so that all error can be reproduced, and all problems with our cli code can be discovered. Unfortunately, that interface in stdlib is still a WIP.
soedirgo: > I think we want to be explicit about invariants and post-conditions of postgres queries

Right, but stubbing the responses and enforcing the order of interaction makes unit tests fail even when the invariant we're interested in still holds. E.g. say I want to change how we check whether or not `supabase_migrations.schema_migrations` table exists. Instead of doing:

```go
	if _, err := conn.Exec(ctx, CHECK_MIGRATION_EXISTS); err != nil {
		// do stuff
	}
```

which is not ideal since `err` can be caused by other things (like cannot connect to db), I want to change it to:

```go
	results, err := conn.Exec(ctx, `select exist (select from information_schema where schema_name = 'supabase_migrations' and table_name = 'schema_migrations');`);
	if err != nil {
		return err
	}
	// if results is a single row and the value is `true`, do stuff
	// do other stuff
```

(probably doesn't execute, iirc I need to `Query` instead of `Exec`, but that's the gist of it)

This would fail a unit test that asserts `creates schema_migrations if it doesn't exist`, and also other tests that rely on the exact order of queries. But the invariant still holds: it would create the `schema_migrations` table all the same.

If we use a real implementation instead, the unit tests become focused: as long as the `schema_migrations` table exists at the end, the test will pass. We test *exactly* what we want to test, and no more.

> it's actually better for the test to fail because the post-conditions that our application is built upon are no longer valid

> We should find out the behaviour of postgres before writing any code

Can you clarify what these post-conditions are? In my mind the only post-conditions I'm interested in is the state in the db, not the order of db queries/responses. The argument for using a real implementation is we don't need to hard code Postgres's behavior - we're already using Postgres.

> If it does uncover a regression, we can write a unit test for that specific version by mocking the reply

This one I do agree is necessary for some cases, e.g. simulating very slow db responses - it's infeasible to reliably reproduce this using a real implementation.

I'm quite bad at explaining myself clearly 😅 but the angle I'm arguing from is basically this: https://abseil.io/resources/swe-book/html/ch13.html#real_implementations.
soedirgo: So I tried running some example tests using `embedded-postgres` and... it's dog slow 😂

```
❯ go test -run Test_GooseMigrations ./...
ok      github.com/fergusstrange/embedded-postgres/examples     4.923s
```

(running all 5 tests takes 24s)

We can pull some tricks like starting the db only once, and doing `CREATE DATABASE` for each unit test, but there's a permanent 5s runtime hit.


sweatybridge: > stubbing the responses and enforcing the order of interaction makes unit tests fail even when the invariant we're interested in still holds

Let me address two separate issues here, stubbing responses and enforcing order.

If the cli needs to issue complex joins with many tables, views, constraints, etc., then I would probably agree to use a real embedded postgres. It makes sense to test the real implementation because the interactions are complex. But in our case, we simply query 1 migration table with 1 version column. Most of our application logic is actually dealing with error from external processes. I believe this falls under an appropriate use case for stubbing:

> stubbing is appropriate when you need a function to return a specific value to get the system under test into a certain state
> Because a function’s behavior is defined inline in the test, stubbing can simulate a wide variety of return values or errors that might not be possible to trigger from a real implementation or a fake.
> A test that requires many functions to be stubbed can be a sign that stubbing is being overused, or that the system under test is too complex and should be refactored.
https://abseil.io/resources/swe-book/html/ch13.html#example_onethree-onefourdot_refactoring

Also notice that the mock class works in a similar way as [gock](https://github.com/h2non/gock) library for http protocol. It's different from a [sql _driver_ mock](https://github.com/DATA-DOG/go-sqlmock) where `sql.Open`, `sql.Query`, `sql.Exec`, etc. needs to be stubbed out individually. I didn't choose `go-sqlmock` because I find its assertions are too fine grained and we don't actually need to support other SQL databases.

Because we mock at the [postgres text protocol](https://www.postgresql.org/docs/current/protocol-flow.html#id-1.10.5.7.4) level, both `pgx.Query` and `pgx.Exec` are serialised to the same "Query" message. This makes our mock class easier to use and less brittle when users call different pgx functions.

> In my mind the only post-conditions I'm interested in is the state in the db, not the order of db queries/responses.

Just to be clear, "enforcing" the execution order is kind of an unintended side effect because it's easier to implement the mock steps in an array. We can generalize with a `map` so it becomes a bag of logical implications. For eg. given A -> B and C -> D, is the program correct? Our query messages are A and C, post-conditions are reply messages B and D, invariants are everything external to the cli process, including db state.

For example, in the original query `_, err := conn.Exec(ctx, CHECK_MIGRATION_EXISTS)`, one possible post-condition is that postgres returns success with 0 rows, indicating that the table exists. We express this post-condition as `Reply("SELECT 0")`.

Another possible post-condition is that postgres replies "ErrorResponse" message of the text protocol, indicating that the table does not exist. We express this post-condition as `ReplyError(code, message)`. Unfortunately pgx groups this error together with unrelated network issues in the return type.

> This would fail a unit test that asserts creates schema_migrations if it doesn't exist, and also other tests that rely on the exact order of queries. But the invariant still holds: it would create the schema_migrations table all the same.
> If we use a real implementation instead, the unit tests become focused: as long as the schema_migrations table exists at the end, the test will pass. We test exactly what we want to test, and no more.

Sounds like the main argument here is that mocked test is brittle. A good rule of thumb is to reduce the number of post-conditions, ie. mock method calls, per unit test. One possible way is to refactor `db remote set` into smaller functions so that each test is more focused.
soedirgo: > If the cli needs to issue complex joins with many tables, views, constraints, etc., then I would probably agree to use a real embedded postgres. It makes sense to test the real implementation because the interactions are complex.

I'd actually argue the inverse - it's noted in the link that real implementations (or failing that, fakes) are preferred if it's fast, deterministic, and has a simple schema (i.e. the dependency). 5s is not fast, though if it happens once I wouldn't say it's a deal breaker. Given how we use the db, it should be deterministic. The schema as you say is simple - just one table, can be set up in 1 query.

> Most of our application logic is actually dealing with error from external processes

This is where I agree mocks would be necessary, e.g. to simulate db connection errors, slow db responses, nondeterministic queries etc. But for most things, e.g. checking if `schema_migrations` is created, that'd be the ideal place to use a real db.

> Also notice that the mock class works in a similar way as [gock](https://github.com/h2non/gock) library for http protocol

The use of gock I think is pretty justified - we can't use the real endpoint (it may 500 and we'd be generating unnecessary load), there's no fake available (I still prefer we create one, but maybe at a later time). But here we have the option to use a real db.

> It's different from a [sql driver mock](https://github.com/DATA-DOG/go-sqlmock)

Yup, agree that the mock struct is better than sqlmock here - though a real db wouldn't have these problems either. And there *might* be actual nuances of `Exec` vs. `Query` that wouldn't be caught if we treat them the same: if the difference manifests when run on a real db, that doesn't mean that the test is brittle, because you *want* it to fail in that case.

> We can generalize with a map so it becomes a bag of logical implications.

If we're only doing reads I think that's fine (like most of our gock use cases), but here we're interested specifically in how the state stored in the db changes. This is demonstrated in [this section](https://abseil.io/resources/swe-book/html/ch13.html#tests_become_less_effective).

> Unfortunately pgx groups this error together with unrelated network issues in the return type

Since we don't need to behave differently between "table doesn't exist" vs. "network error", I think this is fine - in both cases we just need to assert that there's an error. Otherwise, "table doesn't exist" can easily simulated on a real db, "network error" might need a mock (unless it'd suffice to use an invalid db url).

> One possible way is to refactor db remote set into smaller functions so that each test is more focused

We can also refactor the test to avoid stubbing, demonstrated [here](https://abseil.io/resources/swe-book/html/ch13.html#example_onethree-onefourdot_refactoring).

Let me try and set up a wrapper around embedded-postgres based off of this branch - if there are any deal breakers I think we can put this convo behind.
soedirgo: Given that this command will be deprecated, will just merge this in for now - we can try the embedded Postgres approach on other commands that interact with the db (`db remote commit`, `db push`, etc.)

Patch:
diff --git a/go.mod b/go.mod
index fed65af9..b34639c6 100644
--- a/go.mod
+++ b/go.mod
@@ -59,11 +59,13 @@ require (
 	github.com/inconshreveable/mousetrap v1.0.0 // indirect
 	github.com/jackc/chunkreader/v2 v2.0.1 // indirect
 	github.com/jackc/pgconn v1.11.0 // indirect
+	github.com/jackc/pgerrcode v0.0.0-20220416144525-469b46aa5efa
 	github.com/jackc/pgio v1.0.0 // indirect
+	github.com/jackc/pgmock v0.0.0-20210724152146-4ad1a8207f65
 	github.com/jackc/pgpassfile v1.0.0 // indirect
-	github.com/jackc/pgproto3/v2 v2.2.0 // indirect
+	github.com/jackc/pgproto3/v2 v2.2.0
 	github.com/jackc/pgservicefile v0.0.0-20200714003250-2b9c44734f2b // indirect
-	github.com/jackc/pgtype v1.10.0 // indirect
+	github.com/jackc/pgtype v1.10.0
 	github.com/lucasb-eyer/go-colorful v1.2.0 // indirect
 	github.com/magiconair/properties v1.8.6 // indirect
 	github.com/mattn/go-runewidth v0.0.13 // indirect
@@ -94,7 +96,7 @@ require (
 	golang.org/x/term v0.0.0-20210927222741-03fcf44c2211 // indirect
 	golang.org/x/text v0.3.7 // indirect
 	google.golang.org/genproto v0.0.0-20220323144105-ec3c684e5b14 // indirect
-	google.golang.org/grpc v1.45.0 // indirect
+	google.golang.org/grpc v1.45.0
 	google.golang.org/protobuf v1.28.0 // indirect
 	gopkg.in/h2non/gock.v1 v1.1.2
 	gopkg.in/ini.v1 v1.66.4 // indirect
diff --git a/go.sum b/go.sum
index 56b2f848..ec666f68 100644
--- a/go.sum
+++ b/go.sum
@@ -224,6 +224,8 @@ github.com/jackc/pgconn v1.9.0/go.mod h1:YctiPyvzfU11JFxoXokUOOKQXQmDMoJL9vJzHH8
 github.com/jackc/pgconn v1.9.1-0.20210724152538-d89c8390a530/go.mod h1:4z2w8XhRbP1hYxkpTuBjTS3ne3J48K83+u0zoyvg2pI=
 github.com/jackc/pgconn v1.11.0 h1:HiHArx4yFbwl91X3qqIHtUFoiIfLNJXCQRsnzkiwwaQ=
 github.com/jackc/pgconn v1.11.0/go.mod h1:4z2w8XhRbP1hYxkpTuBjTS3ne3J48K83+u0zoyvg2pI=
+github.com/jackc/pgerrcode v0.0.0-20220416144525-469b46aa5efa h1:s+4MhCQ6YrzisK6hFJUX53drDT4UsSW3DEhKn0ifuHw=
+github.com/jackc/pgerrcode v0.0.0-20220416144525-469b46aa5efa/go.mod h1:a/s9Lp5W7n/DD0VrVoyJ00FbP2ytTPDVOivvn2bMlds=
 github.com/jackc/pgio v1.0.0 h1:g12B9UwVnzGhueNavwioyEEpAmqMe1E/BN9ES+8ovkE=
 github.com/jackc/pgio v1.0.0/go.mod h1:oP+2QK2wFfUWgr+gxjoBH9KGBb31Eio69xUb0w5bYf8=
 github.com/jackc/pgmock v0.0.0-20190831213851-13a1b77aafa2/go.mod h1:fGZlG77KXmcq05nJLRkk0+p82V8B8Dw8KN2/V9c/OAE=
diff --git a/internal/db/remote/set/set.go b/internal/db/remote/set/set.go
index 787b94fd..dcb61ab3 100644
--- a/internal/db/remote/set/set.go
+++ b/internal/db/remote/set/set.go
@@ -6,13 +6,29 @@ import (
 	"fmt"
 	u "net/url"
 	"path/filepath"
+	"strconv"
 
-	pgx "github.com/jackc/pgx/v4"
+	"github.com/jackc/pgx/v4"
 	"github.com/spf13/afero"
 	"github.com/supabase/cli/internal/utils"
 )
 
-func Run(url string, fsys afero.Fs) error {
+const (
+	CHECK_MIGRATION_EXISTS = "SELECT 1 FROM supabase_migrations.schema_migrations LIMIT 1"
+	LIST_MIGRATION_VERSION = "SELECT version FROM supabase_migrations.schema_migrations ORDER BY version"
+	CREATE_MIGRATION_TABLE = `CREATE SCHEMA IF NOT EXISTS supabase_migrations;
+CREATE TABLE supabase_migrations.schema_migrations (version text NOT NULL PRIMARY KEY);
+`
+)
+
+func min(a, b int) int {
+	if a < b {
+		return a
+	}
+	return b
+}
+
+func Run(url string, fsys afero.Fs, options ...func(*pgx.ConnConfig)) error {
 	// Sanity checks.
 	{
 		if err := utils.LoadConfigFS(fsys); err != nil {
@@ -26,20 +42,39 @@ func Run(url string, fsys afero.Fs) error {
 	}
 	url = "postgresql://postgres:" + u.QueryEscape(matches[1]) + "@" + u.QueryEscape(matches[2]) + matches[3] + "/postgres"
 
-	ctx := context.Background()
+	config, err := pgx.ParseConfig(url)
+	if err != nil {
+		return err
+	}
+
+	// Simple protocol is preferred over pgx default Parse -> Bind flow because
+	//   1. Using a single command for each query reduces RTT over an Internet connection.
+	//   2. Performance gains from using the alternate binary protocol is negligible because
+	//      we are only selecting from migrations table. Large reads are handled by PostgREST.
+	//   3. Any prepared statements are cleared server side upon closing the TCP connection.
+	//      Since CLI workloads are one-off scripts, we don't use connection pooling and hence
+	//      don't benefit from per connection server side cache.
+	config.PreferSimpleProtocol = true
+	for _, op := range options {
+		op(config)
+	}
 
-	conn, err := pgx.Connect(ctx, url)
+	ctx := context.Background()
+	conn, err := pgx.ConnectConfig(ctx, config)
 	if err != nil {
 		return err
 	}
-	defer conn.Close(context.Background())
+	defer conn.Close(ctx)
 
 	// 1. Assert db.major_version is compatible.
-	var dbMajorVersion uint
-	if err := conn.QueryRow(ctx, "SELECT current_setting('server_version_num')::int8 / 10000").Scan(&dbMajorVersion); err != nil {
+	serverVersion := conn.PgConn().ParameterStatus("server_version")
+	// Safe to assume that supported Postgres version is 10.0 <= n < 100.0
+	dbMajorVersion, err := strconv.ParseUint(serverVersion[:min(2, len(serverVersion))], 10, 7)
+	if err != nil {
 		return err
 	}
-	if dbMajorVersion != utils.Config.Db.MajorVersion {
+
+	if dbMajorVersion != uint64(utils.Config.Db.MajorVersion) {
 		return fmt.Errorf(
 			"Remote database Postgres version %[1]d is incompatible with %[3]s %[2]d. If you are setting up a fresh Supabase CLI project, try changing %[3]s in %[4]s to %[1]d.",
 			dbMajorVersion,
@@ -52,20 +87,15 @@ func Run(url string, fsys afero.Fs) error {
 	// 2. Setup & validate `schema_migrations`.
 
 	// If `schema_migrations` doesn't exist on the remote database, create it.
-	if _, err := conn.Exec(ctx, "SELECT 1 FROM supabase_migrations.schema_migrations"); err != nil {
-		if _, err := conn.Exec(
-			ctx,
-			`CREATE SCHEMA IF NOT EXISTS supabase_migrations;
-CREATE TABLE supabase_migrations.schema_migrations (version text NOT NULL PRIMARY KEY);
-`,
-		); err != nil {
+	if _, err := conn.Exec(ctx, CHECK_MIGRATION_EXISTS); err != nil {
+		if _, err := conn.Exec(ctx, CREATE_MIGRATION_TABLE); err != nil {
 			return err
 		}
 	}
 
 	// If `schema_migrations` is not a "prefix" of list of migrations in repo, fail &
 	// warn user.
-	rows, err := conn.Query(ctx, "SELECT version FROM supabase_migrations.schema_migrations ORDER BY version")
+	rows, err := conn.Query(ctx, LIST_MIGRATION_VERSION)
 	if err != nil {
 		return err
 	} else {
@@ -78,7 +108,7 @@ CREATE TABLE supabase_migrations.schema_migrations (version text NOT NULL PRIMAR
 			remoteMigrations = append(remoteMigrations, version)
 		}
 
-		if err := utils.MkdirIfNotExist("supabase/migrations"); err != nil {
+		if err := utils.MkdirIfNotExistFS(fsys, "supabase/migrations"); err != nil {
 			return err
 		}
 		localMigrations, err := afero.ReadDir(fsys, "supabase/migrations")
@@ -94,8 +124,8 @@ CREATE TABLE supabase_migrations.schema_migrations (version text NOT NULL PRIMAR
 3. Run ` + utils.Aqua("supabase db remote commit") + ".")
 			}
 
-			localTimestamp := utils.MigrateFilePattern.FindStringSubmatch(localMigrations[i].Name())[1]
-			if localTimestamp == remoteTimestamp {
+			localTimestamp := utils.MigrateFilePattern.FindStringSubmatch(localMigrations[i].Name())
+			if len(localTimestamp) > 1 && localTimestamp[1] == remoteTimestamp {
 				continue
 			}
 


Test patch:
diff --git a/internal/db/remote/set/set_test.go b/internal/db/remote/set/set_test.go
new file mode 100644
index 00000000..a7755cb2
--- /dev/null
+++ b/internal/db/remote/set/set_test.go
@@ -0,0 +1,197 @@
+package set
+
+import (
+	"testing"
+
+	"github.com/jackc/pgerrcode"
+	"github.com/spf13/afero"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/supabase/cli/internal/testing/pgtest"
+	"github.com/supabase/cli/internal/utils"
+)
+
+func TestDbRemoteSetCommand(t *testing.T) {
+	const postgresUrl = "postgresql://postgres:password@localhost:5432/postgres"
+
+	t.Run("sets the remote database url", func(t *testing.T) {
+		// Setup in-memory fs
+		fsys := afero.NewMemMapFs()
+		require.NoError(t, utils.WriteConfig(fsys, true))
+		// Setup initial migration
+		version := "20220727064247"
+		_, err := fsys.Create("supabase/migrations/" + version + "_init.sql")
+		require.NoError(t, err)
+		// Setup mock postgres
+		conn := pgtest.NewConn()
+		defer conn.Close()
+		conn.Query(CHECK_MIGRATION_EXISTS).
+			Reply("SELECT 0").
+			Query(LIST_MIGRATION_VERSION).
+			Reply("SELECT 1", map[string]interface{}{"version": version})
+		// Run test
+		assert.NoError(t, Run(postgresUrl, fsys, conn.Intercept))
+		assert.NoError(t, <-conn.ErrChan)
+	})
+
+	t.Run("creates migrations table if absent", func(t *testing.T) {
+		// Setup in-memory fs
+		fsys := afero.NewMemMapFs()
+		require.NoError(t, utils.WriteConfig(fsys, true))
+		// Setup mock postgres
+		conn := pgtest.NewConn()
+		defer conn.Close()
+		conn.Query(CHECK_MIGRATION_EXISTS).
+			ReplyError(pgerrcode.UndefinedTable, "relation \"supabase_migrations.schema_migrations\" does not exist").
+			Query(CREATE_MIGRATION_TABLE).
+			Reply("CREATE TABLE").
+			Query(LIST_MIGRATION_VERSION).
+			Reply("SELECT 0")
+		// Run test
+		assert.NoError(t, Run(postgresUrl, fsys, conn.Intercept))
+		assert.NoError(t, <-conn.ErrChan)
+		// Validate file contents
+		content, err := afero.ReadFile(fsys, utils.RemoteDbPath)
+		assert.NoError(t, err)
+		assert.Equal(t, []byte(postgresUrl), content)
+	})
+
+	t.Run("throws error on missing config file", func(t *testing.T) {
+		assert.Error(t, Run("", afero.NewMemMapFs()))
+	})
+
+	t.Run("throws error on invalid postgres url", func(t *testing.T) {
+		// Setup in-memory fs
+		fsys := afero.NewMemMapFs()
+		require.NoError(t, utils.WriteConfig(fsys, true))
+		// Run test
+		assert.Error(t, Run("invalid", fsys))
+	})
+
+	t.Run("throws error on failture to connect", func(t *testing.T) {
+		// Setup in-memory fs
+		fsys := afero.NewMemMapFs()
+		require.NoError(t, utils.WriteConfig(fsys, true))
+		// Setup mock postgres
+		conn := pgtest.NewWithStatus(map[string]string{})
+		defer conn.Close()
+		// Run test
+		assert.Error(t, Run(postgresUrl, fsys, conn.Intercept))
+		assert.NoError(t, <-conn.ErrChan)
+	})
+
+	t.Run("throws error on missing server version", func(t *testing.T) {
+		// Setup in-memory fs
+		fsys := afero.NewMemMapFs()
+		require.NoError(t, utils.WriteConfig(fsys, true))
+		// Setup mock postgres
+		conn := pgtest.NewWithStatus(map[string]string{
+			"standard_conforming_strings": "on",
+		})
+		defer conn.Close()
+		// Run test
+		assert.Error(t, Run(postgresUrl, fsys, conn.Intercept))
+		assert.NoError(t, <-conn.ErrChan)
+	})
+
+	t.Run("throws error on unsupported server version", func(t *testing.T) {
+		// Setup in-memory fs
+		fsys := afero.NewMemMapFs()
+		require.NoError(t, utils.WriteConfig(fsys, true))
+		// Setup mock postgres
+		conn := pgtest.NewWithStatus(map[string]string{
+			"standard_conforming_strings": "on",
+			"server_version":              "13.1",
+		})
+		defer conn.Close()
+		// Run test
+		assert.Error(t, Run(postgresUrl, fsys, conn.Intercept))
+		assert.NoError(t, <-conn.ErrChan)
+	})
+
+	t.Run("throws error on failure to create table", func(t *testing.T) {
+		// Setup in-memory fs
+		fsys := afero.NewMemMapFs()
+		require.NoError(t, utils.WriteConfig(fsys, true))
+		// Setup mock postgres
+		conn := pgtest.NewConn()
+		defer conn.Close()
+		conn.Query(CHECK_MIGRATION_EXISTS).
+			ReplyError(pgerrcode.UndefinedTable, "relation \"supabase_migrations.schema_migrations\" does not exist").
+			Query(CREATE_MIGRATION_TABLE).
+			ReplyError(pgerrcode.DuplicateTable, "relation \"schema_migrations\" already exists")
+		// Run test
+		assert.Error(t, Run(postgresUrl, fsys, conn.Intercept))
+		assert.NoError(t, <-conn.ErrChan)
+	})
+
+	t.Run("throws error on failure to list migrations", func(t *testing.T) {
+		// Setup in-memory fs
+		fsys := afero.NewMemMapFs()
+		require.NoError(t, utils.WriteConfig(fsys, true))
+		// Setup mock postgres
+		conn := pgtest.NewConn()
+		defer conn.Close()
+		conn.Query(CHECK_MIGRATION_EXISTS).
+			Reply("SELECT 0").
+			Query(LIST_MIGRATION_VERSION).
+			ReplyError(pgerrcode.UndefinedTable, "relation \"supabase_migrations.schema_migrations\" does not exist")
+		// Run test
+		assert.Error(t, Run(postgresUrl, fsys, conn.Intercept))
+		assert.NoError(t, <-conn.ErrChan)
+	})
+
+	t.Run("throws error on migration mismatch", func(t *testing.T) {
+		// Setup in-memory fs
+		fsys := afero.NewMemMapFs()
+		require.NoError(t, utils.WriteConfig(fsys, true))
+		// Setup mock postgres
+		conn := pgtest.NewConn()
+		defer conn.Close()
+		conn.Query(CHECK_MIGRATION_EXISTS).
+			Reply("SELECT 0").
+			Query(LIST_MIGRATION_VERSION).
+			Reply("SELECT 1", map[string]interface{}{"version": "20220727064247"})
+		// Run test
+		assert.Error(t, Run(postgresUrl, fsys, conn.Intercept))
+		assert.NoError(t, <-conn.ErrChan)
+	})
+
+	t.Run("throws error on malformed file name", func(t *testing.T) {
+		// Setup in-memory fs
+		fsys := afero.NewMemMapFs()
+		require.NoError(t, utils.WriteConfig(fsys, true))
+		// Setup initial migration
+		version := "20220727064247"
+		_, err := fsys.Create("supabase/migrations/" + version + ".sql")
+		require.NoError(t, err)
+		// Setup mock postgres
+		conn := pgtest.NewConn()
+		defer conn.Close()
+		conn.Query(CHECK_MIGRATION_EXISTS).
+			Reply("SELECT 0").
+			Query(LIST_MIGRATION_VERSION).
+			Reply("SELECT 1", map[string]interface{}{"version": "20220727064247"})
+		// Run test
+		assert.Error(t, Run(postgresUrl, fsys, conn.Intercept))
+		assert.NoError(t, <-conn.ErrChan)
+	})
+
+	t.Run("throws error on failure to create directory", func(t *testing.T) {
+		// Setup in-memory fs
+		fsys := afero.NewMemMapFs()
+		require.NoError(t, utils.WriteConfig(fsys, true))
+		// Setup mock postgres
+		conn := pgtest.NewConn()
+		defer conn.Close()
+		conn.Query(CHECK_MIGRATION_EXISTS).
+			ReplyError(pgerrcode.UndefinedTable, "relation \"supabase_migrations.schema_migrations\" does not exist").
+			Query(CREATE_MIGRATION_TABLE).
+			Reply("CREATE TABLE").
+			Query(LIST_MIGRATION_VERSION).
+			Reply("SELECT 0")
+		// Run test
+		assert.Error(t, Run(postgresUrl, afero.NewReadOnlyFs(fsys), conn.Intercept))
+		assert.NoError(t, <-conn.ErrChan)
+	})
+}
diff --git a/internal/testing/pgtest/mock.go b/internal/testing/pgtest/mock.go
new file mode 100644
index 00000000..ec82a68a
--- /dev/null
+++ b/internal/testing/pgtest/mock.go
@@ -0,0 +1,220 @@
+package pgtest
+
+import (
+	"context"
+	"net"
+	"reflect"
+	"time"
+
+	"github.com/jackc/pgmock"
+	"github.com/jackc/pgproto3/v2"
+	"github.com/jackc/pgtype"
+	"github.com/jackc/pgx/v4"
+	"google.golang.org/grpc/test/bufconn"
+)
+
+var ci = pgtype.NewConnInfo()
+
+type MockConn struct {
+	// Duplex server listener backed by in-memory buffer
+	server *bufconn.Listener
+
+	// Mock server requests and responses
+	script pgmock.Script
+
+	// Status parameters emitted by postgres on first connect
+	status map[string]string
+
+	// Channel for reporting all server error
+	ErrChan chan error
+}
+
+func (r *MockConn) getStartupMessage(config *pgx.ConnConfig) []pgmock.Step {
+	var steps []pgmock.Step
+	// Add auth message
+	steps = append(
+		steps,
+		pgmock.ExpectMessage(&pgproto3.StartupMessage{
+			ProtocolVersion: pgproto3.ProtocolVersionNumber,
+			Parameters:      map[string]string{"database": config.Database, "user": config.User},
+		}),
+		pgmock.SendMessage(&pgproto3.AuthenticationOk{}),
+	)
+	// Add status message
+	r.status["session_authorization"] = config.User
+	for key, value := range r.status {
+		steps = append(steps, pgmock.SendMessage(&pgproto3.ParameterStatus{Name: key, Value: value}))
+	}
+	// Add ready message
+	steps = append(
+		steps,
+		pgmock.SendMessage(&pgproto3.BackendKeyData{ProcessID: 0, SecretKey: 0}),
+		pgmock.SendMessage(&pgproto3.ReadyForQuery{TxStatus: 'I'}),
+	)
+	return steps
+}
+
+// Configures pgx to use the mock dialer.
+//
+// The mock dialer provides a full duplex net.Conn backed by an in-memory buffer.
+// It is implemented by grcp/test/bufconn package.
+func (r *MockConn) Intercept(config *pgx.ConnConfig) {
+	// TODO: check for config.PreferSimpleProtocol
+	config.DialFunc = func(ctx context.Context, network, addr string) (net.Conn, error) {
+		return r.server.DialContext(ctx)
+	}
+	config.TLSConfig = nil
+	// Add startup message
+	r.script.Steps = append(r.getStartupMessage(config), r.script.Steps...)
+}
+
+// Adds a simple query to the mock connection.
+//
+// TODO: support prepared statements that involve multiple round trips, ie. Parse -> Bind.
+func (r *MockConn) Query(psql string) *MockConn {
+	r.script.Steps = append(r.script.Steps, pgmock.ExpectMessage(&pgproto3.Query{String: psql}))
+	return r
+}
+
+func getDataTypeSize(v interface{}) int16 {
+	t := reflect.TypeOf(v)
+	k := t.Kind()
+	if k < reflect.Int || k > reflect.Complex128 {
+		return -1
+	}
+	return int16(t.Size())
+}
+
+// Adds a server reply using text protocol format.
+//
+// TODO: support binary protocol
+func (r *MockConn) Reply(tag string, rows ...map[string]interface{}) *MockConn {
+	// Add field description
+	if len(rows) > 0 {
+		var desc pgproto3.RowDescription
+		for k, v := range rows[0] {
+			if dt, ok := ci.DataTypeForValue(v); ok {
+				size := getDataTypeSize(v)
+				desc.Fields = append(desc.Fields, pgproto3.FieldDescription{
+					Name:                 []byte(k),
+					TableOID:             17131,
+					TableAttributeNumber: 1,
+					DataTypeOID:          dt.OID,
+					DataTypeSize:         size,
+					TypeModifier:         -1,
+					Format:               pgtype.TextFormatCode,
+				})
+			}
+		}
+		r.script.Steps = append(r.script.Steps, pgmock.SendMessage(&desc))
+	}
+	// Note: Postgres emits field descriptions even if no rows are returned. However,
+	// pgx does not care about it so we do not need to handle the else case.
+
+	// Add row data
+	for _, data := range rows {
+		var dr pgproto3.DataRow
+		for _, v := range data {
+			if dt, ok := ci.DataTypeForValue(v); ok {
+				if err := dt.Value.Set(v); err != nil {
+					continue
+				}
+				if value, err := (dt.Value).(pgtype.TextEncoder).EncodeText(ci, []byte{}); err == nil {
+					dr.Values = append(dr.Values, value)
+				}
+			}
+		}
+		r.script.Steps = append(r.script.Steps, pgmock.SendMessage(&dr))
+	}
+
+	// Add completion message
+	// Subquery emits additional completion messages but pgx doesn't seem to care. For eg.
+	//   Reply("CREATE SCHEMA").
+	//   pgmock.SendMessage(&pgproto3.NoticeResponse{
+	// 	     Severity:            "NOTICE",
+	// 	     SeverityUnlocalized: "NOTICE",
+	// 	     Code:                "42P06",
+	// 	     Message:             "schema \"supabase_migrations\" already exists, skipping",
+	//   }),
+	//   pgmock.SendMessage(&pgproto3.CommandComplete{CommandTag: []byte("CREATE SCHEMA")}),
+	r.script.Steps = append(
+		r.script.Steps,
+		pgmock.SendMessage(&pgproto3.CommandComplete{CommandTag: []byte(tag)}),
+		pgmock.SendMessage(&pgproto3.ReadyForQuery{TxStatus: 'I'}),
+	)
+	return r
+}
+
+// Simulates an error reply from the server.
+//
+// TODO: simulate a notice reply
+func (r *MockConn) ReplyError(code, message string) *MockConn {
+	r.script.Steps = append(
+		r.script.Steps,
+		pgmock.SendMessage(&pgproto3.ErrorResponse{
+			Severity:            "ERROR",
+			SeverityUnlocalized: "ERROR",
+			Code:                code,
+			Message:             message,
+		}),
+		pgmock.SendMessage(&pgproto3.ReadyForQuery{TxStatus: 'I'}),
+	)
+	return r
+}
+
+func (r *MockConn) Close() {
+	_ = r.server.Close()
+}
+
+func NewWithStatus(status map[string]string) *MockConn {
+	const bufSize = 1024 * 1024
+	mock := MockConn{
+		server:  bufconn.Listen(bufSize),
+		status:  status,
+		ErrChan: make(chan error, 1),
+	}
+	// Start server in background
+	const timeout = time.Millisecond * 450
+	go func() {
+		defer close(mock.ErrChan)
+		// Block until we've opened a TCP connection
+		conn, err := mock.server.Accept()
+		if err != nil {
+			mock.ErrChan <- err
+			return
+		}
+		defer conn.Close()
+		// Prevent server from hanging the test
+		err = conn.SetDeadline(time.Now().Add(timeout))
+		if err != nil {
+			mock.ErrChan <- err
+			return
+		}
+		// Always expect clients to terminate the request
+		mock.script.Steps = append(mock.script.Steps, pgmock.ExpectMessage(&pgproto3.Terminate{}))
+		err = mock.script.Run(pgproto3.NewBackend(pgproto3.NewChunkReader(conn), conn))
+		if err != nil {
+			mock.ErrChan <- err
+			return
+		}
+	}()
+
+	return &mock
+}
+
+func NewConn() *MockConn {
+	return NewWithStatus(map[string]string{
+		"application_name":              "",
+		"client_encoding":               "UTF8",
+		"DateStyle":                     "ISO, MDY",
+		"default_transaction_read_only": "off",
+		"in_hot_standby":                "off",
+		"integer_datetimes":             "on",
+		"IntervalStyle":                 "postgres",
+		"is_superuser":                  "on",
+		"server_encoding":               "UTF8",
+		"server_version":                "14.3 (Debian 14.3-1.pgdg110+1)",
+		"standard_conforming_strings":   "on",
+		"TimeZone":                      "UTC",
+	})
+}