From 754d5899f2d59dc4efe36439b63f526c1e0b4795 Mon Sep 17 00:00:00 2001
From: gurkanindibay <gindibay@microsoft.com>
Date: Mon, 29 Jan 2024 17:50:27 +0300
Subject: [PATCH] Refactors MarkObjectDistributed logic

---
 .../distributed/commands/utility_hook.c       | 58 ++++++++++++++-----
 1 file changed, 45 insertions(+), 13 deletions(-)

diff --git a/src/backend/distributed/commands/utility_hook.c b/src/backend/distributed/commands/utility_hook.c
index 435af514f..5bc6426d2 100644
--- a/src/backend/distributed/commands/utility_hook.c
+++ b/src/backend/distributed/commands/utility_hook.c
@@ -110,6 +110,12 @@ typedef struct NonMainDbDistributedStatementInfo
 	bool explicitlyMarkAsDistributed;
 } NonMainDbDistributedStatementInfo;
 
+typedef struct ObjectInfo
+{
+    char *name;
+    Oid id;
+} ObjectInfo;
+
 /*
  * NonMainDbSupportedStatements is an array of statements that are supported
  * from non-main databases.
@@ -151,6 +157,8 @@ static void RunPreprocessMainDBCommand(Node *parsetree, const char *queryString)
 static void RunPostprocessMainDBCommand(Node *parsetree);
 static bool IsStatementSupportedInNonMainDb(Node *parsetree);
 static bool StatementRequiresMarkDistributedFromNonMainDb(Node *parsetree);
+static ObjectInfo GetObjectInfo(Node *parsetree);
+static void MarkObjectDistributedInNonMainDb( Node *parsetree, ObjectInfo objectInfo);
 
 /*
  * ProcessUtilityParseTree is a convenience method to create a PlannedStmt out of
@@ -1662,21 +1670,45 @@ RunPostprocessMainDBCommand(Node *parsetree)
 		return;
 	}
 
-	if (IsA(parsetree, CreateRoleStmt))
-	{
-		StringInfo mainDBQuery = makeStringInfo();
-		CreateRoleStmt *createRoleStmt = castNode(CreateRoleStmt, parsetree);
-		Oid roleOid = get_role_oid(createRoleStmt->role, false);
-		appendStringInfo(mainDBQuery,
-						 MARK_OBJECT_DISTRIBUTED,
-						 AuthIdRelationId,
-						 quote_literal_cstr(createRoleStmt->role),
-						 roleOid,
-						 quote_literal_cstr(CurrentUserName()));
-		RunCitusMainDBQuery(mainDBQuery->data);
-	}
+	ObjectInfo objectInfo = GetObjectInfo(parsetree);
+	MarkObjectDistributedInNonMainDb(parsetree, objectInfo);
 }
 
+/*
+ * GetObjectInfo returns the name and oid of the object in the given parsetree.
+*/
+static ObjectInfo GetObjectInfo(Node *parsetree)
+{
+    ObjectInfo info;
+
+    if (IsA(parsetree, CreateRoleStmt))
+    {
+        CreateRoleStmt *stmt = castNode(CreateRoleStmt, parsetree);
+        info.name = stmt->role;
+        info.id = get_role_oid(stmt->role, false);
+    }
+    // Add else if branches for other statement types
+
+    return info;
+}
+
+/*
+ * MarkObjectDistributedInNonMainDb marks the given object as distributed on the
+ * non-main database.
+*/
+static void MarkObjectDistributedInNonMainDb( Node *parsetree, ObjectInfo objectInfo)
+{
+    StringInfo mainDBQuery = makeStringInfo();
+    appendStringInfo(mainDBQuery,
+                     MARK_OBJECT_DISTRIBUTED,
+                     AuthIdRelationId,
+                     quote_literal_cstr(objectInfo.name),
+                     objectInfo.id,
+                     quote_literal_cstr(CurrentUserName()));
+    RunCitusMainDBQuery(mainDBQuery->data);
+}
+
+
 
 /*
  * IsStatementSupportedIn2Pc returns true if the statement is supported from a