Support GRANT ON SCHEMA commands in CREATE SCHEMA statements (#5789)

* Support GRANT ON SCHEMA commands in CREATE SCHEMA statements

* Add test

* add comment

* Rename to GetGrantCommandsFromCreateSchemaStmt

src/backend/distributed/commands/schema.c

@@ -44,6 +44,7 @@ static ObjectAddress GetObjectAddressBySchemaName(char *schemaName, bool missing
 static List * FilterDistributedSchemas(List *schemas);
 static bool SchemaHasDistributedTableWithFKey(char *schemaName);
 static bool ShouldPropagateCreateSchemaStmt(void);
+static List * GetGrantCommandsFromCreateSchemaStmt(Node *node);
 
 
 /*
@@ -63,13 +64,17 @@ PreprocessCreateSchemaStmt(Node *node, const char *queryString,
 
 	EnsureSequentialMode(OBJECT_SCHEMA);
 
+	/* to prevent recursion with mx we disable ddl propagation */
+	List *commands = list_make1(DISABLE_DDL_PROPAGATION);
+
 	/* deparse sql*/
 	const char *sql = DeparseTreeNode(node);
 
-	/* to prevent recursion with mx we disable ddl propagation */
+	commands = lappend(commands, (void *) sql);
-	List *commands = list_make3(DISABLE_DDL_PROPAGATION,
+
-								(void *) sql,
+	commands = list_concat(commands, GetGrantCommandsFromCreateSchemaStmt(node));
-								ENABLE_DDL_PROPAGATION);
+
+	commands = lappend(commands, ENABLE_DDL_PROPAGATION);
 
 	return NodeDDLTaskList(NON_COORDINATOR_NODES, commands);
 }
@@ -392,3 +397,44 @@ ShouldPropagateCreateSchemaStmt()
 
 	return true;
 }
+
+
+/*
+ * GetGrantCommandsFromCreateSchemaStmt takes a CreateSchemaStmt and returns the
+ * list of deparsed queries of the inner GRANT commands of the given statement.
+ * Ignores commands other than GRANT statements.
+ */
+static List *
+GetGrantCommandsFromCreateSchemaStmt(Node *node)
+{
+	List *commands = NIL;
+	CreateSchemaStmt *stmt = castNode(CreateSchemaStmt, node);
+
+	Node *element = NULL;
+	foreach_ptr(element, stmt->schemaElts)
+	{
+		if (!IsA(element, GrantStmt))
+		{
+			continue;
+		}
+
+		GrantStmt *grantStmt = castNode(GrantStmt, element);
+
+		switch (grantStmt->objtype)
+		{
+			/* we only propagate GRANT ON SCHEMA in community */
+			case OBJECT_SCHEMA:
+			{
+				commands = lappend(commands, DeparseGrantOnSchemaStmt(element));
+				break;
+			}
+
+			default:
+			{
+				break;
+			}
+		}
+	}
+
+	return commands;
+}

src/backend/distributed/deparser/deparse_schema_stmts.c

