From 01ce923aea5774582d0bf4fb52b79b6b50b8c81f Mon Sep 17 00:00:00 2001
From: Hanefi Onaldi <Hanefi.Onaldi@microsoft.com>
Date: Wed, 21 Jun 2023 01:07:39 +0300
Subject: [PATCH] Introduce helper functions for getting lockmode

---
 src/backend/distributed/commands/index.c | 40 +++++++++++++++++++++---
 src/include/distributed/commands.h       |  2 ++
 2 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/src/backend/distributed/commands/index.c b/src/backend/distributed/commands/index.c
index 4edb319f7..782fa71fb 100644
--- a/src/backend/distributed/commands/index.c
+++ b/src/backend/distributed/commands/index.c
@@ -553,11 +553,9 @@ ReindexStmtFindRelationOid(ReindexStmt *reindexStmt, bool missingOk)
 
 	Oid relationId = InvalidOid;
 
-	LOCKMODE lockmode = IsReindexWithParam_compat(reindexStmt, "concurrently") ?
-						ShareUpdateExclusiveLock : AccessExclusiveLock;
-
 	if (reindexStmt->kind == REINDEX_OBJECT_INDEX)
 	{
+		LOCKMODE lockmode = GetReindexIndexRelationLockMode(reindexStmt);
 		struct ReindexIndexCallbackState state;
 		state.concurrent = IsReindexWithParam_compat(reindexStmt,
 													 "concurrently");
@@ -571,6 +569,7 @@ ReindexStmtFindRelationOid(ReindexStmt *reindexStmt, bool missingOk)
 	}
 	else
 	{
+		LOCKMODE lockmode = GetReindexTableRelationLockMode(reindexStmt);
 		relationId = RangeVarGetRelidExtended(reindexStmt->relation, lockmode,
 											  (missingOk) ? RVR_MISSING_OK : 0,
 											  RangeVarCallbackOwnsTable, NULL);
@@ -580,6 +579,35 @@ ReindexStmtFindRelationOid(ReindexStmt *reindexStmt, bool missingOk)
 }
 
 
+/*
+ * GetReindexRelationLockMode returns required lock mode to open the
+ * index that given REINDEX INDEX command operates on.
+ */
+LOCKMODE
+GetReindexIndexRelationLockMode(ReindexStmt *reindexStmt)
+{
+	if (IsReindexWithParam_compat(reindexStmt, "concurrently"))
+	{
+		return ShareUpdateExclusiveLock;
+	}
+	else
+	{
+		return AccessExclusiveLock;
+	}
+}
+
+
+/*
+ * GetReindexTableLockMode returns required lock mode to open the
+ * relation that given REINDEX TABLE command operates on.
+ */
+LOCKMODE
+GetReindexTableRelationLockMode(ReindexStmt *reindexStmt)
+{
+	return ShareLock;
+}
+
+
 /*
  * PreprocessReindexStmt determines whether a given REINDEX statement involves
  * a distributed table. If so (and if the statement does not use unsupported
@@ -605,11 +633,11 @@ PreprocessReindexStmt(Node *node, const char *reindexCommand,
 		Oid relationId = ReindexStmtFindRelationOid(reindexStatement, false);
 		MemoryContext relationContext = NULL;
 		Relation relation = NULL;
-		LOCKMODE lockmode = IsReindexWithParam_compat(reindexStatement, "concurrently") ?
-							ShareUpdateExclusiveLock : AccessExclusiveLock;
 
 		if (reindexStatement->kind == REINDEX_OBJECT_INDEX)
 		{
+			LOCKMODE lockmode = GetReindexIndexRelationLockMode(reindexStatement);
+
 			/*
 			 * Acquire global lock to prevent concurrent writes or DDL. However,
 			 * we do not bother for REINDEX CONCURRENTLY, since we'll have
@@ -626,6 +654,8 @@ PreprocessReindexStmt(Node *node, const char *reindexCommand,
 		}
 		else
 		{
+			LOCKMODE lockmode = GetReindexTableRelationLockMode(reindexStatement);
+
 			AcquireDistributedLockOnRelations(list_make1(reindexStatement->relation),
 											  lockmode, 0);
 
diff --git a/src/include/distributed/commands.h b/src/include/distributed/commands.h
index a013f3977..0ca46f70b 100644
--- a/src/include/distributed/commands.h
+++ b/src/include/distributed/commands.h
@@ -388,6 +388,8 @@ extern char * ChooseIndexName(const char *tabname, Oid namespaceId,
 extern char * ChooseIndexNameAddition(List *colnames);
 extern List * ChooseIndexColumnNames(List *indexElems);
 extern LOCKMODE GetCreateIndexRelationLockMode(IndexStmt *createIndexStatement);
+extern LOCKMODE GetReindexTableRelationLockMode(ReindexStmt *reindexStmt);
+extern LOCKMODE GetReindexIndexRelationLockMode(ReindexStmt *reindexStmt);
 extern List * PreprocessReindexStmt(Node *ReindexStatement,
 									const char *ReindexCommand,
 									ProcessUtilityContext processUtilityContext);