Fixes the issue seen in https://github.com/citusdata/citus-enterprise/issues/745

With this commit, rebalancer backends are identified by application_name = citus_rebalancer and the regular internal backends are identified by application_name = citus_internal
2022-01-27 15:29:51 -08:00 · 2022-01-27 15:29:51 -08:00 · f31bce5b48
parent b072b9235e
commit f31bce5b48
17 changed files with 229 additions and 159 deletions
--- a/src/backend/distributed/commands/call.c
+++ b/src/backend/distributed/commands/call.c
@ -68,7 +68,7 @@ CallDistributedProcedureRemotely(CallStmt *callStmt, DestReceiver *dest)
 		return false;
 	}

-	if (IsCitusInitiatedRemoteBackend())
+	if (IsCitusInternalBackend())
 	{
 		/*
 		 * We are in a citus-initiated backend handling a CALL to a distributed
--- a/src/backend/distributed/commands/utility_hook.c
+++ b/src/backend/distributed/commands/utility_hook.c
@ -897,7 +897,7 @@ ShouldCheckUndistributeCitusLocalTables(void)
 		return false;
 	}

-	if (IsCitusInitiatedRemoteBackend())
+	if (IsCitusInternalBackend() || IsRebalancerInternalBackend())
 	{
 		/* connection from the coordinator operating on a shard */
 		return false;
--- a/src/backend/distributed/connection/connection_management.c
+++ b/src/backend/distributed/connection/connection_management.c
@ -1429,7 +1429,7 @@ ShouldShutdownConnection(MultiConnection *connection, const int cachedConnection
 	 * escalating the number of cached connections. We can recognize such backends
 	 * from their application name.
 	 */
-	return IsCitusInitiatedRemoteBackend() ||
+	return (IsCitusInternalBackend() || IsRebalancerInternalBackend()) ||
 		   connection->initilizationState != POOL_STATE_INITIALIZED ||
 		   cachedConnectionCount >= MaxCachedConnectionsPerWorker ||
 		   connection->forceCloseAtTransactionEnd ||
@ -1441,12 +1441,23 @@ ShouldShutdownConnection(MultiConnection *connection, const int cachedConnection
 }


+/*
+ * IsRebalancerInitiatedBackend returns true if we are in a backend that citus
+ * rebalancer initiated.
+ */
+bool
+IsRebalancerInternalBackend(void)
+{
+	return application_name && strcmp(application_name, CITUS_REBALANCER_NAME) == 0;
+}
+
+
 /*
 * IsCitusInitiatedRemoteBackend returns true if we are in a backend that citus
 * initiated via remote connection.
 */
 bool
-IsCitusInitiatedRemoteBackend(void)
+IsCitusInternalBackend(void)
 {
 	return application_name && strcmp(application_name, CITUS_APPLICATION_NAME) == 0;
 }
--- a/src/backend/distributed/executor/multi_executor.c
+++ b/src/backend/distributed/executor/multi_executor.c
@ -779,7 +779,7 @@ InTaskExecution(void)
 	 * is considered a task execution, but an exception is when we
 	 * are in a delegated function/procedure call.
 	 */
-	return IsCitusInitiatedRemoteBackend() &&
+	return IsCitusInternalBackend() &&
 		   !InTopLevelDelegatedFunctionCall &&
 		   !InDelegatedProcedureCall;
 }
--- a/src/backend/distributed/metadata/metadata_sync.c
+++ b/src/backend/distributed/metadata/metadata_sync.c
@ -2636,7 +2636,8 @@ EnsureCoordinatorInitiatedOperation(void)
 	 * check. The other two checks are to ensure that the operation is initiated
 	 * by the coordinator.
 	 */
-	if (!IsCitusInitiatedRemoteBackend() || !MyBackendIsInDisributedTransaction() ||
+	if (!(IsCitusInternalBackend() || IsRebalancerInternalBackend()) ||
+		!MyBackendIsInDisributedTransaction() ||
 		GetLocalGroupId() == COORDINATOR_GROUP_ID)
 	{
 		ereport(ERROR, (errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
--- a/src/backend/distributed/operations/shard_cleaner.c
+++ b/src/backend/distributed/operations/shard_cleaner.c
@ -96,7 +96,7 @@ isolation_cleanup_orphaned_shards(PG_FUNCTION_ARGS)
 void
 DropOrphanedShardsInSeparateTransaction(void)
 {
-	ExecuteCriticalCommandInSeparateTransaction("CALL citus_cleanup_orphaned_shards()");
+	ExecuteRebalancerCommandInSeparateTransaction("CALL citus_cleanup_orphaned_shards()");
 }


--- a/src/backend/distributed/operations/shard_rebalancer.c
+++ b/src/backend/distributed/operations/shard_rebalancer.c
@ -919,10 +919,10 @@ citus_drain_node(PG_FUNCTION_ARGS)
 	 * This is done in a separate session. This way it's not undone if the
 	 * draining fails midway through.
 	 */
-	ExecuteCriticalCommandInSeparateTransaction(psprintf(
-													"SELECT master_set_node_property(%s, %i, 'shouldhaveshards', false)",
-													quote_literal_cstr(nodeName),
-													nodePort));
+	ExecuteRebalancerCommandInSeparateTransaction(psprintf(
+													  "SELECT master_set_node_property(%s, %i, 'shouldhaveshards', false)",
+													  quote_literal_cstr(nodeName),
+													  nodePort));

 	RebalanceTableShards(&options, shardTransferModeOid);

@ -1696,7 +1696,7 @@ UpdateShardPlacement(PlacementUpdateEvent *placementUpdateEvent,
 	 * In case of failure, we throw an error such that rebalance_table_shards
 	 * fails early.
 	 */
-	ExecuteCriticalCommandInSeparateTransaction(placementUpdateCommand->data);
+	ExecuteRebalancerCommandInSeparateTransaction(placementUpdateCommand->data);

 	UpdateColocatedShardPlacementProgress(shardId,
 										  sourceNode->workerName,
@ -1711,12 +1711,18 @@ UpdateShardPlacement(PlacementUpdateEvent *placementUpdateEvent,
 * don't want to rollback when the current transaction is rolled back.
 */
 void
-ExecuteCriticalCommandInSeparateTransaction(char *command)
+ExecuteRebalancerCommandInSeparateTransaction(char *command)
 {
 	int connectionFlag = FORCE_NEW_CONNECTION;
 	MultiConnection *connection = GetNodeConnection(connectionFlag, LocalHostName,
 													PostPortNumber);
+	StringInfo setApplicationName = makeStringInfo();
+	appendStringInfo(setApplicationName, "SET application_name TO %s",
+					 CITUS_REBALANCER_NAME);
+
+	ExecuteCriticalRemoteCommand(connection, setApplicationName->data);
 	ExecuteCriticalRemoteCommand(connection, command);
+
 	CloseConnection(connection);
 }

--- a/src/backend/distributed/planner/function_call_delegation.c
+++ b/src/backend/distributed/planner/function_call_delegation.c
@ -260,7 +260,7 @@ TryToDelegateFunctionCall(DistributedPlanningContext *planContext)
 		ereport(DEBUG4, (errmsg("function is distributed")));
 	}

-	if (IsCitusInitiatedRemoteBackend())
+	if (IsCitusInternalBackend())
 	{
 		bool isFunctionForceDelegated = procedure->forceDelegation;

--- a/src/backend/distributed/transaction/citus_dist_stat_activity.c
+++ b/src/backend/distributed/transaction/citus_dist_stat_activity.c
@ -188,7 +188,7 @@ FROM \
 	get_all_active_transactions() AS dist_txs(database_id, process_id, initiator_node_identifier, worker_query, transaction_number, transaction_stamp) \
 	ON pg_stat_activity.pid = dist_txs.process_id \
 WHERE \
-	pg_stat_activity.application_name = 'citus' \
+	pg_stat_activity.application_name = 'citus_internal' \
 	AND \
 	pg_stat_activity.query NOT ILIKE '%stat_activity%';"

--- a/src/backend/distributed/transaction/transaction_management.c
+++ b/src/backend/distributed/transaction/transaction_management.c
@ -824,7 +824,7 @@ EnsurePrepareTransactionIsAllowed(void)
 		return;
 	}

-	if (IsCitusInitiatedRemoteBackend())
+	if (IsCitusInternalBackend())
 	{
 		/*
 		 * If this is a Citus-initiated backend.
--- a/src/backend/distributed/worker/worker_shard_visibility.c
+++ b/src/backend/distributed/worker/worker_shard_visibility.c
@ -150,7 +150,8 @@ ErrorIfRelationIsAKnownShard(Oid relationId)
 void
 ErrorIfIllegallyChangingKnownShard(Oid relationId)
 {
-	if (LocalExecutorLevel > 0 || IsCitusInitiatedRemoteBackend() ||
+	if (LocalExecutorLevel > 0 ||
+		(IsCitusInternalBackend() || IsRebalancerInternalBackend()) ||
 		EnableManualChangesToShards)
 	{
 		return;
@ -330,7 +331,7 @@ ResetHideShardsDecision(void)
 static bool
 ShouldHideShardsInternal(void)
 {
-	if (IsCitusInitiatedRemoteBackend())
+	if (IsCitusInternalBackend() || IsRebalancerInternalBackend())
 	{
 		/* we never hide shards from Citus */
 		return false;
--- a/src/include/distributed/connection_management.h
+++ b/src/include/distributed/connection_management.h
@ -29,7 +29,10 @@
 #define ERROR_BUFFER_SIZE 256

 /* application name used for internal connections in Citus */
-#define CITUS_APPLICATION_NAME "citus"
+#define CITUS_APPLICATION_NAME "citus_internal"
+
+/* application name used for internal connections in rebalancer */
+#define CITUS_REBALANCER_NAME "citus_rebalancer"

 /* forward declare, to avoid forcing large headers on everyone */
 struct pg_conn; /* target of the PGconn typedef */
@ -277,7 +280,8 @@ extern void FinishConnectionListEstablishment(List *multiConnectionList);
 extern void FinishConnectionEstablishment(MultiConnection *connection);
 extern void ClaimConnectionExclusively(MultiConnection *connection);
 extern void UnclaimConnection(MultiConnection *connection);
-extern bool IsCitusInitiatedRemoteBackend(void);
+extern bool IsCitusInternalBackend(void);
+extern bool IsRebalancerInternalBackend(void);
 extern void MarkConnectionConnected(MultiConnection *connection);

 /* time utilities */
--- a/src/include/distributed/shard_rebalancer.h
+++ b/src/include/distributed/shard_rebalancer.h
@ -190,7 +190,7 @@ extern List * RebalancePlacementUpdates(List *workerNodeList, List *shardPlaceme
 										RebalancePlanFunctions *rebalancePlanFunctions);
 extern List * ReplicationPlacementUpdates(List *workerNodeList, List *shardPlacementList,
 										  int shardReplicationFactor);
-extern void ExecuteCriticalCommandInSeparateTransaction(char *command);
+extern void ExecuteRebalancerCommandInSeparateTransaction(char *command);


 #endif   /* SHARD_REBALANCER_H */
--- a/src/test/regress/expected/metadata_sync_helpers.out
+++ b/src/test/regress/expected/metadata_sync_helpers.out
@ -36,7 +36,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	SELECT citus_internal_add_partition_metadata ('test'::regclass, 'h', 'col_1', 0, 's');
 ERROR:  This is an internal Citus function can only be used in a distributed transaction
 ROLLBACK;
@ -73,7 +73,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	SELECT citus_internal_add_partition_metadata ('test'::regclass, 'h', 'col_1', 0, 's');
 ERROR:  must be owner of table test
 ROLLBACK;
@ -85,7 +85,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	SELECT citus_internal_update_relation_colocation ('test'::regclass, 10);
 ERROR:  must be owner of table test
 ROLLBACK;
@ -99,7 +99,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	SELECT citus_internal_add_partition_metadata ('test_2'::regclass, 'h', 'col_1', 0, 's');
 citus_internal_add_partition_metadata
 ---------------------------------------------------------------------
@ -121,7 +121,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	SELECT citus_internal_add_partition_metadata ('test_2'::regclass, 'X', 'col_1', 0, 's');
 ERROR:  Metadata syncing is only allowed for hash, reference and local tables:X
 ROLLBACK;
@ -133,7 +133,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	SELECT citus_internal_add_partition_metadata ('test_2'::regclass, 'h', 'non_existing_col', 0, 's');
 ERROR:  column "non_existing_col" of relation "test_2" does not exist
 ROLLBACK;
@ -145,7 +145,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-		SET application_name to 'citus';
+		SET application_name to 'citus_internal';
 		SELECT citus_internal_add_partition_metadata (NULL, 'h', 'non_existing_col', 0, 's');
 ERROR:  relation cannot be NULL
 ROLLBACK;
@ -157,7 +157,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	SELECT citus_internal_add_partition_metadata ('test_2'::regclass, 'h', 'col_1', -1, 's');
 ERROR:  Metadata syncing is only allowed for valid colocation id values.
 ROLLBACK;
@ -169,7 +169,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	SELECT citus_internal_add_partition_metadata ('test_2'::regclass, 'X', 'col_1', 0, 'X');
 ERROR:  Metadata syncing is only allowed for hash, reference and local tables:X
 ROLLBACK;
@ -181,7 +181,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	SELECT citus_internal_add_partition_metadata ('test_2'::regclass, 'h', 'col_1', 0, 's');
 citus_internal_add_partition_metadata
@ -200,7 +200,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	SELECT citus_internal_add_partition_metadata ('test_2'::regclass, 'h', 'col_1', 0, 's');
 citus_internal_add_partition_metadata
@ -219,7 +219,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	SELECT citus_internal_add_partition_metadata ('test_2'::regclass, 'h', NULL, 0, 's');
 ERROR:  Distribution column cannot be NULL for relation "test_2"
@ -252,7 +252,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	SELECT citus_internal_add_partition_metadata ('test_2'::regclass, 'X', 'col_1', 0, 's');
 citus_internal_add_partition_metadata
 ---------------------------------------------------------------------
@ -268,7 +268,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	SELECT citus_internal_update_placement_metadata(1420007, 10000, 11111);
 ERROR:  could not find valid entry for shard xxxxx
@ -298,7 +298,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	SELECT citus_internal_add_partition_metadata ('test_2'::regclass, 'X', 'col_1', 0, 's');
 ERROR:  role "non_existing_user" does not exist
 ROLLBACK;
@ -329,7 +329,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	SELECT citus_internal_add_partition_metadata ('test_ref'::regclass, 'n', 'col_1', 0, 's');
 ERROR:  Reference or local tables cannot have distribution columns
 ROLLBACK;
@ -341,7 +341,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	SELECT citus_internal_add_partition_metadata ('test_ref'::regclass, 'n', NULL, 0, 'A');
 ERROR:  Metadata syncing is only allowed for known replication models.
 ROLLBACK;
@ -353,7 +353,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	SELECT citus_internal_add_partition_metadata ('test_ref'::regclass, 'n', NULL, 0, 'c');
 ERROR:  Local or references tables can only have 's' or 't' as the replication model.
 ROLLBACK;
@ -368,7 +368,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	SELECT citus_internal_add_partition_metadata ('super_user_table'::regclass, 'h', 'col_1', 0, 's');
 citus_internal_add_partition_metadata
 ---------------------------------------------------------------------
@ -387,7 +387,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(relationname, shardid, storagetype, shardminvalue, shardmaxvalue)
 		AS (VALUES ('super_user_table'::regclass, 1420000::bigint, 't'::"char", '-2147483648'::text, '-1610612737'::text))
@ -402,7 +402,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(relationname, shardid, storagetype, shardminvalue, shardmaxvalue)
 		AS (VALUES ('test_2'::regclass, 1420000::bigint, 't'::"char", '-2147483648'::text, '-1610612737'::text))
@ -417,7 +417,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	SELECT citus_internal_add_partition_metadata ('test_2'::regclass, 'h', 'col_1', 250, 's');
 citus_internal_add_partition_metadata
 ---------------------------------------------------------------------
@ -445,7 +445,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	SELECT citus_internal_update_relation_colocation ('test_2'::regclass, 1231231232);
 citus_internal_update_relation_colocation
 ---------------------------------------------------------------------
@ -461,7 +461,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(relationname, shardid, storagetype, shardminvalue, shardmaxvalue)
 		AS (VALUES ('test_2'::regclass, -1, 't'::"char", '-2147483648'::text, '-1610612737'::text))
@ -476,7 +476,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(relationname, shardid, storagetype, shardminvalue, shardmaxvalue)
 		AS (VALUES ('test_2'::regclass, 1420000, 'X'::"char", '-2147483648'::text, '-1610612737'::text))
@ -491,7 +491,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(relationname, shardid, storagetype, shardminvalue, shardmaxvalue)
 		AS (VALUES ('test_2'::regclass, 1420000, 't'::"char", NULL, '-1610612737'::text))
@ -506,7 +506,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(relationname, shardid, storagetype, shardminvalue, shardmaxvalue)
 		AS (VALUES ('test_2'::regclass, 1420000::bigint, 't'::"char", 'non-int'::text, '-1610612737'::text))
@ -521,7 +521,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(relationname, shardid, storagetype, shardminvalue, shardmaxvalue)
 		AS (VALUES ('test_2'::regclass, 1420000::bigint, 't'::"char", '-1610612737'::text, '-2147483648'::text))
@ -536,7 +536,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(relationname, shardid, storagetype, shardminvalue, shardmaxvalue)
 		AS (VALUES ('test_2'::regclass, 1420000::bigint, 't'::"char", '10'::text, '20'::text),
@ -554,7 +554,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

- SET application_name to 'citus';
+ SET application_name to 'citus_internal';
 \set VERBOSITY terse
 WITH distributed_object_data(typetext, objnames, objargs, distargumentindex, colocationid, force_delegation)
   AS (VALUES ('non_existing_type', ARRAY['non_existing_user']::text[], ARRAY[]::text[], -1, 0, false))
@ -569,7 +569,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH distributed_object_data(typetext, objnames, objargs, distargumentindex, colocationid, force_delegation)
 		AS (VALUES ('role', ARRAY['metadata_sync_helper_role']::text[], ARRAY[]::text[], -100, 0, false))
@ -583,7 +583,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH distributed_object_data(typetext, objnames, objargs, distargumentindex, colocationid, force_delegation)
 		AS (VALUES ('role', ARRAY['metadata_sync_helper_role']::text[], ARRAY[]::text[], -1, -1, false))
@ -598,7 +598,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

- SET application_name to 'citus';
+ SET application_name to 'citus_internal';
 \set VERBOSITY terse
 WITH distributed_object_data(typetext, objnames, objargs, distargumentindex, colocationid, force_delegation)
   AS (VALUES ('role', ARRAY['non_existing_user']::text[], ARRAY[]::text[], -1, 0, false))
@ -614,7 +614,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

- SET application_name to 'citus';
+ SET application_name to 'citus_internal';
 \set VERBOSITY terse
 WITH distributed_object_data(typetext, objnames, objargs, distargumentindex, colocationid, force_delegation)
   AS (VALUES ('role', ARRAY['metadata_sync_helper_role']::text[], ARRAY[]::text[], 0, NULL::int, false))
@ -635,7 +635,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	CREATE TABLE publication_test_table(id int);
 	CREATE PUBLICATION publication_test FOR TABLE publication_test_table;
@ -653,7 +653,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-    SET application_name to 'citus';
+    SET application_name to 'citus_internal';
    \set VERBOSITY terse
    CREATE FUNCTION distribution_test_function(int) RETURNS int
    AS $$ SELECT $1 $$
@ -674,7 +674,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	UPDATE pg_dist_partition SET partmethod = 'X';
 	WITH shard_data(relationname, shardid, storagetype, shardminvalue, shardmaxvalue)
@ -693,7 +693,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(relationname, shardid, storagetype, shardminvalue, shardmaxvalue)
 		AS (VALUES ('test_2'::regclass, 1420000::bigint, 't'::"char", '10'::text, '20'::text))
@ -720,7 +720,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(relationname, shardid, storagetype, shardminvalue, shardmaxvalue)
 		AS (VALUES ('test_2'::regclass, 1420000::bigint, 't'::"char", '11'::text, '20'::text),
@ -751,7 +751,7 @@ BEGIN;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	SELECT citus_internal_update_relation_colocation('test_2'::regclass, 251);
 ERROR:  cannot colocate tables test_2 and test_3
 ROLLBACK;
@ -763,7 +763,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(relationname, shardid, storagetype, shardminvalue, shardmaxvalue)
 		AS (VALUES ('test_3'::regclass, 1420009::bigint, 't'::"char", '21'::text, '30'::text),
@ -790,7 +790,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(relationname, shardid, storagetype, shardminvalue, shardmaxvalue)
 		AS (VALUES ('test_ref'::regclass, 1420003::bigint, 't'::"char", '-1610612737'::text, NULL))
@ -805,7 +805,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(relationname, shardid, storagetype, shardminvalue, shardmaxvalue)
 		AS (VALUES ('test_ref'::regclass, 1420006::bigint, 't'::"char", NULL, NULL),
@ -821,7 +821,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(relationname, shardid, storagetype, shardminvalue, shardmaxvalue)
 		AS (VALUES ('test_ref'::regclass, 1420006::bigint, 't'::"char", NULL, NULL))
@ -842,7 +842,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(relationname, shardid, storagetype, shardminvalue, shardmaxvalue)
 		AS (VALUES ('super_user_table'::regclass, 1420007::bigint, 't'::"char", '11'::text, '20'::text))
@ -864,7 +864,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH placement_data(shardid, shardstate, shardlength, groupid, placementid) AS
 		(VALUES (-10, 1, 0::bigint, 1::int, 1500000::bigint))
@ -879,7 +879,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH placement_data(shardid, shardstate, shardlength, groupid, placementid) AS
 		(VALUES (1420000, 1, 0::bigint, 1::int, -10))
@ -894,7 +894,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH placement_data(shardid, shardstate, shardlength, groupid, placementid) AS
 		(VALUES (1430100, 1, 0::bigint, 1::int, 10))
@ -909,7 +909,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH placement_data(shardid, shardstate, shardlength, groupid, placementid) AS
 		(VALUES (1420000, 10, 0::bigint, 1::int, 1500000))
@ -924,7 +924,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH placement_data(shardid, shardstate, shardlength, groupid, placementid) AS
 		(VALUES ( 1420000, 1, 0::bigint, 123123123::int, 1500000))
@ -952,7 +952,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH placement_data(shardid, shardstate, shardlength, groupid, placementid) AS
 		(VALUES (1420000, 1, 0::bigint, get_node_id(), 1500000),
@ -968,7 +968,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH placement_data(shardid, shardstate, shardlength, groupid, placementid) AS
 		(VALUES (1420007, 1, 0::bigint, get_node_id(), 1500000))
@ -983,7 +983,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH placement_data(shardid, shardstate, shardlength, groupid, placementid) AS
 		(VALUES (1420000, 1, 0::bigint, get_node_id(), 1500000),
@ -1024,7 +1024,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	SELECT citus_internal_update_relation_colocation('test_2'::regclass, 251);
 citus_internal_update_relation_colocation
 ---------------------------------------------------------------------
@ -1041,7 +1041,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	SELECT citus_internal_update_placement_metadata(1420000, get_node_id(), get_node_id()+1000);
 ERROR:  Node with group id 1014 for shard placement xxxxx does not exist
@ -1054,7 +1054,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	SELECT citus_internal_update_placement_metadata(1420000, get_node_id()+10000, get_node_id());
 ERROR:  Active placement for shard xxxxx is not found on group:14
@ -1067,7 +1067,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	SELECT citus_internal_update_placement_metadata(0, get_node_id(), get_node_id()+1);
 ERROR:  Shard id does not exists: 0
@ -1080,7 +1080,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	SELECT citus_internal_update_placement_metadata(213123123123, get_node_id(), get_node_id()+1);
 ERROR:  Shard id does not exists: 213123123123
@ -1093,7 +1093,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	SELECT citus_internal_update_placement_metadata(1420007, get_node_id(), get_node_id()+1);
 ERROR:  must be owner of table super_user_table
@ -1106,7 +1106,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(shardid)
 		AS (VALUES (1420007))
@ -1115,7 +1115,7 @@ ERROR:  must be owner of table super_user_table
 ROLLBACK;
 -- the user only allowed to delete shards in a distributed transaction
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(shardid)
 		AS (VALUES (1420007))
@ -1130,7 +1130,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(shardid)
 		AS (VALUES (1420100))
@ -1157,7 +1157,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(shardid)
 		AS (VALUES (1420000))
@ -1191,7 +1191,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	-- with an ugly trick, update the repmodel
 	-- so that making two tables colocated fails
 	UPDATE pg_dist_partition SET repmodel = 't'
@ -1206,7 +1206,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	-- with an ugly trick, update the vartype of table from int to bigint
 	-- so that making two tables colocated fails
 	UPDATE pg_dist_partition SET partkey = '{VAR :varno 1 :varattno 1 :vartype 20 :vartypmod -1 :varcollid 0 :varlevelsup 0 :varnoold 1 :varoattno 1 :location -1}'
@ -1221,7 +1221,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	-- with an ugly trick, update the partmethod of the table to not-valid
 	-- so that making two tables colocated fails
 	UPDATE pg_dist_partition SET partmethod = ''
@ -1236,7 +1236,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	-- with an ugly trick, update the partmethod of the table to not-valid
 	-- so that making two tables colocated fails
 	UPDATE pg_dist_partition SET partmethod = 'a'
@ -1254,7 +1254,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	SELECT citus_internal_add_partition_metadata ('test_5'::regclass, 'h', 'int_col', 500, 's');
 citus_internal_add_partition_metadata
@ -1277,7 +1277,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;

 (1 row)

-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	SELECT citus_internal_add_partition_metadata ('test_7'::regclass, 'h', 'text_col', 500, 's');
 citus_internal_add_partition_metadata
--- a/src/test/regress/expected/shard_rebalancer.out
+++ b/src/test/regress/expected/shard_rebalancer.out
@ -2349,3 +2349,32 @@ WHERE logicalrelid = 'r1'::regclass;
 (1 row)

 DROP TABLE t1, r1;
+-- Test rebalancer with index on a table
+DROP TABLE IF EXISTS test_rebalance_with_index;
+CREATE TABLE test_rebalance_with_index (measureid integer PRIMARY KEY);
+SELECT create_distributed_table('test_rebalance_with_index', 'measureid');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE INDEX rebalance_with_index ON test_rebalance_with_index(measureid);
+INSERT INTO test_rebalance_with_index VALUES(0);
+INSERT INTO test_rebalance_with_index VALUES(1);
+INSERT INTO test_rebalance_with_index VALUES(2);
+SELECT * FROM master_drain_node('localhost', :worker_2_port);
+ master_drain_node
+---------------------------------------------------------------------
+
+(1 row)
+
+CALL citus_cleanup_orphaned_shards();
+UPDATE pg_dist_node SET shouldhaveshards=true WHERE nodeport = :worker_2_port;
+SELECT rebalance_table_shards();
+ rebalance_table_shards
+---------------------------------------------------------------------
+
+(1 row)
+
+CALL citus_cleanup_orphaned_shards();
+DROP TABLE test_rebalance_with_index CASCADE;
--- a/src/test/regress/sql/metadata_sync_helpers.sql
+++ b/src/test/regress/sql/metadata_sync_helpers.sql
@ -28,7 +28,7 @@ ROLLBACK;
 -- but we are on the coordinator, so still not allowed
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	SELECT citus_internal_add_partition_metadata ('test'::regclass, 'h', 'col_1', 0, 's');
 ROLLBACK;

@ -67,14 +67,14 @@ SET search_path TO metadata_sync_helpers;
 -- owner of the table test
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	SELECT citus_internal_add_partition_metadata ('test'::regclass, 'h', 'col_1', 0, 's');
 ROLLBACK;

 -- we do not own the relation
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	SELECT citus_internal_update_relation_colocation ('test'::regclass, 10);
 ROLLBACK;

@ -83,7 +83,7 @@ CREATE TABLE test_2(col_1 int, col_2 int);
 CREATE TABLE test_3(col_1 int, col_2 int);
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	SELECT citus_internal_add_partition_metadata ('test_2'::regclass, 'h', 'col_1', 0, 's');
 	SELECT count(*) FROM pg_dist_partition WHERE logicalrelid = 'metadata_sync_helpers.test_2'::regclass;
 ROLLBACK;
@ -91,42 +91,42 @@ ROLLBACK;
 -- fails because there is no X distribution method
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	SELECT citus_internal_add_partition_metadata ('test_2'::regclass, 'X', 'col_1', 0, 's');
 ROLLBACK;

 -- fails because there is the column does not exist
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	SELECT citus_internal_add_partition_metadata ('test_2'::regclass, 'h', 'non_existing_col', 0, 's');
 ROLLBACK;

 --- fails because we do not allow NULL parameters
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 		SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-		SET application_name to 'citus';
+		SET application_name to 'citus_internal';
 		SELECT citus_internal_add_partition_metadata (NULL, 'h', 'non_existing_col', 0, 's');
 ROLLBACK;

 -- fails because colocationId cannot be negative
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	SELECT citus_internal_add_partition_metadata ('test_2'::regclass, 'h', 'col_1', -1, 's');
 ROLLBACK;

 -- fails because there is no X replication model
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	SELECT citus_internal_add_partition_metadata ('test_2'::regclass, 'X', 'col_1', 0, 'X');
 ROLLBACK;

 -- the same table cannot be added twice, that is enforced by a primary key
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	SELECT citus_internal_add_partition_metadata ('test_2'::regclass, 'h', 'col_1', 0, 's');
 	SELECT citus_internal_add_partition_metadata ('test_2'::regclass, 'h', 'col_1', 0, 's');
@ -135,7 +135,7 @@ ROLLBACK;
 -- the same table cannot be added twice, that is enforced by a primary key even if distribution key changes
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	SELECT citus_internal_add_partition_metadata ('test_2'::regclass, 'h', 'col_1', 0, 's');
 	SELECT citus_internal_add_partition_metadata ('test_2'::regclass, 'h', 'col_2', 0, 's');
@ -144,7 +144,7 @@ ROLLBACK;
 -- hash distributed table cannot have NULL distribution key
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	SELECT citus_internal_add_partition_metadata ('test_2'::regclass, 'h', NULL, 0, 's');
 ROLLBACK;
@ -165,14 +165,14 @@ SET search_path TO metadata_sync_helpers;

 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	SELECT citus_internal_add_partition_metadata ('test_2'::regclass, 'X', 'col_1', 0, 's');
 ROLLBACK;

 -- should throw error even if we skip the checks, there are no such nodes
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	SELECT citus_internal_update_placement_metadata(1420007, 10000, 11111);
 ROLLBACK;
@ -189,7 +189,7 @@ SET search_path TO metadata_sync_helpers;

 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	SELECT citus_internal_add_partition_metadata ('test_2'::regclass, 'X', 'col_1', 0, 's');
 ROLLBACK;

@ -207,21 +207,21 @@ SET search_path TO metadata_sync_helpers;
 CREATE TABLE test_ref(col_1 int, col_2 int);
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	SELECT citus_internal_add_partition_metadata ('test_ref'::regclass, 'n', 'col_1', 0, 's');
 ROLLBACK;

 -- non-valid replication model
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	SELECT citus_internal_add_partition_metadata ('test_ref'::regclass, 'n', NULL, 0, 'A');
 ROLLBACK;

 -- not-matching replication model for reference table
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	SELECT citus_internal_add_partition_metadata ('test_ref'::regclass, 'n', NULL, 0, 'c');
 ROLLBACK;

@ -231,7 +231,7 @@ SET search_path TO metadata_sync_helpers;
 CREATE TABLE super_user_table(col_1 int);
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	SELECT citus_internal_add_partition_metadata ('super_user_table'::regclass, 'h', 'col_1', 0, 's');
 COMMIT;

@ -244,7 +244,7 @@ SET search_path TO metadata_sync_helpers;

 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(relationname, shardid, storagetype, shardminvalue, shardmaxvalue)
 		AS (VALUES ('super_user_table'::regclass, 1420000::bigint, 't'::"char", '-2147483648'::text, '-1610612737'::text))
@ -254,7 +254,7 @@ ROLLBACK;
 -- the user is only allowed to add a shard for add a table which is in pg_dist_partition
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(relationname, shardid, storagetype, shardminvalue, shardmaxvalue)
 		AS (VALUES ('test_2'::regclass, 1420000::bigint, 't'::"char", '-2147483648'::text, '-1610612737'::text))
@ -264,7 +264,7 @@ ROLLBACK;
 -- ok, now add the table to the pg_dist_partition
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	SELECT citus_internal_add_partition_metadata ('test_2'::regclass, 'h', 'col_1', 250, 's');
 	SELECT citus_internal_add_partition_metadata ('test_3'::regclass, 'h', 'col_1', 251, 's');
 	SELECT citus_internal_add_partition_metadata ('test_ref'::regclass, 'n', NULL, 0, 't');
@ -273,14 +273,14 @@ COMMIT;
 -- we can update to a non-existing colocation group (e.g., colocate_with:=none)
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	SELECT citus_internal_update_relation_colocation ('test_2'::regclass, 1231231232);
 ROLLBACK;

 -- invalid shard ids are not allowed
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(relationname, shardid, storagetype, shardminvalue, shardmaxvalue)
 		AS (VALUES ('test_2'::regclass, -1, 't'::"char", '-2147483648'::text, '-1610612737'::text))
@ -290,7 +290,7 @@ ROLLBACK;
 -- invalid storage types are not allowed
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(relationname, shardid, storagetype, shardminvalue, shardmaxvalue)
 		AS (VALUES ('test_2'::regclass, 1420000, 'X'::"char", '-2147483648'::text, '-1610612737'::text))
@ -300,7 +300,7 @@ ROLLBACK;
 -- NULL shard ranges are not allowed for hash distributed tables
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(relationname, shardid, storagetype, shardminvalue, shardmaxvalue)
 		AS (VALUES ('test_2'::regclass, 1420000, 't'::"char", NULL, '-1610612737'::text))
@ -310,7 +310,7 @@ ROLLBACK;
 -- non-integer shard ranges are not allowed
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(relationname, shardid, storagetype, shardminvalue, shardmaxvalue)
 		AS (VALUES ('test_2'::regclass, 1420000::bigint, 't'::"char", 'non-int'::text, '-1610612737'::text))
@ -320,7 +320,7 @@ ROLLBACK;
 -- shardMinValue should be smaller than shardMaxValue
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(relationname, shardid, storagetype, shardminvalue, shardmaxvalue)
 		AS (VALUES ('test_2'::regclass, 1420000::bigint, 't'::"char", '-1610612737'::text, '-2147483648'::text))
@ -330,7 +330,7 @@ ROLLBACK;
 -- we do not allow overlapping shards for the same table
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(relationname, shardid, storagetype, shardminvalue, shardmaxvalue)
 		AS (VALUES ('test_2'::regclass, 1420000::bigint, 't'::"char", '10'::text, '20'::text),
@ -344,7 +344,7 @@ ROLLBACK;
 -- check with non-existing object type
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH distributed_object_data(typetext, objnames, objargs, distargumentindex, colocationid, force_delegation)
 		AS (VALUES ('non_existing_type', ARRAY['non_existing_user']::text[], ARRAY[]::text[], -1, 0, false))
@ -354,7 +354,7 @@ ROLLBACK;
 -- check the sanity of distributionArgumentIndex and colocationId
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH distributed_object_data(typetext, objnames, objargs, distargumentindex, colocationid, force_delegation)
 		AS (VALUES ('role', ARRAY['metadata_sync_helper_role']::text[], ARRAY[]::text[], -100, 0, false))
@ -363,7 +363,7 @@ ROLLBACK;

 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH distributed_object_data(typetext, objnames, objargs, distargumentindex, colocationid, force_delegation)
 		AS (VALUES ('role', ARRAY['metadata_sync_helper_role']::text[], ARRAY[]::text[], -1, -1, false))
@ -373,7 +373,7 @@ ROLLBACK;
 -- check with non-existing object
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH distributed_object_data(typetext, objnames, objargs, distargumentindex, colocationid, force_delegation)
 		AS (VALUES ('role', ARRAY['non_existing_user']::text[], ARRAY[]::text[], -1, 0, false))
@ -384,7 +384,7 @@ ROLLBACK;
 -- if any parameter is NULL
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH distributed_object_data(typetext, objnames, objargs, distargumentindex, colocationid, force_delegation)
 		AS (VALUES ('role', ARRAY['metadata_sync_helper_role']::text[], ARRAY[]::text[], 0, NULL::int, false))
@ -397,7 +397,7 @@ ROLLBACK;
 -- which is known how to distribute
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse

 	CREATE TABLE publication_test_table(id int);
@ -412,7 +412,7 @@ ROLLBACK;
 -- Show that citus_internal_add_object_metadata checks the priviliges
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
    SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-    SET application_name to 'citus';
+    SET application_name to 'citus_internal';
    \set VERBOSITY terse

    CREATE FUNCTION distribution_test_function(int) RETURNS int
@ -430,7 +430,7 @@ ROLLBACK;
 SET search_path TO metadata_sync_helpers;
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	UPDATE pg_dist_partition SET partmethod = 'X';
 	WITH shard_data(relationname, shardid, storagetype, shardminvalue, shardmaxvalue)
@ -444,7 +444,7 @@ ROLLBACK;
 SET search_path TO metadata_sync_helpers;
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(relationname, shardid, storagetype, shardminvalue, shardmaxvalue)
 		AS (VALUES ('test_2'::regclass, 1420000::bigint, 't'::"char", '10'::text, '20'::text))
@ -462,7 +462,7 @@ SET search_path TO metadata_sync_helpers;
 -- now, add few shards
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(relationname, shardid, storagetype, shardminvalue, shardmaxvalue)
 		AS (VALUES ('test_2'::regclass, 1420000::bigint, 't'::"char", '11'::text, '20'::text),
@ -478,14 +478,14 @@ COMMIT;
 -- we cannot mark these two tables colocated because they are not colocated
 BEGIN;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	SELECT citus_internal_update_relation_colocation('test_2'::regclass, 251);
 ROLLBACK;

 -- now, add few more shards for test_3 to make it colocated with test_2
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(relationname, shardid, storagetype, shardminvalue, shardmaxvalue)
 		AS (VALUES ('test_3'::regclass, 1420009::bigint, 't'::"char", '21'::text, '30'::text),
@ -499,7 +499,7 @@ COMMIT;
 -- shardMin/MaxValues should be NULL for reference tables
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(relationname, shardid, storagetype, shardminvalue, shardmaxvalue)
 		AS (VALUES ('test_ref'::regclass, 1420003::bigint, 't'::"char", '-1610612737'::text, NULL))
@ -509,7 +509,7 @@ ROLLBACK;
 -- reference tables cannot have multiple shards
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(relationname, shardid, storagetype, shardminvalue, shardmaxvalue)
 		AS (VALUES ('test_ref'::regclass, 1420006::bigint, 't'::"char", NULL, NULL),
@ -520,7 +520,7 @@ ROLLBACK;
 -- finally, add a shard for reference tables
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(relationname, shardid, storagetype, shardminvalue, shardmaxvalue)
 		AS (VALUES ('test_ref'::regclass, 1420006::bigint, 't'::"char", NULL, NULL))
@ -533,7 +533,7 @@ SET search_path TO metadata_sync_helpers;
 -- and a shard for the superuser table
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(relationname, shardid, storagetype, shardminvalue, shardmaxvalue)
 		AS (VALUES ('super_user_table'::regclass, 1420007::bigint, 't'::"char", '11'::text, '20'::text))
@ -548,7 +548,7 @@ SET search_path TO metadata_sync_helpers;
 -- shard does not exist
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH placement_data(shardid, shardstate, shardlength, groupid, placementid) AS
 		(VALUES (-10, 1, 0::bigint, 1::int, 1500000::bigint))
@ -558,7 +558,7 @@ ROLLBACK;
 -- invalid placementid
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH placement_data(shardid, shardstate, shardlength, groupid, placementid) AS
 		(VALUES (1420000, 1, 0::bigint, 1::int, -10))
@ -568,7 +568,7 @@ ROLLBACK;
 -- non-existing shard
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH placement_data(shardid, shardstate, shardlength, groupid, placementid) AS
 		(VALUES (1430100, 1, 0::bigint, 1::int, 10))
@ -578,7 +578,7 @@ ROLLBACK;
 -- invalid shard state
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH placement_data(shardid, shardstate, shardlength, groupid, placementid) AS
 		(VALUES (1420000, 10, 0::bigint, 1::int, 1500000))
@ -588,7 +588,7 @@ ROLLBACK;
 -- non-existing node with non-existing node-id 123123123
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH placement_data(shardid, shardstate, shardlength, groupid, placementid) AS
 		(VALUES ( 1420000, 1, 0::bigint, 123123123::int, 1500000))
@ -612,7 +612,7 @@ END; $$ language plpgsql;
 -- fails because we ingest more placements for the same shards to the same worker node
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH placement_data(shardid, shardstate, shardlength, groupid, placementid) AS
 		(VALUES (1420000, 1, 0::bigint, get_node_id(), 1500000),
@ -623,7 +623,7 @@ ROLLBACK;
 -- shard is not owned by us
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH placement_data(shardid, shardstate, shardlength, groupid, placementid) AS
 		(VALUES (1420007, 1, 0::bigint, get_node_id(), 1500000))
@ -633,7 +633,7 @@ ROLLBACK;
 -- sucessfully add placements
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH placement_data(shardid, shardstate, shardlength, groupid, placementid) AS
 		(VALUES (1420000, 1, 0::bigint, get_node_id(), 1500000),
@ -654,7 +654,7 @@ COMMIT;
 -- we should be able to colocate both tables now
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	SELECT citus_internal_update_relation_colocation('test_2'::regclass, 251);
 ROLLBACK;

@ -663,7 +663,7 @@ ROLLBACK;
 -- fails because we are trying to update it to non-existing node
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	SELECT citus_internal_update_placement_metadata(1420000, get_node_id(), get_node_id()+1000);
 COMMIT;
@ -671,7 +671,7 @@ COMMIT;
 -- fails because the source node doesn't contain the shard
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	SELECT citus_internal_update_placement_metadata(1420000, get_node_id()+10000, get_node_id());
 COMMIT;
@ -679,7 +679,7 @@ COMMIT;
 -- fails because shard does not exist
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	SELECT citus_internal_update_placement_metadata(0, get_node_id(), get_node_id()+1);
 COMMIT;
@ -687,7 +687,7 @@ COMMIT;
 -- fails because none-existing shard
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	SELECT citus_internal_update_placement_metadata(213123123123, get_node_id(), get_node_id()+1);
 COMMIT;
@ -695,7 +695,7 @@ COMMIT;
 -- fails because we do not own the shard
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	SELECT citus_internal_update_placement_metadata(1420007, get_node_id(), get_node_id()+1);
 COMMIT;
@ -703,7 +703,7 @@ COMMIT;
 -- the user only allowed to delete their own shards
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(shardid)
 		AS (VALUES (1420007))
@ -712,7 +712,7 @@ ROLLBACK;

 -- the user only allowed to delete shards in a distributed transaction
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(shardid)
 		AS (VALUES (1420007))
@ -722,7 +722,7 @@ ROLLBACK;
 -- the user cannot delete non-existing shards
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(shardid)
 		AS (VALUES (1420100))
@ -737,7 +737,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT count(*) FROM pg_dist_placement WHERE shardid = 1420000;

 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	WITH shard_data(shardid)
 		AS (VALUES (1420000))
@ -754,7 +754,7 @@ ROLLBACK;
 SET search_path TO metadata_sync_helpers;
 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	-- with an ugly trick, update the repmodel
 	-- so that making two tables colocated fails
 	UPDATE pg_dist_partition SET repmodel = 't'
@ -765,7 +765,7 @@ ROLLBACK;

 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	-- with an ugly trick, update the vartype of table from int to bigint
 	-- so that making two tables colocated fails
 	UPDATE pg_dist_partition SET partkey = '{VAR :varno 1 :varattno 1 :vartype 20 :vartypmod -1 :varcollid 0 :varlevelsup 0 :varnoold 1 :varoattno 1 :location -1}'
@ -775,7 +775,7 @@ ROLLBACK;

 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	-- with an ugly trick, update the partmethod of the table to not-valid
 	-- so that making two tables colocated fails
 	UPDATE pg_dist_partition SET partmethod = ''
@ -785,7 +785,7 @@ ROLLBACK;

 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	-- with an ugly trick, update the partmethod of the table to not-valid
 	-- so that making two tables colocated fails
 	UPDATE pg_dist_partition SET partmethod = 'a'
@ -799,7 +799,7 @@ CREATE TABLE test_6(int_col int, text_col text);

 BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	SELECT citus_internal_add_partition_metadata ('test_5'::regclass, 'h', 'int_col', 500, 's');
 	SELECT citus_internal_add_partition_metadata ('test_6'::regclass, 'h', 'text_col', 500, 's');
@ -815,7 +815,7 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	CREATE TABLE test_8(int_col int, text_col text COLLATE "caseinsensitive");

 	SELECT assign_distributed_transaction_id(0, 8, '2021-07-09 15:41:55.542377+02');
-	SET application_name to 'citus';
+	SET application_name to 'citus_internal';
 	\set VERBOSITY terse
 	SELECT citus_internal_add_partition_metadata ('test_7'::regclass, 'h', 'text_col', 500, 's');
 	SELECT citus_internal_add_partition_metadata ('test_8'::regclass, 'h', 'text_col', 500, 's');
--- a/src/test/regress/sql/shard_rebalancer.sql
+++ b/src/test/regress/sql/shard_rebalancer.sql
@ -1399,3 +1399,21 @@ WHERE logicalrelid = 'r1'::regclass;

 DROP TABLE t1, r1;

+-- Test rebalancer with index on a table
+
+DROP TABLE IF EXISTS test_rebalance_with_index;
+CREATE TABLE test_rebalance_with_index (measureid integer PRIMARY KEY);
+SELECT create_distributed_table('test_rebalance_with_index', 'measureid');
+CREATE INDEX rebalance_with_index ON test_rebalance_with_index(measureid);
+
+INSERT INTO test_rebalance_with_index VALUES(0);
+INSERT INTO test_rebalance_with_index VALUES(1);
+INSERT INTO test_rebalance_with_index VALUES(2);
+
+SELECT * FROM master_drain_node('localhost', :worker_2_port);
+CALL citus_cleanup_orphaned_shards();
+UPDATE pg_dist_node SET shouldhaveshards=true WHERE nodeport = :worker_2_port;
+
+SELECT rebalance_table_shards();
+CALL citus_cleanup_orphaned_shards();
+DROP TABLE test_rebalance_with_index CASCADE;