feat: add db test command using pgtap

cmd/db.go

@@ -17,6 +17,7 @@ import (
 	"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/db/test"
 	"github.com/supabase/cli/internal/utils"
 )
 
@@ -149,7 +150,17 @@ var (
 		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())
+			ctx, _ := signal.NotifyContext(cmd.Context(), os.Interrupt)
+			return lint.Run(ctx, schema, level.Value, afero.NewOsFs())
+		},
+	}
+
+	dbTestCmd = &cobra.Command{
+		Use:   "test",
+		Short: "Tests local database with pgTAP.",
+		RunE: func(cmd *cobra.Command, args []string) error {
+			ctx, _ := signal.NotifyContext(cmd.Context(), os.Interrupt)
+			return test.Run(ctx, afero.NewOsFs())
 		},
 	}
 )
@@ -187,5 +198,7 @@ func init() {
 	lintFlags.StringSliceVarP(&schema, "schema", "s", []string{"public"}, "List of schema to include.")
 	lintFlags.Var(&level, "level", "Error level to emit.")
 	dbCmd.AddCommand(dbLintCmd)
+	// Build test command
+	dbCmd.AddCommand(dbTestCmd)
 	rootCmd.AddCommand(dbCmd)
 }

internal/db/branch/switch_/switch_.go

