support grant statement propagation for aggregates (#6132)

src/backend/distributed/commands/distribute_object_ops.c

@@ -76,6 +76,15 @@ static DistributeObjectOps Aggregate_Rename = {
 	.address = RenameFunctionStmtObjectAddress,
 	.markDistributed = false,
 };
+static DistributeObjectOps Aggregate_Grant = {
+	.deparse = DeparseGrantOnFunctionStmt,
+	.qualify = NULL,
+	.preprocess = PreprocessGrantOnFunctionStmt,
+	.postprocess = PostprocessGrantOnFunctionStmt,
+	.operationType = DIST_OPS_ALTER,
+	.address = NULL,
+	.markDistributed = false,
+};
 static DistributeObjectOps Any_AlterEnum = {
 	.deparse = DeparseAlterEnumStmt,
 	.qualify = QualifyAlterEnumStmt,
@@ -1745,6 +1754,11 @@ GetDistributeObjectOps(Node *node)
 					return &Function_Grant;
 				}
 
+				case OBJECT_AGGREGATE:
+				{
+					return &Aggregate_Grant;
+				}
+
 				case OBJECT_PROCEDURE:
 				{
 					return &Procedure_Grant;

src/backend/distributed/deparser/deparse_function_stmts.c

@@ -789,7 +789,15 @@ static void
 AppendGrantOnFunctionFunctions(StringInfo buf, GrantStmt *stmt)
 {
 	ListCell *cell = NULL;
-	appendStringInfo(buf, " ON %s ", ObjectTypeToKeyword(stmt->objtype));
+
+	/*
+	 * The FUNCTION syntax works for plain functions, aggregate functions, and window
+	 * functions, but not for procedures; use PROCEDURE for those. Alternatively, use
+	 * ROUTINE to refer to a function, aggregate function, window function, or procedure
+	 * regardless of its precise type.
+	 * https://www.postgresql.org/docs/current/sql-grant.html
+	 */
+	appendStringInfo(buf, " ON ROUTINE ");
 
 	foreach(cell, stmt->objects)
 	{
@@ -815,12 +823,12 @@ AppendGrantOnFunctionFunctions(StringInfo buf, GrantStmt *stmt)
 
 
 /*
- * isFunction returns true if the given ObjectType is a function, a procedure or a routine
+ * isFunction returns true if the given ObjectType is a function, a procedure, a routine
- * otherwise returns false
+ * or an aggregate otherwise returns false
  */
 bool
 isFunction(ObjectType objectType)
 {
 	return (objectType == OBJECT_FUNCTION || objectType == OBJECT_PROCEDURE ||
-			objectType == OBJECT_ROUTINE);
+			objectType == OBJECT_ROUTINE || objectType == OBJECT_AGGREGATE);
 }

src/backend/distributed/metadata/metadata_sync.c

@@ -2070,6 +2070,7 @@ GetObjectsForGrantStmt(ObjectType objectType, Oid objectId)
 
 		/* enterprise supported object types */
 		case OBJECT_FUNCTION:
+		case OBJECT_AGGREGATE:
 		case OBJECT_PROCEDURE:
 		{
 			ObjectWithArgs *owa = ObjectWithArgsFromOid(objectId);
@@ -2252,11 +2253,15 @@ GenerateGrantOnFunctionQueriesFromAclItem(Oid functionOid, AclItem *aclItem)
 		{
 			objectType = OBJECT_PROCEDURE;
 		}
+		else if (prokind == PROKIND_AGGREGATE)
+		{
+			objectType = OBJECT_AGGREGATE;
+		}
 		else
 		{
 			ereport(ERROR, (errmsg("unsupported prokind"),
 							errdetail("GRANT commands on procedures are propagated only "
-									  "for procedures and functions.")));
+									  "for procedures, functions, and aggregates.")));
 		}
 
 		char *query = DeparseTreeNode((Node *) GenerateGrantStmtForRights(

src/test/regress/expected/grant_on_function_propagation.out

@@ -274,6 +274,27 @@ GRANT EXECUTE ON FUNCTION function_notice() TO function_user_1, non_existent_use
 ERROR:  role "non_existent_user" does not exist
 GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA grant_on_function, non_existent_schema TO function_user_1;
 ERROR:  schema "non_existent_schema" does not exist
+-- ROUTINE syntax in GRANT should support functions.
+CREATE USER function_user_4;
+GRANT EXECUTE ON ROUTINE function_notice() TO function_user_4;
+-- check if the grants are propagated correctly
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname IN ('function_notice') ORDER BY 1, 2;
+     proname     | pronargs |                                                                            proacl
+---------------------------------------------------------------------
+ function_notice |        0 | {=X/postgres,postgres=X/postgres,function_user_2=X*/postgres,function_user_1=X/function_user_2,function_user_3=X/function_user_2,function_user_4=X/postgres}
+ function_notice |        1 | {=X/postgres,postgres=X/postgres,function_user_2=X*/postgres,function_user_1=X/function_user_2,function_user_3=X/function_user_2}
+(2 rows)
+
+\c - - - :worker_2_port
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname IN ('function_notice') ORDER BY 1, 2;
+     proname     | pronargs |                                                                            proacl
+---------------------------------------------------------------------
+ function_notice |        0 | {=X/postgres,postgres=X/postgres,function_user_2=X*/postgres,function_user_1=X/function_user_2,function_user_3=X/function_user_2,function_user_4=X/postgres}
+ function_notice |        1 | {=X/postgres,postgres=X/postgres,function_user_2=X*/postgres,function_user_1=X/function_user_2,function_user_3=X/function_user_2}
+(2 rows)
+
+\c - - - :master_port
+SET search_path TO grant_on_function, public;
 DROP FUNCTION function_notice(), function_notice(text), function_hello, not_distributed_function, not_propagated_function_user_test;
 -- add similar tests for procedures
 -- remove one of the worker nodes to test adding a new node later
@@ -532,7 +553,281 @@ GRANT EXECUTE ON PROCEDURE procedure_notice() TO procedure_user_1, non_existent_
 ERROR:  role "non_existent_user" does not exist
 GRANT EXECUTE ON ALL PROCEDURES IN SCHEMA grant_on_function, non_existent_schema TO procedure_user_1;
 ERROR:  schema "non_existent_schema" does not exist
+-- ROUTINE syntax in GRANT should support procedures.
+CREATE USER procedure_user_4;
+GRANT EXECUTE ON ROUTINE procedure_notice() TO procedure_user_4;
+-- check if the grants are propagated correctly
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname IN ('procedure_notice') ORDER BY 1, 2;
+     proname      | pronargs |                                                                               proacl
+---------------------------------------------------------------------
+ procedure_notice |        0 | {=X/postgres,postgres=X/postgres,procedure_user_2=X*/postgres,procedure_user_1=X/procedure_user_2,procedure_user_3=X/procedure_user_2,procedure_user_4=X/postgres}
+ procedure_notice |        1 | {=X/postgres,postgres=X/postgres,procedure_user_2=X*/postgres,procedure_user_1=X/procedure_user_2,procedure_user_3=X/procedure_user_2}
+(2 rows)
+
+\c - - - :worker_2_port
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname IN ('procedure_notice') ORDER BY 1, 2;
+     proname      | pronargs |                                                                               proacl
+---------------------------------------------------------------------
+ procedure_notice |        0 | {=X/postgres,postgres=X/postgres,procedure_user_2=X*/postgres,procedure_user_1=X/procedure_user_2,procedure_user_3=X/procedure_user_2,procedure_user_4=X/postgres}
+ procedure_notice |        1 | {=X/postgres,postgres=X/postgres,procedure_user_2=X*/postgres,procedure_user_1=X/procedure_user_2,procedure_user_3=X/procedure_user_2}
+(2 rows)
+
+\c - - - :master_port
+SET search_path TO grant_on_function, public;
 DROP PROCEDURE procedure_notice(), procedure_notice(text), procedure_hello, not_distributed_procedure, not_propagated_procedure_user_test;
+-- add similar tests for aggregates
+-- remove one of the worker nodes to test adding a new node later
+SELECT 1 FROM citus_remove_node('localhost', :worker_2_port);
+ ?column?
+---------------------------------------------------------------------
+        1
+(1 row)
+
+-- create some simple aggregates
+CREATE AGGREGATE dist_float_avg (float8)
+(
+    sfunc = float8_accum,
+    stype = float8[],
+    finalfunc = float8_avg,
+    initcond = '{0,0,0}'
+);
+CREATE AGGREGATE dist_float_avg_2 (float8)
+(
+    sfunc = float8_accum,
+    stype = float8[],
+    finalfunc = float8_avg,
+    initcond = '{0,0,0}'
+);
+SET citus.enable_metadata_sync TO OFF;
+CREATE AGGREGATE no_dist_float_avg (float8)
+(
+    sfunc = float8_accum,
+    stype = float8[],
+    finalfunc = float8_avg,
+    initcond = '{0,0,0}'
+);
+RESET citus.enable_metadata_sync;
+-- create some users and grant them permission on grant_on_function schema
+CREATE USER aggregate_user_1;
+CREATE USER aggregate_user_2;
+CREATE USER aggregate_user_3;
+GRANT ALL ON SCHEMA grant_on_function TO aggregate_user_1, aggregate_user_2, aggregate_user_3;
+-- do some varying grants
+GRANT EXECUTE ON ROUTINE dist_float_avg(float8) TO aggregate_user_1;
+GRANT EXECUTE ON ROUTINE dist_float_avg(float8) TO aggregate_user_2 WITH GRANT OPTION;
+SET ROLE aggregate_user_2;
+GRANT EXECUTE ON ROUTINE dist_float_avg(float8) TO aggregate_user_3;
+RESET ROLE;
+SELECT create_distributed_function('dist_float_avg(float8)');
+NOTICE:  procedure grant_on_function.dist_float_avg is already distributed
+DETAIL:  Citus distributes procedures with CREATE [PROCEDURE|FUNCTION|AGGREGATE] commands
+ create_distributed_function
+---------------------------------------------------------------------
+
+(1 row)
+
+-- re-distributing the same aggregate with GRANTs should be fine
+SELECT create_distributed_function('dist_float_avg(float8)');
+NOTICE:  procedure grant_on_function.dist_float_avg is already distributed
+DETAIL:  Citus distributes procedures with CREATE [PROCEDURE|FUNCTION|AGGREGATE] commands
+ create_distributed_function
+---------------------------------------------------------------------
+
+(1 row)
+
+-- check grants propagated correctly after create_distributed_function
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'dist_float_avg' ORDER BY 1, 2;
+    proname     | pronargs |                                                             proacl
+---------------------------------------------------------------------
+ dist_float_avg |        1 | {=X/postgres,postgres=X/postgres,aggregate_user_1=X/postgres,aggregate_user_2=X*/postgres,aggregate_user_3=X/aggregate_user_2}
+(1 row)
+
+\c - - - :worker_1_port
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'dist_float_avg' ORDER BY 1, 2;
+    proname     | pronargs |                                                             proacl
+---------------------------------------------------------------------
+ dist_float_avg |        1 | {=X/postgres,postgres=X/postgres,aggregate_user_1=X/postgres,aggregate_user_2=X*/postgres,aggregate_user_3=X/aggregate_user_2}
+(1 row)
+
+\c - - - :master_port
+-- do some varying revokes
+REVOKE EXECUTE ON ROUTINE grant_on_function.dist_float_avg(float8) FROM aggregate_user_1, aggregate_user_3;
+REVOKE GRANT OPTION FOR EXECUTE ON ROUTINE grant_on_function.dist_float_avg(float8) FROM aggregate_user_2 CASCADE;
+-- check revokes propagated correctly for the distributed aggregate dist_float_avg(float8)
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'dist_float_avg' ORDER BY 1, 2;
+    proname     | pronargs |                            proacl
+---------------------------------------------------------------------
+ dist_float_avg |        1 | {=X/postgres,postgres=X/postgres,aggregate_user_2=X/postgres}
+(1 row)
+
+\c - - - :worker_1_port
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'dist_float_avg' ORDER BY 1, 2;
+    proname     | pronargs |                            proacl
+---------------------------------------------------------------------
+ dist_float_avg |        1 | {=X/postgres,postgres=X/postgres,aggregate_user_2=X/postgres}
+(1 row)
+
+\c - - - :master_port
+REVOKE EXECUTE ON ROUTINE grant_on_function.dist_float_avg(float8) FROM aggregate_user_2;
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'dist_float_avg' ORDER BY 1, 2;
+    proname     | pronargs |              proacl
+---------------------------------------------------------------------
+ dist_float_avg |        1 | {=X/postgres,postgres=X/postgres}
+(1 row)
+
+\c - - - :worker_1_port
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'dist_float_avg' ORDER BY 1, 2;
+    proname     | pronargs |              proacl
+---------------------------------------------------------------------
+ dist_float_avg |        1 | {=X/postgres,postgres=X/postgres}
+(1 row)
+
+\c - - - :master_port
+-- GRANT .. ON ALL ROUTINES IN SCHEMA .. with multiple roles
+GRANT EXECUTE ON ALL ROUTINES IN SCHEMA grant_on_function TO aggregate_user_1, aggregate_user_3;
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'dist_float_avg' ORDER BY 1, 2;
+    proname     | pronargs |                                          proacl
+---------------------------------------------------------------------
+ dist_float_avg |        1 | {=X/postgres,postgres=X/postgres,aggregate_user_1=X/postgres,aggregate_user_3=X/postgres}
+(1 row)
+
+\c - - - :worker_1_port
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'dist_float_avg' ORDER BY 1, 2;
+    proname     | pronargs |                                          proacl
+---------------------------------------------------------------------
+ dist_float_avg |        1 | {=X/postgres,postgres=X/postgres,aggregate_user_1=X/postgres,aggregate_user_3=X/postgres}
+(1 row)
+
+\c - - - :master_port
+-- REVOKE .. ON ALL PROCEDURES IN SCHEMA .. with multiple roles
+REVOKE EXECUTE ON ALL ROUTINES IN SCHEMA grant_on_function FROM aggregate_user_1, aggregate_user_3;
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'dist_float_avg' ORDER BY 1, 2;
+    proname     | pronargs |              proacl
+---------------------------------------------------------------------
+ dist_float_avg |        1 | {=X/postgres,postgres=X/postgres}
+(1 row)
+
+\c - - - :worker_1_port
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'dist_float_avg' ORDER BY 1, 2;
+    proname     | pronargs |              proacl
+---------------------------------------------------------------------
+ dist_float_avg |        1 | {=X/postgres,postgres=X/postgres}
+(1 row)
+
+\c - - - :master_port
+-- distribute another aggregate
+SET search_path TO grant_on_function, public;
+SELECT create_distributed_function('dist_float_avg_2(float8)');
+NOTICE:  procedure grant_on_function.dist_float_avg_2 is already distributed
+DETAIL:  Citus distributes procedures with CREATE [PROCEDURE|FUNCTION|AGGREGATE] commands
+ create_distributed_function
+---------------------------------------------------------------------
+
+(1 row)
+
+-- GRANT with multiple aggregates and multiple roles
+GRANT EXECUTE ON ROUTINE grant_on_function.dist_float_avg(float8), grant_on_function.dist_float_avg_2(float8),
+grant_on_function.no_dist_float_avg(float8) TO aggregate_user_2 WITH GRANT OPTION;
+SET ROLE aggregate_user_2;
+GRANT EXECUTE ON ROUTINE grant_on_function.dist_float_avg(float8), grant_on_function.dist_float_avg_2(float8),
+grant_on_function.no_dist_float_avg(float8) TO aggregate_user_1, aggregate_user_3;
+RESET ROLE;
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname IN ('dist_float_avg', 'dist_float_avg_2', 'no_dist_float_avg') ORDER BY 1, 2;
+      proname      | pronargs |                                                                 proacl
+---------------------------------------------------------------------
+ dist_float_avg    |        1 | {=X/postgres,postgres=X/postgres,aggregate_user_2=X*/postgres,aggregate_user_1=X/aggregate_user_2,aggregate_user_3=X/aggregate_user_2}
+ dist_float_avg_2  |        1 | {=X/postgres,postgres=X/postgres,aggregate_user_2=X*/postgres,aggregate_user_1=X/aggregate_user_2,aggregate_user_3=X/aggregate_user_2}
+ no_dist_float_avg |        1 | {=X/postgres,postgres=X/postgres,aggregate_user_2=X*/postgres,aggregate_user_1=X/aggregate_user_2,aggregate_user_3=X/aggregate_user_2}
+(3 rows)
+
+\c - - - :worker_1_port
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname IN ('dist_float_avg', 'dist_float_avg_2', 'no_dist_float_avg') ORDER BY 1, 2;
+     proname      | pronargs |                                                                 proacl
+---------------------------------------------------------------------
+ dist_float_avg   |        1 | {=X/postgres,postgres=X/postgres,aggregate_user_2=X*/postgres,aggregate_user_1=X/aggregate_user_2,aggregate_user_3=X/aggregate_user_2}
+ dist_float_avg_2 |        1 | {=X/postgres,postgres=X/postgres,aggregate_user_2=X*/postgres,aggregate_user_1=X/aggregate_user_2,aggregate_user_3=X/aggregate_user_2}
+(2 rows)
+
+\c - - - :master_port
+-- add the previously removed node
+SELECT 1 FROM citus_add_node('localhost', :worker_2_port);
+ ?column?
+---------------------------------------------------------------------
+        1
+(1 row)
+
+-- check if the grants are propagated correctly
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname IN ('dist_float_avg', 'dist_float_avg_2') ORDER BY 1, 2;
+     proname      | pronargs |                                                                 proacl
+---------------------------------------------------------------------
+ dist_float_avg   |        1 | {=X/postgres,postgres=X/postgres,aggregate_user_2=X*/postgres,aggregate_user_1=X/aggregate_user_2,aggregate_user_3=X/aggregate_user_2}
+ dist_float_avg_2 |        1 | {=X/postgres,postgres=X/postgres,aggregate_user_2=X*/postgres,aggregate_user_1=X/aggregate_user_2,aggregate_user_3=X/aggregate_user_2}
+(2 rows)
+
+\c - - - :worker_2_port
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname IN ('dist_float_avg', 'dist_float_avg_2') ORDER BY 1, 2;
+     proname      | pronargs |                                                                 proacl
+---------------------------------------------------------------------
+ dist_float_avg   |        1 | {=X/postgres,postgres=X/postgres,aggregate_user_2=X*/postgres,aggregate_user_1=X/aggregate_user_2,aggregate_user_3=X/aggregate_user_2}
+ dist_float_avg_2 |        1 | {=X/postgres,postgres=X/postgres,aggregate_user_2=X*/postgres,aggregate_user_1=X/aggregate_user_2,aggregate_user_3=X/aggregate_user_2}
+(2 rows)
+
+\c - - - :master_port
+-- check that it works correctly with a user that is not distributed
+CREATE AGGREGATE not_propagated_aggregate_user_test (float8)
+(
+    sfunc = float8_accum,
+    stype = float8[],
+    finalfunc = float8_avg,
+    initcond = '{0,0,0}'
+);
+SET citus.enable_ddl_propagation TO off;
+CREATE USER not_propagated_aggregate_user_4;
+SET citus.enable_ddl_propagation TO on;
+GRANT EXECUTE ON ROUTINE not_propagated_aggregate_user_test TO aggregate_user_1, not_propagated_aggregate_user_4;
+-- check if the grants are propagated correctly
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname IN ('not_propagated_aggregate_user_test') ORDER BY 1, 2;
+              proname               | pronargs |                                                  proacl
+---------------------------------------------------------------------
+ not_propagated_aggregate_user_test |        1 | {=X/postgres,postgres=X/postgres,aggregate_user_1=X/postgres,not_propagated_aggregate_user_4=X/postgres}
+(1 row)
+
+\c - - - :worker_2_port
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname IN ('not_propagated_aggregate_user_test') ORDER BY 1, 2;
+              proname               | pronargs |                                                  proacl
+---------------------------------------------------------------------
+ not_propagated_aggregate_user_test |        1 | {=X/postgres,postgres=X/postgres,aggregate_user_1=X/postgres,not_propagated_aggregate_user_4=X/postgres}
+(1 row)
+
+\c - - - :master_port
+SET search_path TO grant_on_function, public;
+-- the following should fail is in plain PG
+GRANT EXECUTE ON ROUTINE dist_float_avg(float8), non_existent_aggregate TO aggregate_user_1;
+ERROR:  could not find a function named "non_existent_aggregate"
+GRANT EXECUTE ON ROUTINE dist_float_avg(float8) TO aggregate_user_1, non_existent_user;
+ERROR:  role "non_existent_user" does not exist
+GRANT EXECUTE ON ALL ROUTINES IN SCHEMA grant_on_function, non_existent_schema TO aggregate_user_1;
+ERROR:  schema "non_existent_schema" does not exist
+-- ROUTINE syntax in GRANT should support aggregates.
+CREATE USER aggregate_user_4;
+GRANT EXECUTE ON ROUTINE dist_float_avg(float8) TO aggregate_user_4;
+-- check if the grants are propagated correctly
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname IN ('dist_float_avg') ORDER BY 1, 2;
+    proname     | pronargs |                                                                               proacl
+---------------------------------------------------------------------
+ dist_float_avg |        1 | {=X/postgres,postgres=X/postgres,aggregate_user_2=X*/postgres,aggregate_user_1=X/aggregate_user_2,aggregate_user_3=X/aggregate_user_2,aggregate_user_4=X/postgres}
+(1 row)
+
+\c - - - :worker_2_port
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname IN ('dist_float_avg') ORDER BY 1, 2;
+    proname     | pronargs |                                                                               proacl
+---------------------------------------------------------------------
+ dist_float_avg |        1 | {=X/postgres,postgres=X/postgres,aggregate_user_2=X*/postgres,aggregate_user_1=X/aggregate_user_2,aggregate_user_3=X/aggregate_user_2,aggregate_user_4=X/postgres}
+(1 row)
+
+\c - - - :master_port
+SET search_path TO grant_on_function, public;
+DROP AGGREGATE dist_float_avg(float8), dist_float_avg_2(float8), no_dist_float_avg(float8), not_propagated_aggregate_user_test(float8);
 SELECT run_command_on_coordinator_and_workers('DROP SCHEMA grant_on_function CASCADE');
  run_command_on_coordinator_and_workers
 ---------------------------------------------------------------------

src/test/regress/sql/grant_on_function_propagation.sql

@@ -159,6 +159,18 @@ GRANT EXECUTE ON FUNCTION function_notice(), non_existent_function TO function_u
 GRANT EXECUTE ON FUNCTION function_notice() TO function_user_1, non_existent_user;
 GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA grant_on_function, non_existent_schema TO function_user_1;
 
+-- ROUTINE syntax in GRANT should support functions.
+CREATE USER function_user_4;
+GRANT EXECUTE ON ROUTINE function_notice() TO function_user_4;
+
+-- check if the grants are propagated correctly
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname IN ('function_notice') ORDER BY 1, 2;
+\c - - - :worker_2_port
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname IN ('function_notice') ORDER BY 1, 2;
+\c - - - :master_port
+
+SET search_path TO grant_on_function, public;
+
 DROP FUNCTION function_notice(), function_notice(text), function_hello, not_distributed_function, not_propagated_function_user_test;
 
 -- add similar tests for procedures
@@ -308,8 +320,178 @@ GRANT EXECUTE ON PROCEDURE procedure_notice(), non_existent_procedure TO procedu
 GRANT EXECUTE ON PROCEDURE procedure_notice() TO procedure_user_1, non_existent_user;
 GRANT EXECUTE ON ALL PROCEDURES IN SCHEMA grant_on_function, non_existent_schema TO procedure_user_1;
 
+-- ROUTINE syntax in GRANT should support procedures.
+CREATE USER procedure_user_4;
+GRANT EXECUTE ON ROUTINE procedure_notice() TO procedure_user_4;
+
+-- check if the grants are propagated correctly
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname IN ('procedure_notice') ORDER BY 1, 2;
+\c - - - :worker_2_port
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname IN ('procedure_notice') ORDER BY 1, 2;
+\c - - - :master_port
+
+SET search_path TO grant_on_function, public;
+
 DROP PROCEDURE procedure_notice(), procedure_notice(text), procedure_hello, not_distributed_procedure, not_propagated_procedure_user_test;
 
+-- add similar tests for aggregates
+
+-- remove one of the worker nodes to test adding a new node later
+SELECT 1 FROM citus_remove_node('localhost', :worker_2_port);
+
+-- create some simple aggregates
+CREATE AGGREGATE dist_float_avg (float8)
+(
+    sfunc = float8_accum,
+    stype = float8[],
+    finalfunc = float8_avg,
+    initcond = '{0,0,0}'
+);
+
+CREATE AGGREGATE dist_float_avg_2 (float8)
+(
+    sfunc = float8_accum,
+    stype = float8[],
+    finalfunc = float8_avg,
+    initcond = '{0,0,0}'
+);
+
+SET citus.enable_metadata_sync TO OFF;
+CREATE AGGREGATE no_dist_float_avg (float8)
+(
+    sfunc = float8_accum,
+    stype = float8[],
+    finalfunc = float8_avg,
+    initcond = '{0,0,0}'
+);
+
+RESET citus.enable_metadata_sync;
+
+-- create some users and grant them permission on grant_on_function schema
+CREATE USER aggregate_user_1;
+CREATE USER aggregate_user_2;
+CREATE USER aggregate_user_3;
+GRANT ALL ON SCHEMA grant_on_function TO aggregate_user_1, aggregate_user_2, aggregate_user_3;
+
+-- do some varying grants
+GRANT EXECUTE ON ROUTINE dist_float_avg(float8) TO aggregate_user_1;
+GRANT EXECUTE ON ROUTINE dist_float_avg(float8) TO aggregate_user_2 WITH GRANT OPTION;
+SET ROLE aggregate_user_2;
+GRANT EXECUTE ON ROUTINE dist_float_avg(float8) TO aggregate_user_3;
+RESET ROLE;
+
+SELECT create_distributed_function('dist_float_avg(float8)');
+
+-- re-distributing the same aggregate with GRANTs should be fine
+SELECT create_distributed_function('dist_float_avg(float8)');
+
+-- check grants propagated correctly after create_distributed_function
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'dist_float_avg' ORDER BY 1, 2;
+\c - - - :worker_1_port
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'dist_float_avg' ORDER BY 1, 2;
+\c - - - :master_port
+
+-- do some varying revokes
+REVOKE EXECUTE ON ROUTINE grant_on_function.dist_float_avg(float8) FROM aggregate_user_1, aggregate_user_3;
+REVOKE GRANT OPTION FOR EXECUTE ON ROUTINE grant_on_function.dist_float_avg(float8) FROM aggregate_user_2 CASCADE;
+
+-- check revokes propagated correctly for the distributed aggregate dist_float_avg(float8)
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'dist_float_avg' ORDER BY 1, 2;
+\c - - - :worker_1_port
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'dist_float_avg' ORDER BY 1, 2;
+\c - - - :master_port
+
+REVOKE EXECUTE ON ROUTINE grant_on_function.dist_float_avg(float8) FROM aggregate_user_2;
+
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'dist_float_avg' ORDER BY 1, 2;
+\c - - - :worker_1_port
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'dist_float_avg' ORDER BY 1, 2;
+\c - - - :master_port
+
+-- GRANT .. ON ALL ROUTINES IN SCHEMA .. with multiple roles
+GRANT EXECUTE ON ALL ROUTINES IN SCHEMA grant_on_function TO aggregate_user_1, aggregate_user_3;
+
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'dist_float_avg' ORDER BY 1, 2;
+\c - - - :worker_1_port
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'dist_float_avg' ORDER BY 1, 2;
+\c - - - :master_port
+
+-- REVOKE .. ON ALL PROCEDURES IN SCHEMA .. with multiple roles
+REVOKE EXECUTE ON ALL ROUTINES IN SCHEMA grant_on_function FROM aggregate_user_1, aggregate_user_3;
+
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'dist_float_avg' ORDER BY 1, 2;
+\c - - - :worker_1_port
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'dist_float_avg' ORDER BY 1, 2;
+\c - - - :master_port
+
+-- distribute another aggregate
+SET search_path TO grant_on_function, public;
+SELECT create_distributed_function('dist_float_avg_2(float8)');
+
+-- GRANT with multiple aggregates and multiple roles
+GRANT EXECUTE ON ROUTINE grant_on_function.dist_float_avg(float8), grant_on_function.dist_float_avg_2(float8),
+grant_on_function.no_dist_float_avg(float8) TO aggregate_user_2 WITH GRANT OPTION;
+SET ROLE aggregate_user_2;
+GRANT EXECUTE ON ROUTINE grant_on_function.dist_float_avg(float8), grant_on_function.dist_float_avg_2(float8),
+grant_on_function.no_dist_float_avg(float8) TO aggregate_user_1, aggregate_user_3;
+RESET ROLE;
+
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname IN ('dist_float_avg', 'dist_float_avg_2', 'no_dist_float_avg') ORDER BY 1, 2;
+\c - - - :worker_1_port
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname IN ('dist_float_avg', 'dist_float_avg_2', 'no_dist_float_avg') ORDER BY 1, 2;
+\c - - - :master_port
+
+-- add the previously removed node
+SELECT 1 FROM citus_add_node('localhost', :worker_2_port);
+
+-- check if the grants are propagated correctly
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname IN ('dist_float_avg', 'dist_float_avg_2') ORDER BY 1, 2;
+\c - - - :worker_2_port
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname IN ('dist_float_avg', 'dist_float_avg_2') ORDER BY 1, 2;
+\c - - - :master_port
+
+-- check that it works correctly with a user that is not distributed
+CREATE AGGREGATE not_propagated_aggregate_user_test (float8)
+(
+    sfunc = float8_accum,
+    stype = float8[],
+    finalfunc = float8_avg,
+    initcond = '{0,0,0}'
+);
+
+SET citus.enable_ddl_propagation TO off;
+CREATE USER not_propagated_aggregate_user_4;
+SET citus.enable_ddl_propagation TO on;
+
+GRANT EXECUTE ON ROUTINE not_propagated_aggregate_user_test TO aggregate_user_1, not_propagated_aggregate_user_4;
+
+-- check if the grants are propagated correctly
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname IN ('not_propagated_aggregate_user_test') ORDER BY 1, 2;
+\c - - - :worker_2_port
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname IN ('not_propagated_aggregate_user_test') ORDER BY 1, 2;
+\c - - - :master_port
+
+SET search_path TO grant_on_function, public;
+
+-- the following should fail is in plain PG
+GRANT EXECUTE ON ROUTINE dist_float_avg(float8), non_existent_aggregate TO aggregate_user_1;
+GRANT EXECUTE ON ROUTINE dist_float_avg(float8) TO aggregate_user_1, non_existent_user;
+GRANT EXECUTE ON ALL ROUTINES IN SCHEMA grant_on_function, non_existent_schema TO aggregate_user_1;
+
+-- ROUTINE syntax in GRANT should support aggregates.
+CREATE USER aggregate_user_4;
+GRANT EXECUTE ON ROUTINE dist_float_avg(float8) TO aggregate_user_4;
+
+-- check if the grants are propagated correctly
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname IN ('dist_float_avg') ORDER BY 1, 2;
+\c - - - :worker_2_port
+SELECT proname, pronargs, proacl FROM pg_proc WHERE proname IN ('dist_float_avg') ORDER BY 1, 2;
+\c - - - :master_port
+
+SET search_path TO grant_on_function, public;
+
+DROP AGGREGATE dist_float_avg(float8), dist_float_avg_2(float8), no_dist_float_avg(float8), not_propagated_aggregate_user_test(float8);
+
 SELECT run_command_on_coordinator_and_workers('DROP SCHEMA grant_on_function CASCADE');
 DROP USER function_user_1, function_user_2, function_user_3, not_propagated_function_user_4;
 DROP USER procedure_user_1, procedure_user_2, procedure_user_3, not_propagated_procedure_user_4;