Refactor utility hook global state changes (#3990)

src/backend/distributed/commands/utility_hook.c

@@ -76,6 +76,9 @@ static int activeDropSchemaOrDBs = 0;
 static void ExecuteDistributedDDLJob(DDLJob *ddlJob);
 static char * SetSearchPathToCurrentSearchPathCommand(void);
 static char * CurrentSearchPath(void);
+static void IncrementUtilityHookCountersIfNecessary(Node *parsetree);
+static void PostStandardProcessUtility(Node *parsetree);
+static void DecrementUtilityHookCountersIfNecessary(Node *parsetree);
 static bool IsDropSchemaOrDB(Node *parsetree);
 
 
@@ -467,15 +470,7 @@ multi_ProcessUtility(PlannedStmt *pstmt,
 
 	PG_TRY();
 	{
-		if (IsA(parsetree, AlterTableStmt))
+		IncrementUtilityHookCountersIfNecessary(parsetree);
-		{
-			activeAlterTables++;
-		}
-
-		if (IsDropSchemaOrDB(parsetree))
-		{
-			activeDropSchemaOrDBs++;
-		}
 
 		/*
 		 * Check if we are running ALTER EXTENSION citus UPDATE (TO "<version>") command and
@@ -521,27 +516,11 @@ multi_ProcessUtility(PlannedStmt *pstmt,
 		 */
 		CommandCounterIncrement();
 
-		if (IsA(parsetree, AlterTableStmt))
+		PostStandardProcessUtility(parsetree);
-		{
-			activeAlterTables--;
-		}
-
-		if (IsDropSchemaOrDB(parsetree))
-		{
-			activeDropSchemaOrDBs--;
-		}
 	}
 	PG_CATCH();
 	{
-		if (IsA(parsetree, AlterTableStmt))
+		PostStandardProcessUtility(parsetree);
-		{
-			activeAlterTables--;
-		}
-
-		if (IsDropSchemaOrDB(parsetree))
-		{
-			activeDropSchemaOrDBs--;
-		}
 
 		PG_RE_THROW();
 	}
@@ -595,16 +574,6 @@ multi_ProcessUtility(PlannedStmt *pstmt,
 		PostprocessAlterTableStmtAttachPartition(alterTableStatement, queryString);
 	}
 
-	/*
-	 * Re-forming the foreign key graph relies on the command being executed
-	 * on the local table first. However, in order to decide whether the
-	 * command leads to an invalidation, we need to check before the command
-	 * is being executed since we read pg_constraint table. Thus, we maintain a
-	 * local flag and do the invalidation after multi_ProcessUtility,
-	 * before ExecuteDistributedDDLJob().
-	 */
-	InvalidateForeignKeyGraphForDDL();
-
 	/* after local command has completed, finish by executing worker DDLJobs, if any */
 	if (ddlJobs != NIL)
 	{
@@ -826,6 +795,68 @@ CurrentSearchPath(void)
 }
 
 
+/*
+ * IncrementUtilityHookCountersIfNecessary increments activeAlterTables and
+ * activeDropSchemaOrDBs counters if utility command being processed implies
+ * to do so.
+ */
+static void
+IncrementUtilityHookCountersIfNecessary(Node *parsetree)
+{
+	if (IsA(parsetree, AlterTableStmt))
+	{
+		activeAlterTables++;
+	}
+
+	if (IsDropSchemaOrDB(parsetree))
+	{
+		activeDropSchemaOrDBs++;
+	}
+}
+
+
+/*
+ * PostStandardProcessUtility performs operations to alter (backend) global
+ * state of citus utility hook. Those operations should be done after standard
+ * process utility executes even if it errors out.
+ */
+static void
+PostStandardProcessUtility(Node *parsetree)
+{
+	DecrementUtilityHookCountersIfNecessary(parsetree);
+
+	/*
+	 * Re-forming the foreign key graph relies on the command being executed
+	 * on the local table first. However, in order to decide whether the
+	 * command leads to an invalidation, we need to check before the command
+	 * is being executed since we read pg_constraint table. Thus, we maintain a
+	 * local flag and do the invalidation after multi_ProcessUtility,
+	 * before ExecuteDistributedDDLJob().
+	 */
+	InvalidateForeignKeyGraphForDDL();
+}
+
+
+/*
+ * DecrementUtilityHookCountersIfNecessary decrements activeAlterTables and
+ * activeDropSchemaOrDBs counters if utility command being processed implies
+ * to do so.
+ */
+static void
+DecrementUtilityHookCountersIfNecessary(Node *parsetree)
+{
+	if (IsA(parsetree, AlterTableStmt))
+	{
+		activeAlterTables--;
+	}
+
+	if (IsDropSchemaOrDB(parsetree))
+	{
+		activeDropSchemaOrDBs--;
+	}
+}
+
+
 /*
  * MarkInvalidateForeignKeyGraph marks whether the foreign key graph should be
  * invalidated due to a DDL.