From f62d4a6712c3ffac442e278379e160c37d2d0858 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Philip=20Dub=C3=A9?= <gunkmute@gmail.com>
Date: Mon, 19 Aug 2019 17:01:16 +0000
Subject: [PATCH] citus_rm_job_directory for multi_query_directory_cleanup

---
 src/backend/distributed/test/file_utils.c     | 30 +++++++++++++++++++
 .../multi_query_directory_cleanup.out         | 15 ++++++++++
 .../multi_query_directory_cleanup_0.out       | 15 ++++++++++
 .../sql/multi_query_directory_cleanup.sql     | 11 +++++++
 4 files changed, 71 insertions(+)
 create mode 100644 src/backend/distributed/test/file_utils.c

diff --git a/src/backend/distributed/test/file_utils.c b/src/backend/distributed/test/file_utils.c
new file mode 100644
index 000000000..02ecaa0bd
--- /dev/null
+++ b/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();
+}
diff --git a/src/test/regress/expected/multi_query_directory_cleanup.out b/src/test/regress/expected/multi_query_directory_cleanup.out
index 12d9c6cd3..5d3e38f5e 100644
--- a/src/test/regress/expected/multi_query_directory_cleanup.out
+++ b/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
diff --git a/src/test/regress/expected/multi_query_directory_cleanup_0.out b/src/test/regress/expected/multi_query_directory_cleanup_0.out
index 8547afd99..ed98c36b4 100644
--- a/src/test/regress/expected/multi_query_directory_cleanup_0.out
+++ b/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
diff --git a/src/test/regress/sql/multi_query_directory_cleanup.sql b/src/test/regress/sql/multi_query_directory_cleanup.sql
index cde747bf8..c024670f3 100644
--- a/src/test/regress/sql/multi_query_directory_cleanup.sql
+++ b/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;