Fixes create user queries from Citus non-main databases with other users (#7442)

This PR makes the connections to other nodes for
`mark_object_distributed` use the same user as
`execute_command_on_remote_nodes_as_user` so they'll use the same
connection.

src/backend/distributed/commands/utility_hook.c

@@ -92,7 +92,7 @@
 #define START_MANAGEMENT_TRANSACTION \
 	"SELECT citus_internal.start_management_transaction('%lu')"
 #define MARK_OBJECT_DISTRIBUTED \
-	"SELECT citus_internal.mark_object_distributed(%d, %s, %d)"
+	"SELECT citus_internal.mark_object_distributed(%d, %s, %d, %s)"
 
 
 bool EnableDDLPropagation = true; /* ddl propagation is enabled */
@@ -1636,7 +1636,8 @@ RunPostprocessMainDBCommand(Node *parsetree)
 						 MARK_OBJECT_DISTRIBUTED,
 						 AuthIdRelationId,
 						 quote_literal_cstr(createRoleStmt->role),
-						 roleOid);
+						 roleOid,
+						 quote_literal_cstr(CurrentUserName()));
 		RunCitusMainDBQuery(mainDBQuery->data);
 	}
 }

src/backend/distributed/metadata/distobject.c

@@ -67,7 +67,8 @@ PG_FUNCTION_INFO_V1(master_unmark_object_distributed);
 
 /*
  * mark_object_distributed adds an object to pg_dist_object
- * in all of the nodes.
+ * in all of the nodes, for the connections to the other nodes this function
+ * uses the user passed.
  */
 Datum
 mark_object_distributed(PG_FUNCTION_ARGS)
@@ -81,6 +82,8 @@ mark_object_distributed(PG_FUNCTION_ARGS)
 	Oid objectId = PG_GETARG_OID(2);
 	ObjectAddress *objectAddress = palloc0(sizeof(ObjectAddress));
 	ObjectAddressSet(*objectAddress, classId, objectId);
+	text *connectionUserText = PG_GETARG_TEXT_P(3);
+	char *connectionUser = text_to_cstring(connectionUserText);
 
 	/*
 	 * This function is called when a query is run from a Citus non-main database.
@@ -88,7 +91,8 @@ mark_object_distributed(PG_FUNCTION_ARGS)
 	 * 2PC still works.
 	 */
 	bool useConnectionForLocalQuery = true;
-	MarkObjectDistributedWithName(objectAddress, objectName, useConnectionForLocalQuery);
+	MarkObjectDistributedWithName(objectAddress, objectName, useConnectionForLocalQuery,
+								  connectionUser);
 	PG_RETURN_VOID();
 }
 
@@ -193,7 +197,8 @@ void
 MarkObjectDistributed(const ObjectAddress *distAddress)
 {
 	bool useConnectionForLocalQuery = false;
-	MarkObjectDistributedWithName(distAddress, "", useConnectionForLocalQuery);
+	MarkObjectDistributedWithName(distAddress, "", useConnectionForLocalQuery,
+								  CurrentUserName());
 }
 
 
