From f3b7e53734b46917234a5b47c647b1fb64cfd675 Mon Sep 17 00:00:00 2001
From: Jelte Fennema <github-tech@jeltef.nl>
Date: Thu, 22 Jul 2021 11:26:01 +0200
Subject: [PATCH] Simplify marking a sequence list as distributed

---
 .../citus_add_local_table_to_metadata.c       | 11 +------
 .../commands/create_distributed_table.c       | 29 +++++++++++++------
 src/backend/distributed/commands/table.c      | 24 ++++-----------
 src/include/distributed/metadata_utility.h    |  1 +
 4 files changed, 28 insertions(+), 37 deletions(-)

diff --git a/src/backend/distributed/commands/citus_add_local_table_to_metadata.c b/src/backend/distributed/commands/citus_add_local_table_to_metadata.c
index 3552ccdb4..1cf533535 100644
--- a/src/backend/distributed/commands/citus_add_local_table_to_metadata.c
+++ b/src/backend/distributed/commands/citus_add_local_table_to_metadata.c
@@ -1063,16 +1063,7 @@ FinalizeCitusLocalTableCreation(Oid relationId, List *dependentSequenceList)
 			PropagateSequenceListDependencies(dependentSequenceList);
 		}
 		CreateTableMetadataOnWorkers(relationId);
-
-		Oid sequenceOid;
-		foreach_oid(sequenceOid, dependentSequenceList)
-		{
-			ObjectAddress address;
-
-			ObjectAddressSet(address, RelationRelationId, sequenceOid);
-			bool shouldSyncMetadata = true;
-			MarkObjectDistributed(&address, shouldSyncMetadata);
-		}
+		MarkSequenceListDistributed(dependentSequenceList);
 	}
 
 	/*
diff --git a/src/backend/distributed/commands/create_distributed_table.c b/src/backend/distributed/commands/create_distributed_table.c
index af1648859..c95e3568c 100644
--- a/src/backend/distributed/commands/create_distributed_table.c
+++ b/src/backend/distributed/commands/create_distributed_table.c
@@ -557,15 +557,7 @@ CreateDistributedTable(Oid relationId, Var *distributionColumn, char distributio
 	{
 		if (ClusterHasKnownMetadataWorkers())
 		{
-			Oid sequenceOid;
-			foreach_oid(sequenceOid, dependentSequenceList)
-			{
-				ObjectAddress address;
-
-				ObjectAddressSet(address, RelationRelationId, sequenceOid);
-				bool shouldSyncMetadata = true;
-				MarkObjectDistributed(&address, shouldSyncMetadata);
-			}
+			MarkSequenceListDistributed(dependentSequenceList);
 		}
 	}
 
@@ -695,6 +687,25 @@ PropagateSequenceDependencies(Oid sequenceOid)
 }
 
 
+/*
+ * MarkSequenceListDistributed marks all sequences in the list as distributed.
+ * NOTE: The sequences should have already been created on the worker nodes.
+ */
+void
+MarkSequenceListDistributed(List *sequenceList)
+{
+	Oid sequenceOid = InvalidOid;
+	foreach_oid(sequenceOid, sequenceList)
+	{
+		ObjectAddress address;
+
+		ObjectAddressSet(address, RelationRelationId, sequenceOid);
+		bool shouldSyncMetadata = true;
+		MarkObjectDistributed(&address, shouldSyncMetadata);
+	}
+}
+
+
 /*
  * EnsureDistributedSequencesHaveOneType first ensures that the type of the column
  * in which the sequence is used as default is supported for each sequence in input
diff --git a/src/backend/distributed/commands/table.c b/src/backend/distributed/commands/table.c
index 659c9269a..5b20ebc41 100644
--- a/src/backend/distributed/commands/table.c
+++ b/src/backend/distributed/commands/table.c
@@ -1600,7 +1600,7 @@ PostprocessAlterTableStmt(AlterTableStmt *alterTableStatement)
 		EnsureDependenciesExistOnAllNodes(&tableAddress);
 	}
 
-	List *newDistributedObjects = NIL;
+	List *newDependentSequences = NIL;
 
 	List *commandList = alterTableStatement->cmds;
 	AlterTableCmd *command = NULL;
@@ -1692,12 +1692,8 @@ PostprocessAlterTableStmt(AlterTableStmt *alterTableStatement)
 								{
 									PropagateSequenceDependencies(
 										seqOid);
-									ObjectAddress *sequenceAddress = palloc(
-										sizeof(ObjectAddress));
-									ObjectAddressSet(*sequenceAddress, RelationRelationId,
-													 seqOid);
-									newDistributedObjects = lappend(newDistributedObjects,
-																	sequenceAddress);
+									newDependentSequences = lappend_oid(
+										newDependentSequences, seqOid);
 								}
 							}
 						}
@@ -1730,10 +1726,8 @@ PostprocessAlterTableStmt(AlterTableStmt *alterTableStatement)
 						ClusterHasKnownMetadataWorkers())
 					{
 						PropagateSequenceDependencies(seqOid);
-						ObjectAddress *sequenceAddress = palloc(sizeof(ObjectAddress));
-						ObjectAddressSet(*sequenceAddress, RelationRelationId, seqOid);
-						newDistributedObjects = lappend(newDistributedObjects,
-														sequenceAddress);
+						newDependentSequences = lappend_oid(newDependentSequences,
+															seqOid);
 					}
 				}
 			}
@@ -1760,13 +1754,7 @@ PostprocessAlterTableStmt(AlterTableStmt *alterTableStatement)
 		}
 
 		SendCommandToWorkersWithMetadata(ENABLE_DDL_PROPAGATION);
-
-		ObjectAddress *address;
-		foreach_ptr(address, newDistributedObjects)
-		{
-			bool shouldSyncMetadata = true;
-			MarkObjectDistributed(address, shouldSyncMetadata);
-		}
+		MarkSequenceListDistributed(newDependentSequences);
 	}
 }
 
diff --git a/src/include/distributed/metadata_utility.h b/src/include/distributed/metadata_utility.h
index d70207a72..3f9fbd090 100644
--- a/src/include/distributed/metadata_utility.h
+++ b/src/include/distributed/metadata_utility.h
@@ -295,6 +295,7 @@ extern void EnsureSequenceTypeSupported(Oid seqOid, Oid seqTypId);
 extern void AlterSequenceType(Oid seqOid, Oid typeOid);
 extern void PropagateSequenceListDependencies(List *sequenceList);
 extern void PropagateSequenceDependencies(Oid sequenceOid);
+extern void MarkSequenceListDistributed(List *sequenceList);
 extern void EnsureDistributedSequencesHaveOneType(Oid relationId,
 												  List *dependentSequenceList,
 												  List *attnumList);