nocfbot.changes.diff

diff --git a/src/backend/commands/extension.c b/src/backend/commands/extension.c
index 77bc6a82987..03b3de357bd 100644
--- a/src/backend/commands/extension.c
+++ b/src/backend/commands/extension.c
@@ -1243,6 +1243,44 @@ extension_is_trusted(ExtensionControlFile *control)
 	return false;
 }
 
+/*
+ * Enforce superuser requirements for extension operations.
+ *
+ * Returns true if we should switch to superuser for trusted extensions.
+ * Throws an error if superuser is required but not available.
+ *
+ * This function should only be called after choosing the appropriate
+ * control file (including any secondary control files for updates).
+ */
+static bool
+check_extension_superuser_requirements(ExtensionControlFile *control,
+									   const char *from_version)
+{
+	if (control->superuser && !superuser())
+	{
+		if (extension_is_trusted(control))
+			return true;
+		else if (from_version == NULL)
+			ereport(ERROR,
+					(errcode(ERRCODE_INSUFFICIENT_PRIVILEGE),
+					 errmsg("permission denied to create extension \"%s\"",
+							control->name),
+					 control->trusted
+					 ? errhint("Must have CREATE privilege on current database to create this extension.")
+					 : errhint("Must be superuser to create this extension.")));
+		else
+			ereport(ERROR,
+					(errcode(ERRCODE_INSUFFICIENT_PRIVILEGE),
+					 errmsg("permission denied to update extension \"%s\"",
+							control->name),
+					 control->trusted
+					 ? errhint("Must have CREATE privilege on current database to update this extension.")
+					 : errhint("Must be superuser to update this extension.")));
+	}
+
+	return false;
+}
+
 /*
  * Execute the appropriate script file for installing or updating the extension
  *
@@ -1271,27 +1309,7 @@ execute_extension_script(Oid extensionOid, ExtensionControlFile *control,
 	 * here so that the control flags are correctly associated with the right
 	 * script(s) if they happen to be set in secondary control files.
 	 */
-	if (control->superuser && !superuser())
-	{
-		if (extension_is_trusted(control))
-			switch_to_superuser = true;
-		else if (from_version == NULL)
-			ereport(ERROR,
-					(errcode(ERRCODE_INSUFFICIENT_PRIVILEGE),
-					 errmsg("permission denied to create extension \"%s\"",
-							control->name),
-					 control->trusted
-					 ? errhint("Must have CREATE privilege on current database to create this extension.")
-					 : errhint("Must be superuser to create this extension.")));
-		else
-			ereport(ERROR,
-					(errcode(ERRCODE_INSUFFICIENT_PRIVILEGE),
-					 errmsg("permission denied to update extension \"%s\"",
-							control->name),
-					 control->trusted
-					 ? errhint("Must have CREATE privilege on current database to update this extension.")
-					 : errhint("Must be superuser to update this extension.")));
-	}
+	switch_to_superuser = check_extension_superuser_requirements(control, from_version);
 
 	filename = get_extension_script_filename(control, from_version, version);
 
@@ -1854,7 +1872,7 @@ CreateSchemaForExtension(char *schemaName)
 }
 
 /*
- * Create a owned schema with the given name, as part of CREATE EXTENSION, and
+ * Create an owned schema with the given name, as part of CREATE EXTENSION, and
  * fails if the schema already exist.
  */
 static Oid