@@ -65,7 +65,7 @@ func switchDatabase(ctx context.Context, source, target string, options ...func(
 	if err != nil {
 		return err
 	}
-	defer conn.Close(ctx)
+	defer conn.Close(context.Background())
 	if err := reset.DisconnectClients(ctx, conn); err != nil {
 		return err
 	}

internal/db/diff/migra.go

@@ -103,7 +103,7 @@ func ApplyMigrations(ctx context.Context, url string, fsys afero.Fs, options ...
 	if err != nil {
 		return err
 	}
-	defer conn.Close(ctx)
+	defer conn.Close(context.Background())
 	// Apply migrations
 	if migrations, err := afero.ReadDir(fsys, utils.MigrationsDir); err == nil {
 		for i, migration := range migrations {

internal/db/lint/lint.go

@@ -51,7 +51,7 @@ func Run(ctx context.Context, schema []string, level string, fsys afero.Fs, opts
 	if err != nil {
 		return err
 	}
-	defer conn.Close(ctx)
+	defer conn.Close(context.Background())
 	result, err := LintDatabase(ctx, conn, schema)
 	if err != nil {
 		return err

internal/db/reset/reset.go

@@ -95,7 +95,7 @@ func SeedDatabase(ctx context.Context, url string, fsys afero.Fs, options ...fun
 	if err != nil {
 		return err
 	}
-	defer conn.Close(ctx)
+	defer conn.Close(context.Background())
 	// Batch seed commands, safe to use statement cache
 	batch := pgx.Batch{}
 	lines, err := parser.Split(sql)
@@ -130,7 +130,7 @@ func ActivateDatabase(ctx context.Context, branch string, options ...func(*pgx.C
 	if err != nil {
 		return err
 	}
-	defer conn.Close(ctx)
+	defer conn.Close(context.Background())
 	if err := DisconnectClients(ctx, conn); err != nil {
 		return err
 	}

internal/db/test/templates/test.sh

@@ -0,0 +1,29 @@
+#!/bin/bash
+set -euo pipefail
+
+# TODO: move this dependency to base image
+if ! command -v pg_prove &> /dev/null; then
+    apt-get -qq update && apt-get -qq install postgresql-14-pgtap
+fi
+
+# temporarily enable pgtap
+enable="create extension if not exists pgtap with schema extensions"
+notice=$(psql -h localhost -U postgres -p 5432 -d postgres -c "$enable" 2>&1 >/dev/null)
+
+# run postgres unit tests
+pg_prove -h localhost -U postgres -r "$@"
+
+cleanup() {
+    # save the return code of the script
+    status=$?
+    # disable pgtap if previously not enabled
+    if [ -z "$notice" ]; then
+        psql -h localhost -U postgres -p 5432 -d postgres -c "drop extension if exists pgtap"
+    fi
+    # clean up test files
+    rm -rf "$@"
+    # actually quit
+    exit $status
+}
+
+trap cleanup EXIT

internal/db/test/test.go

@@ -0,0 +1,114 @@
+package test
+
+import (
+	"archive/tar"
+	"bytes"
+	"context"
+	_ "embed"
+	"fmt"
+	"io"
+	"os"
+	"path/filepath"
+
+	"github.com/docker/docker/api/types"
+	"github.com/spf13/afero"
+	"github.com/supabase/cli/internal/utils"
+)
+
+var (
+	//go:embed templates/test.sh
+	testScript string
+)
+
+func Run(ctx context.Context, fsys afero.Fs) error {
+	// Sanity checks.
+	{
+		if err := utils.LoadConfigFS(fsys); err != nil {
+			return err
+		}
+
+		if err := utils.AssertSupabaseDbIsRunning(); err != nil {
+			return err
+		}
+	}
+
+	var buf bytes.Buffer
+	if err := compress(utils.DbTestsDir, &buf, fsys); err != nil {
+		return err
+	}
+	dstPath := "/tmp"
+	if err := utils.Docker.CopyToContainer(ctx, utils.DbId, dstPath, &buf, types.CopyToContainerOptions{}); err != nil {
+		return err
+	}
+
+	// Requires unix path inside container
+	cmd := []string{"/bin/bash", "-c", testScript, dstPath + "/" + utils.DbTestsDir}
+	out, err := utils.DockerExecOnce(ctx, utils.DbId, nil, cmd)
+	if err != nil {
+		return err
+	}
+
+	fmt.Print(out)
+	return nil
+}
+
+// Ref 1: https://medium.com/@skdomino/taring-untaring-files-in-go-6b07cf56bc07
+// Ref 2: https://gist.github.com/mimoo/25fc9716e0f1353791f5908f94d6e726
+func compress(src string, buf io.Writer, fsys afero.Fs) error {
+	tw := tar.NewWriter(buf)
+
+	// walk through every file in the folder
+	if err := afero.Walk(fsys, src, func(file string, fi os.FileInfo, err error) error {
+		// return on any error
+		if err != nil {
+			fmt.Println(err)
+			return err
+		}
+
+		// return on non-regular files
+		if !fi.Mode().IsRegular() {
+			return nil
+		}
+
+		// create a new dir/file header
+		header, err := tar.FileInfoHeader(fi, fi.Name())
+		if err != nil {
+			return err
+		}
+
+		// must provide real name
+		header.Name = filepath.ToSlash(file)
+
+		// write the header
+		if err := tw.WriteHeader(header); err != nil {
+			return err
+		}
+
+		// open files for taring
+		f, err := fsys.Open(file)
+		if err != nil {
+			return err
+		}
+
+		// copy file data into tar writer
+		if _, err := io.Copy(tw, f); err != nil {
+			return err
+		}
+
+		// manually close here after each file operation; defering would cause each file close
+		// to wait until all operations have completed.
+		if err := f.Close(); err != nil {
+			return err
+		}
+
+		return nil
+	}); err != nil {
+		return err
+	}
+
+	// produce tar
+	if err := tw.Close(); err != nil {
+		return err
+	}
+	return nil
+}

internal/db/test/test_test.go

@@ -0,0 +1,23 @@
+package test
+
+import (
+	"bytes"
+	"path/filepath"
+	"testing"
+
+	"github.com/spf13/afero"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestTarDirectory(t *testing.T) {
+	t.Run("tars a given directory", func(t *testing.T) {
+		// Setup in-memory fs
+		fsys := afero.NewMemMapFs()
+		dir := "/tests"
+		require.NoError(t, afero.WriteFile(fsys, filepath.Join(dir, "order_test.pg"), []byte("SELECT 0;"), 0644))
+		// Run test
+		var buf bytes.Buffer
+		assert.NoError(t, compress(dir, &buf, fsys))
+	})
+}

internal/utils/misc.go

@@ -75,6 +75,7 @@ DO 'BEGIN WHILE (
 	CurrBranchPath = "supabase/.branches/_current_branch"
 	MigrationsDir  = "supabase/migrations"
 	FunctionsDir   = "supabase/functions"
+	DbTestsDir     = "supabase/tests"
 	SeedDataPath   = "supabase/seed.sql"
 )