diff --git a/src/backend/distributed/executor/repartition_join_execution.c b/src/backend/distributed/executor/repartition_join_execution.c
index 696c1ad9f..7809008de 100644
--- a/src/backend/distributed/executor/repartition_join_execution.c
+++ b/src/backend/distributed/executor/repartition_join_execution.c
@@ -104,7 +104,7 @@ CreateTemporarySchemasForMergeTasks(Job *topLeveLJob)
 {
 	List *jobIds = ExtractJobsInJobTree(topLeveLJob);
 	char *createSchemasCommand = GenerateCreateSchemasCommand(jobIds, CurrentUserName());
-	SendCommandToWorkersInParallel(NON_COORDINATOR_NODES, createSchemasCommand,
+	SendCommandToWorkersInParallel(ALL_DATA_NODES, createSchemasCommand,
 								   CitusExtensionOwnerName());
 	return jobIds;
 }
@@ -191,8 +191,8 @@ GenerateJobCommands(List *jobIds, char *templateCommand)
 void
 DoRepartitionCleanup(List *jobIds)
 {
-	SendCommandToWorkersOptionalInParallel(NON_COORDINATOR_NODES,
-										   GenerateDeleteJobsCommand(jobIds),
+	SendCommandToWorkersOptionalInParallel(ALL_DATA_NODES, GenerateDeleteJobsCommand(
+											   jobIds),
 										   CitusExtensionOwnerName());
 }
 
diff --git a/src/backend/distributed/transaction/worker_transaction.c b/src/backend/distributed/transaction/worker_transaction.c
index 1415e1d4a..31a8e9073 100644
--- a/src/backend/distributed/transaction/worker_transaction.c
+++ b/src/backend/distributed/transaction/worker_transaction.c
@@ -130,7 +130,15 @@ SendCommandToWorkersWithMetadata(const char *command)
 List *
 TargetWorkerSetNodeList(TargetWorkerSet targetWorkerSet, LOCKMODE lockMode)
 {
-	List *workerNodeList = ActivePrimaryNonCoordinatorNodeList(lockMode);
+	List *workerNodeList = NIL;
+	if (targetWorkerSet == ALL_DATA_NODES)
+	{
+		workerNodeList = ActivePrimaryNodeList(lockMode);
+	}
+	else
+	{
+		workerNodeList = ActivePrimaryNonCoordinatorNodeList(lockMode);
+	}
 	List *result = NIL;
 
 	WorkerNode *workerNode = NULL;
diff --git a/src/include/distributed/worker_transaction.h b/src/include/distributed/worker_transaction.h
index 7ea24fd97..21aa47e34 100644
--- a/src/include/distributed/worker_transaction.h
+++ b/src/include/distributed/worker_transaction.h
@@ -24,7 +24,8 @@
 typedef enum TargetWorkerSet
 {
 	NON_COORDINATOR_METADATA_NODES,
-	NON_COORDINATOR_NODES
+	NON_COORDINATOR_NODES,
+	ALL_DATA_NODES
 } TargetWorkerSet;
 
 
diff --git a/src/test/regress/expected/coordinator_shouldhaveshards.out b/src/test/regress/expected/coordinator_shouldhaveshards.out
index 58ad5cb09..40f239547 100644
--- a/src/test/regress/expected/coordinator_shouldhaveshards.out
+++ b/src/test/regress/expected/coordinator_shouldhaveshards.out
@@ -198,13 +198,23 @@ NOTICE:  executing the command locally: SELECT count(*) AS count FROM coordinato
 (1 row)
 
 ROLLBACK;
--- Commented out since it currently does not clean up task files on the
--- coordinator
--- See this issue for details: https://github.com/citusdata/citus/issues/3996
--- BEGIN;
--- SET citus.enable_repartition_joins TO ON;
--- SELECT count(*) FROM test t1, test t2 WHERE t1.x = t2.y;
--- ROLLBACK;
+-- repartition queries should work fine
+SET citus.enable_repartition_joins TO ON;
+SELECT count(*) FROM test t1, test t2 WHERE t1.x = t2.y;
+ count
+---------------------------------------------------------------------
+   100
+(1 row)
+
+BEGIN;
+SET citus.enable_repartition_joins TO ON;
+SELECT count(*) FROM test t1, test t2 WHERE t1.x = t2.y;
+ count
+---------------------------------------------------------------------
+   100
+(1 row)
+
+END;
 BEGIN;
 SET citus.enable_repartition_joins TO ON;
 -- trigger local execution
diff --git a/src/test/regress/sql/coordinator_shouldhaveshards.sql b/src/test/regress/sql/coordinator_shouldhaveshards.sql
index 8816579d2..0707be471 100644
--- a/src/test/regress/sql/coordinator_shouldhaveshards.sql
+++ b/src/test/regress/sql/coordinator_shouldhaveshards.sql
@@ -87,13 +87,14 @@ SELECT create_distributed_table('dist_table', 'a', colocate_with := 'none');
 SELECT count(*) FROM dist_table;
 ROLLBACK;
 
--- Commented out since it currently does not clean up task files on the
--- coordinator
--- See this issue for details: https://github.com/citusdata/citus/issues/3996
--- BEGIN;
--- SET citus.enable_repartition_joins TO ON;
--- SELECT count(*) FROM test t1, test t2 WHERE t1.x = t2.y;
--- ROLLBACK;
+-- repartition queries should work fine
+SET citus.enable_repartition_joins TO ON;
+SELECT count(*) FROM test t1, test t2 WHERE t1.x = t2.y;
+
+BEGIN;
+SET citus.enable_repartition_joins TO ON;
+SELECT count(*) FROM test t1, test t2 WHERE t1.x = t2.y;
+END;
 
 BEGIN;
 SET citus.enable_repartition_joins TO ON;