From 0fab75ea55ee2ed88df099d9f073b271c1e67e83 Mon Sep 17 00:00:00 2001
From: Ahmet Gedemenli <afgedemenli@gmail.com>
Date: Wed, 22 Dec 2021 13:01:40 +0300
Subject: [PATCH] Disable dropping multiple servers in a single command

---
 .../distributed/commands/foreign_server.c     | 79 +++++++++++--------
 1 file changed, 45 insertions(+), 34 deletions(-)

diff --git a/src/backend/distributed/commands/foreign_server.c b/src/backend/distributed/commands/foreign_server.c
index f6c58186162..c8bce71b50b 100644
--- a/src/backend/distributed/commands/foreign_server.c
+++ b/src/backend/distributed/commands/foreign_server.c
@@ -24,6 +24,7 @@
 #include "nodes/primnodes.h"
 
 static Node * RecreateForeignServerStmt(Oid serverId);
+static bool NameListHasDistributedServer(List *serverNames);
 
 
 /*
@@ -164,29 +165,18 @@ PreprocessDropForeignServerStmt(Node *node, const char *queryString,
 	DropStmt *stmt = castNode(DropStmt, node);
 	Assert(stmt->removeType == OBJECT_FOREIGN_SERVER);
 
-	List *allServerNamesToDrop = stmt->objects;
-	List *distributedServerAddresses = NIL;
-	List *distributedServerNames = NIL;
-	Value *serverValue = NULL;
-	foreach_ptr(serverValue, stmt->objects)
-	{
-		char *serverString = strVal(serverValue);
-		ForeignServer *server = GetForeignServerByName(serverString, false);
-
-		ObjectAddress address = { 0 };
-		ObjectAddressSet(address, ForeignServerRelationId, server->serverid);
+	bool includesDistributedServer = NameListHasDistributedServer(stmt->objects);
 
-		/* filter distributed servers */
-		if (IsObjectDistributed(&address))
-		{
-			distributedServerAddresses = lappend(distributedServerAddresses, &address);
-			distributedServerNames = lappend(distributedServerNames, serverValue);
-		}
+	if (!includesDistributedServer)
+	{
+		return NIL;
 	}
 
-	if (list_length(distributedServerNames) <= 0)
+	if (list_length(stmt->objects) > 1)
 	{
-		return NIL;
+		ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+						errmsg("cannot drop distributed server with other servers"),
+						errhint("Try dropping each object in a separate DROP command")));
 	}
 
 	if (!ShouldPropagate())
@@ -196,23 +186,18 @@ PreprocessDropForeignServerStmt(Node *node, const char *queryString,
 
 	EnsureCoordinator();
 
-	/* unmark each distributed server */
-	ObjectAddress *address = NULL;
-	foreach_ptr(address, distributedServerAddresses)
-	{
-		UnmarkObjectDistributed(address);
-	}
+	Assert(list_length(stmt->objects) == 1);
+
+	Value *serverValue = linitial(stmt->objects);
+	ForeignServer *server = GetForeignServerByName(strVal(serverValue), false);
+
+	ObjectAddress address = { 0 };
+	ObjectAddressSet(address, ForeignServerRelationId, server->serverid);
+
+	/* unmark distributed server */
+	UnmarkObjectDistributed(&address);
 
-	/*
-	 * Temporary swap the lists of objects to delete with the distributed
-	 * objects and deparse to an sql statement for the workers.
-	 * Then switch back to allServerNamesToDrop to drop all specified
-	 * servers in coordinator after PreprocessDropForeignServerStmt completes
-	 * its execution.
-	 */
-	stmt->objects = distributedServerNames;
 	const char *deparsedStmt = DeparseTreeNode((Node *) stmt);
-	stmt->objects = allServerNamesToDrop;
 
 	/*
 	 * To prevent recursive propagation in mx architecture, we disable ddl
@@ -323,3 +308,29 @@ RecreateForeignServerStmt(Oid serverId)
 
 	return (Node *) createStmt;
 }
+
+
+/*
+ * NameListHasDistributedServer takes a namelist of servers and returns true if at least
+ * one of them is distributed. Returns false otherwise.
+ */
+static bool
+NameListHasDistributedServer(List *serverNames)
+{
+	Value *serverValue = NULL;
+	foreach_ptr(serverValue, serverNames)
+	{
+		char *serverNameString = strVal(serverValue);
+		ForeignServer *server = GetForeignServerByName(serverNameString, false);
+
+		ObjectAddress address = { 0 };
+		ObjectAddressSet(address, ForeignServerRelationId, server->serverid);
+
+		if (IsObjectDistributed(&address))
+		{
+			return true;
+		}
+	}
+
+	return false;
+}