From 2cfac213825101f0d5ac0195a8d9dfc81077638b Mon Sep 17 00:00:00 2001
From: Ahmet Gedemenli <afgedemenli@gmail.com>
Date: Wed, 22 Sep 2021 17:41:40 +0300
Subject: [PATCH] Dont auto-undistribute user-added citus local tables -
 Refactor

---
 .../distributed/commands/utility_hook.c        | 15 ++++++++-------
 .../utils/foreign_key_relationship.c           | 18 ++++++++++++++----
 .../distributed/foreign_key_relationship.h     |  2 +-
 3 files changed, 23 insertions(+), 12 deletions(-)

diff --git a/src/backend/distributed/commands/utility_hook.c b/src/backend/distributed/commands/utility_hook.c
index 29bc2df01..e89a1a7ea 100644
--- a/src/backend/distributed/commands/utility_hook.c
+++ b/src/backend/distributed/commands/utility_hook.c
@@ -95,7 +95,7 @@ static void IncrementUtilityHookCountersIfNecessary(Node *parsetree);
 static void PostStandardProcessUtility(Node *parsetree);
 static void DecrementUtilityHookCountersIfNecessary(Node *parsetree);
 static bool IsDropSchemaOrDB(Node *parsetree);
-static bool ShouldUndistributeCitusLocalTables(void);
+static bool ShouldCheckUndistributeCitusLocalTables(void);
 
 
 /*
@@ -271,7 +271,7 @@ multi_ProcessUtility(PlannedStmt *pstmt,
 			 * can happen due to various kinds of drop commands, we immediately
 			 * undistribute them at the end of the command.
 			 */
-			if (ShouldUndistributeCitusLocalTables())
+			if (ShouldCheckUndistributeCitusLocalTables())
 			{
 				UndistributeDisconnectedCitusLocalTables();
 			}
@@ -687,7 +687,8 @@ ProcessUtilityInternal(PlannedStmt *pstmt,
 /*
  * UndistributeDisconnectedCitusLocalTables undistributes citus local tables that
  * are not connected to any reference tables via their individual foreign key
- * subgraphs.
+ * subgraphs. Note that this function undistributes only the auto-converted tables,
+ * i.e the ones that are converted by Citus by cascading through foreign keys.
  */
 void
 UndistributeDisconnectedCitusLocalTables(void)
@@ -720,7 +721,7 @@ UndistributeDisconnectedCitusLocalTables(void)
 			continue;
 		}
 
-		if (ConnectedToReferenceTableViaFKey(citusLocalTableId))
+		if (!ShouldUndistributeCitusLocalTable(citusLocalTableId))
 		{
 			/* still connected to a reference table, skip it */
 			UnlockRelationOid(citusLocalTableId, lockMode);
@@ -751,11 +752,11 @@ UndistributeDisconnectedCitusLocalTables(void)
 
 
 /*
- * ShouldUndistributeCitusLocalTables returns true if we might need to check
- * citus local tables for their connectivity to reference tables.
+ * ShouldCheckUndistributeCitusLocalTables returns true if we might need to check
+ * citus local tables for undistributing automatically.
  */
 static bool
-ShouldUndistributeCitusLocalTables(void)
+ShouldCheckUndistributeCitusLocalTables(void)
 {
 	if (!ConstraintDropped)
 	{
diff --git a/src/backend/distributed/utils/foreign_key_relationship.c b/src/backend/distributed/utils/foreign_key_relationship.c
index 8996feb15..dd085191d 100644
--- a/src/backend/distributed/utils/foreign_key_relationship.c
+++ b/src/backend/distributed/utils/foreign_key_relationship.c
@@ -139,11 +139,12 @@ GetForeignKeyConnectedRelationIdList(Oid relationId)
 
 
 /*
- * ConnectedToReferenceTableViaFKey returns true if given relationId is
- * connected to a reference table via its foreign key subgraph.
+ * ShouldUndistributeCitusLocalTable returns true if given relationId needs
+ * to be undistributed. Here we do not undistribute table if it's converted by the user,
+ * or connected to a table converted by the user, or a reference table, via foreign keys.
  */
 bool
-ConnectedToReferenceTableViaFKey(Oid relationId)
+ShouldUndistributeCitusLocalTable(Oid relationId)
 {
 	/*
 	 * As we will operate on foreign key connected relations, here we
@@ -152,7 +153,16 @@ ConnectedToReferenceTableViaFKey(Oid relationId)
 	InvalidateForeignKeyGraph();
 
 	List *fkeyConnectedRelations = GetForeignKeyConnectedRelationIdList(relationId);
-	return RelationIdListHasReferenceTable(fkeyConnectedRelations);
+	Oid relationOid = InvalidOid;
+	foreach_oid(relationOid, fkeyConnectedRelations)
+	{
+		if (IsCitusTableType(relationOid, REFERENCE_TABLE))
+		{
+			return false;
+		}
+	}
+
+	return true;
 }
 
 
diff --git a/src/include/distributed/foreign_key_relationship.h b/src/include/distributed/foreign_key_relationship.h
index 491142d13..6e476162e 100644
--- a/src/include/distributed/foreign_key_relationship.h
+++ b/src/include/distributed/foreign_key_relationship.h
@@ -16,7 +16,7 @@
 #include "nodes/primnodes.h"
 
 extern List * GetForeignKeyConnectedRelationIdList(Oid relationId);
-extern bool ConnectedToReferenceTableViaFKey(Oid relationId);
+extern bool ShouldUndistributeCitusLocalTable(Oid relationId);
 extern List * ReferencedRelationIdList(Oid relationId);
 extern List * ReferencingRelationIdList(Oid relationId);
 extern void SetForeignConstraintRelationshipGraphInvalid(void);