addressing reviews

src/backend/distributed/commands/role.c

@@ -570,7 +570,7 @@ GenerateCreateOrAlterRoleCommand(Oid roleOid)
 			completeRoleList = lappend(completeRoleList, DeparseTreeNode(stmt));
 		}
 
-		/* append SECURITY LABEL ON ROLE commands fot this specific user */
+		/* append SECURITY LABEL ON ROLE commands for this specific user */
 		List *secLabelOnRoleStmts = GenerateSecLabelOnRoleStmts(roleOid, rolename);
 		stmt = NULL;
 		foreach_ptr(stmt, secLabelOnRoleStmts)
@@ -913,20 +913,20 @@ GenerateGrantRoleStmtsOfRole(Oid roleid)
 static List *
 GenerateSecLabelOnRoleStmts(Oid roleid, char *rolename)
 {
-	ScanKeyData skey[1];
-	HeapTuple tuple = NULL;
 	List *secLabelStmts = NIL;
 
 	/*
 	 * Note that roles are shared database objects, therefore their
 	 * security labels are stored in pg_shseclabel instead of pg_seclabel.
 	 */
 	Relation pg_shseclabel = table_open(SharedSecLabelRelationId, AccessShareLock);
+	ScanKeyData skey[1];
 	ScanKeyInit(&skey[0], Anum_pg_shseclabel_objoid, BTEqualStrategyNumber, F_OIDEQ,
 				ObjectIdGetDatum(roleid));
 	SysScanDesc scan = systable_beginscan(pg_shseclabel, SharedSecLabelObjectIndexId,
 										  true, NULL, 1, &skey[0]);
 
+	HeapTuple tuple = NULL;
 	while (HeapTupleIsValid(tuple = systable_getnext(scan)))
 	{
 		SecLabelStmt *secLabelStmt = makeNode(SecLabelStmt);

src/backend/distributed/commands/seclabel.c

@@ -11,7 +11,6 @@
 
 #include "postgres.h"
 
-#include "commands/seclabel.h"
 #include "distributed/commands.h"
 #include "distributed/commands/utility_hook.h"
 #include "distributed/coordinator_protocol.h"
@@ -20,22 +19,6 @@
 #include "distributed/metadata_sync.h"
 #include "distributed/metadata/distobject.h"
 
-
-PG_FUNCTION_INFO_V1(citus_test_register_label_provider);
-
-
-/*
- * citus_test_register_label_provider registers a dummy label provider
- * named 'citus_tests_label_provider'. This is aimed to be used for testing.
- */
-Datum
-citus_test_register_label_provider(PG_FUNCTION_ARGS)
-{
-	register_label_provider("citus_tests_label_provider", citus_test_object_relabel);
-	PG_RETURN_VOID();
-}
-
-
 /*
  * PreprocessSecLabelStmt is executed before the statement is applied to the local
  * postgres instance.
@@ -54,7 +37,7 @@ PreprocessSecLabelStmt(Node *node, const char *queryString,
 
 	SecLabelStmt *secLabelStmt = castNode(SecLabelStmt, node);
 
-	List *objectAddresses = GetObjectAddressListFromParseTree(node, false, false);
+	List *objectAddresses = GetObjectAddressListFromParseTree(node, false, true);
 	if (!IsAnyObjectDistributed(objectAddresses))
 	{
 		return NIL;
@@ -64,11 +47,11 @@ PreprocessSecLabelStmt(Node *node, const char *queryString,
 	{
 		if (EnableUnsupportedFeatureMessages)
 		{
-			ereport(NOTICE, (errmsg("not propagating SECURITY LABEL commands whose"
-									" object type is not role"),
-							 errhint(
-								 "Connect to worker nodes directly to manually run the same"
-								 " SECURITY LABEL command after disabling DDL propagation.")));
+			ereport(NOTICE, (errmsg("not propagating SECURITY LABEL commands whose "
+									"object type is not role"),
+							 errhint("Connect to worker nodes directly to manually "
+									 "run the same SECURITY LABEL command after "
+									 "disabling DDL propagation.")));
 		}
 		return NIL;
 	}

src/backend/distributed/operations/shard_rebalancer.c

@@ -317,7 +317,7 @@ PG_FUNCTION_INFO_V1(citus_rebalance_start);
 PG_FUNCTION_INFO_V1(citus_rebalance_stop);
 PG_FUNCTION_INFO_V1(citus_rebalance_wait);
 
-bool RunningUnderIsolationTest = false;
+bool RunningUnderCitusTestSuite = false;
 int MaxRebalancerLoggedIgnoredMoves = 5;
 int RebalancerByDiskSizeBaseCost = 100 * 1024 * 1024;
 bool PropagateSessionSettingsForLoopbackConnection = false;

src/backend/distributed/replication/multi_logical_replication.c

@@ -1143,7 +1143,7 @@ ConflictWithIsolationTestingBeforeCopy(void)
 	const bool sessionLock = false;
 	const bool dontWait = false;
 
-	if (RunningUnderIsolationTest)
+	if (RunningUnderCitusTestSuite)
 	{
 		SET_LOCKTAG_ADVISORY(tag, MyDatabaseId,
 							 SHARD_MOVE_ADVISORY_LOCK_SECOND_KEY,
@@ -1177,7 +1177,7 @@ ConflictWithIsolationTestingAfterCopy(void)
 	const bool sessionLock = false;
 	const bool dontWait = false;
 
-	if (RunningUnderIsolationTest)
+	if (RunningUnderCitusTestSuite)
 	{
 		SET_LOCKTAG_ADVISORY(tag, MyDatabaseId,
 							 SHARD_MOVE_ADVISORY_LOCK_FIRST_KEY,

src/backend/distributed/shared_library_init.c

@@ -29,6 +29,7 @@
 #include "citus_version.h"
 #include "commands/explain.h"
 #include "commands/extension.h"
+#include "commands/seclabel.h"
 #include "common/string.h"
 #include "executor/executor.h"
 #include "distributed/backend_data.h"
@@ -574,6 +575,16 @@ _PG_init(void)
 	INIT_COLUMNAR_SYMBOL(PGFunction, columnar_storage_info);
 	INIT_COLUMNAR_SYMBOL(PGFunction, columnar_store_memory_stats);
 	INIT_COLUMNAR_SYMBOL(PGFunction, test_columnar_storage_write_new_page);
+
+	/*
+	 * This part is only for SECURITY LABEL tests
+	 * mimicking what an actual security label provider would do
+	 */
+	if (RunningUnderCitusTestSuite)
+	{
+		register_label_provider("citus_tests_label_provider",
+								citus_test_object_relabel);
+	}
 }
 
 
@@ -2305,13 +2316,14 @@ RegisterCitusConfigVariables(void)
 		WarnIfReplicationModelIsSet, NULL, NULL);
 
 	DefineCustomBoolVariable(
-		"citus.running_under_isolation_test",
+		"citus.running_under_citus_test_suite",
 		gettext_noop(
 			"Only useful for testing purposes, when set to true, Citus does some "
-			"tricks to implement useful isolation tests with rebalancing. Should "
+			"tricks to implement useful isolation tests with rebalancing. It also "
+			"registers a dummy label provider for SECURITY LABEL tests. Should "
 			"never be set to true on production systems "),
 		gettext_noop("for details of the tricks implemented, refer to the source code"),
-		&RunningUnderIsolationTest,
+		&RunningUnderCitusTestSuite,
 		false,
 		PGC_SUSET,
 		GUC_SUPERUSER_ONLY | GUC_NO_SHOW_ALL | GUC_NOT_IN_SAMPLE,

src/include/distributed/shard_rebalancer.h

@@ -189,7 +189,7 @@ typedef struct RebalancePlanFunctions
 extern char *VariablesToBePassedToNewConnections;
 extern int MaxRebalancerLoggedIgnoredMoves;
 extern int RebalancerByDiskSizeBaseCost;
-extern bool RunningUnderIsolationTest;
+extern bool RunningUnderCitusTestSuite;
 extern bool PropagateSessionSettingsForLoopbackConnection;
 extern int MaxBackgroundTaskExecutorsPerNode;
 

src/test/regress/Makefile

@@ -101,7 +101,7 @@ check-base: all
 
 # check-minimal only sets up the cluster
 check-minimal: all
-	$(pg_regress_multi_check) --load-extension=citus \
+	$(pg_regress_multi_check) --load-extension=citus --seclabel \
 	-- $(MULTI_REGRESS_OPTS) --schedule=$(citus_abs_srcdir)/minimal_schedule $(EXTRA_TESTS)
 
 check-base-vg: all
@@ -159,7 +159,7 @@ check-enterprise: all
 	$(pg_regress_multi_check) --load-extension=citus \
 	-- $(MULTI_REGRESS_OPTS) --schedule=$(citus_abs_srcdir)/enterprise_schedule $(EXTRA_TESTS)
 check-multi-1: all
-	$(pg_regress_multi_check) --load-extension=citus \
+	$(pg_regress_multi_check) --load-extension=citus --seclabel-test \
 	-- $(MULTI_REGRESS_OPTS) --schedule=$(citus_abs_srcdir)/multi_1_schedule $(EXTRA_TESTS)
 
 check-multi-hyperscale: all

src/test/regress/expected/seclabel.out

@@ -11,17 +11,6 @@ SELECT citus_remove_node('localhost', :worker_2_port);
 
 (1 row)
 
--- now we register a label provider
-CREATE FUNCTION citus_test_register_label_provider()
-  RETURNS void
-  LANGUAGE C
-  AS 'citus', $$citus_test_register_label_provider$$;
-SELECT citus_test_register_label_provider();
- citus_test_register_label_provider
----------------------------------------------------------------------
-
-(1 row)
-
 CREATE ROLE user1;
 -- check an invalid label for our current dummy hook citus_test_object_relabel
 SECURITY LABEL ON ROLE user1 IS 'invalid_label';
@@ -42,12 +31,6 @@ SELECT objtype, objname, provider, label FROM pg_seclabels;
 (0 rows)
 
 \c - - - :master_port
-SELECT citus_test_register_label_provider();
- citus_test_register_label_provider
----------------------------------------------------------------------
-
-(1 row)
-
 -- check that we only support propagating for roles
 SET citus.shard_replication_factor to 1;
 CREATE TABLE a (a int);
@@ -61,51 +44,6 @@ SECURITY LABEL ON TABLE a IS 'citus_classified';
 NOTICE:  not propagating SECURITY LABEL commands whose object type is not role
 HINT:  Connect to worker nodes directly to manually run the same SECURITY LABEL command after disabling DDL propagation.
 DROP TABLE a;
--- the registered label provider is per session only
--- this means that we need to maintain the same connection to the worker node
--- in order for the label provider to be visible there
--- hence here we create the necessary session_level_connection_to_node functions
-SET citus.enable_metadata_sync TO off;
-CREATE OR REPLACE FUNCTION start_session_level_connection_to_node(text, integer)
-        RETURNS void
-        LANGUAGE C STRICT VOLATILE
-        AS 'citus', $$start_session_level_connection_to_node$$;
-CREATE OR REPLACE FUNCTION override_backend_data_gpid(bigint)
-        RETURNS void
-        LANGUAGE C STRICT IMMUTABLE
-        AS 'citus', $$override_backend_data_gpid$$;
-SELECT run_command_on_workers($$SET citus.enable_metadata_sync TO off;CREATE OR REPLACE FUNCTION override_backend_data_gpid(bigint)
-        RETURNS void
-        LANGUAGE C STRICT IMMUTABLE
-        AS 'citus'$$);
- run_command_on_workers
----------------------------------------------------------------------
- (localhost,57637,t,SET)
-(1 row)
-
-CREATE OR REPLACE FUNCTION run_commands_on_session_level_connection_to_node(text)
-        RETURNS void
-        LANGUAGE C STRICT VOLATILE
-        AS 'citus', $$run_commands_on_session_level_connection_to_node$$;
-CREATE OR REPLACE FUNCTION stop_session_level_connection_to_node()
-        RETURNS void
-        LANGUAGE C STRICT VOLATILE
-        AS 'citus', $$stop_session_level_connection_to_node$$;
-RESET citus.enable_metadata_sync;
--- now we establish a connection to the worker node
-SELECT start_session_level_connection_to_node('localhost', :worker_1_port);
- start_session_level_connection_to_node
----------------------------------------------------------------------
-
-(1 row)
-
--- with that same connection, we register the label provider in the worker node
-SELECT run_commands_on_session_level_connection_to_node('SELECT citus_test_register_label_provider()');
- run_commands_on_session_level_connection_to_node
----------------------------------------------------------------------
-
-(1 row)
-
 SET citus.log_remote_commands TO on;
 SET citus.grep_remote_commands = '%SECURITY LABEL%';
 -- then we run a security label statement which will use the same connection to the worker node
@@ -120,12 +58,6 @@ SECURITY LABEL for citus_tests_label_provider ON ROLE user1 IS 'citus_unclassifi
 NOTICE:  issuing SECURITY LABEL FOR citus_tests_label_provider ON ROLE user1 IS 'citus_unclassified'
 DETAIL:  on server postgres@localhost:xxxxx connectionId: xxxxxxx
 RESET citus.log_remote_commands;
-SELECT stop_session_level_connection_to_node();
- stop_session_level_connection_to_node
----------------------------------------------------------------------
-
-(1 row)
-
 \c - - - :worker_1_port
 SELECT objtype, objname, provider, label FROM pg_seclabels;
  objtype | objname |          provider          |       label
@@ -142,24 +74,11 @@ SET citus.grep_remote_commands = '%SECURITY LABEL%';
 SELECT 1 FROM citus_add_node('localhost', :worker_2_port);
 NOTICE:  issuing SELECT worker_create_or_alter_role('user1', 'CREATE ROLE user1 NOSUPERUSER NOCREATEDB NOCREATEROLE INHERIT NOLOGIN NOREPLICATION NOBYPASSRLS CONNECTION LIMIT -1 PASSWORD NULL VALID UNTIL ''infinity''', 'ALTER ROLE user1 NOSUPERUSER NOCREATEDB NOCREATEROLE INHERIT NOLOGIN NOREPLICATION NOBYPASSRLS CONNECTION LIMIT -1 PASSWORD NULL VALID UNTIL ''infinity''');SECURITY LABEL FOR citus_tests_label_provider ON ROLE user1 IS 'citus_unclassified'
 DETAIL:  on server postgres@localhost:xxxxx connectionId: xxxxxxx
-WARNING:  security label provider "citus_tests_label_provider" is not loaded
-CONTEXT:  while executing command on localhost:xxxxx
-ERROR:  failure on connection marked as essential: localhost:xxxxx
--- cleanup
-RESET citus.log_remote_commands;
-DROP FUNCTION stop_session_level_connection_to_node, run_commands_on_session_level_connection_to_node,
-              override_backend_data_gpid, start_session_level_connection_to_node;
-SELECT run_command_on_workers($$ DROP FUNCTION override_backend_data_gpid $$);
-       run_command_on_workers
----------------------------------------------------------------------
- (localhost,57637,t,"DROP FUNCTION")
-(1 row)
-
-DROP FUNCTION citus_test_register_label_provider;
-DROP ROLE user1;
-SELECT 1 FROM citus_add_node('localhost', :worker_2_port);
  ?column?
 ---------------------------------------------------------------------
         1
 (1 row)
 
+-- cleanup
+RESET citus.log_remote_commands;
+DROP ROLE user1;

src/test/regress/pg_regress_multi.pl

@@ -50,6 +50,7 @@ ()
     print "  --connection-timeout	Timeout for connecting to worker nodes\n";
     print "  --mitmproxy        	Start a mitmproxy for one of the workers\n";
     print "  --worker-count         Number of workers in Citus cluster (default: 2)\n";
+	print "  --seclabel-test        This test runs seclabel propagation tests";
     exit 1;
 }
 
