fix: simplify query for user defined schemas (#2206)

internal/db/diff/diff.go

@@ -83,24 +83,8 @@ func loadSchema(ctx context.Context, config pgconn.Config, options ...func(*pgx.
 		return nil, err
 	}
 	defer conn.Close(context.Background())
-	return LoadUserSchemas(ctx, conn)
-}
-
-func LoadUserSchemas(ctx context.Context, conn *pgx.Conn) ([]string, error) {
 	// RLS policies in auth and storage schemas can be included with -s flag
-	exclude := append([]string{
-		"auth",
-		// "extensions",
-		"pgbouncer",
-		"realtime",
-		"_realtime",
-		"storage",
-		"_analytics",
-		// Exclude functions because Webhooks support is early alpha
-		"supabase_functions",
-		"supabase_migrations",
-	}, utils.SystemSchemas...)
-	return reset.ListSchemas(ctx, conn, exclude...)
+	return reset.LoadUserSchemas(ctx, conn)
 }
 
 func CreateShadowDatabase(ctx context.Context) (string, error) {

internal/db/diff/diff_test.go

@@ -35,28 +35,10 @@ var dbConfig = pgconn.Config{
 }
 
 var escapedSchemas = []string{
-	"auth",
 	"pgbouncer",
-	"realtime",
-	`\_realtime`,
-	"storage",
-	`\_analytics`,
-	`supabase\_functions`,
-	`supabase\_migrations`,
-	"cron",
-	"dbdev",
-	"graphql",
-	`graphql\_public`,
-	"net",
 	"pgsodium",
-	`pgsodium\_masks`,
 	"pgtle",
-	"repack",
-	"tiger",
-	`tiger\_data`,
-	`timescaledb\_%`,
-	`\_timescaledb\_%`,
-	"topology",
+	`supabase\_migrations`,
 	"vault",
 	`information\_schema`,
 	`pg\_%`,
@@ -130,7 +112,7 @@ func TestRun(t *testing.T) {
 		// Setup mock postgres
 		conn := pgtest.NewConn()
 		defer conn.Close(t)
-		conn.Query(reset.LIST_SCHEMAS, escapedSchemas).
+		conn.Query(reset.ListSchemas, escapedSchemas).
 			ReplyError(pgerrcode.DuplicateTable, `relation "test" already exists`)
 		// Run test
 		err := Run(context.Background(), []string{}, "", dbConfig, DiffSchemaMigra, fsys, conn.Intercept)
@@ -356,24 +338,6 @@ At statement 0: create schema public`)
 	})
 }
 
-func TestUserSchema(t *testing.T) {
-	// Setup mock postgres
-	conn := pgtest.NewConn()
-	defer conn.Close(t)
-	conn.Query(reset.LIST_SCHEMAS, escapedSchemas).
-		Reply("SELECT 1", []interface{}{"test"})
-	// Connect to mock
-	ctx := context.Background()
-	mock, err := utils.ConnectByConfig(ctx, dbConfig, conn.Intercept)
-	require.NoError(t, err)
-	defer mock.Close(ctx)
-	// Run test
-	schemas, err := LoadUserSchemas(ctx, mock)
-	// Check error
-	assert.NoError(t, err)
-	assert.ElementsMatch(t, []string{"test"}, schemas)
-}
-
 func TestDropStatements(t *testing.T) {
 	drops := findDropStatements("create table t(); drop table t; alter table t drop column c")
 	assert.Equal(t, []string{"drop table t", "alter table t drop column c"}, drops)

internal/db/dump/templates/dump_schema.sh

@@ -18,7 +18,7 @@ export PGDATABASE="$PGDATABASE"
 #   - do not include ACL changes on internal schemas
 #   - do not include RLS policies on cron extension schema
 #   - do not include event triggers
-#   - do not include creating publication "supabase_realtime"
+#   - do not create publication "supabase_realtime"
 pg_dump \
     --schema-only \
     --quote-all-identifier \

internal/db/lint/lint.go

@@ -93,7 +93,7 @@ func LintDatabase(ctx context.Context, conn *pgx.Conn, schema []string) ([]Resul
 		return nil, errors.Errorf("failed to begin transaction: %w", err)
 	}
 	if len(schema) == 0 {
-		schema, err = reset.ListSchemas(ctx, conn, utils.InternalSchemas...)
+		schema, err = reset.LoadUserSchemas(ctx, conn)
 		if err != nil {
 			return nil, err
 		}

internal/db/pull/pull.go

@@ -14,6 +14,7 @@ import (
 	"github.com/spf13/afero"
 	"github.com/supabase/cli/internal/db/diff"
 	"github.com/supabase/cli/internal/db/dump"
+	"github.com/supabase/cli/internal/db/reset"
 	"github.com/supabase/cli/internal/migration/list"
 	"github.com/supabase/cli/internal/migration/new"
 	"github.com/supabase/cli/internal/migration/repair"
@@ -78,7 +79,7 @@ func run(p utils.Program, ctx context.Context, schema []string, path string, con
 	defaultSchema := len(schema) == 0
 	if defaultSchema {
 		var err error
-		schema, err = diff.LoadUserSchemas(ctx, conn)
+		schema, err = reset.LoadUserSchemas(ctx, conn)
 		if err != nil {
 			return err
 		}

internal/db/pull/pull_test.go

@@ -31,28 +31,10 @@ var dbConfig = pgconn.Config{
 }
 
 var escapedSchemas = []string{
-	"auth",
 	"pgbouncer",
-	"realtime",
-	`\_realtime`,
-	"storage",
-	`\_analytics`,
-	`supabase\_functions`,
-	`supabase\_migrations`,
-	"cron",
-	"dbdev",
-	"graphql",
-	`graphql\_public`,
-	"net",
 	"pgsodium",
-	`pgsodium\_masks`,
 	"pgtle",
-	"repack",
-	"tiger",
-	`tiger\_data`,
-	`timescaledb\_%`,
-	`\_timescaledb\_%`,
-	"topology",
+	`supabase\_migrations`,
 	"vault",
 	`information\_schema`,
 	`pg\_%`,
@@ -196,7 +178,7 @@ func TestPullSchema(t *testing.T) {
 		defer conn.Close(t)
 		conn.Query(list.LIST_MIGRATION_VERSION).
 			Reply("SELECT 1", []interface{}{"0"}).
-			Query(reset.LIST_SCHEMAS, escapedSchemas).
+			Query(reset.ListSchemas, escapedSchemas).
 			ReplyError(pgerrcode.DuplicateTable, `relation "test" already exists`)
 		// Connect to mock
 		ctx := context.Background()

internal/db/remote/changes/changes.go

@@ -7,6 +7,7 @@ import (
 	"github.com/jackc/pgconn"
 	"github.com/spf13/afero"
 	"github.com/supabase/cli/internal/db/diff"
+	"github.com/supabase/cli/internal/db/reset"
 	"github.com/supabase/cli/internal/utils"
 )
 
@@ -53,5 +54,5 @@ func loadSchema(ctx context.Context, config pgconn.Config, w io.Writer) ([]strin
 		return nil, err
 	}
 	defer conn.Close(context.Background())
-	return diff.LoadUserSchemas(ctx, conn)
+	return reset.LoadUserSchemas(ctx, conn)
 }

internal/db/remote/commit/commit.go

@@ -12,6 +12,7 @@ import (
 	"github.com/spf13/afero"
 	"github.com/supabase/cli/internal/db/diff"
 	"github.com/supabase/cli/internal/db/dump"
+	"github.com/supabase/cli/internal/db/reset"
 	"github.com/supabase/cli/internal/migration/list"
 	"github.com/supabase/cli/internal/migration/repair"
 	"github.com/supabase/cli/internal/utils"
@@ -52,7 +53,7 @@ func run(p utils.Program, ctx context.Context, schema []string, config pgconn.Co
 
 	// 2. Fetch remote schema changes
 	if len(schema) == 0 {
-		schema, err = diff.LoadUserSchemas(ctx, conn)
+		schema, err = reset.LoadUserSchemas(ctx, conn)
 		if err != nil {
 			return err
 		}

internal/db/reset/reset.go

@@ -28,16 +28,13 @@ import (
 	"github.com/supabase/cli/internal/utils/pgxv5"
 )
 
-const (
-	SET_POSTGRES_ROLE = "SET ROLE postgres;"
-	LIST_SCHEMAS      = "SELECT schema_name FROM information_schema.schemata WHERE NOT schema_name LIKE ANY($1) ORDER BY schema_name"
-)
-
 var (
 	ErrUnhealthy   = errors.New("service not healthy")
 	serviceTimeout = 30 * time.Second
 	//go:embed templates/drop.sql
 	dropObjects string
+	//go:embed templates/list.sql
+	ListSchemas string
 )
 
 func Run(ctx context.Context, version string, config pgconn.Config, fsys afero.Fs, options ...func(*pgx.ConnConfig)) error {
@@ -264,39 +261,41 @@ func resetRemote(ctx context.Context, version string, config pgconn.Config, fsys
 		return err
 	}
 	defer conn.Close(context.Background())
-	// List user defined schemas
-	excludes := []string{"public"}
-	for _, schema := range utils.InternalSchemas {
-		if schema != "supabase_migrations" {
-			excludes = append(excludes, schema)
-		}
-	}
-	userSchemas, err := ListSchemas(ctx, conn, excludes...)
+	// Only drop objects in extensions and public schema
+	excludes := append([]string{
+		"extensions",
+		"public",
+	}, utils.ManagedSchemas...)
+	userSchemas, err := LoadUserSchemas(ctx, conn, excludes...)
 	if err != nil {
 		return err
 	}
-	// Drop user defined objects
+	// Drop all user defined schemas
 	migration := repair.MigrationFile{}
 	for _, schema := range userSchemas {
 		sql := fmt.Sprintf("DROP SCHEMA IF EXISTS %s CASCADE", schema)
 		migration.Lines = append(migration.Lines, sql)
 	}
+	// If an extension uses a schema it doesn't create, dropping the schema will cascade to also
+	// drop the extension. But if an extension creates its own schema, dropping the schema will
+	// throw an error. Hence, we drop the extension instead so it cascades to its own schema.
 	migration.Lines = append(migration.Lines, dropObjects)
 	if err := migration.ExecBatch(ctx, conn); err != nil {
 		return err
 	}
 	return apply.MigrateAndSeed(ctx, version, conn, fsys)
 }
 
-func ListSchemas(ctx context.Context, conn *pgx.Conn, exclude ...string) ([]string, error) {
-	exclude = LikeEscapeSchema(exclude)
+func LoadUserSchemas(ctx context.Context, conn *pgx.Conn, exclude ...string) ([]string, error) {
 	if len(exclude) == 0 {
-		exclude = append(exclude, "")
+		exclude = utils.ManagedSchemas
 	}
-	rows, err := conn.Query(ctx, LIST_SCHEMAS, exclude)
+	exclude = LikeEscapeSchema(exclude)
+	rows, err := conn.Query(ctx, ListSchemas, exclude)
 	if err != nil {
 		return nil, errors.Errorf("failed to list schemas: %w", err)
 	}
+	// TODO: show detail and hint from pgconn.PgError
 	return pgxv5.CollectStrings(rows)
 }
 

internal/db/reset/reset_test.go

@@ -304,29 +304,12 @@ func TestRestartDatabase(t *testing.T) {
 }
 
 var escapedSchemas = []string{
-	"public",
-	"auth",
 	"extensions",
+	"public",
 	"pgbouncer",
-	"realtime",
-	`\_realtime`,
-	"storage",
-	`\_analytics`,
-	`supabase\_functions`,
-	"cron",
-	"dbdev",
-	"graphql",
-	`graphql\_public`,
-	"net",
 	"pgsodium",
-	`pgsodium\_masks`,
 	"pgtle",
-	"repack",
-	"tiger",
-	`tiger\_data`,
-	`timescaledb\_%`,
-	`\_timescaledb\_%`,
-	"topology",
+	`supabase\_migrations`,
 	"vault",
 	`information\_schema`,
 	`pg\_%`,
@@ -347,7 +330,7 @@ func TestResetRemote(t *testing.T) {
 		// Setup mock postgres
 		conn := pgtest.NewConn()
 		defer conn.Close(t)
-		conn.Query(LIST_SCHEMAS, escapedSchemas).
+		conn.Query(ListSchemas, escapedSchemas).
 			Reply("SELECT 1", []interface{}{"private"}).
 			Query("DROP SCHEMA IF EXISTS private CASCADE").
 			Reply("DROP SCHEMA").
@@ -374,7 +357,7 @@ func TestResetRemote(t *testing.T) {
 		// Setup mock postgres
 		conn := pgtest.NewConn()
 		defer conn.Close(t)
-		conn.Query(LIST_SCHEMAS, escapedSchemas).
+		conn.Query(ListSchemas, escapedSchemas).
 			ReplyError(pgerrcode.InsufficientPrivilege, "permission denied for relation information_schema")
 		// Run test
 		err := resetRemote(context.Background(), "", dbConfig, fsys, conn.Intercept)
@@ -388,7 +371,7 @@ func TestResetRemote(t *testing.T) {
 		// Setup mock postgres
 		conn := pgtest.NewConn()
 		defer conn.Close(t)
-		conn.Query(LIST_SCHEMAS, escapedSchemas).
+		conn.Query(ListSchemas, escapedSchemas).
 			Reply("SELECT 0").
 			Query(dropObjects).
 			ReplyError(pgerrcode.InsufficientPrivilege, "permission denied for relation supabase_migrations")

internal/db/reset/templates/drop.sql

@@ -33,12 +33,12 @@ begin
     execute format('drop table if exists %I.%I cascade', rec.relnamespace::regnamespace::name, rec.relname);
   end loop;
 
-  -- truncate auth tables
+  -- truncate tables in auth and migrations schema
   for rec in
     select *
     from pg_class c
     where
-      c.relnamespace::regnamespace::name = 'auth'
+      c.relnamespace::regnamespace::name in ('auth', 'supabase_migrations')
       and c.relkind = 'r'
   loop
     execute format('truncate %I.%I restart identity cascade', rec.relnamespace::regnamespace::name, rec.relname);

internal/db/reset/templates/list.sql

@@ -0,0 +1,13 @@
+-- List user defined schemas, excluding
+--  Extension created schemas
+--  Supabase managed schemas
+select pn.nspname
+from pg_namespace pn
+left join pg_depend pd
+  on pd.objid = pn.oid
+join pg_roles r 
+  on pn.nspowner = r.oid
+where pd.deptype is null
+  and not pn.nspname like any($1)
+  and r.rolname != 'supabase_admin'
+order by pn.nspname