Succesfully create CitusLocalTables from partitioned tables

src/backend/distributed/commands/citus_add_local_table_to_metadata.c

@@ -47,6 +47,8 @@
 
 static void citus_add_local_table_to_metadata_internal(Oid relationId,
 													   bool cascadeViaForeignKeys);
+static void CreateChildLocalTablesIfRelationIsPartitioned(Oid shellRelationId,
+														  Oid shardRelationId);
 static void ErrorIfUnsupportedCreateCitusLocalTable(Relation relation);
 static void ErrorIfUnsupportedCitusLocalTableKind(Oid relationId);
 static void ErrorIfUnsupportedCitusLocalColumnDefinition(Relation relation);
@@ -331,14 +333,40 @@ CreateCitusLocalTable(Oid relationId, bool cascadeViaForeignKeys)
 
 	FinalizeCitusLocalTableCreation(shellRelationId, dependentSequenceList);
 
+	CreateChildLocalTablesIfRelationIsPartitioned(shellRelationId, shardRelationId);
+}
+
+
+/*
+ * CreateChildLocalTablesIfRelationIsPartitioned takes a relation id and creates
+ * Citus Local Tables for its partitions, if it's a partitioned table.
+ */
+static void
+CreateChildLocalTablesIfRelationIsPartitioned(Oid shellRelationId, Oid shardRelationId)
+{
 	/* if this table is partitioned table, add its partitions to metadata too */
-	if (PartitionedTable(relationId))
+	if (PartitionedTable(shardRelationId))
 	{
-		List *partitionList = PartitionList(relationId);
+		List *partitionList = PartitionList(shardRelationId);
 		Oid partitionRelationId = InvalidOid;
 		foreach_oid(partitionRelationId, partitionList)
 		{
+			/*
+			 * We get the parent shell relation name before creating Citus Local Table,
+			 * since it will be renamed later. Then we generate the command in form of
+			 * ALTER TABLE .. ATTACH PARTITION .., for attaching the shell child to the
+			 * shell parent later.
+			 */
+			char *qualifiedShellRelationName = generate_qualified_relation_name(
+				shellRelationId);
+			char *attachToParentCommand =
+				GenerateAlterTableAttachPartitionToParentCommand(partitionRelationId,
+																 qualifiedShellRelationName);
 			CreateCitusLocalTable(partitionRelationId, false);
+			char *detachPartitionCommand = GenerateDetachPartitionCommand(
+				partitionRelationId);
+			ExecuteAndLogUtilityCommandList(list_make2(detachPartitionCommand,
+													   attachToParentCommand));
 		}
 	}
 }

src/backend/distributed/commands/table.c

@@ -78,10 +78,6 @@ static List * GetRelationIdListFromRangeVarList(List *rangeVarList, LOCKMODE loc
 static bool AlterTableCommandTypeIsTrigger(AlterTableType alterTableType);
 static bool AlterTableDropsForeignKey(AlterTableStmt *alterTableStatement);
 static void ErrorIfUnsupportedAlterTableStmt(AlterTableStmt *alterTableStatement);
-static void ErrorIfCitusLocalTablePartitionCommand(AlterTableCmd *alterTableCmd,
-												   Oid parentRelationId);
-static Oid GetPartitionCommandChildRelationId(AlterTableCmd *alterTableCmd,
-											  bool missingOk);
 static List * InterShardDDLTaskList(Oid leftRelationId, Oid rightRelationId,
 									const char *commandString);
 static bool AlterInvolvesPartitionColumn(AlterTableStmt *alterTableStatement,
@@ -362,7 +358,10 @@ PostprocessCreateTableStmtPartitionOf(CreateStmt *createStatement, const
 	{
 		if (IsCitusTableType(parentRelationId, CITUS_LOCAL_TABLE))
 		{
-			/* if it's a citus local table, we don't need distribution column */
+			/*
+			 * If the parent is a citus local table, we don't need distribution column.
+			 * We can do create a Citus Local Table with current table and early return.
+			 */
 			CreateCitusLocalTable(relationId, false);
 			return;
 		}
@@ -2484,8 +2483,16 @@ ErrorIfUnsupportedAlterTableStmt(AlterTableStmt *alterTableStatement)
 											"separately.")));
 				}
 
