fix the problem #5763 (#6519)

Co-authored-by: TsinghuaLucky912 <tsinghualucky912@foxmail.com>
Fixes #5763

src/backend/distributed/commands/distribute_object_ops.c

@@ -254,6 +254,15 @@ static DistributeObjectOps Any_CreateRole = {
 	.address = CreateRoleStmtObjectAddress,
 	.markDistributed = true,
 };
+static DistributeObjectOps Any_DropOwned = {
+	.deparse = DeparseDropOwnedStmt,
+	.qualify = NULL,
+	.preprocess = PreprocessDropOwnedStmt,
+	.postprocess = NULL,
+	.operationType = DIST_OPS_DROP,
+	.address = NULL,
+	.markDistributed = false,
+};
 static DistributeObjectOps Any_DropRole = {
 	.deparse = DeparseDropRoleStmt,
 	.qualify = NULL,
@@ -1658,6 +1667,11 @@ GetDistributeObjectOps(Node *node)
 			return &Any_DropRole;
 		}
 
+		case T_DropOwnedStmt:
+		{
+			return &Any_DropOwned;
+		}
+
 		case T_DropStmt:
 		{
 			DropStmt *stmt = castNode(DropStmt, node);

src/backend/distributed/commands/owned.c

@@ -0,0 +1,90 @@
+/*-------------------------------------------------------------------------
+ *
+ * owned.c
+ *    Commands for DROP OWNED statements.
+ *
+ * Copyright (c) Citus Data, Inc.
+ *
+ *-------------------------------------------------------------------------
+ */
+
+#include "postgres.h"
+
+#include "access/heapam.h"
+#include "access/htup_details.h"
+#include "access/genam.h"
+#include "access/table.h"
+#include "access/xact.h"
+#include "catalog/catalog.h"
+#include "catalog/pg_auth_members.h"
+#include "catalog/pg_authid.h"
+#include "catalog/pg_db_role_setting.h"
+#include "catalog/pg_type.h"
+#include "catalog/objectaddress.h"
+#include "commands/dbcommands.h"
+#include "distributed/citus_ruleutils.h"
+#include "distributed/citus_safe_lib.h"
+#include "distributed/commands.h"
+#include "distributed/commands/utility_hook.h"
+#include "distributed/deparser.h"
+#include "distributed/listutils.h"
+#include "distributed/coordinator_protocol.h"
+#include "distributed/metadata/distobject.h"
+#include "distributed/metadata_sync.h"
+#include "distributed/metadata/distobject.h"
+#include "distributed/multi_executor.h"
+#include "distributed/relation_access_tracking.h"
+#include "distributed/version_compat.h"
+#include "distributed/worker_transaction.h"
+#include "miscadmin.h"
+#include "nodes/makefuncs.h"
+#include "nodes/parsenodes.h"
+#include "nodes/pg_list.h"
+#include "parser/scansup.h"
+#include "utils/builtins.h"
+#include "utils/fmgroids.h"
+#include "utils/rel.h"
+#include "utils/varlena.h"
+#include "utils/syscache.h"
+
+/*
+ * PreprocessDropOwnedStmt finds the distributed role out of the ones
+ * being dropped and unmarks them distributed and creates the drop statements
+ * for the workers.
+ */
+List *
+PreprocessDropOwnedStmt(Node *node, const char *queryString,
+						ProcessUtilityContext processUtilityContext)
+{
+	DropOwnedStmt *stmt = castNode(DropOwnedStmt, node);
+	List *allDropRoles = stmt->roles;
+
+	List *distributedDropRoles = FilterDistributedRoles(allDropRoles);
+	if (list_length(distributedDropRoles) <= 0)
+	{
+		return NIL;
+	}
+
+	if (!ShouldPropagate())
+	{
+		return NIL;
+	}
+
+	/* check creation against multi-statement transaction policy */
+	if (!ShouldPropagateCreateInCoordinatedTransction())
+	{
+		return NIL;
+	}
+
+	EnsureCoordinator();
+
+	stmt->roles = distributedDropRoles;
+	char *sql = DeparseTreeNode((Node *) stmt);
+	stmt->roles = allDropRoles;
+
+	List *commands = list_make3(DISABLE_DDL_PROPAGATION,
+								sql,
+								ENABLE_DDL_PROPAGATION);
+
+	return NodeDDLTaskList(NON_COORDINATOR_NODES, commands);
+}

src/backend/distributed/deparser/deparse_owned_stmts.c

@@ -0,0 +1,84 @@
+/*-------------------------------------------------------------------------
+ *
+ * deparse_owned_stmts.c
+ *	  Functions to turn all Statement structures related to owned back
+ *    into sql.
+ *
+ * Copyright (c), Citus Data, Inc.
+ *
+ *-------------------------------------------------------------------------
+ */
+
+#include "postgres.h"
+
+#include "pg_version_compat.h"
+
+#include "distributed/citus_ruleutils.h"
+#include "distributed/deparser.h"
+#include "lib/stringinfo.h"
+#include "nodes/parsenodes.h"
+#include "utils/builtins.h"
+
+static void AppendDropOwnedStmt(StringInfo buf, DropOwnedStmt *stmt);
+static void AppendRoleList(StringInfo buf, List *roleList);
+
+/*
+ * DeparseDropOwnedStmt builds and returns a string representing of the
+ * DropOwnedStmt for application on a remote server.
+ */
+char *
+DeparseDropOwnedStmt(Node *node)
+{
+	DropOwnedStmt *stmt = castNode(DropOwnedStmt, node);
+
+	StringInfoData buf = { 0 };
+	initStringInfo(&buf);
+
+	AppendDropOwnedStmt(&buf, stmt);
+
+	return buf.data;
+}
+
+
+/*
+ * AppendDropOwnedStmt generates the string representation of the
+ * DropOwnedStmt and appends it to the buffer.
+ */
+static void
+AppendDropOwnedStmt(StringInfo buf, DropOwnedStmt *stmt)
+{
+	appendStringInfo(buf, "DROP OWNED BY ");
+
+	AppendRoleList(buf, stmt->roles);
+
+	if (stmt->behavior == DROP_RESTRICT)
+	{
+		appendStringInfo(buf, " RESTRICT");
+	}
+	else if (stmt->behavior == DROP_CASCADE)
+	{
+		appendStringInfo(buf, " CASCADE");
+	}
+}
+
+
+static void
+AppendRoleList(StringInfo buf, List *roleList)
+{
+	ListCell *cell = NULL;
+	foreach(cell, roleList)
+	{
+		Node *roleNode = (Node *) lfirst(cell);
+		Assert(IsA(roleNode, RoleSpec) || IsA(roleNode, AccessPriv));
+		char const *rolename = NULL;
+		if (IsA(roleNode, RoleSpec))
+		{
+			rolename = RoleSpecString((RoleSpec *) roleNode, true);
+		}
+		appendStringInfoString(buf, rolename);
+		if (cell != list_tail(roleList))
+		{
+			appendStringInfo(buf, ", ");
+		}
+	}
+}

src/include/distributed/commands.h

@@ -385,6 +385,9 @@ extern bool IsReindexWithParam_compat(ReindexStmt *stmt, char *paramName);
 extern List * CreateExtensionStmtObjectAddress(Node *stmt, bool missing_ok, bool
 											   isPostprocess);
 
+/* owned.c -  forward declarations */
+extern List * PreprocessDropOwnedStmt(Node *node, const char *queryString,
+									  ProcessUtilityContext processUtilityContext);
 
 /* policy.c -  forward declarations */
 extern List * CreatePolicyCommands(Oid relationId);

src/include/distributed/deparser.h

@@ -196,6 +196,9 @@ extern char * DeparseCreateRoleStmt(Node *stmt);
 extern char * DeparseDropRoleStmt(Node *stmt);
 extern char * DeparseGrantRoleStmt(Node *stmt);
 
+/* forward declarations for deparse_owned_stmts.c */
+extern char * DeparseDropOwnedStmt(Node *node);
+
 /* forward declarations for deparse_extension_stmts.c */
 extern DefElem * GetExtensionOption(List *extensionOptions,
 									const char *defname);

src/test/regress/expected/issue_5763.out

@@ -0,0 +1,38 @@
+--
+-- ISSUE_5763
+--
+-- Issue: DROP OWNED BY fails to drop the schemas on the workers
+-- Link: https://github.com/citusdata/citus/issues/5763
+--
+CREATE USER issue_5763_1 WITH SUPERUSER;
+CREATE USER issue_5763_2 WITH SUPERUSER;
+\c - issue_5763_1 - :master_port
+CREATE SCHEMA issue_5763_sc_1;
+\c - issue_5763_2 - :master_port
+CREATE SCHEMA issue_5763_sc_2;
+\c - postgres - :master_port
+DROP OWNED BY issue_5763_1, issue_5763_2;
+\c - issue_5763_1 - :master_port
+CREATE SCHEMA issue_5763_sc_1;
+\c - postgres - :master_port
+DROP SCHEMA issue_5763_sc_1;
+DROP USER issue_5763_1, issue_5763_2;
+-- test CASCADE options
+CREATE USER issue_5763_3 WITH SUPERUSER;
+\c - issue_5763_3 - :master_port
+CREATE SCHEMA issue_5763_sc_3;
+CREATE TABLE issue_5763_sc_3.tb1(id int);
+\c - postgres - :master_port
+DROP OWNED BY issue_5763_3 CASCADE;
+DROP USER issue_5763_3;
+-- test non-distributed role
+SET citus.enable_create_role_propagation TO off;
+CREATE USER issue_5763_4 WITH SUPERUSER;
+NOTICE:  not propagating CREATE ROLE/USER commands to worker nodes
+HINT:  Connect to worker nodes directly to manually create all necessary users and roles.
+\c - issue_5763_4 - :master_port
+set citus.enable_ddl_propagation = off;
+CREATE SCHEMA issue_5763_sc_4;
+\c - postgres - :master_port
+DROP OWNED BY issue_5763_4 RESTRICT;
+DROP USER issue_5763_4;

src/test/regress/expected/multi_multiuser.out

@@ -545,27 +545,6 @@ DROP TABLE
     test,
     test_coloc,
     colocation_table;
-SELECT run_command_on_workers($$DROP OWNED BY full_access$$);
-      run_command_on_workers
----------------------------------------------------------------------
- (localhost,57637,t,"DROP OWNED")
- (localhost,57638,t,"DROP OWNED")
-(2 rows)
-
-SELECT run_command_on_workers($$DROP OWNED BY some_role$$);
-      run_command_on_workers
----------------------------------------------------------------------
- (localhost,57637,t,"DROP OWNED")
- (localhost,57638,t,"DROP OWNED")
-(2 rows)
-
-SELECT run_command_on_workers($$DROP OWNED BY read_access$$);
-      run_command_on_workers
----------------------------------------------------------------------
- (localhost,57637,t,"DROP OWNED")
- (localhost,57638,t,"DROP OWNED")
-(2 rows)
-
 DROP USER full_access;
 DROP USER read_access;
 DROP USER no_access;

src/test/regress/expected/multi_schema_support.out

@@ -1154,13 +1154,6 @@ SET citus.next_shard_id TO 1197000;
 -- we do not use run_command_on_coordinator_and_workers here because when there is CASCADE, it causes deadlock
 DROP OWNED BY "test-user" CASCADE;
 NOTICE:  drop cascades to table schema_with_user.test_table
-SELECT run_command_on_workers('DROP OWNED BY "test-user" CASCADE');
-      run_command_on_workers
----------------------------------------------------------------------
- (localhost,57637,t,"DROP OWNED")
- (localhost,57638,t,"DROP OWNED")
-(2 rows)
-
 DROP USER "test-user";
 DROP FUNCTION run_command_on_coordinator_and_workers(p_sql text);
 -- test run_command_on_* UDFs with schema

src/test/regress/expected/single_node_enterprise.out

@@ -489,18 +489,6 @@ SET client_min_messages TO WARNING;
 DROP SCHEMA single_node_ent CASCADE;
 DROP OWNED BY full_access_single_node;
 DROP OWNED BY read_access_single_node;
-SELECT run_command_on_workers($$DROP OWNED BY full_access_single_node$$);
-      run_command_on_workers
----------------------------------------------------------------------
- (localhost,57637,t,"DROP OWNED")
-(1 row)
-
-SELECT run_command_on_workers($$DROP OWNED BY read_access_single_node$$);
-      run_command_on_workers
----------------------------------------------------------------------
- (localhost,57637,t,"DROP OWNED")
-(1 row)
-
 DROP ROLE full_access_single_node;
 DROP ROLE read_access_single_node;
 -- remove the nodes for next tests

src/test/regress/multi_schedule

@@ -95,7 +95,7 @@ test: multi_dropped_column_aliases foreign_key_restriction_enforcement
 test: binary_protocol
 test: alter_table_set_access_method
 test: alter_distributed_table
-test: issue_5248 issue_5099
+test: issue_5248 issue_5099 issue_5763
 test: object_propagation_debug
 test: undistribute_table
 test: run_command_on_all_nodes
@@ -113,3 +113,4 @@ test: ensure_no_intermediate_data_leak
 # --------
 test: ensure_no_shared_connection_leak
 test: check_mx
+

src/test/regress/sql/issue_5763.sql

@@ -0,0 +1,50 @@
+--
+-- ISSUE_5763
+--
+-- Issue: DROP OWNED BY fails to drop the schemas on the workers
+-- Link: https://github.com/citusdata/citus/issues/5763
+--
+
+CREATE USER issue_5763_1 WITH SUPERUSER;
+CREATE USER issue_5763_2 WITH SUPERUSER;
+
+\c - issue_5763_1 - :master_port
+CREATE SCHEMA issue_5763_sc_1;
+
+\c - issue_5763_2 - :master_port
+CREATE SCHEMA issue_5763_sc_2;
+
+\c - postgres - :master_port
+DROP OWNED BY issue_5763_1, issue_5763_2;
+
+\c - issue_5763_1 - :master_port
+CREATE SCHEMA issue_5763_sc_1;
+
+\c - postgres - :master_port
+DROP SCHEMA issue_5763_sc_1;
+DROP USER issue_5763_1, issue_5763_2;
+
+-- test CASCADE options
+CREATE USER issue_5763_3 WITH SUPERUSER;
+
+\c - issue_5763_3 - :master_port
+CREATE SCHEMA issue_5763_sc_3;
+CREATE TABLE issue_5763_sc_3.tb1(id int);
+
+\c - postgres - :master_port
+DROP OWNED BY issue_5763_3 CASCADE;
+
+DROP USER issue_5763_3;
+
+-- test non-distributed role
+SET citus.enable_create_role_propagation TO off;
+CREATE USER issue_5763_4 WITH SUPERUSER;
+
+\c - issue_5763_4 - :master_port
+set citus.enable_ddl_propagation = off;
+CREATE SCHEMA issue_5763_sc_4;
+
+\c - postgres - :master_port
+DROP OWNED BY issue_5763_4 RESTRICT;
+
+DROP USER issue_5763_4;

src/test/regress/sql/multi_multiuser.sql

@@ -328,9 +328,6 @@ DROP TABLE
     test,
     test_coloc,
     colocation_table;
-SELECT run_command_on_workers($$DROP OWNED BY full_access$$);
-SELECT run_command_on_workers($$DROP OWNED BY some_role$$);
-SELECT run_command_on_workers($$DROP OWNED BY read_access$$);
 DROP USER full_access;
 DROP USER read_access;
 DROP USER no_access;