@@ -87,11 +87,6 @@ DeparseAlterSchemaRenameStmt(Node *node)
 static void
 AppendCreateSchemaStmt(StringInfo buf, CreateSchemaStmt *stmt)
 {
-	if (stmt->schemaElts != NIL)
-	{
-		elog(ERROR, "schema creating is not supported with other create commands");
-	}
-
 	appendStringInfoString(buf, "CREATE SCHEMA ");
 
 	if (stmt->if_not_exists)

src/test/regress/expected/multi_mx_schema_support.out

@@ -538,9 +538,9 @@ SELECT run_command_on_workers($$DROP SCHEMA localschema;$$);
  (localhost,57638,f,"ERROR:  schema ""localschema"" does not exist")
 (2 rows)
 
+SET client_min_messages TO ERROR;
 CREATE ROLE schema_owner WITH LOGIN;
-NOTICE:  not propagating CREATE ROLE/USER commands to worker nodes
+RESET client_min_messages;
-HINT:  Connect to worker nodes directly to manually create all necessary users and roles.
 SELECT run_command_on_workers($$CREATE ROLE schema_owner WITH LOGIN;$$);
       run_command_on_workers
 ---------------------------------------------------------------------
@@ -560,6 +560,59 @@ SELECT run_command_on_workers($$SELECT COUNT(*) FROM pg_namespace WHERE nspname=
 (2 rows)
 
 DROP SCHEMA schema_owner;
+-- test CREATE SCHEMA .. GRANT ON SCHEMA commands
+-- first create the role to be granted
+SET citus.enable_ddl_propagation TO OFF;
+SET client_min_messages TO ERROR;
+CREATE ROLE role_to_be_granted WITH LOGIN;
+RESET client_min_messages;
+SELECT run_command_on_workers($$CREATE ROLE role_to_be_granted WITH LOGIN;$$);
+      run_command_on_workers
+---------------------------------------------------------------------
+ (localhost,57637,t,"CREATE ROLE")
+ (localhost,57638,t,"CREATE ROLE")
+(2 rows)
+
+RESET citus.enable_ddl_propagation;
+CREATE SCHEMA old_schema;
+CREATE SCHEMA new_schema
+    CREATE TABLE t1 (a int)
+    GRANT ALL ON SCHEMA old_schema TO role_to_be_granted
+    GRANT ALL ON SCHEMA new_schema TO role_to_be_granted;
+-- the role should be granted on both the new and the old schema
+SELECT nspacl FROM pg_namespace WHERE nspname='old_schema' OR nspname='new_schema';
+                        nspacl
+---------------------------------------------------------------------
+ {postgres=UC/postgres,role_to_be_granted=UC/postgres}
+ {postgres=UC/postgres,role_to_be_granted=UC/postgres}
+(2 rows)
+
+-- verify on workers
+SELECT run_command_on_workers($$SELECT nspacl FROM pg_namespace WHERE nspname='new_schema';$$);
+                           run_command_on_workers
+---------------------------------------------------------------------
+ (localhost,57637,t,"{postgres=UC/postgres,role_to_be_granted=UC/postgres}")
+ (localhost,57638,t,"{postgres=UC/postgres,role_to_be_granted=UC/postgres}")
+(2 rows)
+
+SELECT run_command_on_workers($$SELECT nspacl FROM pg_namespace WHERE nspname='old_schema';$$);
+                           run_command_on_workers
+---------------------------------------------------------------------
+ (localhost,57637,t,"{postgres=UC/postgres,role_to_be_granted=UC/postgres}")
+ (localhost,57638,t,"{postgres=UC/postgres,role_to_be_granted=UC/postgres}")
+(2 rows)
+
+-- verify the table t1 is created as a local pg table
+-- this might be changed after some improvements on use_citus_managed_tables
+-- if so, please verify that t1 is added to metadata
+SELECT COUNT(*)=0 FROM pg_dist_partition WHERE logicalrelid='new_schema.t1'::regclass;
+ ?column?
+---------------------------------------------------------------------
+ t
+(1 row)
+
+DROP SCHEMA old_schema, new_schema CASCADE;
+NOTICE:  drop cascades to table new_schema.t1
 DROP SCHEMA mx_old_schema CASCADE;
 DROP SCHEMA mx_new_schema CASCADE;
 NOTICE:  drop cascades to table mx_new_schema.table_set_schema

src/test/regress/sql/multi_mx_schema_support.sql

@@ -351,7 +351,9 @@ CREATE SCHEMA localschema;
 -- should error out
 SELECT run_command_on_workers($$DROP SCHEMA localschema;$$);
 
+SET client_min_messages TO ERROR;
 CREATE ROLE schema_owner WITH LOGIN;
+RESET client_min_messages;
 SELECT run_command_on_workers($$CREATE ROLE schema_owner WITH LOGIN;$$);
 RESET citus.enable_ddl_propagation;
 -- create schema with the name of the owner
@@ -361,5 +363,32 @@ SELECT run_command_on_workers($$SELECT COUNT(*) FROM pg_namespace WHERE nspname=
 
 DROP SCHEMA schema_owner;
 
+-- test CREATE SCHEMA .. GRANT ON SCHEMA commands
+-- first create the role to be granted
+SET citus.enable_ddl_propagation TO OFF;
+SET client_min_messages TO ERROR;
+CREATE ROLE role_to_be_granted WITH LOGIN;
+RESET client_min_messages;
+SELECT run_command_on_workers($$CREATE ROLE role_to_be_granted WITH LOGIN;$$);
+RESET citus.enable_ddl_propagation;
+
+CREATE SCHEMA old_schema;
+CREATE SCHEMA new_schema
+    CREATE TABLE t1 (a int)
+    GRANT ALL ON SCHEMA old_schema TO role_to_be_granted
+    GRANT ALL ON SCHEMA new_schema TO role_to_be_granted;
+
+-- the role should be granted on both the new and the old schema
+SELECT nspacl FROM pg_namespace WHERE nspname='old_schema' OR nspname='new_schema';
+-- verify on workers
+SELECT run_command_on_workers($$SELECT nspacl FROM pg_namespace WHERE nspname='new_schema';$$);
+SELECT run_command_on_workers($$SELECT nspacl FROM pg_namespace WHERE nspname='old_schema';$$);
+
+-- verify the table t1 is created as a local pg table
+-- this might be changed after some improvements on use_citus_managed_tables
+-- if so, please verify that t1 is added to metadata
+SELECT COUNT(*)=0 FROM pg_dist_partition WHERE logicalrelid='new_schema.t1'::regclass;
+
+DROP SCHEMA old_schema, new_schema CASCADE;
 DROP SCHEMA mx_old_schema CASCADE;
 DROP SCHEMA mx_new_schema CASCADE;