fix: diff schemas using pgkit migra (#3994)

* fix: diff schemas using pgkit migra

* chore: update exclusion rules for pg schema diff

Author: sweatybridge

cmd/db.go

@@ -100,7 +100,7 @@ var (
 			differ := diff.DiffSchemaMigra
 			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 RLS policies, enums, and grants.")
+				fmt.Fprintln(os.Stderr, utils.Yellow("WARNING:"), "--use-pg-schema flag is experimental and may not include all entities, such as views and grants.")
 			}
 			return diff.Run(cmd.Context(), schema, file, flags.DbConfig, differ, afero.NewOsFs())
 		},

internal/db/diff/diff.go

@@ -89,16 +89,6 @@ func findDropStatements(out string) []string {
 	return drops
 }
 
-func loadSchema(ctx context.Context, config pgconn.Config, options ...func(*pgx.ConnConfig)) ([]string, error) {
-	conn, err := utils.ConnectByConfig(ctx, config, 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 CreateShadowDatabase(ctx context.Context, port uint16) (string, error) {
 	// Disable background workers in shadow database
 	config := start.NewContainerConfig("-c", "max_worker_processes=0")
@@ -178,12 +168,11 @@ func DiffDatabase(ctx context.Context, schema []string, config pgconn.Config, w
 		}
 	}
 	// Load all user defined schemas
-	if len(schema) == 0 {
-		if schema, err = loadSchema(ctx, config, options...); err != nil {
-			return "", err
-		}
+	if len(schema) > 0 {
+		fmt.Fprintln(w, "Diffing schemas:", strings.Join(schema, ","))
+	} else {
+		fmt.Fprintln(w, "Diffing schemas...")
 	}
-	fmt.Fprintln(w, "Diffing schemas:", strings.Join(schema, ","))
 	source := utils.ToPostgresURL(shadowConfig)
 	target := utils.ToPostgresURL(config)
 	return differ(ctx, source, target, schema)

internal/db/diff/diff_test.go

@@ -24,7 +24,6 @@ import (
 	"github.com/supabase/cli/internal/testing/helper"
 	"github.com/supabase/cli/internal/utils"
 	"github.com/supabase/cli/internal/utils/flags"
-	"github.com/supabase/cli/pkg/config"
 	"github.com/supabase/cli/pkg/migration"
 	"github.com/supabase/cli/pkg/pgtest"
 )
@@ -64,7 +63,7 @@ func TestRun(t *testing.T) {
 		require.NoError(t, apitest.MockDockerLogs(utils.Docker, "test-shadow-storage", ""))
 		apitest.MockDockerStart(utils.Docker, utils.GetRegistryImageUrl(utils.Config.Auth.Image), "test-shadow-auth")
 		require.NoError(t, apitest.MockDockerLogs(utils.Docker, "test-shadow-auth", ""))
-		apitest.MockDockerStart(utils.Docker, utils.GetRegistryImageUrl(config.Images.Migra), "test-migra")
+		apitest.MockDockerStart(utils.Docker, utils.GetRegistryImageUrl(utils.Config.EdgeRuntime.Image), "test-migra")
 		diff := "create table test();"
 		require.NoError(t, apitest.MockDockerLogs(utils.Docker, "test-migra", diff))
 		// Setup mock postgres
@@ -285,7 +284,7 @@ create schema public`)
 		gock.New(utils.Docker.DaemonHost()).
 			Delete("/v" + utils.Docker.ClientVersion() + "/containers/test-shadow-db").
 			Reply(http.StatusOK)
-		apitest.MockDockerStart(utils.Docker, utils.GetRegistryImageUrl(config.Images.Migra), "test-migra")
+		apitest.MockDockerStart(utils.Docker, utils.GetRegistryImageUrl(utils.Config.EdgeRuntime.Image), "test-migra")
 		gock.New(utils.Docker.DaemonHost()).
 			Get("/v" + utils.Docker.ClientVersion() + "/containers/test-migra/logs").
 			ReplyError(errors.New("network error"))

internal/db/diff/migra.go

@@ -9,15 +9,52 @@ import (
 	"github.com/docker/docker/api/types/container"
 	"github.com/docker/docker/api/types/network"
 	"github.com/go-errors/errors"
+	"github.com/spf13/viper"
 	"github.com/supabase/cli/internal/utils"
 	"github.com/supabase/cli/pkg/config"
 )
 
-//go:embed templates/migra.sh
-var diffSchemaScript string
+var (
+	//go:embed templates/migra.sh
+	diffSchemaScript string
+	//go:embed templates/migra.ts
+	diffSchemaTypeScript string
+
+	managedSchemas = []string{
+		// Local development
+		"_analytics",
+		"_realtime",
+		"_supavisor",
+		// Owned by extensions
+		"cron",
+		"graphql",
+		"graphql_public",
+		"net",
+		"pgroonga",
+		"pgtle",
+		"repack",
+		"tiger_data",
+		"vault",
+		// Deprecated extensions
+		"pgsodium",
+		"pgsodium_masks",
+		"timescaledb_experimental",
+		"timescaledb_information",
+		"_timescaledb_cache",
+		"_timescaledb_catalog",
+		"_timescaledb_config",
+		"_timescaledb_debug",
+		"_timescaledb_functions",
+		"_timescaledb_internal",
+		// Managed by Supabase
+		"pgbouncer",
+		"supabase_functions",
+		"supabase_migrations",
+	}
+)
 
 // Diffs local database schema against shadow, dumps output to stdout.
-func DiffSchemaMigra(ctx context.Context, source, target string, schema []string) (string, error) {
+func DiffSchemaMigraBash(ctx context.Context, source, target string, schema []string) (string, error) {
 	env := []string{"SOURCE=" + source, "TARGET=" + target}
 	// Passing in script string means command line args must be set manually, ie. "$@"
 	args := "set -- " + strings.Join(schema, " ") + ";"
@@ -42,3 +79,41 @@ func DiffSchemaMigra(ctx context.Context, source, target string, schema []string
 	}
 	return out.String(), nil
 }
+
+func DiffSchemaMigra(ctx context.Context, source, target string, schema []string) (string, error) {
+	env := []string{"SOURCE=" + source, "TARGET=" + target}
+	if len(schema) > 0 {
+		env = append(env, "INCLUDED_SCHEMAS="+strings.Join(schema, ","))
+	} else {
+		env = append(env, "EXCLUDED_SCHEMAS="+strings.Join(managedSchemas, ","))
+	}
+	cmd := []string{"edge-runtime", "start", "--main-service=."}
+	if viper.GetBool("DEBUG") {
+		cmd = append(cmd, "--verbose")
+	}
+	cmdString := strings.Join(cmd, " ")
+	entrypoint := []string{"sh", "-c", `cat <<'EOF' > index.ts && ` + cmdString + `
+` + diffSchemaTypeScript + `
+EOF
+`}
+	var out, stderr bytes.Buffer
+	if err := utils.DockerRunOnceWithConfig(
+		ctx,
+		container.Config{
+			Image:      utils.Config.EdgeRuntime.Image,
+			Env:        env,
+			Entrypoint: entrypoint,
+		},
+		container.HostConfig{
+			Binds:       []string{utils.EdgeRuntimeId + ":/root/.cache/deno:rw"},
+			NetworkMode: network.NetworkHost,
+		},
+		network.NetworkingConfig{},
+		"",
+		&out,
+		&stderr,
+	); err != nil && !strings.HasPrefix(stderr.String(), "main worker has been destroyed") {
+		return "", errors.Errorf("error diffing schema: %w:\n%s", err, stderr.String())
+	}
+	return out.String(), nil
+}

internal/db/diff/pgschema.go

@@ -22,12 +22,24 @@ func DiffPgSchema(ctx context.Context, source, target string, schema []string) (
 	}
 	defer dbDst.Close()
 	// Generate DDL based on schema plan
+	opts := []pgschema.PlanOpt{pgschema.WithDoNotValidatePlan()}
+	if len(schema) > 0 {
+		opts = append(opts, pgschema.WithIncludeSchemas(schema...))
+	} else {
+		opts = append(opts,
+			pgschema.WithExcludeSchemas(managedSchemas...),
+			pgschema.WithExcludeSchemas(
+				"topology", // unsupported due to views
+				"realtime", // unsupported due to partitioned table
+				"storage",  // unsupported due to unique index
+			),
+		)
+	}
 	plan, err := pgschema.Generate(
 		ctx,
 		pgschema.DBSchemaSource(dbSrc),
 		pgschema.DBSchemaSource(dbDst),
-		pgschema.WithDoNotValidatePlan(),
-		pgschema.WithIncludeSchemas(schema...),
+		opts...,
 	)
 	if err != nil {
 		return "", errors.Errorf("failed to generate plan: %w", err)

internal/db/diff/templates/migra.ts

@@ -0,0 +1,73 @@
+import { createClient } from "npm:@pgkit/client";
+import { Migration } from "npm:@pgkit/migra";
+
+const clientBase = createClient(Deno.env.get("SOURCE"));
+const clientHead = createClient(Deno.env.get("TARGET"));
+const includedSchemas = Deno.env.get("INCLUDED_SCHEMAS")?.split(",") ?? [];
+const excludedSchemas = Deno.env.get("EXCLUDED_SCHEMAS")?.split(",") ?? [];
+
+const managedSchemas = ["auth", "realtime", "storage"];
+const extensionSchemas = [
+  "pg_catalog",
+  "extensions",
+  "pgmq",
+  "tiger",
+  "topology",
+];
+
+try {
+  let sql = "";
+  for (const schema of includedSchemas) {
+    const m = await Migration.create(clientBase, clientHead, {
+      schema,
+      ignore_extension_versions: true,
+    });
+    m.set_safety(false);
+    m.add_all_changes(true);
+    sql += m.sql;
+  }
+  if (includedSchemas.length === 0) {
+    // Migra does not ignore custom types and triggers created by extensions, so we diff
+    // them separately. This workaround only applies to a known list of managed schemas.
+    for (const schema of extensionSchemas) {
+      const e = await Migration.create(clientBase, clientHead, {
+        schema,
+        ignore_extension_versions: true,
+      });
+      e.set_safety(false);
+      e.add(e.changes.schemas({ creations_only: true }));
+      e.add_extension_changes();
+      sql += e.sql;
+    }
+    // Diff user defined entities in non-managed schemas, including extensions.
+    const m = await Migration.create(clientBase, clientHead, {
+      exclude_schema: [
+        ...managedSchemas,
+        ...extensionSchemas,
+        ...excludedSchemas,
+      ],
+      ignore_extension_versions: true,
+    });
+    m.set_safety(false);
+    m.add_all_changes(true);
+    sql += m.sql;
+    // For managed schemas, we want to include triggers and RLS policies only.
+    for (const schema of managedSchemas) {
+      const s = await Migration.create(clientBase, clientHead, {
+        schema,
+        ignore_extension_versions: true,
+      });
+      s.set_safety(false);
+      s.add(s.changes.triggers({ drops_only: true }));
+      s.add(s.changes.rlspolicies({ drops_only: true }));
+      s.add(s.changes.rlspolicies({ creations_only: true }));
+      s.add(s.changes.triggers({ creations_only: true }));
+      sql += s.sql;
+    }
+  }
+  console.log(sql);
+} catch (e) {
+  console.error(e);
+} finally {
+  await Promise.all([clientHead.end(), clientBase.end()]);
+}