@@ -204,7 +209,7 @@ MarkObjectDistributed(const ObjectAddress *distAddress)
  */
 void
 MarkObjectDistributedWithName(const ObjectAddress *distAddress, char *objectName,
-							  bool useConnectionForLocalQuery)
+							  bool useConnectionForLocalQuery, char *connectionUser)
 {
 	if (!CitusHasBeenLoaded())
 	{
@@ -234,7 +239,8 @@ MarkObjectDistributedWithName(const ObjectAddress *distAddress, char *objectName
 	{
 		char *workerPgDistObjectUpdateCommand =
 			CreatePgDistObjectEntryCommand(distAddress, objectName);
-		SendCommandToRemoteNodesWithMetadata(workerPgDistObjectUpdateCommand);
+		SendCommandToRemoteMetadataNodesParams(workerPgDistObjectUpdateCommand,
+											   connectionUser, 0, NULL, NULL);
 	}
 }
 

src/backend/distributed/sql/downgrades/citus--12.2-1--12.1-1.sql

@@ -15,7 +15,7 @@ DROP FUNCTION citus_internal.execute_command_on_remote_nodes_as_user(
 );
 
 DROP FUNCTION citus_internal.mark_object_distributed(
-    classId Oid, objectName text, objectId Oid
+    classId Oid, objectName text, objectId Oid, connectionUser text
 );
 
 DROP FUNCTION citus_internal.commit_management_command_2pc();

src/backend/distributed/sql/udfs/mark_object_distributed/12.2-1.sql

@@ -1,7 +1,7 @@
-CREATE OR REPLACE FUNCTION citus_internal.mark_object_distributed(classId Oid, objectName text, objectId Oid)
+CREATE OR REPLACE FUNCTION citus_internal.mark_object_distributed(classId Oid, objectName text, objectId Oid, connectionUser text)
     RETURNS VOID
     LANGUAGE C
 AS 'MODULE_PATHNAME', $$mark_object_distributed$$;
 
-COMMENT ON FUNCTION citus_internal.mark_object_distributed(classId Oid, objectName text, objectId Oid)
+COMMENT ON FUNCTION citus_internal.mark_object_distributed(classId Oid, objectName text, objectId Oid, connectionUser text)
     IS 'adds an object to pg_dist_object on all nodes';

src/backend/distributed/sql/udfs/mark_object_distributed/latest.sql

@@ -1,7 +1,7 @@
-CREATE OR REPLACE FUNCTION citus_internal.mark_object_distributed(classId Oid, objectName text, objectId Oid)
+CREATE OR REPLACE FUNCTION citus_internal.mark_object_distributed(classId Oid, objectName text, objectId Oid, connectionUser text)
     RETURNS VOID
     LANGUAGE C
 AS 'MODULE_PATHNAME', $$mark_object_distributed$$;
 
-COMMENT ON FUNCTION citus_internal.mark_object_distributed(classId Oid, objectName text, objectId Oid)
+COMMENT ON FUNCTION citus_internal.mark_object_distributed(classId Oid, objectName text, objectId Oid, connectionUser text)
     IS 'adds an object to pg_dist_object on all nodes';

src/backend/distributed/transaction/worker_transaction.c

@@ -36,10 +36,6 @@
 #include "distributed/worker_manager.h"
 #include "distributed/worker_transaction.h"
 
-static void SendCommandToRemoteMetadataNodesParams(const char *command,
-												   const char *user, int parameterCount,
-												   const Oid *parameterTypes,
-												   const char *const *parameterValues);
 static void SendBareCommandListToMetadataNodesInternal(List *commandList,
 													   TargetWorkerSet targetWorkerSet);
 static void SendCommandToMetadataWorkersParams(const char *command,
@@ -209,7 +205,7 @@ SendCommandListToRemoteNodesWithMetadata(List *commands)
  * SendCommandToWorkersParamsInternal() that can be used to send commands
  * to remote metadata nodes.
  */
-static void
+void
 SendCommandToRemoteMetadataNodesParams(const char *command,
 									   const char *user, int parameterCount,
 									   const Oid *parameterTypes,

src/include/distributed/metadata/distobject.h

@@ -24,7 +24,8 @@ extern bool IsAnyObjectDistributed(const List *addresses);
 extern bool ClusterHasDistributedFunctionWithDistArgument(void);
 extern void MarkObjectDistributed(const ObjectAddress *distAddress);
 extern void MarkObjectDistributedWithName(const ObjectAddress *distAddress, char *name,
-										  bool useConnectionForLocalQuery);
+										  bool useConnectionForLocalQuery,
+										  char *connectionUser);
 extern void MarkObjectDistributedViaSuperUser(const ObjectAddress *distAddress);
 extern void MarkObjectDistributedLocally(const ObjectAddress *distAddress);
 extern void UnmarkObjectDistributed(const ObjectAddress *address);

src/include/distributed/worker_transaction.h

@@ -68,6 +68,10 @@ extern void SendCommandToWorkersAsUser(TargetWorkerSet targetWorkerSet,
 									   const char *nodeUser, const char *command);
 extern void SendCommandToWorkerAsUser(const char *nodeName, int32 nodePort,
 									  const char *nodeUser, const char *command);
+extern void SendCommandToRemoteMetadataNodesParams(const char *command,
+												   const char *user, int parameterCount,
+												   const Oid *parameterTypes,
+												   const char *const *parameterValues);
 extern bool SendOptionalCommandListToWorkerOutsideTransaction(const char *nodeName,
 															  int32 nodePort,
 															  const char *nodeUser,

src/test/regress/citus_tests/test/test_other_databases.py

@@ -22,7 +22,7 @@ def test_main_commited_outer_not_yet(cluster):
             "SELECT citus_internal.execute_command_on_remote_nodes_as_user('CREATE USER u1;', 'postgres')"
         )
         cur2.execute(
-            "SELECT citus_internal.mark_object_distributed(1260, 'u1', 123123)"
+            "SELECT citus_internal.mark_object_distributed(1260, 'u1', 123123, 'postgres')"
         )
         cur2.execute("COMMIT")
 
@@ -133,7 +133,7 @@ def test_main_commited_outer_aborted(cluster):
             "SELECT citus_internal.execute_command_on_remote_nodes_as_user('CREATE USER u2;', 'postgres')"
         )
         cur2.execute(
-            "SELECT citus_internal.mark_object_distributed(1260, 'u2', 321321)"
+            "SELECT citus_internal.mark_object_distributed(1260, 'u2', 321321, 'postgres')"
         )
         cur2.execute("COMMIT")
 

src/test/regress/expected/multi_extension.out

@@ -1424,7 +1424,7 @@ SELECT * FROM multi_extension.print_extension_changes();
 ---------------------------------------------------------------------
                  | function citus_internal.commit_management_command_2pc() void
                  | function citus_internal.execute_command_on_remote_nodes_as_user(text,text) void
-                 | function citus_internal.mark_object_distributed(oid,text,oid) void
+                 | function citus_internal.mark_object_distributed(oid,text,oid,text) void
                  | function citus_internal.start_management_transaction(xid8) void
                  | function citus_internal_acquire_citus_advisory_object_class_lock(integer,cstring) void
                  | function citus_internal_database_command(text) void

src/test/regress/expected/other_databases.out

@@ -71,6 +71,7 @@ SELECT citus_internal.execute_command_on_remote_nodes_as_user($$SELECT 'dangerou
 ERROR:  operation is not allowed
 HINT:  Run the command with a superuser.
 \c other_db1
+SET ROLE nonsuperuser;
 CREATE USER other_db_user9;
 RESET ROLE;
 \c regression

src/test/regress/expected/upgrade_list_citus_objects.out

@@ -59,7 +59,7 @@ ORDER BY 1;
  function citus_internal.commit_management_command_2pc()
  function citus_internal.execute_command_on_remote_nodes_as_user(text,text)
  function citus_internal.find_groupid_for_node(text,integer)
- function citus_internal.mark_object_distributed(oid,text,oid)
+ function citus_internal.mark_object_distributed(oid,text,oid,text)
  function citus_internal.pg_dist_node_trigger_func()
  function citus_internal.pg_dist_rebalance_strategy_trigger_func()
  function citus_internal.pg_dist_shard_placement_trigger_func()

src/test/regress/sql/other_databases.sql

@@ -51,6 +51,7 @@ SET ROLE nonsuperuser;
 SELECT citus_internal.execute_command_on_remote_nodes_as_user($$SELECT 'dangerous query'$$, 'postgres');
 
 \c other_db1
+SET ROLE nonsuperuser;
 CREATE USER other_db_user9;
 
 RESET ROLE;