feat: allow changing extension schema

Current logic doesn't capture the `ALTER EXTENSION .. SET SCHEMA` statement. This uses the `pageinspect` contrib extension for testing. Closes #88.
supabase · Oct 12, 2024 · 9543b47 · 9543b47
1 parent 6e40398
commit 9543b47
Show file tree

Hide file tree

Showing 6 changed files with 84 additions and 5 deletions.
diff --git a/nix/withTmpDb.sh.in b/nix/withTmpDb.sh.in
@@ -17,7 +17,7 @@ options="-F -c listen_addresses=\"\" -k $PGDATA -c shared_preload_libraries=\"pg
 
 reserved_roles="supabase_storage_admin, anon, reserved_but_not_yet_created, authenticator*"
 reserved_memberships="pg_read_server_files, pg_write_server_files, pg_execute_server_program, role_with_reserved_membership"
-privileged_extensions="autoinc, citext, hstore, sslinfo, pg_tle, postgres_fdw"
+privileged_extensions="autoinc, citext, hstore, sslinfo, pg_tle, postgres_fdw, pageinspect"
 privileged_extensions_custom_scripts_path="$tmpdir/privileged_extensions_custom_scripts"
 privileged_role="privileged_role"
 privileged_role_allowed_configs="session_replication_role, pgrst.*, other.nested.*"

diff --git a/src/privileged_extensions.c b/src/privileged_extensions.c
@@ -272,11 +272,11 @@ void handle_create_extension(
 
 void handle_alter_extension(
     void (*process_utility_hook)(PROCESS_UTILITY_PARAMS),
-    PROCESS_UTILITY_PARAMS, const char *privileged_extensions,
+    PROCESS_UTILITY_PARAMS,
+    const char *extname, const char *privileged_extensions,
     const char *privileged_extensions_superuser) {
-    AlterExtensionStmt *stmt = (AlterExtensionStmt *)pstmt->utilityStmt;
 
-    if (is_string_in_comma_delimited_string(stmt->extname,
+    if (is_string_in_comma_delimited_string(extname,
                                             privileged_extensions)) {
         bool already_switched_to_superuser = false;
         switch_to_superuser(privileged_extensions_superuser,

diff --git a/src/privileged_extensions.h b/src/privileged_extensions.h
@@ -14,6 +14,7 @@ extern void handle_create_extension(
 extern void
 handle_alter_extension(void (*process_utility_hook)(PROCESS_UTILITY_PARAMS),
                        PROCESS_UTILITY_PARAMS,
+                       const char *extname,
                        const char *privileged_extensions,
                        const char *privileged_extensions_superuser);
 

diff --git a/src/supautils.c b/src/supautils.c
@@ -623,7 +623,7 @@ supautils_hook(PROCESS_UTILITY_PARAMS)
         }
 
 		/*
-		 * ALTER EXTENSION <extension> [ UPDATE | SET SCHEMA ]
+		 * ALTER EXTENSION <extension> [ ADD | DROP | UPDATE ]
 		 */
 		case T_AlterExtensionStmt:
 		{
@@ -634,13 +634,41 @@ supautils_hook(PROCESS_UTILITY_PARAMS)
 				break;
 			}
 
+			AlterExtensionStmt *stmt = (AlterExtensionStmt *)pstmt->utilityStmt;
+
 			handle_alter_extension(prev_hook,
 								   PROCESS_UTILITY_ARGS,
+								   stmt->extname,
 								   privileged_extensions,
 								   privileged_extensions_superuser);
 			return;
 		}
 
+		/*
+		 * ALTER EXTENSION <extension> SET SCHEMA
+		 */
+		case T_AlterObjectSchemaStmt:
+		{
+			if (superuser()) {
+				break;
+			}
+			if (privileged_extensions == NULL) {
+				break;
+			}
+
+			AlterObjectSchemaStmt *stmt = (AlterObjectSchemaStmt *)pstmt->utilityStmt;
+
+			if (stmt->objectType == OBJECT_EXTENSION){
+				handle_alter_extension(prev_hook,
+									   PROCESS_UTILITY_ARGS,
+									   strVal(stmt->object),
+									   privileged_extensions,
+									   privileged_extensions_superuser);
+			}
+
+			return;
+		}
+
 		/*
 		 * ALTER EXTENSION <extension> [ ADD | DROP ]
 		 *

diff --git a/test/expected/privileged_extensions.out b/test/expected/privileged_extensions.out
@@ -104,3 +104,31 @@ drop role another_superuser;
 set role extensions_role;
 \echo
 
+-- can change extensions schema
+create extension pageinspect;
+select count(*) = 3 as extensions_in_public_schema
+from information_schema.routines
+where routine_name in ('page_header', 'heap_page_items', 'bt_metap')
+and routine_schema = 'public';
+ extensions_in_public_schema 
+-----------------------------
+ t
+(1 row)
+
+-- go back to postgres role for creating a new schema and switch to extensions_role again
+reset role;
+create schema xtens;
+set role extensions_role;
+\echo
+
+-- now alter extension schema
+alter extension pageinspect set schema xtens;
+select count(*) = 3 as extensions_in_xtens_schema
+from information_schema.routines
+where routine_name in ('page_header', 'heap_page_items', 'bt_metap')
+and routine_schema = 'xtens';
+ extensions_in_xtens_schema 
+----------------------------
+ t
+(1 row)
+
diff --git a/test/sql/privileged_extensions.sql b/test/sql/privileged_extensions.sql
@@ -67,3 +67,25 @@ drop extension sslinfo;
 drop role another_superuser;
 set role extensions_role;
 \echo
+
+-- can change extensions schema
+create extension pageinspect;
+
+select count(*) = 3 as extensions_in_public_schema
+from information_schema.routines
+where routine_name in ('page_header', 'heap_page_items', 'bt_metap')
+and routine_schema = 'public';
+
+-- go back to postgres role for creating a new schema and switch to extensions_role again
+reset role;
+create schema xtens;
+set role extensions_role;
+\echo
+
+-- now alter extension schema
+alter extension pageinspect set schema xtens;
+
+select count(*) = 3 as extensions_in_xtens_schema
+from information_schema.routines
+where routine_name in ('page_header', 'heap_page_items', 'bt_metap')
+and routine_schema = 'xtens';