fix: use pgkit migra with experimental flag

Author: sweatybridge

cmd/db.go

@@ -101,6 +101,8 @@ var (
 			if usePgSchema {
 				differ = diff.DiffPgSchema
 				fmt.Fprintln(os.Stderr, utils.Yellow("WARNING:"), "--use-pg-schema flag is experimental and may not include all entities, such as views and grants.")
+			} else if !viper.GetBool("EXPERIMENTAL") {
+				differ = diff.DiffSchemaMigraBash
 			}
 			return diff.Run(cmd.Context(), schema, file, flags.DbConfig, differ, afero.NewOsFs())
 		},

internal/db/diff/diff.go

@@ -28,7 +28,7 @@ import (
 	"github.com/supabase/cli/pkg/parser"
 )
 
-type DiffFunc func(context.Context, string, string, []string) (string, error)
+type DiffFunc func(context.Context, string, string, []string, ...func(*pgx.ConnConfig)) (string, error)
 
 func Run(ctx context.Context, schema []string, file string, config pgconn.Config, differ DiffFunc, fsys afero.Fs, options ...func(*pgx.ConnConfig)) (err error) {
 	out, err := DiffDatabase(ctx, schema, config, os.Stderr, fsys, differ, options...)
@@ -136,7 +136,7 @@ func MigrateShadowDatabase(ctx context.Context, container string, fsys afero.Fs,
 	return migration.ApplyMigrations(ctx, migrations, conn, afero.NewIOFS(fsys))
 }
 
-func DiffDatabase(ctx context.Context, schema []string, config pgconn.Config, w io.Writer, fsys afero.Fs, differ func(context.Context, string, string, []string) (string, error), options ...func(*pgx.ConnConfig)) (string, error) {
+func DiffDatabase(ctx context.Context, schema []string, config pgconn.Config, w io.Writer, fsys afero.Fs, differ DiffFunc, options ...func(*pgx.ConnConfig)) (string, error) {
 	fmt.Fprintln(w, "Creating shadow database...")
 	shadow, err := CreateShadowDatabase(ctx, utils.Config.Db.ShadowPort)
 	if err != nil {
@@ -175,7 +175,7 @@ func DiffDatabase(ctx context.Context, schema []string, config pgconn.Config, w
 	}
 	source := utils.ToPostgresURL(shadowConfig)
 	target := utils.ToPostgresURL(config)
-	return differ(ctx, source, target, schema)
+	return differ(ctx, source, target, schema, options...)
 }
 
 func migrateBaseDatabase(ctx context.Context, config pgconn.Config, migrations []string, fsys afero.Fs, options ...func(*pgx.ConnConfig)) error {

internal/db/diff/diff_test.go

@@ -15,6 +15,7 @@ import (
 	"github.com/h2non/gock"
 	"github.com/jackc/pgconn"
 	"github.com/jackc/pgerrcode"
+	"github.com/jackc/pgx/v4"
 	"github.com/spf13/afero"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -72,7 +73,17 @@ func TestRun(t *testing.T) {
 			Reply("CREATE DATABASE")
 		defer conn.Close(t)
 		// Run test
-		err := Run(context.Background(), []string{"public"}, "file", dbConfig, DiffSchemaMigra, fsys, conn.Intercept)
+		err := Run(context.Background(), []string{"public"}, "file", dbConfig, DiffSchemaMigra, fsys, func(cc *pgx.ConnConfig) {
+			if cc.Host == dbConfig.Host {
+				// Fake a SSL error when connecting to target database
+				cc.LookupFunc = func(ctx context.Context, host string) (addrs []string, err error) {
+					return nil, errors.New("server refused TLS connection")
+				}
+			} else {
+				// Hijack connection to shadow database
+				conn.Intercept(cc)
+			}
+		})
 		// Check error
 		assert.NoError(t, err)
 		assert.Empty(t, apitest.ListUnmatchedRequests())

internal/db/diff/migra.go

@@ -9,10 +9,12 @@ import (
 	"github.com/docker/docker/api/types/container"
 	"github.com/docker/docker/api/types/network"
 	"github.com/go-errors/errors"
+	"github.com/jackc/pgx/v4"
 	"github.com/spf13/viper"
 	"github.com/supabase/cli/internal/gen/types"
 	"github.com/supabase/cli/internal/utils"
 	"github.com/supabase/cli/pkg/config"
+	"github.com/supabase/cli/pkg/migration"
 )
 
 var (
@@ -60,7 +62,14 @@ var (
 )
 
 // Diffs local database schema against shadow, dumps output to stdout.
-func DiffSchemaMigraBash(ctx context.Context, source, target string, schema []string) (string, error) {
+func DiffSchemaMigraBash(ctx context.Context, source, target string, schema []string, options ...func(*pgx.ConnConfig)) (string, error) {
+	// Load all user defined schemas
+	if len(schema) == 0 {
+		var err error
+		if schema, err = loadSchema(ctx, target, options...); err != nil {
+			return "", err
+		}
+	}
 	env := []string{"SOURCE=" + source, "TARGET=" + target}
 	// Passing in script string means command line args must be set manually, ie. "$@"
 	args := "set -- " + strings.Join(schema, " ") + ";"
@@ -86,10 +95,20 @@ func DiffSchemaMigraBash(ctx context.Context, source, target string, schema []st
 	return out.String(), nil
 }
 
-func DiffSchemaMigra(ctx context.Context, source, target string, schema []string) (string, error) {
+func loadSchema(ctx context.Context, dbURL string, options ...func(*pgx.ConnConfig)) ([]string, error) {
+	conn, err := utils.ConnectByUrl(ctx, dbURL, options...)
+	if err != nil {
+		return nil, err
+	}
+	defer conn.Close(context.Background())
+	// RLS policies in auth and storage schemas can be included with -s flag
+	return migration.ListUserSchemas(ctx, conn)
+}
+
+func DiffSchemaMigra(ctx context.Context, source, target string, schema []string, options ...func(*pgx.ConnConfig)) (string, error) {
 	env := []string{"SOURCE=" + source, "TARGET=" + target}
 	// node-postgres does not support sslmode=prefer
-	if require, err := types.IsRequireSSL(ctx, target); err != nil {
+	if require, err := types.IsRequireSSL(ctx, target, options...); err != nil {
 		return "", err
 	} else if require {
 		rootCA := caStaging + caProd

internal/db/diff/pgschema.go

@@ -7,10 +7,11 @@ import (
 	"strings"
 
 	"github.com/go-errors/errors"
+	"github.com/jackc/pgx/v4"
 	pgschema "github.com/stripe/pg-schema-diff/pkg/diff"
 )
 
-func DiffPgSchema(ctx context.Context, source, target string, schema []string) (string, error) {
+func DiffPgSchema(ctx context.Context, source, target string, schema []string, _ ...func(*pgx.ConnConfig)) (string, error) {
 	dbSrc, err := sql.Open("pgx", source)
 	if err != nil {
 		return "", errors.Errorf("failed to open source database: %w", err)