-				ErrorIfCitusLocalTablePartitionCommand(command, relationId);
-
+				if (IsCitusTableType(partitionRelationId, CITUS_LOCAL_TABLE) &&
+					IsCitusTableType(relationId, CITUS_LOCAL_TABLE))
+				{
+					/*
+					 * Citus Local Tables cannot be colocated with other tables.
+					 * If both the parent and child tables are Citus Local Tables,
+					 * we don't need to check colocation.
+					 */
+					break;
+				}
 				if (IsCitusTable(partitionRelationId) &&
 					!TablesColocated(relationId, partitionRelationId))
 				{
@@ -2603,52 +2610,6 @@ ErrorIfUnsupportedAlterTableStmt(AlterTableStmt *alterTableStatement)
 }
 
 
-/*
- * ErrorIfCitusLocalTablePartitionCommand errors out if given alter table subcommand is
- * an ALTER TABLE ATTACH / DETACH PARTITION command run for a citus local table.
- */
-static void
-ErrorIfCitusLocalTablePartitionCommand(AlterTableCmd *alterTableCmd, Oid parentRelationId)
-{
-	AlterTableType alterTableType = alterTableCmd->subtype;
-	if (alterTableType != AT_AttachPartition && alterTableType != AT_DetachPartition)
-	{
-		return;
-	}
-
-	bool missingOK = false;
-	Oid childRelationId = GetPartitionCommandChildRelationId(alterTableCmd, missingOK);
-	if (!IsCitusTableType(parentRelationId, CITUS_LOCAL_TABLE) &&
-		!IsCitusTableType(childRelationId, CITUS_LOCAL_TABLE))
-	{
-		return;
-	}
-
-	ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-					errmsg("cannot execute ATTACH/DETACH PARTITION command as "
-						   "local tables added to metadata cannot be involved in "
-						   "partition relationships with other tables")));
-}
-
-
-/*
- * GetPartitionCommandChildRelationId returns child relationId for given
- * ALTER TABLE ATTACH / DETACH PARTITION subcommand.
- */
-static Oid
-GetPartitionCommandChildRelationId(AlterTableCmd *alterTableCmd, bool missingOk)
-{
-	AlterTableType alterTableType PG_USED_FOR_ASSERTS_ONLY = alterTableCmd->subtype;
-	Assert(alterTableType == AT_AttachPartition || alterTableType == AT_DetachPartition);
-
-	PartitionCmd *partitionCommand = (PartitionCmd *) alterTableCmd->def;
-	RangeVar *childRelationRangeVar = partitionCommand->name;
-	Oid childRelationId = RangeVarGetRelid(childRelationRangeVar, AccessExclusiveLock,
-										   missingOk);
-	return childRelationId;
-}
-
-
 /*
  * SetupExecutionModeForAlterTable is the function that is responsible
  * for two things for practical purpose for not doing the same checks

src/backend/distributed/utils/multi_partitioning_utils.c

@@ -731,6 +731,36 @@ GenerateAlterTableAttachPartitionCommand(Oid partitionTableId)
 }
 
 
+/*
+ * GenerateAlterTableAttachPartitionCommand returns the necessary command to
+ * attach the given partition to its parent.
+ */
+char *
+GenerateAlterTableAttachPartitionToParentCommand(Oid partitionTableId,
+												 char *parentTableQualifiedName)
+{
+	StringInfo createPartitionCommand = makeStringInfo();
+
+
+	if (!PartitionTable(partitionTableId))
+	{
+		char *relationName = get_rel_name(partitionTableId);
+
+		ereport(ERROR, (errmsg("\"%s\" is not a partition", relationName)));
+	}
+
+	char *tableQualifiedName = generate_qualified_relation_name(partitionTableId);
+
+	char *partitionBoundCString = PartitionBound(partitionTableId);
+
+	appendStringInfo(createPartitionCommand, "ALTER TABLE %s ATTACH PARTITION %s %s;",
+					 parentTableQualifiedName, tableQualifiedName,
+					 partitionBoundCString);
+
+	return createPartitionCommand->data;
+}
+
+
 /*
  * This function heaviliy inspired from RelationBuildPartitionDesc()
  * which is avaliable in src/backend/catalog/partition.c.

src/include/distributed/multi_partitioning_utils.h

@@ -24,6 +24,9 @@ extern List * PartitionList(Oid parentRelationId);
 extern char * GenerateDetachPartitionCommand(Oid partitionTableId);
 extern char * GenerateAttachShardPartitionCommand(ShardInterval *shardInterval);
 extern char * GenerateAlterTableAttachPartitionCommand(Oid partitionTableId);
+extern char * GenerateAlterTableAttachPartitionToParentCommand(Oid partitionTableId,
+															   char *
+															   parentTableQualifiedName);
 extern char * GeneratePartitioningInformation(Oid tableId);
 extern void FixPartitionConstraintsOnWorkers(Oid relationId);
 extern void FixLocalPartitionConstraints(Oid relationId, int64 shardId);