From 16c8706136940538e0689e1f73385469d1601b44 Mon Sep 17 00:00:00 2001
From: gindibay <gindibay@microsoft.com>
Date: Tue, 25 Jul 2023 18:26:25 +0300
Subject: [PATCH] Adds initial code

---
 .gitignore                                    |  1 +
 .../distributed/metadata/node_metadata.c      | 27 +++++++++++++++++++
 .../distributed/sql/citus--12.0-1--12.1-1.sql |  4 +++
 .../sql/downgrades/citus--12.1-1--12.0-1.sql  |  2 ++
 .../sql/udfs/citus_pause_node/12.1-1.sql      |  9 +++++++
 .../sql/udfs/citus_pause_node/latest.sql      |  9 +++++++
 .../regress/sql/multi_cluster_management.sql  |  3 +++
 7 files changed, 55 insertions(+)
 create mode 100644 src/backend/distributed/sql/udfs/citus_pause_node/12.1-1.sql
 create mode 100644 src/backend/distributed/sql/udfs/citus_pause_node/latest.sql

diff --git a/.gitignore b/.gitignore
index df447746a..44846b224 100644
--- a/.gitignore
+++ b/.gitignore
@@ -55,3 +55,4 @@ lib*.pc
 # style related temporary outputs
 *.uncrustify
 .venv
+test-cluster/*
diff --git a/src/backend/distributed/metadata/node_metadata.c b/src/backend/distributed/metadata/node_metadata.c
index 60a5ab92b..b00737add 100644
--- a/src/backend/distributed/metadata/node_metadata.c
+++ b/src/backend/distributed/metadata/node_metadata.c
@@ -152,6 +152,7 @@ PG_FUNCTION_INFO_V1(master_disable_node);
 PG_FUNCTION_INFO_V1(citus_activate_node);
 PG_FUNCTION_INFO_V1(master_activate_node);
 PG_FUNCTION_INFO_V1(citus_update_node);
+PG_FUNCTION_INFO_V1(citus_pause_node);
 PG_FUNCTION_INFO_V1(master_update_node);
 PG_FUNCTION_INFO_V1(get_shard_id_for_distribution_column);
 PG_FUNCTION_INFO_V1(citus_nodename_for_nodeid);
@@ -1329,6 +1330,32 @@ citus_update_node(PG_FUNCTION_ARGS)
 	PG_RETURN_VOID();
 }
 
+Datum
+citus_pause_node(PG_FUNCTION_ARGS)
+{
+	CheckCitusVersion(ERROR);
+
+	int32 nodeId = PG_GETARG_INT32(0);
+	List *placementList = NIL;
+
+	WorkerNode *workerNode = FindNodeAnyClusterByNodeId(nodeId);
+	if (workerNode == NULL)
+	{
+		ereport(ERROR, (errcode(ERRCODE_NO_DATA_FOUND),
+						errmsg("node %u not found", nodeId)));
+	}
+
+
+	if (NodeIsPrimary(workerNode))
+	{
+		placementList = AllShardPlacementsOnNodeGroup(workerNode->groupId);
+		LockShardsInPlacementListMetadata(placementList, AccessExclusiveLock);
+	}
+
+	PG_RETURN_VOID();
+
+}
+
 
 /*
  * master_update_node is a wrapper function for old UDF name.
diff --git a/src/backend/distributed/sql/citus--12.0-1--12.1-1.sql b/src/backend/distributed/sql/citus--12.0-1--12.1-1.sql
index 685915ebd..d314d0b90 100644
--- a/src/backend/distributed/sql/citus--12.0-1--12.1-1.sql
+++ b/src/backend/distributed/sql/citus--12.0-1--12.1-1.sql
@@ -1,3 +1,7 @@
 -- citus--12.0-1--12.1-1
 
+#include "udfs/citus_pause_node/12.1-1.sql"
+
+
+
 -- bump version to 12.1-1
diff --git a/src/backend/distributed/sql/downgrades/citus--12.1-1--12.0-1.sql b/src/backend/distributed/sql/downgrades/citus--12.1-1--12.0-1.sql
index c498a5059..6088c9677 100644
--- a/src/backend/distributed/sql/downgrades/citus--12.1-1--12.0-1.sql
+++ b/src/backend/distributed/sql/downgrades/citus--12.1-1--12.0-1.sql
@@ -1,2 +1,4 @@
 -- citus--12.1-1--12.0-1
 -- this is an empty downgrade path since citus--12.0-1--12.1-1.sql is empty for now
+
+DROP FUNCTION pg_catalog.citus_pause_node(int);
diff --git a/src/backend/distributed/sql/udfs/citus_pause_node/12.1-1.sql b/src/backend/distributed/sql/udfs/citus_pause_node/12.1-1.sql
new file mode 100644
index 000000000..b34961492
--- /dev/null
+++ b/src/backend/distributed/sql/udfs/citus_pause_node/12.1-1.sql
@@ -0,0 +1,9 @@
+CREATE FUNCTION pg_catalog.citus_pause_node(node_id int)
+  RETURNS void
+  LANGUAGE C STRICT
+  AS 'MODULE_PATHNAME', $$citus_pause_node$$;
+
+COMMENT ON FUNCTION pg_catalog.citus_pause_node(node_id int)
+  IS 'pauses node with given id which leads to add lock in tables and prevent any queries to be executed on that node';
+
+REVOKE ALL ON FUNCTION pg_catalog.citus_pause_node(int) FROM PUBLIC;
diff --git a/src/backend/distributed/sql/udfs/citus_pause_node/latest.sql b/src/backend/distributed/sql/udfs/citus_pause_node/latest.sql
new file mode 100644
index 000000000..b34961492
--- /dev/null
+++ b/src/backend/distributed/sql/udfs/citus_pause_node/latest.sql
@@ -0,0 +1,9 @@
+CREATE FUNCTION pg_catalog.citus_pause_node(node_id int)
+  RETURNS void
+  LANGUAGE C STRICT
+  AS 'MODULE_PATHNAME', $$citus_pause_node$$;
+
+COMMENT ON FUNCTION pg_catalog.citus_pause_node(node_id int)
+  IS 'pauses node with given id which leads to add lock in tables and prevent any queries to be executed on that node';
+
+REVOKE ALL ON FUNCTION pg_catalog.citus_pause_node(int) FROM PUBLIC;
diff --git a/src/test/regress/sql/multi_cluster_management.sql b/src/test/regress/sql/multi_cluster_management.sql
index 9ec0eb28e..d36545cfa 100644
--- a/src/test/regress/sql/multi_cluster_management.sql
+++ b/src/test/regress/sql/multi_cluster_management.sql
@@ -376,6 +376,9 @@ SELECT citus_update_node(:worker_1_node, 'localhost', 9992);
 SELECT citus_nodename_for_nodeid(:worker_1_node);
 SELECT citus_nodeport_for_nodeid(:worker_1_node);
 
+--citus_pause_node allows pausing a node from the non-default cluster
+--TODO add test cases here
+
 
 SELECT nodeid AS worker_1_node FROM pg_dist_node WHERE nodeport=:worker_1_port \gset