@@ -2017,6 +2035,9 @@ CreateExtensionInternal(char *extensionName,
 	Oid			extensionOid;
 	ObjectAddress address;
 	ListCell   *lc;
+	bool		switch_to_superuser = false;
+	Oid			save_userid = 0;
+	int			save_sec_context = 0;
 
 	/*
 	 * Read the primary control file.  Note we assume that it does not contain
@@ -2084,8 +2105,27 @@ CreateExtensionInternal(char *extensionName,
 	 */
 	control = read_extension_aux_control_file(pcontrol, versionName);
 
+	/*
+	 * For trusted extensions with owned schemas, we need to create the schema
+	 * as superuser to ensure proper ownership.
+	 */
+	if (control->owned_schema)
+	{
+		switch_to_superuser = check_extension_superuser_requirements(control, NULL);
+		if (switch_to_superuser)
+		{
+			GetUserIdAndSecContext(&save_userid, &save_sec_context);
+			SetUserIdAndSecContext(BOOTSTRAP_SUPERUSERID,
+								   save_sec_context | SECURITY_LOCAL_USERID_CHANGE);
+		}
+	}
+
 	schemaOid = GetOrCreateSchemaForExtension(&schemaName, control, cascade);
 
+	/* Restore authentication state after schema creation if needed */
+	if (switch_to_superuser)
+		SetUserIdAndSecContext(save_userid, save_sec_context);
+
 	/*
 	 * Make note if a temporary namespace has been accessed in this
 	 * transaction.
diff --git a/src/test/modules/test_extensions/Makefile b/src/test/modules/test_extensions/Makefile
index a6594c19d7e..0a7cf692214 100644
--- a/src/test/modules/test_extensions/Makefile
+++ b/src/test/modules/test_extensions/Makefile
@@ -10,7 +10,8 @@ EXTENSION = test_ext1 test_ext2 test_ext3 test_ext4 test_ext5 test_ext6 \
             test_ext_evttrig \
             test_ext_set_schema \
             test_ext_req_schema1 test_ext_req_schema2 test_ext_req_schema3 \
-            test_ext_owned_schema test_ext_owned_schema_relocatable
+            test_ext_owned_schema test_ext_owned_schema_nosuperuser \
+            test_ext_owned_schema_trusted test_ext_owned_schema_relocatable
 
 DATA = test_ext1--1.0.sql test_ext2--1.0.sql test_ext3--1.0.sql \
        test_ext4--1.0.sql test_ext5--1.0.sql test_ext6--1.0.sql \
@@ -28,6 +29,11 @@ DATA = test_ext1--1.0.sql test_ext2--1.0.sql test_ext3--1.0.sql \
        test_ext_req_schema2--1.0.sql \
        test_ext_req_schema3--1.0.sql \
        test_ext_owned_schema--1.0.sql \
+       test_ext_owned_schema--1.0--1.1.sql \
+       test_ext_owned_schema_nosuperuser--1.0.sql \
+       test_ext_owned_schema_nosuperuser--1.0--1.1.sql \
+       test_ext_owned_schema_trusted--1.0.sql \
+       test_ext_owned_schema_trusted--1.0--1.1.sql \
        test_ext_owned_schema_relocatable--1.0.sql
 
 REGRESS = test_extensions test_extdepend
diff --git a/src/test/modules/test_extensions/expected/test_extensions.out b/src/test/modules/test_extensions/expected/test_extensions.out
index b82b3979a90..1820e76d746 100644
--- a/src/test/modules/test_extensions/expected/test_extensions.out
+++ b/src/test/modules/test_extensions/expected/test_extensions.out
@@ -690,7 +690,36 @@ Objects in extension "test_ext_owned_schema"
  schema test_ext_owned_schema
 (2 rows)
 
+-- Test ALTER EXTENSION UPDATE with owned schema
+ALTER EXTENSION test_ext_owned_schema UPDATE TO '1.1';
+\dx+ test_ext_owned_schema;
+Objects in extension "test_ext_owned_schema"
+           Object description            
+-----------------------------------------
+ function test_ext_owned_schema.owned1()
+ function test_ext_owned_schema.owned2()
+ schema test_ext_owned_schema
+(3 rows)
+
+-- Verify that the owned schema is dropped together with the extension
 DROP EXTENSION test_ext_owned_schema;
+SELECT COUNT(*) FROM pg_namespace WHERE nspname = 'test_ext_owned_schema';
+ count 
+-------
+     0
+(1 row)
+
+-- Test drop behavior with non-extension objects in the owned schema
+CREATE EXTENSION test_ext_owned_schema;
+CREATE FUNCTION test_ext_owned_schema.non_ext_func() RETURNS int LANGUAGE SQL AS $$ SELECT 1 $$;
+-- Fails because of the non-extension object
+DROP EXTENSION test_ext_owned_schema;
+ERROR:  cannot drop extension test_ext_owned_schema because other objects depend on it
+DETAIL:  function test_ext_owned_schema.non_ext_func() depends on schema test_ext_owned_schema
+HINT:  Use DROP ... CASCADE to drop the dependent objects too.
+-- Succeeds and cascades to the non-extension object
+DROP EXTENSION test_ext_owned_schema CASCADE;
+NOTICE:  drop cascades to function test_ext_owned_schema.non_ext_func()
 CREATE SCHEMA already_existing;
 -- Fails for an already existing schema to be provided
 CREATE EXTENSION test_ext_owned_schema_relocatable SCHEMA already_existing;
@@ -721,3 +750,77 @@ Objects in extension "test_ext_owned_schema_relocatable"
 (2 rows)
 
 DROP EXTENSION test_ext_owned_schema_relocatable;
+-- Verify the schema was dropped with the extension
+DROP SCHEMA some_other_name;
+ERROR:  schema "some_other_name" does not exist
+-- Test owned_schema + superuser=false extension
+CREATE USER test_ext_user;
+GRANT CREATE ON DATABASE regression_test_extensions TO test_ext_user;
+SET SESSION AUTHORIZATION test_ext_user;
+CREATE EXTENSION test_ext_owned_schema_nosuperuser;
+\dx+ test_ext_owned_schema_nosuperuser;
+Objects in extension "test_ext_owned_schema_nosuperuser"
+               Object description                
+-------------------------------------------------
+ function test_owned_schema_nosuperuser.owned1()
+ schema test_owned_schema_nosuperuser
+(2 rows)
+
+-- Check that schema is owned by the creating user (not bootstrap superuser)
+SELECT n.nspname, n.nspowner = current_user::regrole as owned_by_current_user
+FROM pg_namespace n
+WHERE n.nspname = 'test_owned_schema_nosuperuser';
+            nspname            | owned_by_current_user 
+-------------------------------+-----------------------
+ test_owned_schema_nosuperuser | t
+(1 row)
+
+-- Upgrades should work for superuser=false extensions
+ALTER EXTENSION test_ext_owned_schema_nosuperuser UPDATE TO '1.1';
+\dx+ test_ext_owned_schema_nosuperuser;
+Objects in extension "test_ext_owned_schema_nosuperuser"
+               Object description                
+-------------------------------------------------
+ function test_owned_schema_nosuperuser.owned1()
+ function test_owned_schema_nosuperuser.owned2()
+ schema test_owned_schema_nosuperuser
+(3 rows)
+
+DROP EXTENSION test_ext_owned_schema_nosuperuser;
+RESET SESSION AUTHORIZATION;
+-- Test owned_schema + trusted=true extension
+SET SESSION AUTHORIZATION test_ext_user;
+CREATE EXTENSION test_ext_owned_schema_trusted;
+\dx+ test_ext_owned_schema_trusted;
+Objects in extension "test_ext_owned_schema_trusted"
+             Object description              
+---------------------------------------------
+ function test_owned_schema_trusted.owned1()
+ schema test_owned_schema_trusted
+(2 rows)
+
+-- Check that schema is owned by bootstrap superuser for trusted extension,
+-- even though the extension is owned by test_ext_user
+SELECT n.nspname, n.nspowner = 10 as owned_by_bootstrap_superuser
+FROM pg_namespace n
+WHERE n.nspname = 'test_owned_schema_trusted';
+          nspname          | owned_by_bootstrap_superuser 
+---------------------------+------------------------------
+ test_owned_schema_trusted | t
+(1 row)
+
+-- Updating trusted extensions should work normally
+ALTER EXTENSION test_ext_owned_schema_trusted UPDATE TO '1.1';
+\dx+ test_ext_owned_schema_trusted;
+Objects in extension "test_ext_owned_schema_trusted"
+             Object description              
+---------------------------------------------
+ function test_owned_schema_trusted.owned1()
+ function test_owned_schema_trusted.owned2()
+ schema test_owned_schema_trusted
+(3 rows)
+
+DROP EXTENSION test_ext_owned_schema_trusted;
+RESET SESSION AUTHORIZATION;
+REVOKE CREATE ON DATABASE regression_test_extensions FROM test_ext_user;
+DROP USER test_ext_user;
diff --git a/src/test/modules/test_extensions/meson.build b/src/test/modules/test_extensions/meson.build
index 3e25c4466ec..019e0d2a243 100644
--- a/src/test/modules/test_extensions/meson.build
+++ b/src/test/modules/test_extensions/meson.build
@@ -45,7 +45,14 @@ test_install_data += files(
   'test_ext_set_schema--1.0.sql',
   'test_ext_set_schema.control',
   'test_ext_owned_schema--1.0.sql',
+  'test_ext_owned_schema--1.0--1.1.sql',
   'test_ext_owned_schema.control',
+  'test_ext_owned_schema_nosuperuser--1.0.sql',
+  'test_ext_owned_schema_nosuperuser--1.0--1.1.sql',
+  'test_ext_owned_schema_nosuperuser.control',
+  'test_ext_owned_schema_trusted--1.0.sql',
+  'test_ext_owned_schema_trusted--1.0--1.1.sql',
+  'test_ext_owned_schema_trusted.control',
   'test_ext_owned_schema_relocatable--1.0.sql',
   'test_ext_owned_schema_relocatable.control',
 )
diff --git a/src/test/modules/test_extensions/sql/test_extensions.sql b/src/test/modules/test_extensions/sql/test_extensions.sql
index a97866d00ea..700033fc13b 100644
--- a/src/test/modules/test_extensions/sql/test_extensions.sql
+++ b/src/test/modules/test_extensions/sql/test_extensions.sql
@@ -316,7 +316,20 @@ CREATE EXTENSION test_ext_owned_schema SCHEMA test_schema;
 DROP SCHEMA test_ext_owned_schema;
 CREATE EXTENSION test_ext_owned_schema;
 \dx+ test_ext_owned_schema;
+-- Test ALTER EXTENSION UPDATE with owned schema
+ALTER EXTENSION test_ext_owned_schema UPDATE TO '1.1';
+\dx+ test_ext_owned_schema;
+-- Verify that the owned schema is dropped together with the extension
+DROP EXTENSION test_ext_owned_schema;
+SELECT COUNT(*) FROM pg_namespace WHERE nspname = 'test_ext_owned_schema';
+
+-- Test drop behavior with non-extension objects in the owned schema
+CREATE EXTENSION test_ext_owned_schema;
+CREATE FUNCTION test_ext_owned_schema.non_ext_func() RETURNS int LANGUAGE SQL AS $$ SELECT 1 $$;
+-- Fails because of the non-extension object
 DROP EXTENSION test_ext_owned_schema;
+-- Succeeds and cascades to the non-extension object
+DROP EXTENSION test_ext_owned_schema CASCADE;
 
 CREATE SCHEMA already_existing;
 -- Fails for an already existing schema to be provided
@@ -330,3 +343,38 @@ ALTER EXTENSION test_ext_owned_schema_relocatable SET SCHEMA already_existing;
 ALTER EXTENSION test_ext_owned_schema_relocatable SET SCHEMA some_other_name;
 \dx+ test_ext_owned_schema_relocatable
 DROP EXTENSION test_ext_owned_schema_relocatable;
+-- Verify the schema was dropped with the extension
+DROP SCHEMA some_other_name;
+
+-- Test owned_schema + superuser=false extension
+CREATE USER test_ext_user;
+GRANT CREATE ON DATABASE regression_test_extensions TO test_ext_user;
+SET SESSION AUTHORIZATION test_ext_user;
+CREATE EXTENSION test_ext_owned_schema_nosuperuser;
+\dx+ test_ext_owned_schema_nosuperuser;
+-- Check that schema is owned by the creating user (not bootstrap superuser)
+SELECT n.nspname, n.nspowner = current_user::regrole as owned_by_current_user
+FROM pg_namespace n
+WHERE n.nspname = 'test_owned_schema_nosuperuser';
+-- Upgrades should work for superuser=false extensions
+ALTER EXTENSION test_ext_owned_schema_nosuperuser UPDATE TO '1.1';
+\dx+ test_ext_owned_schema_nosuperuser;
+DROP EXTENSION test_ext_owned_schema_nosuperuser;
+RESET SESSION AUTHORIZATION;
+
+-- Test owned_schema + trusted=true extension
+SET SESSION AUTHORIZATION test_ext_user;
+CREATE EXTENSION test_ext_owned_schema_trusted;
+\dx+ test_ext_owned_schema_trusted;
+-- Check that schema is owned by bootstrap superuser for trusted extension,
+-- even though the extension is owned by test_ext_user
+SELECT n.nspname, n.nspowner = 10 as owned_by_bootstrap_superuser
+FROM pg_namespace n
+WHERE n.nspname = 'test_owned_schema_trusted';
+-- Updating trusted extensions should work normally
+ALTER EXTENSION test_ext_owned_schema_trusted UPDATE TO '1.1';
+\dx+ test_ext_owned_schema_trusted;
+DROP EXTENSION test_ext_owned_schema_trusted;
+RESET SESSION AUTHORIZATION;
+REVOKE CREATE ON DATABASE regression_test_extensions FROM test_ext_user;
+DROP USER test_ext_user;
diff --git a/src/test/modules/test_extensions/test_ext_owned_schema--1.0--1.1.sql b/src/test/modules/test_extensions/test_ext_owned_schema--1.0--1.1.sql
new file mode 100644
index 00000000000..a87be1ac967
--- /dev/null
+++ b/src/test/modules/test_extensions/test_ext_owned_schema--1.0--1.1.sql
@@ -0,0 +1,3 @@
+-- Test extension upgrade script
+CREATE FUNCTION owned2() RETURNS text
+  LANGUAGE SQL AS $$ SELECT 'owned2'::text $$;
diff --git a/src/test/modules/test_extensions/test_ext_owned_schema_nosuperuser--1.0--1.1.sql b/src/test/modules/test_extensions/test_ext_owned_schema_nosuperuser--1.0--1.1.sql
new file mode 100644
index 00000000000..e9ba788d751
--- /dev/null
+++ b/src/test/modules/test_extensions/test_ext_owned_schema_nosuperuser--1.0--1.1.sql
@@ -0,0 +1,5 @@
+-- complain if script is sourced in psql, rather than via CREATE EXTENSION
+\echo Use "ALTER EXTENSION test_ext_owned_schema_nosuperuser UPDATE TO '1.1'" to load this file. \quit
+
+CREATE FUNCTION owned2() RETURNS text
+  LANGUAGE SQL AS $$ SELECT 'owned2'::text $$;
diff --git a/src/test/modules/test_extensions/test_ext_owned_schema_nosuperuser--1.0.sql b/src/test/modules/test_extensions/test_ext_owned_schema_nosuperuser--1.0.sql
new file mode 100644
index 00000000000..6c22c777519
--- /dev/null
+++ b/src/test/modules/test_extensions/test_ext_owned_schema_nosuperuser--1.0.sql
@@ -0,0 +1,3 @@
+-- Test extension with owned schema (superuser=false)
+CREATE FUNCTION owned1() RETURNS text
+  LANGUAGE SQL AS $$ SELECT 'owned1'::text $$;
diff --git a/src/test/modules/test_extensions/test_ext_owned_schema_nosuperuser.control b/src/test/modules/test_extensions/test_ext_owned_schema_nosuperuser.control
new file mode 100644
index 00000000000..caad0960996
--- /dev/null
+++ b/src/test/modules/test_extensions/test_ext_owned_schema_nosuperuser.control
@@ -0,0 +1,5 @@
+comment = 'Test extension with owned schema (superuser=false)'
+default_version = '1.0'
+superuser = false
+owned_schema = true
+schema = 'test_owned_schema_nosuperuser'
diff --git a/src/test/modules/test_extensions/test_ext_owned_schema_trusted--1.0--1.1.sql b/src/test/modules/test_extensions/test_ext_owned_schema_trusted--1.0--1.1.sql
new file mode 100644
index 00000000000..a87be1ac967
--- /dev/null
+++ b/src/test/modules/test_extensions/test_ext_owned_schema_trusted--1.0--1.1.sql
@@ -0,0 +1,3 @@
+-- Test extension upgrade script
+CREATE FUNCTION owned2() RETURNS text
+  LANGUAGE SQL AS $$ SELECT 'owned2'::text $$;
diff --git a/src/test/modules/test_extensions/test_ext_owned_schema_trusted--1.0.sql b/src/test/modules/test_extensions/test_ext_owned_schema_trusted--1.0.sql
new file mode 100644
index 00000000000..4e5410c184d
--- /dev/null
+++ b/src/test/modules/test_extensions/test_ext_owned_schema_trusted--1.0.sql
@@ -0,0 +1,3 @@
+-- Test extension with owned schema (trusted=true)
+CREATE FUNCTION owned1() RETURNS text
+  LANGUAGE SQL AS $$ SELECT 'owned1'::text $$;
diff --git a/src/test/modules/test_extensions/test_ext_owned_schema_trusted.control b/src/test/modules/test_extensions/test_ext_owned_schema_trusted.control
new file mode 100644
index 00000000000..277b3f1603b
--- /dev/null
+++ b/src/test/modules/test_extensions/test_ext_owned_schema_trusted.control
@@ -0,0 +1,5 @@
+comment = 'Test extension with owned schema (trusted=true)'
+default_version = '1.0'
+trusted = true
+owned_schema = true
+schema = 'test_owned_schema_trusted'