fix: escape drop schema statements (#3877)

Author: sweatybridge

internal/migration/down/down_test.go

@@ -76,8 +76,6 @@ func TestMigrationsDown(t *testing.T) {
 	})
 }
 
-var escapedSchemas = append(migration.ManagedSchemas, "extensions", "public")
-
 func TestResetRemote(t *testing.T) {
 	t.Run("resets remote database", func(t *testing.T) {
 		// Setup in-memory fs
@@ -87,11 +85,7 @@ func TestResetRemote(t *testing.T) {
 		// Setup mock postgres
 		conn := pgtest.NewConn()
 		defer conn.Close(t)
-		conn.Query(migration.ListSchemas, escapedSchemas).
-			Reply("SELECT 1", []interface{}{"private"}).
-			Query("DROP SCHEMA IF EXISTS private CASCADE").
-			Reply("DROP SCHEMA").
-			Query(migration.DropObjects).
+		conn.Query(migration.DropObjects).
 			Reply("INSERT 0")
 		helper.MockMigrationHistory(conn).
 			Query(migration.INSERT_MIGRATION_VERSION, "0", "schema", nil).
@@ -113,11 +107,7 @@ func TestResetRemote(t *testing.T) {
 		// Setup mock postgres
 		conn := pgtest.NewConn()
 		defer conn.Close(t)
-		conn.Query(migration.ListSchemas, escapedSchemas).
-			Reply("SELECT 1", []interface{}{"private"}).
-			Query("DROP SCHEMA IF EXISTS private CASCADE").
-			Reply("DROP SCHEMA").
-			Query(migration.DropObjects).
+		conn.Query(migration.DropObjects).
 			Reply("INSERT 0")
 		helper.MockMigrationHistory(conn).
 			Query(migration.INSERT_MIGRATION_VERSION, "0", "schema", nil).
@@ -135,9 +125,7 @@ func TestResetRemote(t *testing.T) {
 		// Setup mock postgres
 		conn := pgtest.NewConn()
 		defer conn.Close(t)
-		conn.Query(migration.ListSchemas, escapedSchemas).
-			Reply("SELECT 0").
-			Query(migration.DropObjects).
+		conn.Query(migration.DropObjects).
 			ReplyError(pgerrcode.InsufficientPrivilege, "permission denied for relation supabase_migrations")
 		// Run test
 		err := ResetAll(context.Background(), "", conn.MockClient(t), fsys)

pkg/migration/drop.go

@@ -3,7 +3,6 @@ package migration
 import (
 	"context"
 	_ "embed"
-	"fmt"
 
 	"github.com/go-errors/errors"
 	"github.com/jackc/pgx/v4"
@@ -33,24 +32,7 @@ var (
 )
 
 func DropUserSchemas(ctx context.Context, conn *pgx.Conn) error {
-	// Only drop objects in extensions and public schema
-	excludes := append(ManagedSchemas,
-		"extensions",
-		"public",
-	)
-	userSchemas, err := ListUserSchemas(ctx, conn, excludes...)
-	if err != nil {
-		return err
-	}
-	// Drop all user defined schemas
 	migration := MigrationFile{}
-	for _, schema := range userSchemas {
-		sql := fmt.Sprintf("DROP SCHEMA IF EXISTS %s CASCADE", schema)
-		migration.Statements = append(migration.Statements, 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.Statements = append(migration.Statements, DropObjects)
 	return migration.ExecBatch(ctx, conn)
 }

pkg/migration/drop_test.go

@@ -9,44 +9,24 @@ import (
 	"github.com/supabase/cli/pkg/pgtest"
 )
 
-var escapedSchemas = append(ManagedSchemas, "extensions", "public")
-
 func TestDropSchemas(t *testing.T) {
 	t.Run("resets remote database", func(t *testing.T) {
 		// Setup mock postgres
 		conn := pgtest.NewConn()
 		defer conn.Close(t)
-		conn.Query(ListSchemas, escapedSchemas).
-			Reply("SELECT 1", []interface{}{"private"}).
-			Query("DROP SCHEMA IF EXISTS private CASCADE").
-			Reply("DROP SCHEMA").
-			Query(DropObjects).
+		conn.Query(DropObjects).
 			Reply("INSERT 0")
 		// Run test
 		err := DropUserSchemas(context.Background(), conn.MockClient(t))
 		// Check error
 		assert.NoError(t, err)
 	})
 
-	t.Run("throws error on list schema failure", func(t *testing.T) {
-		// Setup mock postgres
-		conn := pgtest.NewConn()
-		defer conn.Close(t)
-		conn.Query(ListSchemas, escapedSchemas).
-			ReplyError(pgerrcode.InsufficientPrivilege, "permission denied for relation information_schema")
-		// Run test
-		err := DropUserSchemas(context.Background(), conn.MockClient(t))
-		// Check error
-		assert.ErrorContains(t, err, "ERROR: permission denied for relation information_schema (SQLSTATE 42501)")
-	})
-
 	t.Run("throws error on drop schema failure", func(t *testing.T) {
 		// Setup mock postgres
 		conn := pgtest.NewConn()
 		defer conn.Close(t)
-		conn.Query(ListSchemas, escapedSchemas).
-			Reply("SELECT 0").
-			Query(DropObjects).
+		conn.Query(DropObjects).
 			ReplyError(pgerrcode.InsufficientPrivilege, "permission denied for relation supabase_migrations")
 		// Run test
 		err := DropUserSchemas(context.Background(), conn.MockClient(t))

pkg/migration/queries/drop.sql

@@ -1,6 +1,21 @@
 do $$ declare
   rec record;
 begin
+  -- schemas
+  for rec in
+    select pn.*
+    from pg_namespace pn
+    left join pg_depend pd on pd.objid = pn.oid
+    where pd.deptype is null
+      and not pn.nspname like any(array['information\_schema', 'pg\_%', '\_analytics', '\_realtime', '\_supavisor', 'pgbouncer', 'pgmq', 'pgsodium', 'pgtle', 'supabase\_migrations', 'vault', 'extensions', 'public'])
+      and pn.nspowner::regrole::text != 'supabase_admin'
+  loop
+    -- 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 schemas first while excluding those created by extensions.
+    execute format('drop schema if exists %I cascade', rec.nspname);
+  end loop;
+
   -- extensions
   for rec in
     select *

pkg/migration/queries/list.sql

@@ -3,11 +3,8 @@
 --  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
+left join pg_depend pd on pd.objid = pn.oid
 where pd.deptype is null
   and not pn.nspname like any($1)
-  and r.rolname != 'supabase_admin'
+  and pn.nspowner::regrole::text != 'supabase_admin'
 order by pn.nspname