@@ -86,6 +87,7 @@ ()
 my $publicWorker1Host = "localhost";
 my $publicWorker2Host = "localhost";
 my $workerCount = 2;
+my $seclabelTest = 0;
 
 my $serversAreShutdown = "TRUE";
 my $usingWindows = 0;
@@ -120,6 +122,7 @@ ()
     'worker-1-public-hostname=s' => \$publicWorker1Host,
     'worker-2-public-hostname=s' => \$publicWorker2Host,
     'worker-count=i' => \$workerCount,
+	'seclabel-test' => \$seclabelTest,
     'help' => sub { Usage() });
 
 my $fixopen = "$bindir/postgres.fixopen";
@@ -560,7 +563,7 @@ sub generate_hba
    push(@pgOptions, "citus.metadata_sync_interval=1000");
    push(@pgOptions, "citus.metadata_sync_retry_interval=100");
    push(@pgOptions, "client_min_messages='warning'"); # pg12 introduced notice showing during isolation tests
-   push(@pgOptions, "citus.running_under_isolation_test=true");
+   push(@pgOptions, "citus.running_under_citus_test_suite=true");
 
    # Disable all features of the maintenance daemon. Otherwise queries might
    # randomly show temporarily as "waiting..." because they are waiting for the
@@ -573,6 +576,14 @@ sub generate_hba
    push(@pgOptions, "citus.background_task_queue_interval=-1");
 }
 
+# if the security label propagation tests are running in the suite
+# we need to load the label provider in PG_init by setting
+# running_under_citus_test_suite GUC to true
+if($seclabelTest)
+{
+	push(@pgOptions, "citus.running_under_citus_test_suite=true");
+}
+
 # Add externally added options last, so they overwrite the default ones above
 for my $option (@userPgOptions)
 {