Add support for SECURITY LABEL on ROLE propagation from non-main databases (#7525)

DESCRIPTION: Adds support for distributed "SECURITY LABEL on ROLE"
commands from the databases where Citus is not installed.

src/backend/distributed/commands/utility_hook.c

@@ -175,6 +175,7 @@ static MarkObjectDistributedParams GetMarkObjectDistributedParams(Node *parsetre
  * NonMainDbDistributedStatementInfo objects.
  */
 static bool NonMainDbCheckSupportedObjectTypeForGrant(Node *node);
+static bool NonMainDbCheckSupportedObjectTypeForSecLabel(Node *node);
 
 
 /*
@@ -188,6 +189,7 @@ static const NonMainDbDistributedStatementInfo NonMainDbSupportedStatements[] =
 	{ T_GrantStmt, false, NonMainDbCheckSupportedObjectTypeForGrant },
 	{ T_CreatedbStmt, false, NULL },
 	{ T_DropdbStmt, false, NULL },
+	{ T_SecLabelStmt, false, NonMainDbCheckSupportedObjectTypeForSecLabel },
 };
 
 
@@ -1862,3 +1864,15 @@ NonMainDbCheckSupportedObjectTypeForGrant(Node *node)
 	GrantStmt *stmt = castNode(GrantStmt, node);
 	return stmt->objtype == OBJECT_DATABASE;
 }
+
+
+/*
+ * NonMainDbCheckSupportedObjectTypeForSecLabel implements checkSupportedObjectTypes
+ * callback for SecLabel.
+ */
+static bool
+NonMainDbCheckSupportedObjectTypeForSecLabel(Node *node)
+{
+	SecLabelStmt *stmt = castNode(SecLabelStmt, node);
+	return stmt->objtype == OBJECT_ROLE;
+}

src/test/regress/expected/metadata_sync_from_non_maindb.out

@@ -63,7 +63,25 @@ select check_database_privileges('grant_role2pc''_user3','metadata_sync_2pc_db',
  (TEMPORARY,t)
 (8 rows)
 
+-- test for security label on role
+\c metadata_sync_2pc_db - - :master_port
+SECURITY LABEL FOR "citus '!tests_label_provider" ON ROLE grant_role2pc_user4 IS 'citus_unclassified';
+SECURITY LABEL FOR "citus '!tests_label_provider" ON ROLE "grant_role2pc'_user1" IS 'citus_classified';
 \c regression
+SELECT node_type, result FROM get_citus_tests_label_provider_labels('grant_role2pc_user4') ORDER BY node_type;
+  node_type  |                                             result
+---------------------------------------------------------------------
+ coordinator | {"label": "citus_unclassified", "objtype": "role", "provider": "citus '!tests_label_provider"}
+ worker_1    | {"label": "citus_unclassified", "objtype": "role", "provider": "citus '!tests_label_provider"}
+(2 rows)
+
+SELECT node_type, result FROM get_citus_tests_label_provider_labels($$"grant_role2pc''_user1"$$) ORDER BY node_type;
+  node_type  |                                            result
+---------------------------------------------------------------------
+ coordinator | {"label": "citus_classified", "objtype": "role", "provider": "citus '!tests_label_provider"}
+ worker_1    | {"label": "citus_classified", "objtype": "role", "provider": "citus '!tests_label_provider"}
+(2 rows)
+
 set citus.enable_create_database_propagation to on;
 select 1 from citus_add_node('localhost', :worker_2_port);
  ?column?
@@ -121,6 +139,22 @@ select check_database_privileges('grant_role2pc''_user3','metadata_sync_2pc_db',
  (TEMPORARY,t)
 (12 rows)
 
+SELECT node_type, result FROM get_citus_tests_label_provider_labels('grant_role2pc_user4') ORDER BY node_type;
+  node_type  |                                             result
+---------------------------------------------------------------------
+ coordinator | {"label": "citus_unclassified", "objtype": "role", "provider": "citus '!tests_label_provider"}
+ worker_1    | {"label": "citus_unclassified", "objtype": "role", "provider": "citus '!tests_label_provider"}
+ worker_2    | {"label": "citus_unclassified", "objtype": "role", "provider": "citus '!tests_label_provider"}
+(3 rows)
+
+SELECT node_type, result FROM get_citus_tests_label_provider_labels($$"grant_role2pc''_user1"$$) ORDER BY node_type;
+  node_type  |                                            result
+---------------------------------------------------------------------
+ coordinator | {"label": "citus_classified", "objtype": "role", "provider": "citus '!tests_label_provider"}
+ worker_1    | {"label": "citus_classified", "objtype": "role", "provider": "citus '!tests_label_provider"}
+ worker_2    | {"label": "citus_classified", "objtype": "role", "provider": "citus '!tests_label_provider"}
+(3 rows)
+
 \c metadata_sync_2pc_db
 revoke "grant_role2pc'_user1","grant_role2pc'_user2" from grant_role2pc_user4,grant_role2pc_user5 ;
 revoke admin option for "grant_role2pc'_user1","grant_role2pc'_user2" from "grant_role2pc'_user3";

src/test/regress/expected/seclabel_non_maindb.out

@@ -0,0 +1,111 @@
+-- SECLABEL
+--
+-- Test suite for running SECURITY LABEL ON ROLE statements from non-main databases
+SET citus.enable_create_database_propagation to ON;
+CREATE DATABASE database1;
+CREATE DATABASE database2;
+\c - - - :worker_1_port
+SET citus.enable_create_database_propagation to ON;
+CREATE DATABASE database_w1;
+\c - - - :master_port
+CREATE ROLE user1;
+\c database1
+SHOW citus.main_db;
+ citus.main_db
+---------------------------------------------------------------------
+ regression
+(1 row)
+
+SHOW citus.superuser;
+ citus.superuser
+---------------------------------------------------------------------
+ postgres
+(1 row)
+
+CREATE ROLE "user 2";
+-- Set a SECURITY LABEL on a role from a non-main database
+SECURITY LABEL FOR "citus '!tests_label_provider" ON ROLE user1 IS 'citus_classified';
+SECURITY LABEL FOR "citus '!tests_label_provider" ON ROLE "user 2" IS 'citus_unclassified';
+-- Check the result
+\c regression
+SELECT node_type, result FROM get_citus_tests_label_provider_labels('user1') ORDER BY node_type;
+  node_type  |                                            result
+---------------------------------------------------------------------
+ coordinator | {"label": "citus_classified", "objtype": "role", "provider": "citus '!tests_label_provider"}
+ worker_1    | {"label": "citus_classified", "objtype": "role", "provider": "citus '!tests_label_provider"}
+ worker_2    | {"label": "citus_classified", "objtype": "role", "provider": "citus '!tests_label_provider"}
+(3 rows)
+
+SELECT node_type, result FROM get_citus_tests_label_provider_labels('"user 2"') ORDER BY node_type;
+  node_type  |                                             result
+---------------------------------------------------------------------
+ coordinator | {"label": "citus_unclassified", "objtype": "role", "provider": "citus '!tests_label_provider"}
+ worker_1    | {"label": "citus_unclassified", "objtype": "role", "provider": "citus '!tests_label_provider"}
+ worker_2    | {"label": "citus_unclassified", "objtype": "role", "provider": "citus '!tests_label_provider"}
+(3 rows)
+
+\c database1
+-- Set a SECURITY LABEL on database, it should not be propagated
+SECURITY LABEL FOR "citus '!tests_label_provider" ON DATABASE database1 IS 'citus_classified';
+-- Set a SECURITY LABEL on a table, it should not be propagated
+CREATE TABLE a (i int);
+SECURITY LABEL ON TABLE a IS 'citus_classified';
+\c regression
+SELECT node_type, result FROM get_citus_tests_label_provider_labels('database1') ORDER BY node_type;
+  node_type  |                                              result
+---------------------------------------------------------------------
+ coordinator | {"label": "citus_classified", "objtype": "database", "provider": "citus '!tests_label_provider"}
+ worker_1    |
+ worker_2    |
+(3 rows)
+
+-- Check that only the SECURITY LABEL for ROLES is propagated to the non-main databases on other nodes
+\c database_w1 - - :worker_1_port
+SELECT provider, objtype, label, objname FROM pg_seclabels ORDER BY objname;
+           provider           | objtype |       label        | objname
+---------------------------------------------------------------------
+ citus '!tests_label_provider | role    | citus_unclassified | "user 2"
+ citus '!tests_label_provider | role    | citus_classified   | user1
+(2 rows)
+
+-- Check the result after a transaction
+BEGIN;
+SECURITY LABEL FOR "citus '!tests_label_provider" ON ROLE user1 IS 'citus_unclassified';
+SECURITY LABEL FOR "citus '!tests_label_provider" ON DATABASE database_w1 IS 'citus_classified';
+COMMIT;
+\c regression
+SELECT node_type, result FROM get_citus_tests_label_provider_labels('database_w1') ORDER BY node_type;
+  node_type  |                                              result
+---------------------------------------------------------------------
+ coordinator |
+ worker_1    | {"label": "citus_classified", "objtype": "database", "provider": "citus '!tests_label_provider"}
+ worker_2    |
+(3 rows)
+
+SELECT node_type, result FROM get_citus_tests_label_provider_labels('user1') ORDER BY node_type;
+  node_type  |                                             result
+---------------------------------------------------------------------
+ coordinator | {"label": "citus_unclassified", "objtype": "role", "provider": "citus '!tests_label_provider"}
+ worker_1    | {"label": "citus_unclassified", "objtype": "role", "provider": "citus '!tests_label_provider"}
+ worker_2    | {"label": "citus_unclassified", "objtype": "role", "provider": "citus '!tests_label_provider"}
+(3 rows)
+
+BEGIN;
+SECURITY LABEL FOR "citus '!tests_label_provider" ON ROLE "user 2" IS 'citus_classified';
+ROLLBACK;
+SELECT node_type, result FROM get_citus_tests_label_provider_labels('"user 2"') ORDER BY node_type;
+  node_type  |                                             result
+---------------------------------------------------------------------
+ coordinator | {"label": "citus_unclassified", "objtype": "role", "provider": "citus '!tests_label_provider"}
+ worker_1    | {"label": "citus_unclassified", "objtype": "role", "provider": "citus '!tests_label_provider"}
+ worker_2    | {"label": "citus_unclassified", "objtype": "role", "provider": "citus '!tests_label_provider"}
+(3 rows)
+
+-- clean up
+SET citus.enable_create_database_propagation to ON;
+DROP DATABASE database1;
+DROP DATABASE database2;
+DROP DATABASE database_w1;
+DROP ROLE user1;
+DROP ROLE "user 2";
+RESET citus.enable_create_database_propagation;

src/test/regress/multi_schedule

@@ -108,7 +108,7 @@ test: object_propagation_debug
 test: undistribute_table
 test: run_command_on_all_nodes
 test: background_task_queue_monitor
-test: other_databases grant_role_from_non_maindb
+test: other_databases grant_role_from_non_maindb seclabel_non_maindb
 test: citus_internal_access
 
 # Causal clock test

src/test/regress/sql/metadata_sync_from_non_maindb.sql

@@ -37,7 +37,15 @@ select check_database_privileges('grant_role2pc''_user1','metadata_sync_2pc_db',
 select check_database_privileges('grant_role2pc''_user2','metadata_sync_2pc_db',ARRAY['CONNECT']);
 select check_database_privileges('grant_role2pc''_user3','metadata_sync_2pc_db',ARRAY['CREATE','CONNECT','TEMP','TEMPORARY']);
 
+-- test for security label on role
+\c metadata_sync_2pc_db - - :master_port
+SECURITY LABEL FOR "citus '!tests_label_provider" ON ROLE grant_role2pc_user4 IS 'citus_unclassified';
+SECURITY LABEL FOR "citus '!tests_label_provider" ON ROLE "grant_role2pc'_user1" IS 'citus_classified';
+
 \c regression
+SELECT node_type, result FROM get_citus_tests_label_provider_labels('grant_role2pc_user4') ORDER BY node_type;
+SELECT node_type, result FROM get_citus_tests_label_provider_labels($$"grant_role2pc''_user1"$$) ORDER BY node_type;
+
 set citus.enable_create_database_propagation to on;
 select 1 from citus_add_node('localhost', :worker_2_port);
 
@@ -56,6 +64,9 @@ select check_database_privileges('grant_role2pc''_user1','metadata_sync_2pc_db',
 select check_database_privileges('grant_role2pc''_user2','metadata_sync_2pc_db',ARRAY['CONNECT']);
 select check_database_privileges('grant_role2pc''_user3','metadata_sync_2pc_db',ARRAY['CREATE','CONNECT','TEMP','TEMPORARY']);
 
+SELECT node_type, result FROM get_citus_tests_label_provider_labels('grant_role2pc_user4') ORDER BY node_type;
+SELECT node_type, result FROM get_citus_tests_label_provider_labels($$"grant_role2pc''_user1"$$) ORDER BY node_type;
+
 \c metadata_sync_2pc_db
 revoke "grant_role2pc'_user1","grant_role2pc'_user2" from grant_role2pc_user4,grant_role2pc_user5 ;
 

src/test/regress/sql/seclabel_non_maindb.sql

@@ -0,0 +1,71 @@
+-- SECLABEL
+--
+-- Test suite for running SECURITY LABEL ON ROLE statements from non-main databases
+
+SET citus.enable_create_database_propagation to ON;
+
+CREATE DATABASE database1;
+CREATE DATABASE database2;
+
+\c - - - :worker_1_port
+SET citus.enable_create_database_propagation to ON;
+CREATE DATABASE database_w1;
+
+
+\c - - - :master_port
+CREATE ROLE user1;
+\c database1
+SHOW citus.main_db;
+SHOW citus.superuser;
+
+CREATE ROLE "user 2";
+
+-- Set a SECURITY LABEL on a role from a non-main database
+SECURITY LABEL FOR "citus '!tests_label_provider" ON ROLE user1 IS 'citus_classified';
+SECURITY LABEL FOR "citus '!tests_label_provider" ON ROLE "user 2" IS 'citus_unclassified';
+
+-- Check the result
+\c regression
+SELECT node_type, result FROM get_citus_tests_label_provider_labels('user1') ORDER BY node_type;
+SELECT node_type, result FROM get_citus_tests_label_provider_labels('"user 2"') ORDER BY node_type;
+
+\c database1
+-- Set a SECURITY LABEL on database, it should not be propagated
+SECURITY LABEL FOR "citus '!tests_label_provider" ON DATABASE database1 IS 'citus_classified';
+
+-- Set a SECURITY LABEL on a table, it should not be propagated
+CREATE TABLE a (i int);
+SECURITY LABEL ON TABLE a IS 'citus_classified';
+
+\c regression
+SELECT node_type, result FROM get_citus_tests_label_provider_labels('database1') ORDER BY node_type;
+
+-- Check that only the SECURITY LABEL for ROLES is propagated to the non-main databases on other nodes
+\c database_w1 - - :worker_1_port
+SELECT provider, objtype, label, objname FROM pg_seclabels ORDER BY objname;
+
+
+-- Check the result after a transaction
+BEGIN;
+SECURITY LABEL FOR "citus '!tests_label_provider" ON ROLE user1 IS 'citus_unclassified';
+SECURITY LABEL FOR "citus '!tests_label_provider" ON DATABASE database_w1 IS 'citus_classified';
+COMMIT;
+
+\c regression
+SELECT node_type, result FROM get_citus_tests_label_provider_labels('database_w1') ORDER BY node_type;
+SELECT node_type, result FROM get_citus_tests_label_provider_labels('user1') ORDER BY node_type;
+
+BEGIN;
+SECURITY LABEL FOR "citus '!tests_label_provider" ON ROLE "user 2" IS 'citus_classified';
+ROLLBACK;
+
+SELECT node_type, result FROM get_citus_tests_label_provider_labels('"user 2"') ORDER BY node_type;
+
+-- clean up
+SET citus.enable_create_database_propagation to ON;
+DROP DATABASE database1;
+DROP DATABASE database2;
+DROP DATABASE database_w1;
+DROP ROLE user1;
+DROP ROLE "user 2";
+RESET citus.enable_create_database_propagation;