citus_rm_job_directory for multi_query_directory_cleanup

src/backend/distributed/test/file_utils.c

@@ -0,0 +1,30 @@
+#include "postgres.h"
+
+#include "distributed/worker_protocol.h"
+#include "distributed/transmit.h"
+#include "distributed/master_metadata_utility.h"
+#include "fmgr.h"
+#include "lib/stringinfo.h"
+
+PG_FUNCTION_INFO_V1(citus_rm_job_directory);
+
+/*
+ * citus_rm_job_directory removes the job directory for the given job id.
+ * Used before beginning multi_query_directory_cleanup.
+ */
+Datum
+citus_rm_job_directory(PG_FUNCTION_ARGS)
+{
+	uint64 jobId = PG_GETARG_INT64(0);
+	StringInfo jobCacheDirectory = makeStringInfo();
+
+	EnsureSuperUser();
+
+	appendStringInfo(jobCacheDirectory, "base/%s/%s%0*" INT64_MODIFIER "u",
+					 PG_JOB_CACHE_DIR, JOB_DIRECTORY_PREFIX,
+					 MIN_JOB_DIRNAME_WIDTH, jobId);
+	CitusRemoveDirectory(jobCacheDirectory);
+	FreeStringInfo(jobCacheDirectory);
+
+	PG_RETURN_VOID();
+}

src/test/regress/expected/multi_query_directory_cleanup.out

@@ -7,6 +7,21 @@
 -- result files.
 SET citus.next_shard_id TO 810000;
 SET citus.enable_unique_job_ids TO off;
+CREATE FUNCTION citus_rm_job_directory(bigint)
+	RETURNS void
+	AS 'citus'
+	LANGUAGE C STRICT;
+with silence as (
+	SELECT citus_rm_job_directory(split_part(f, '_', 2)::bigint)
+	from pg_ls_dir('base/pgsql_job_cache') f
+)
+select count(*) * 0 zero
+from silence;
+ zero 
+------
+    0
+(1 row)
+
 BEGIN;
 -- pg_ls_dir() displays jobids. We explicitly set the jobId sequence
 -- here so that the regression output becomes independent of the

src/test/regress/expected/multi_query_directory_cleanup_0.out

@@ -7,6 +7,21 @@
 -- result files.
 SET citus.next_shard_id TO 810000;
 SET citus.enable_unique_job_ids TO off;
+CREATE FUNCTION citus_rm_job_directory(bigint)
+	RETURNS void
+	AS 'citus'
+	LANGUAGE C STRICT;
+with silence as (
+	SELECT citus_rm_job_directory(split_part(f, '_', 2)::bigint)
+	from pg_ls_dir('base/pgsql_job_cache') f
+)
+select count(*) * 0 zero
+from silence;
+ zero 
+------
+    0
+(1 row)
+
 BEGIN;
 -- pg_ls_dir() displays jobids. We explicitly set the jobId sequence
 -- here so that the regression output becomes independent of the

src/test/regress/sql/multi_query_directory_cleanup.sql

@@ -11,6 +11,17 @@
 SET citus.next_shard_id TO 810000;
 SET citus.enable_unique_job_ids TO off;
 
+CREATE FUNCTION citus_rm_job_directory(bigint)
+	RETURNS void
+	AS 'citus'
+	LANGUAGE C STRICT;
+
+with silence as (
+	SELECT citus_rm_job_directory(split_part(f, '_', 2)::bigint)
+	from pg_ls_dir('base/pgsql_job_cache') f
+)
+select count(*) * 0 zero
+from silence;
 
 BEGIN;