From f596827e7e13c17cc8ef224b02610f369719115d Mon Sep 17 00:00:00 2001
From: Jelte Fennema <github-tech@jeltef.nl>
Date: Tue, 8 Jun 2021 15:53:14 +0200
Subject: [PATCH] BuildCitusTableCacheEntry

---
 .../commands/create_distributed_table.c       | 18 ++++++--
 .../distributed/metadata/metadata_cache.c     | 43 ++++++++++++++++++-
 .../distributed/metadata/metadata_utility.c   |  3 ++
 src/include/distributed/metadata_cache.h      |  2 +
 4 files changed, 62 insertions(+), 4 deletions(-)

diff --git a/src/backend/distributed/commands/create_distributed_table.c b/src/backend/distributed/commands/create_distributed_table.c
index e768c6e2f..5cb059fab 100644
--- a/src/backend/distributed/commands/create_distributed_table.c
+++ b/src/backend/distributed/commands/create_distributed_table.c
@@ -518,8 +518,14 @@ CreateDistributedTable(Oid relationId, List *distributionColumnList,
 	 * ColocationIdForNewTable assumes caller acquires lock on relationId. In our case,
 	 * our caller already acquired lock on relationId.
 	 */
-	uint32 colocationId = ColocationIdForNewTable(relationId, linitial(
-													  distributionColumnList),
+
+	/* TODO: Use full column list */
+	Var *distributionColumn = NULL;
+	if (list_length(distributionColumnList) > 0)
+	{
+		distributionColumn = linitial(distributionColumnList);
+	}
+	uint32 colocationId = ColocationIdForNewTable(relationId, distributionColumn,
 												  distributionMethod, replicationModel,
 												  shardCount, shardCountIsStrict,
 												  colocateWithTableName,
@@ -1222,8 +1228,14 @@ EnsureRelationCanBeDistributed(Oid relationId, List *distributionColumnList,
 		}
 	}
 
+	/* TODO: Use full column list */
+	distributionColumn = NULL;
+	if (list_length(distributionColumnList) > 0)
+	{
+		distributionColumn = linitial(distributionColumnList);
+	}
 	ErrorIfUnsupportedConstraint(relation, distributionMethod, replicationModel,
-								 linitial(distributionColumnList), colocationId);
+								 distributionColumn, colocationId);
 
 
 	ErrorIfUnsupportedPolicy(relation);
diff --git a/src/backend/distributed/metadata/metadata_cache.c b/src/backend/distributed/metadata/metadata_cache.c
index 2ca01f449..c4e24d523 100644
--- a/src/backend/distributed/metadata/metadata_cache.c
+++ b/src/backend/distributed/metadata/metadata_cache.c
@@ -1240,11 +1240,39 @@ BuildCitusTableCacheEntry(Oid relationId)
 	cacheEntry->partitionMethod = datumArray[Anum_pg_dist_partition_partmethod - 1];
 	Datum partitionKeyDatum = datumArray[Anum_pg_dist_partition_partkey - 1];
 	bool partitionKeyIsNull = isNullArray[Anum_pg_dist_partition_partkey - 1];
+	Datum partitionKeysDatum = datumArray[Anum_pg_dist_partition_partkeys - 1];
+	bool partitionKeysIsNull = isNullArray[Anum_pg_dist_partition_partkeys - 1];
 
-	/* note that for reference tables partitionKeyisNull is true */
+	/* note that for reference tables partitionKeyIsNull is true */
 	if (!partitionKeyIsNull)
 	{
 		oldContext = MemoryContextSwitchTo(MetadataCacheMemoryContext);
+		if (!partitionKeysIsNull)
+		{
+			ArrayType *partitionKeysArray = DatumGetArrayTypeP(partitionKeysDatum);
+			int partitionKeysCount = ArrayObjectCount(partitionKeysArray);
+			Datum *partitionKeysArrayDatum = DeconstructArrayObject(partitionKeysArray);
+			for (int i = 0; i < partitionKeysCount; i++)
+			{
+				/* get the string representation of the partition column Var */
+				char *partitionKeyString = TextDatumGetCString(
+					partitionKeysArrayDatum[i]);
+
+				/* convert the string to a Node and ensure it is a Var */
+				Node *partitionNode = stringToNode(partitionKeyString);
+				Assert(IsA(partitionNode, Var));
+
+				cacheEntry->partitionKeyStrings =
+					lappend(cacheEntry->partitionKeyStrings, partitionKeyString);
+				cacheEntry->partitionColumns =
+					lappend(cacheEntry->partitionColumns, partitionNode);
+			}
+
+			/* TODO: uncomment once fixed
+			 * cacheEntry->partitionColumn = linitial(cacheEntry->partitionColumns);
+			 * cacheEntry->partitionKeyString = linitial(cacheEntry->partitionKeyStrings);
+			 */
+		}
 
 		/* get the string representation of the partition column Var */
 		cacheEntry->partitionKeyString = TextDatumGetCString(partitionKeyDatum);
@@ -1255,6 +1283,7 @@ BuildCitusTableCacheEntry(Oid relationId)
 
 		cacheEntry->partitionColumn = (Var *) partitionNode;
 
+
 		MemoryContextSwitchTo(oldContext);
 	}
 	else
@@ -3483,6 +3512,18 @@ ResetCitusTableCacheEntry(CitusTableCacheEntry *cacheEntry)
 		return;
 	}
 
+	if (cacheEntry->partitionKeyStrings != NIL)
+	{
+		list_free(cacheEntry->partitionKeyStrings);
+		cacheEntry->partitionKeyStrings = NIL;
+	}
+
+	if (cacheEntry->partitionColumns != NIL)
+	{
+		list_free(cacheEntry->partitionColumns);
+		cacheEntry->partitionColumns = NIL;
+	}
+
 	/* clean up ShardIdCacheHash */
 	RemoveStaleShardIdCacheEntries(cacheEntry);
 
diff --git a/src/backend/distributed/metadata/metadata_utility.c b/src/backend/distributed/metadata/metadata_utility.c
index 17598b528..2b3875f39 100644
--- a/src/backend/distributed/metadata/metadata_utility.c
+++ b/src/backend/distributed/metadata/metadata_utility.c
@@ -1756,6 +1756,9 @@ InsertIntoPgDistPartition(Oid relationId, char distributionMethod,
 	/* set partkey column to NULL for reference tables */
 	if (distributionMethod != DISTRIBUTE_BY_NONE)
 	{
+		Assert(list_length(distributionColumnList) > 0);
+		Assert(linitial(distributionColumnList) != NULL);
+
 		Datum *distributionColumnDatumArray =
 			palloc0(list_length(distributionColumnList) * sizeof(Datum));
 
diff --git a/src/include/distributed/metadata_cache.h b/src/include/distributed/metadata_cache.h
index 43af99a27..295fd07f9 100644
--- a/src/include/distributed/metadata_cache.h
+++ b/src/include/distributed/metadata_cache.h
@@ -60,6 +60,8 @@ typedef struct
 	/* pg_dist_partition metadata for this table */
 	char *partitionKeyString;
 	Var *partitionColumn;
+	List *partitionKeyStrings;
+	List *partitionColumns;
 	char partitionMethod;
 	uint32 colocationId;
 	char replicationModel;