From 4ef57cb67af37f4a7673bb485cc73c68f043d294 Mon Sep 17 00:00:00 2001
From: Brian Cloutier <github@briancloutier.com>
Date: Fri, 23 Sep 2016 15:53:47 +0300
Subject: [PATCH] Fail cluster_remove_node if node has any shard placements

---
 .../master/master_metadata_utility.c          | 37 +++++++++++++++++++
 .../distributed/utils/metadata_cache.c        | 12 ++++++
 src/backend/distributed/utils/node_metadata.c | 12 +++---
 .../distributed/master_metadata_utility.h     |  1 +
 src/include/distributed/metadata_cache.h      |  1 +
 5 files changed, 57 insertions(+), 6 deletions(-)

diff --git a/src/backend/distributed/master/master_metadata_utility.c b/src/backend/distributed/master/master_metadata_utility.c
index 7244119f7..b63584b9a 100644
--- a/src/backend/distributed/master/master_metadata_utility.c
+++ b/src/backend/distributed/master/master_metadata_utility.c
@@ -199,6 +199,43 @@ ShardLength(uint64 shardId)
 }
 
 
+/*
+ * NodeHasShardPlacements returns whether any shards are placed on this node
+ */
+bool
+NodeHasShardPlacements(char *nodeName, int32 nodePort)
+{
+	const int scanKeyCount = 2;
+	const bool indexOK = true;
+
+	bool hasPlacements = false;
+
+	HeapTuple heapTuple = NULL;
+	SysScanDesc scanDescriptor = NULL;
+	ScanKeyData scanKey[scanKeyCount];
+
+	Relation pgShardPlacement = heap_open(DistShardPlacementRelationId(),
+										  AccessShareLock);
+
+	ScanKeyInit(&scanKey[0], Anum_pg_dist_shard_placement_nodename,
+				BTEqualStrategyNumber, F_TEXTEQ, CStringGetTextDatum(nodeName));
+	ScanKeyInit(&scanKey[1], Anum_pg_dist_shard_placement_nodeport,
+				BTEqualStrategyNumber, F_INT4EQ, Int32GetDatum(nodePort));
+
+	scanDescriptor = systable_beginscan(pgShardPlacement,
+										DistShardPlacementNodeidIndexId(), indexOK,
+										NULL, scanKeyCount, scanKey);
+
+	heapTuple = systable_getnext(scanDescriptor);
+	hasPlacements = HeapTupleIsValid(heapTuple);
+
+	systable_endscan(scanDescriptor);
+	heap_close(pgShardPlacement, AccessShareLock);
+
+	return hasPlacements;
+}
+
+
 /*
  * FinalizedShardPlacementList finds shard placements for the given shardId from
  * system catalogs, chooses placements that are in finalized state, and returns
diff --git a/src/backend/distributed/utils/metadata_cache.c b/src/backend/distributed/utils/metadata_cache.c
index 5c65a57cd..5fcdc0dc1 100644
--- a/src/backend/distributed/utils/metadata_cache.c
+++ b/src/backend/distributed/utils/metadata_cache.c
@@ -56,6 +56,7 @@ static Oid distPartitionLogicalRelidIndexId = InvalidOid;
 static Oid distShardLogicalRelidIndexId = InvalidOid;
 static Oid distShardShardidIndexId = InvalidOid;
 static Oid distShardPlacementShardidIndexId = InvalidOid;
+static Oid distShardPlacementNodeidIndexId = InvalidOid;
 static Oid extraDataContainerFuncId = InvalidOid;
 
 /* Hash table for informations about each partition */
@@ -667,6 +668,17 @@ DistShardPlacementShardidIndexId(void)
 }
 
 
+/* return oid of pg_dist_shard_placement_nodeid_index */
+Oid
+DistShardPlacementNodeidIndexId(void)
+{
+	CachedRelationLookup("pg_dist_shard_placement_nodeid_index",
+						 &distShardPlacementNodeidIndexId);
+
+	return distShardPlacementNodeidIndexId;
+}
+
+
 /* return oid of the citus_extradata_container(internal) function */
 Oid
 CitusExtraDataContainerFuncId(void)
diff --git a/src/backend/distributed/utils/node_metadata.c b/src/backend/distributed/utils/node_metadata.c
index 75f469681..dcec491c7 100644
--- a/src/backend/distributed/utils/node_metadata.c
+++ b/src/backend/distributed/utils/node_metadata.c
@@ -132,13 +132,13 @@ cluster_remove_node(PG_FUNCTION_ARGS)
 	int32 nodePort = PG_GETARG_INT32(1);
 	char *nodeNameString = text_to_cstring(nodeName);
 
-	DeleteNodeRow(nodeNameString, nodePort);
+	bool hasShardPlacements = NodeHasShardPlacements(nodeNameString, nodePort);
+	if (hasShardPlacements)
+	{
+		ereport(ERROR, (errmsg("you cannot remove a node which has shard placements")));
+	}
 
-	/*
-	 * 1) lookup the node
-	 * 2) ensure there are no existing shard placements for this node
-	 * 3) remove the row
-	 */
+	DeleteNodeRow(nodeNameString, nodePort);
 
 	PG_RETURN_VOID();
 }
diff --git a/src/include/distributed/master_metadata_utility.h b/src/include/distributed/master_metadata_utility.h
index dd38ae22b..2b4b10e68 100644
--- a/src/include/distributed/master_metadata_utility.h
+++ b/src/include/distributed/master_metadata_utility.h
@@ -62,6 +62,7 @@ extern int ShardIntervalCount(Oid relationId);
 extern List * LoadShardList(Oid relationId);
 extern void CopyShardInterval(ShardInterval *srcInterval, ShardInterval *destInterval);
 extern uint64 ShardLength(uint64 shardId);
+extern bool NodeHasShardPlacements(char *nodeName, int32 nodePort);
 extern List * FinalizedShardPlacementList(uint64 shardId);
 extern List * ShardPlacementList(uint64 shardId);
 extern ShardPlacement * TupleToShardPlacement(TupleDesc tupleDesc,
diff --git a/src/include/distributed/metadata_cache.h b/src/include/distributed/metadata_cache.h
index b19c48b5e..d66981437 100644
--- a/src/include/distributed/metadata_cache.h
+++ b/src/include/distributed/metadata_cache.h
@@ -74,6 +74,7 @@ extern Oid DistPartitionLogicalRelidIndexId(void);
 extern Oid DistShardLogicalRelidIndexId(void);
 extern Oid DistShardShardidIndexId(void);
 extern Oid DistShardPlacementShardidIndexId(void);
+extern Oid DistShardPlacementNodeidIndexId(void);
 
 /* function oids */
 extern Oid CitusExtraDataContainerFuncId(void);