Repository: supabase/cli
Pull request #398 (2022-08-22 07:59:13)

feat: add db lint command
## What kind of change does this PR introduce?

feature #298 

## What is the current behavior?

no linting

## What is the new behavior?

- adds `supabase db lint` with `--schema` flag
- uses [plpgsql_check](https://github.com/okbob/plpgsql_check)  extension
- uses prepared statement to support multiple schema

TODO:
- [x] format output (json for now)
- [x] add unit tests
- [x] separate warning vs error
- [ ] create ignore file for different types of error (comment might be hard)
- [x] custom lint check (might be able to use pg_tap)

use case
- run check on anything that touches supabase schema
- useful to have the line of code where error occurs (body of function)
- want to check columns are created, ie. created_by, modified_by, etc

## Additional context

Add `f1()` example as a migration script, then run

```
$ supabase db lint -s public,private
Linting schema: public
Linting schema: private
[
  {
    "function": "public.f1",
    "issues": [
      {
        "level": "error",
        "message": "relation \"t1\" does not exist",
        "statement": {
          "lineNumber": "4",
          "text": "FOR over SELECT rows"
        },
        "query": {
          "position": "15",
          "text": "SELECT * FROM t1"
        },
        "sqlState": "42P01"
      }
    ]
  }
]
```


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

* **9** of **74**   **(12.16%)**  changed or added relevant lines in **2** files are covered.
* No unchanged relevant lines lost coverage.
* Overall coverage decreased (**-0.3%**) to **29.305%**

---

|  Changes Missing Coverage | Covered Lines | Changed/Added Lines | % |
| :-----|--------------|--------|---: |
| [cmd/db.go](https://coveralls.io/builds/51828958/source?filename=cmd%2Fdb.go#L144) | 9 | 12 | 75.0%
| [internal/db/lint/lint.go](https://coveralls.io/builds/51828958/source?filename=internal%2Fdb%2Flint%2Flint.go#L21) | 0 | 62 | 0.0%
<!-- | **Total:** | **9** | **74** | **12.16%** | -->


|  Totals | [![Coverage Status](https://coveralls.io/builds/51828958/badge)](https://coveralls.io/builds/51828958) |
| :-- | --: |
| Change from base [Build 2901881545](https://coveralls.io/builds/51827440): |  -0.3% |
| Covered Lines: | 1460 |
| Relevant Lines: | 4982 |

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


Patch:
diff --git a/cmd/db.go b/cmd/db.go
index 6b907d47..aca4d610 100644
--- a/cmd/db.go
+++ b/cmd/db.go
@@ -12,10 +12,12 @@ import (
 	"github.com/supabase/cli/internal/db/branch/list"
 	"github.com/supabase/cli/internal/db/branch/switch_"
 	"github.com/supabase/cli/internal/db/diff"
+	"github.com/supabase/cli/internal/db/lint"
 	"github.com/supabase/cli/internal/db/push"
 	"github.com/supabase/cli/internal/db/remote/changes"
 	"github.com/supabase/cli/internal/db/remote/commit"
 	"github.com/supabase/cli/internal/db/reset"
+	"github.com/supabase/cli/internal/utils"
 )
 
 var (
@@ -136,6 +138,19 @@ var (
 			return reset.Run(ctx, afero.NewOsFs())
 		},
 	}
+
+	level = utils.EnumFlag{
+		Allowed: lint.AllowedLevels,
+		Value:   lint.AllowedLevels[0],
+	}
+
+	dbLintCmd = &cobra.Command{
+		Use:   "lint",
+		Short: "Checks local database for typing error",
+		RunE: func(cmd *cobra.Command, args []string) error {
+			return lint.Run(cmd.Context(), schema, level.Value, afero.NewOsFs())
+		},
+	}
 )
 
 func init() {
@@ -146,9 +161,10 @@ func init() {
 	dbBranchCmd.AddCommand(dbSwitchCmd)
 	dbCmd.AddCommand(dbBranchCmd)
 	// Build diff command
-	dbDiffCmd.Flags().BoolVar(&useMigra, "use-migra", false, "Use migra to generate schema diff.")
-	dbDiffCmd.Flags().StringVarP(&file, "file", "f", "", "Saves schema diff to a file.")
-	dbDiffCmd.Flags().StringSliceVarP(&schema, "schema", "s", []string{"public"}, "List of schema to include.")
+	diffFlags := dbDiffCmd.Flags()
+	diffFlags.BoolVar(&useMigra, "use-migra", false, "Use migra to generate schema diff.")
+	diffFlags.StringVarP(&file, "file", "f", "", "Saves schema diff to a file.")
+	diffFlags.StringSliceVarP(&schema, "schema", "s", []string{"public"}, "List of schema to include.")
 	dbCmd.AddCommand(dbDiffCmd)
 	// Build push command
 	pushFlags := dbPushCmd.Flags()
@@ -163,7 +179,12 @@ func init() {
 	dbRemoteCmd.AddCommand(dbRemoteChangesCmd)
 	dbRemoteCmd.AddCommand(dbRemoteCommitCmd)
 	dbCmd.AddCommand(dbRemoteCmd)
-	// Buidl reset command
+	// Build reset command
 	dbCmd.AddCommand(dbResetCmd)
+	// Build lint command
+	lintFlags := dbLintCmd.Flags()
+	lintFlags.StringSliceVarP(&schema, "schema", "s", []string{"public"}, "List of schema to include.")
+	lintFlags.Var(&level, "level", "Error level to emit.")
+	dbCmd.AddCommand(dbLintCmd)
 	rootCmd.AddCommand(dbCmd)
 }
diff --git a/internal/db/diff/migra.go b/internal/db/diff/migra.go
index 42d14eea..8744b2be 100644
--- a/internal/db/diff/migra.go
+++ b/internal/db/diff/migra.go
@@ -131,7 +131,7 @@ func ApplyMigrations(ctx context.Context, url string, fsys afero.Fs, options ...
 			for _, line := range parser.Split(sql) {
 				trim := strings.TrimSpace(strings.TrimRight(line, ";"))
 				if len(trim) > 0 {
-					batch.ExecParams(line, nil, nil, nil, nil)
+					batch.ExecParams(trim, nil, nil, nil, nil)
 				}
 			}
 			if err := conn.PgConn().ExecBatch(ctx, &batch).Close(); err != nil {
diff --git a/internal/db/lint/lint.go b/internal/db/lint/lint.go
new file mode 100644
index 00000000..269ea1bc
--- /dev/null
+++ b/internal/db/lint/lint.go
@@ -0,0 +1,182 @@
+package lint
+
+import (
+	"context"
+	_ "embed"
+	"encoding/json"
+	"fmt"
+	"io"
+	"os"
+	"strings"
+
+	"github.com/jackc/pgx/v4"
+	"github.com/spf13/afero"
+	"github.com/spf13/viper"
+	"github.com/supabase/cli/internal/debug"
+	"github.com/supabase/cli/internal/utils"
+)
+
+const ENABLE_PGSQL_CHECK = "CREATE EXTENSION IF NOT EXISTS plpgsql_check"
+
+var (
+	AllowedLevels = []string{
+		"warning",
+		"error",
+	}
+	//go:embed templates/check.sql
+	checkSchemaScript string
+)
+
+type LintLevel int
+
+func toEnum(level string) LintLevel {
+	for i, curr := range AllowedLevels {
+		if strings.HasPrefix(level, curr) {
+			return LintLevel(i)
+		}
+	}
+	return -1
+}
+
+func Run(ctx context.Context, schema []string, level string, fsys afero.Fs, opts ...func(*pgx.ConnConfig)) error {
+	// Sanity checks.
+	if err := utils.LoadConfigFS(fsys); err != nil {
+		return err
+	}
+	if err := utils.AssertSupabaseDbIsRunning(); err != nil {
+		return err
+	}
+	// Run lint script
+	conn, err := ConnectLocalPostgres(ctx, "localhost", utils.Config.Db.Port, "postgres", opts...)
+	if err != nil {
+		return err
+	}
+	defer conn.Close(ctx)
+	result, err := LintDatabase(ctx, conn, schema)
+	if err != nil {
+		return err
+	}
+	return printResultJSON(result, toEnum(level), os.Stdout)
+}
+
+func filterResult(result []Result, minLevel LintLevel) (filtered []Result) {
+	for _, r := range result {
+		out := Result{Function: r.Function}
+		for _, issue := range r.Issues {
+			if toEnum(issue.Level) >= minLevel {
+				out.Issues = append(out.Issues, issue)
+			}
+		}
+		if len(out.Issues) > 0 {
+			filtered = append(filtered, out)
+		}
+	}
+	return filtered
+}
+
+func printResultJSON(result []Result, minLevel LintLevel, stdout io.Writer) error {
+	filtered := filterResult(result, minLevel)
+	if len(filtered) == 0 {
+		return nil
+	}
+	// Pretty print output
+	enc := json.NewEncoder(stdout)
+	enc.SetIndent("", "  ")
+	return enc.Encode(filtered)
+}
+
+// Connnect to local Postgres with optimised settings. The caller is responsible for closing the connection returned.
+func ConnectLocalPostgres(ctx context.Context, host string, port uint, database string, options ...func(*pgx.ConnConfig)) (*pgx.Conn, error) {
+	url := fmt.Sprintf("postgresql://postgres:postgres@%s:%d/%s", host, port, database)
+	// Parse connection url
+	config, err := pgx.ParseConfig(url)
+	if err != nil {
+		return nil, err
+	}
+	// Apply config overrides
+	for _, op := range options {
+		op(config)
+	}
+	if viper.GetBool("DEBUG") {
+		debug.SetupPGX(config)
+	}
+	// Connect to database
+	return pgx.ConnectConfig(ctx, config)
+}
+
+func LintDatabase(ctx context.Context, conn *pgx.Conn, schema []string) ([]Result, error) {
+	tx, err := conn.Begin(ctx)
+	if err != nil {
+		return nil, err
+	}
+	// Always rollback since lint should not have side effects
+	defer func() {
+		if err := tx.Rollback(context.Background()); err != nil {
+			fmt.Fprintln(os.Stderr, err)
+		}
+	}()
+	if _, err := conn.Exec(ctx, ENABLE_PGSQL_CHECK); err != nil {
+		return nil, err
+	}
+	// Batch prepares statements
+	batch := pgx.Batch{}
+	for _, s := range schema {
+		batch.Queue(checkSchemaScript, s)
+	}
+	br := conn.SendBatch(ctx, &batch)
+	defer br.Close()
+	var result []Result
+	for _, s := range schema {
+		fmt.Fprintln(os.Stderr, "Linting schema:", s)
+		rows, err := br.Query()
+		if err != nil {
+			return nil, err
+		}
+		// Parse result row
+		for rows.Next() {
+			var name string
+			var data []byte
+			if err := rows.Scan(&name, &data); err != nil {
+				return nil, err
+			}
+			var r Result
+			if err := json.Unmarshal(data, &r); err != nil {
+				return nil, err
+			}
+			// Update function name
+			r.Function = s + "." + name
+			result = append(result, r)
+		}
+		err = rows.Err()
+		if err != nil {
+			return nil, err
+		}
+	}
+	return result, nil
+}
+
+type Query struct {
+	Position string `json:"position"`
+	Text     string `json:"text"`
+}
+
+type Statement struct {
+	LineNumber string `json:"lineNumber"`
+	Text       string `json:"text"`
+}
+
+type Issue struct {
+	Level     string     `json:"level"`
+	Message   string     `json:"message"`
+	Statement *Statement `json:"statement,omitempty"`
+	Query     *Query     `json:"query,omitempty"`
+	Hint      string     `json:"hint,omitempty"`
+	Detail    string     `json:"detail,omitempty"`
+	Context   string     `json:"context,omitempty"`
+	SQLState  string     `json:"sqlState,omitempty"`
+}
+
+type Result struct {
+	Function string  `json:"function"`
+	Issues   []Issue `json:"issues"`
+}
diff --git a/internal/db/lint/templates/check.sql b/internal/db/lint/templates/check.sql
new file mode 100644
index 00000000..aa382ff4
--- /dev/null
+++ b/internal/db/lint/templates/check.sql
@@ -0,0 +1,6 @@
+-- Ref: https://github.com/okbob/plpgsql_check#mass-check
+SELECT p.proname, plpgsql_check_function(p.oid, format:='json')
+FROM pg_catalog.pg_namespace n
+JOIN pg_catalog.pg_proc p ON pronamespace = n.oid
+JOIN pg_catalog.pg_language l ON p.prolang = l.oid
+WHERE l.lanname = 'plpgsql' AND p.prorettype <> 2279 AND n.nspname = $1::text;
diff --git a/internal/db/remote/commit/commit.go b/internal/db/remote/commit/commit.go
index b3e6442c..d2f4bd71 100644
--- a/internal/db/remote/commit/commit.go
+++ b/internal/db/remote/commit/commit.go
@@ -500,11 +500,7 @@ func ConnectRemotePostgres(ctx context.Context, username, password, database, ho
 	if viper.GetBool("DEBUG") {
 		debug.SetupPGX(config)
 	}
-	conn, err := pgx.ConnectConfig(ctx, config)
-	if err != nil {
-		return nil, err
-	}
-	return conn, nil
+	return pgx.ConnectConfig(ctx, config)
 }
 
 func AssertRemoteInSync(ctx context.Context, conn *pgx.Conn, fsys afero.Fs) error {
diff --git a/internal/db/reset/reset.go b/internal/db/reset/reset.go
index de1bfd18..e0dfe868 100644
--- a/internal/db/reset/reset.go
+++ b/internal/db/reset/reset.go
@@ -51,12 +51,12 @@ func Run(ctx context.Context, fsys afero.Fs) error {
 		return err
 	}
 
-	if err := SeedDatabase(ctx, url, fsys); err != nil && !errors.Is(err, os.ErrNotExist) {
+	if err := SeedDatabase(ctx, url, fsys, opts...); err != nil && !errors.Is(err, os.ErrNotExist) {
 		return err
 	}
 
 	fmt.Fprintln(os.Stderr, "Activating branch...")
-	if err := ActivateDatabase(ctx, branch); err != nil {
+	if err := ActivateDatabase(ctx, branch, opts...); err != nil {
 		return err
 	}
 


Test patch:
diff --git a/internal/db/diff/migra_test.go b/internal/db/diff/migra_test.go
index 412dca05..556feed7 100644
--- a/internal/db/diff/migra_test.go
+++ b/internal/db/diff/migra_test.go
@@ -17,7 +17,6 @@ func TestApplyMigrations(t *testing.T) {
 	const postgresUrl = "postgresql://postgres:password@localhost:5432/postgres"
 
 	t.Run("applies migrations from local directory", func(t *testing.T) {
-		t.Skip("pgmock does not support batch query")
 		// Setup in-memory fs
 		fsys := afero.NewMemMapFs()
 		// Setup initial migration
@@ -33,8 +32,10 @@ func TestApplyMigrations(t *testing.T) {
 		defer conn.Close(t)
 		conn.Query("create table test").
 			Reply("SELECT 0").
-			Query("drop table test;-- ignore me").
-			Reply("SELECT 0")
+			Query("drop table test").
+			Reply("SELECT 0").
+			Query("-- ignore me").
+			Reply("")
 		// Run test
 		assert.NoError(t, ApplyMigrations(context.Background(), postgresUrl, fsys, conn.Intercept))
 	})
@@ -73,7 +74,6 @@ func TestApplyMigrations(t *testing.T) {
 	})
 
 	t.Run("throws error on failture to send batch", func(t *testing.T) {
-		t.Skip("pgmock does not support batch query")
 		// Setup in-memory fs
 		fsys := afero.NewMemMapFs()
 		// Setup initial migration
diff --git a/internal/db/lint/lint_test.go b/internal/db/lint/lint_test.go
new file mode 100644
index 00000000..4f175d79
--- /dev/null
+++ b/internal/db/lint/lint_test.go
@@ -0,0 +1,269 @@
+package lint
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"net/http"
+	"testing"
+
+	"github.com/docker/docker/api/types"
+	"github.com/docker/docker/client"
+	"github.com/jackc/pgerrcode"
+	"github.com/spf13/afero"
+	"github.com/spf13/viper"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/supabase/cli/internal/testing/pgtest"
+	"github.com/supabase/cli/internal/utils"
+	"gopkg.in/h2non/gock.v1"
+)
+
+func TestLintCommand(t *testing.T) {
+	const version = "1.41"
+
+	// 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{})
+	// Setup db response
+	expected := Result{
+		Function: "22751",
+		Issues: []Issue{{
+			Level:   AllowedLevels[1],
+			Message: `record "r" has no field "c"`,
+			Statement: &Statement{
+				LineNumber: "6",
+				Text:       "RAISE",
+			},
+			Context:  `SQL expression "r.c"`,
+			SQLState: pgerrcode.UndefinedColumn,
+		}},
+	}
+	data, err := json.Marshal(expected)
+	require.NoError(t, err)
+	// Setup mock postgres
+	conn := pgtest.NewConn()
+	defer conn.Close(t)
+	conn.Query("begin").Reply("BEGIN").
+		Query(ENABLE_PGSQL_CHECK).
+		Reply("CREATE EXTENSION").
+		Query(checkSchemaScript, "public").
+		Reply("SELECT 1", []interface{}{"f1", string(data)}).
+		Query("rollback").Reply("ROLLBACK")
+	// Run test
+	assert.NoError(t, Run(context.Background(), []string{"public"}, "warning", fsys, conn.Intercept))
+}
+
+func TestLintDatabase(t *testing.T) {
+	t.Run("parses lint results", func(t *testing.T) {
+		expected := []Result{{
+			Function: "public.f1",
+			Issues: []Issue{{
+				Level:   AllowedLevels[1],
+				Message: `record "r" has no field "c"`,
+				Statement: &Statement{
+					LineNumber: "6",
+					Text:       "RAISE",
+				},
+				Context:  `SQL expression "r.c"`,
+				SQLState: pgerrcode.UndefinedColumn,
+			}, {
+				Level:    "warning extra",
+				Message:  `never read variable "entity"`,
+				SQLState: pgerrcode.SuccessfulCompletion,
+			}},
+		}, {
+			Function: "public.f2",
+			Issues: []Issue{{
+				Level:   AllowedLevels[1],
+				Message: `relation "t2" does not exist`,
+				Statement: &Statement{
+					LineNumber: "4",
+					Text:       "FOR over SELECT rows",
+				},
+				Query: &Query{
+					Position: "15",
+					Text:     "SELECT * FROM t2",
+				},
+				SQLState: pgerrcode.UndefinedTable,
+			}},
+		}}
+		r1, err := json.Marshal(expected[0])
+		require.NoError(t, err)
+		r2, err := json.Marshal(expected[1])
+		require.NoError(t, err)
+		// Setup mock postgres
+		conn := pgtest.NewConn()
+		defer conn.Close(t)
+		conn.Query("begin").Reply("BEGIN").
+			Query(ENABLE_PGSQL_CHECK).
+			Reply("CREATE EXTENSION").
+			Query(checkSchemaScript, "public").
+			Reply("SELECT 2",
+				[]interface{}{"f1", string(r1)},
+				[]interface{}{"f2", string(r2)},
+			).
+			Query("rollback").Reply("ROLLBACK")
+		// Connect to mock
+		ctx := context.Background()
+		mock, err := ConnectLocalPostgres(ctx, "localhost", 5432, "postgres", conn.Intercept)
+		require.NoError(t, err)
+		defer mock.Close(ctx)
+		// Run test
+		result, err := LintDatabase(ctx, mock, []string{"public"})
+		assert.NoError(t, err)
+		// Validate result
+		assert.ElementsMatch(t, expected, result)
+	})
+
+	t.Run("supports multiple schema", func(t *testing.T) {
+		expected := []Result{{
+			Function: "public.where_clause",
+			Issues: []Issue{{
+				Level:   AllowedLevels[0],
+				Message: "target type is different type than source type",
+				Statement: &Statement{
+					LineNumber: "32",
+					Text:       "statement block",
+				},
+				Hint:     "The input expression type does not have an assignment cast to the target type.",
+				Detail:   `cast "text" value to "text[]" type`,
+				Context:  `during statement block local variable "clause_arr" initialization on line 3`,
+				SQLState: pgerrcode.DatatypeMismatch,
+			}},
+		}, {
+			Function: "private.f2",
+			Issues:   []Issue{},
+		}}
+		r1, err := json.Marshal(expected[0])
+		require.NoError(t, err)
+		r2, err := json.Marshal(expected[1])
+		require.NoError(t, err)
+		// Setup mock postgres
+		conn := pgtest.NewConn()
+		defer conn.Close(t)
+		conn.Query("begin").Reply("BEGIN").
+			Query(ENABLE_PGSQL_CHECK).
+			Reply("CREATE EXTENSION").
+			Query(checkSchemaScript, "public").
+			Reply("SELECT 1", []interface{}{"where_clause", string(r1)}).
+			Query(checkSchemaScript, "private").
+			Reply("SELECT 1", []interface{}{"f2", string(r2)}).
+			Query("rollback").Reply("ROLLBACK")
+		// Connect to mock
+		ctx := context.Background()
+		mock, err := ConnectLocalPostgres(ctx, "localhost", 5432, "postgres", conn.Intercept)
+		require.NoError(t, err)
+		defer mock.Close(ctx)
+		// Run test
+		result, err := LintDatabase(ctx, mock, []string{"public", "private"})
+		assert.NoError(t, err)
+		// Validate result
+		assert.ElementsMatch(t, expected, result)
+	})
+
+	t.Run("throws error on missing extension", func(t *testing.T) {
+		// Setup mock postgres
+		conn := pgtest.NewConn()
+		defer conn.Close(t)
+		conn.Query("begin").Reply("BEGIN").
+			Query(ENABLE_PGSQL_CHECK).
+			ReplyError(pgerrcode.UndefinedFile, `could not open extension control file "/usr/share/postgresql/14/extension/plpgsql_check.control": No such file or directory"`).
+			Query("rollback").Reply("ROLLBACK")
+		// Connect to mock
+		ctx := context.Background()
+		mock, err := ConnectLocalPostgres(ctx, "localhost", 5432, "postgres", conn.Intercept)
+		require.NoError(t, err)
+		defer mock.Close(ctx)
+		// Run test
+		_, err = LintDatabase(ctx, mock, []string{"public"})
+		assert.Error(t, err)
+	})
+
+	t.Run("throws error on malformed json", func(t *testing.T) {
+		// Setup mock postgres
+		conn := pgtest.NewConn()
+		defer conn.Close(t)
+		conn.Query("begin").Reply("BEGIN").
+			Query(ENABLE_PGSQL_CHECK).
+			Reply("CREATE EXTENSION").
+			Query(checkSchemaScript, "public").
+			Reply("SELECT 1", []interface{}{"f1", "malformed"}).
+			Query("rollback").Reply("ROLLBACK")
+		// Connect to mock
+		ctx := context.Background()
+		mock, err := ConnectLocalPostgres(ctx, "localhost", 5432, "postgres", conn.Intercept)
+		require.NoError(t, err)
+		defer mock.Close(ctx)
+		// Run test
+		_, err = LintDatabase(ctx, mock, []string{"public"})
+		assert.Error(t, err)
+	})
+}
+
+func TestConnectLocal(t *testing.T) {
+	t.Run("connects with debug log", func(t *testing.T) {
+		viper.Set("DEBUG", true)
+		_, err := ConnectLocalPostgres(context.Background(), "localhost", 5432, "postgres")
+		assert.Error(t, err)
+	})
+
+	t.Run("throws error on invalid port", func(t *testing.T) {
+		_, err := ConnectLocalPostgres(context.Background(), "localhost", 0, "postgres")
+		assert.Error(t, err)
+	})
+}
+
+func TestPrintResult(t *testing.T) {
+	result := []Result{{
+		Function: "public.f1",
+		Issues: []Issue{{
+			Level:   "warning",
+			Message: "test 1a",
+		}, {
+			Level:   "error",
+			Message: "test 1b",
+		}},
+	}, {
+		Function: "private.f2",
+		Issues: []Issue{{
+			Level:   "warning extra",
+			Message: "test 2",
+		}},
+	}}
+
+	t.Run("filters warning level", func(t *testing.T) {
+		// Run test
+		var out bytes.Buffer
+		assert.NoError(t, printResultJSON(result, toEnum("warning"), &out))
+		// Validate output
+		var actual []Result
+		assert.NoError(t, json.Unmarshal(out.Bytes(), &actual))
+		assert.ElementsMatch(t, result, actual)
+	})
+
+	t.Run("filters error level", func(t *testing.T) {
+		// Run test
+		var out bytes.Buffer
+		assert.NoError(t, printResultJSON(result, toEnum("error"), &out))
+		// Validate output
+		var actual []Result
+		assert.NoError(t, json.Unmarshal(out.Bytes(), &actual))
+		assert.ElementsMatch(t, []Result{{
+			Function: result[0].Function,
+			Issues:   []Issue{result[0].Issues[1]},
+		}}, actual)
+	})
+}
diff --git a/internal/db/remote/set/set_test.go b/internal/db/remote/set/set_test.go
index 5475f567..26508333 100644
--- a/internal/db/remote/set/set_test.go
+++ b/internal/db/remote/set/set_test.go
@@ -28,7 +28,7 @@ func TestDbRemoteSetCommand(t *testing.T) {
 		conn.Query(CHECK_MIGRATION_EXISTS).
 			Reply("SELECT 0").
 			Query(LIST_MIGRATION_VERSION).
-			Reply("SELECT 1", map[string]interface{}{"version": version})
+			Reply("SELECT 1", []interface{}{version})
 		// Run test
 		assert.NoError(t, Run(postgresUrl, fsys, conn.Intercept))
 	})
@@ -141,7 +141,7 @@ func TestDbRemoteSetCommand(t *testing.T) {
 		conn.Query(CHECK_MIGRATION_EXISTS).
 			Reply("SELECT 0").
 			Query(LIST_MIGRATION_VERSION).
-			Reply("SELECT 1", map[string]interface{}{"version": "20220727064247"})
+			Reply("SELECT 1", []interface{}{"20220727064247"})
 		// Run test
 		assert.Error(t, Run(postgresUrl, fsys, conn.Intercept))
 	})
@@ -160,7 +160,7 @@ func TestDbRemoteSetCommand(t *testing.T) {
 		conn.Query(CHECK_MIGRATION_EXISTS).
 			Reply("SELECT 0").
 			Query(LIST_MIGRATION_VERSION).
-			Reply("SELECT 1", map[string]interface{}{"version": "20220727064247"})
+			Reply("SELECT 1", []interface{}{"20220727064247"})
 		// Run test
 		assert.Error(t, Run(postgresUrl, fsys, conn.Intercept))
 	})
diff --git a/internal/testing/pgtest/mock.go b/internal/testing/pgtest/mock.go
index 348b7c3f..564866c1 100644
--- a/internal/testing/pgtest/mock.go
+++ b/internal/testing/pgtest/mock.go
@@ -2,6 +2,7 @@ package pgtest
 
 import (
 	"context"
+	"fmt"
 	"net"
 	"reflect"
 	"testing"
@@ -64,17 +65,29 @@ func (r *MockConn) Intercept(config *pgx.ConnConfig) {
 	config.DialFunc = func(ctx context.Context, network, addr string) (net.Conn, error) {
 		return r.server.DialContext(ctx)
 	}
-	config.PreferSimpleProtocol = true
 	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(sql string) *MockConn {
-	r.script.Steps = append(r.script.Steps, pgmock.ExpectMessage(&pgproto3.Query{String: sql}))
+// Adds a simple query or prepared statement to the mock connection.
+func (r *MockConn) Query(sql string, args ...interface{}) *MockConn {
+	var oids []uint32
+	var params [][]byte
+	for _, v := range args {
+		if dt, ok := ci.DataTypeForValue(v); ok {
+			if err := dt.Value.Set(v); err != nil {
+				continue
+			}
+			value, err := (dt.Value).(pgtype.TextEncoder).EncodeText(ci, []byte{})
+			if err != nil {
+				continue
+			}
+			params = append(params, value)
+			oids = append(oids, dt.OID)
+		}
+	}
+	r.script.Steps = append(r.script.Steps, ExpectQuery(sql, params, oids))
 	return r
 }
 
@@ -87,18 +100,24 @@ func getDataTypeSize(v interface{}) int16 {
 	return int16(t.Size())
 }
 
+func (r *MockConn) lastQuery() *extendedQueryStep {
+	return r.script.Steps[len(r.script.Steps)-1].(*extendedQueryStep)
+}
+
 // Adds a server reply using text protocol format.
 //
 // TODO: support binary protocol
-func (r *MockConn) Reply(tag string, rows ...map[string]interface{}) *MockConn {
+func (r *MockConn) Reply(tag string, rows ...[]interface{}) *MockConn {
+	q := r.lastQuery()
 	// Add field description
 	if len(rows) > 0 {
 		var desc pgproto3.RowDescription
-		for k, v := range rows[0] {
+		for i, v := range rows[0] {
+			name := fmt.Sprintf("c_%02d", i)
 			if dt, ok := ci.DataTypeForValue(v); ok {
 				size := getDataTypeSize(v)
 				desc.Fields = append(desc.Fields, pgproto3.FieldDescription{
-					Name:                 []byte(k),
+					Name:                 []byte(name),
 					TableOID:             17131,
 					TableAttributeNumber: 1,
 					DataTypeOID:          dt.OID,
@@ -108,11 +127,11 @@ func (r *MockConn) Reply(tag string, rows ...map[string]interface{}) *MockConn {
 				})
 			}
 		}
-		r.script.Steps = append(r.script.Steps, pgmock.SendMessage(&desc))
+		q.reply.Steps = append(q.reply.Steps, pgmock.SendMessage(&desc))
+	} else {
+		// No data is optional, but we add for completeness
+		q.reply.Steps = append(q.reply.Steps, pgmock.SendMessage(&pgproto3.NoData{}))
 	}
-	// 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
@@ -126,15 +145,16 @@ func (r *MockConn) Reply(tag string, rows ...map[string]interface{}) *MockConn {
 				}
 			}
 		}
-		r.script.Steps = append(r.script.Steps, pgmock.SendMessage(&dr))
+		q.reply.Steps = append(q.reply.Steps, pgmock.SendMessage(&dr))
 	}
-
 	// Add completion message
-	r.script.Steps = append(
-		r.script.Steps,
-		pgmock.SendMessage(&pgproto3.CommandComplete{CommandTag: []byte(tag)}),
-		pgmock.SendMessage(&pgproto3.ReadyForQuery{TxStatus: 'I'}),
-	)
+	var complete pgproto3.BackendMessage
+	if tag == "" {
+		complete = &pgproto3.EmptyQueryResponse{}
+	} else {
+		complete = &pgproto3.CommandComplete{CommandTag: []byte(tag)}
+	}
+	q.reply.Steps = append(q.reply.Steps, pgmock.SendMessage(complete))
 	return r
 }
 
@@ -142,15 +162,15 @@ func (r *MockConn) Reply(tag string, rows ...map[string]interface{}) *MockConn {
 //
 // TODO: simulate a notice reply
 func (r *MockConn) ReplyError(code, message string) *MockConn {
-	r.script.Steps = append(
-		r.script.Steps,
+	q := r.lastQuery()
+	q.reply.Steps = append(
+		q.reply.Steps,
 		pgmock.SendMessage(&pgproto3.ErrorResponse{
 			Severity:            "ERROR",
 			SeverityUnlocalized: "ERROR",
 			Code:                code,
 			Message:             message,
 		}),
-		pgmock.SendMessage(&pgproto3.ReadyForQuery{TxStatus: 'I'}),
 	)
 	return r
 }
@@ -189,7 +209,7 @@ func NewWithStatus(status map[string]string) *MockConn {
 			return
 		}
 		// Always expect clients to terminate the request
-		mock.script.Steps = append(mock.script.Steps, pgmock.ExpectMessage(&pgproto3.Terminate{}))
+		mock.script.Steps = append(mock.script.Steps, ExpectTerminate())
 		err = mock.script.Run(pgproto3.NewBackend(pgproto3.NewChunkReader(conn), conn))
 		if err != nil {
 			mock.errChan <- err
diff --git a/internal/testing/pgtest/step.go b/internal/testing/pgtest/step.go
new file mode 100644
index 00000000..2dd37a2e
--- /dev/null
+++ b/internal/testing/pgtest/step.go
@@ -0,0 +1,130 @@
+package pgtest
+
+import (
+	"fmt"
+	"reflect"
+
+	"github.com/jackc/pgmock"
+	"github.com/jackc/pgproto3/v2"
+	"github.com/jackc/pgtype"
+)
+
+type extendedQueryStep struct {
+	sql    string
+	params [][]byte
+	oids   []uint32
+	reply  pgmock.Script
+}
+
+func (e *extendedQueryStep) Step(backend *pgproto3.Backend) error {
+	msg, err := getFrontendMessage(backend)
+	if err != nil {
+		return err
+	}
+
+	// Handle prepared statements, name can be dynamic: lrupsc_5_0
+	if m, ok := msg.(*pgproto3.Parse); ok {
+		want := &pgproto3.Parse{Name: m.Name, Query: e.sql, ParameterOIDs: m.ParameterOIDs}
+		if !reflect.DeepEqual(m, want) {
+			return fmt.Errorf("msg => %#v, e.want => %#v", m, want)
+		}
+		// Anonymous ps falls through
+		if m.Name != "" {
+			script := pgmock.Script{Steps: []pgmock.Step{
+				pgmock.ExpectMessage(&pgproto3.Describe{ObjectType: 'S', Name: m.Name}),
+				pgmock.ExpectMessage(&pgproto3.Sync{}),
+				pgmock.SendMessage(&pgproto3.ParseComplete{}),
+				pgmock.SendMessage(&pgproto3.ParameterDescription{ParameterOIDs: e.oids}),
+				// Postgres responds pgproto3.RowDescription but it's optional for pgx
+				pgmock.SendMessage(&pgproto3.ReadyForQuery{TxStatus: 'I'}),
+			}}
+			if err := script.Run(backend); err != nil {
+				return err
+			}
+		}
+		// Expect bind command next
+		msg, err = backend.Receive()
+		if err != nil {
+			return err
+		}
+	}
+
+	if m, ok := msg.(*pgproto3.Bind); ok {
+		var codes []int16
+		for range e.oids {
+			codes = append(codes, pgtype.TextFormatCode)
+		}
+		want := &pgproto3.Bind{
+			ParameterFormatCodes: codes,
+			Parameters:           e.params,
+			ResultFormatCodes:    []int16{},
+			DestinationPortal:    m.DestinationPortal,
+			PreparedStatement:    m.PreparedStatement,
+		}
+		if !reflect.DeepEqual(m, want) {
+			return fmt.Errorf("msg => %#v, e.want => %#v", msg, want)
+		}
+		e.reply.Steps = append([]pgmock.Step{
+			pgmock.ExpectMessage(&pgproto3.Describe{ObjectType: 'P'}),
+			pgmock.ExpectMessage(&pgproto3.Execute{}),
+			pgmock.SendMessage(&pgproto3.ParseComplete{}),
+			pgmock.SendMessage(&pgproto3.BindComplete{}),
+		}, e.reply.Steps...)
+		return e.reply.Run(backend)
+	}
+
+	// Handle simple query
+	want := &pgproto3.Query{String: e.sql}
+	if m, ok := msg.(*pgproto3.Query); ok && reflect.DeepEqual(m, want) {
+		e.reply.Steps = append(e.reply.Steps, pgmock.SendMessage(&pgproto3.ReadyForQuery{TxStatus: 'I'}))
+		return e.reply.Run(backend)
+	}
+
+	return fmt.Errorf("msg => %#v, e.want => %#v", msg, want)
+}
+
+// Expects a SQL query in any form: simple, prepared, or anonymous.
+func ExpectQuery(sql string, params [][]byte, oids []uint32) pgmock.Step {
+	return &extendedQueryStep{sql: sql, params: params, oids: oids}
+}
+
+type terminateStep struct{}
+
+func (e *terminateStep) Step(backend *pgproto3.Backend) error {
+	msg, err := getFrontendMessage(backend)
+	if err != nil {
+		return err
+	}
+
+	// Handle simple query
+	if _, ok := msg.(*pgproto3.Terminate); ok {
+		return nil
+	}
+
+	return fmt.Errorf("msg => %#v, e.want => %#v", msg, &pgproto3.Terminate{})
+}
+
+func ExpectTerminate() pgmock.Step {
+	return &terminateStep{}
+}
+
+func getFrontendMessage(backend *pgproto3.Backend) (pgproto3.FrontendMessage, error) {
+	msg, err := backend.Receive()
+	if err != nil {
+		return nil, err
+	}
+
+	// Sync signals end of batch statements
+	if _, ok := msg.(*pgproto3.Sync); ok {
+		reply := pgmock.SendMessage(&pgproto3.ReadyForQuery{TxStatus: 'I'})
+		if err := reply.Step(backend); err != nil {
+			return nil, err
+		}
+		msg, err = backend.Receive()
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	return msg, nil
+}