Refactor distributed table creation logic

This commit is preperation for introducing distributed partitioned
table support. We want to clean and refactor some code in distributed
table creation logic so that we can handle partitioned tables in more
robust way.

src/backend/distributed/master/master_expire_table_cache.c

@@ -183,7 +183,7 @@ DropShardsFromWorker(WorkerNode *workerNode, Oid relationId, List *shardInterval
 		}
 	}
 
-	if (relationKind == RELKIND_RELATION)
+	if (RegularTable(relationId))
 	{
 		appendStringInfo(workerCommand, DROP_REGULAR_TABLE_COMMAND, shardNames->data);
 	}

src/backend/distributed/master/master_node_protocol.c

@@ -721,7 +721,7 @@ ShardStorageType(Oid relationId)
 	char shardStorageType = 0;
 
 	char relationType = get_rel_relkind(relationId);
-	if (relationType == RELKIND_RELATION)
+	if (RegularTable(relationId))
 	{
 		shardStorageType = SHARD_STORAGE_TABLE;
 	}

src/backend/distributed/master/master_repair_shards.c

@@ -389,7 +389,7 @@ RecreateTableDDLCommandList(Oid relationId)
 	bool includeSequenceDefaults = false;
 
 	/* build appropriate DROP command based on relation kind */
-	if (relationKind == RELKIND_RELATION)
+	if (RegularTable(relationId))
 	{
 		appendStringInfo(dropCommand, DROP_REGULAR_TABLE_COMMAND,
 						 qualifiedRelationName);

src/backend/distributed/utils/citus_ruleutils.c

@@ -36,6 +36,7 @@
 #include "distributed/citus_ruleutils.h"
 #include "distributed/multi_partitioning_utils.h"
 #include "distributed/relay_utility.h"
+#include "distributed/master_metadata_utility.h"
 #include "foreign/foreign.h"
 #include "lib/stringinfo.h"
 #include "nodes/nodes.h"
@@ -59,7 +60,6 @@
 
 
 static void AppendOptionListToString(StringInfo stringData, List *options);
-static bool SupportedRelationKind(Relation relation);
 static const char * convert_aclright_to_string(int aclright);
 
 
@@ -291,7 +291,6 @@ pg_get_tableschemadef_string(Oid tableRelationId, bool includeSequenceDefaults)
 	AttrNumber constraintIndex = 0;
 	AttrNumber constraintCount = 0;
 	StringInfoData buffer = { NULL, 0, 0, 0 };
-	bool supportedRelationKind = false;
 
 	/*
 	 * Instead of retrieving values from system catalogs as other functions in
@@ -304,13 +303,7 @@ pg_get_tableschemadef_string(Oid tableRelationId, bool includeSequenceDefaults)
 	relation = relation_open(tableRelationId, AccessShareLock);
 	relationName = generate_relation_name(tableRelationId, NIL);
 
-	supportedRelationKind = SupportedRelationKind(relation);
-	if (!supportedRelationKind)
-	{
-		ereport(ERROR, (errcode(ERRCODE_WRONG_OBJECT_TYPE),
-						errmsg("%s is not a regular, foreign, or partitioned table",
-							   relationName)));
-	}
+	EnsureRelationKindSupported(tableRelationId);
 
 	initStringInfo(&buffer);
 
@@ -491,25 +484,33 @@ pg_get_tableschemadef_string(Oid tableRelationId, bool includeSequenceDefaults)
 
 
 /*
- * SupportedRelationKind returns true if the given relation is supported as a
- * distributed relation.
+ * EnsureRelationKindSupported errors out if the given relation is not supported
+ * as a distributed relation.
  */
-static bool
-SupportedRelationKind(Relation relation)
+void
+EnsureRelationKindSupported(Oid relationId)
 {
-	char relationKind = relation->rd_rel->relkind;
-	bool supportedRelationKind = (relationKind == RELKIND_RELATION || relationKind ==
-								  RELKIND_FOREIGN_TABLE);
-#if (PG_VERSION_NUM >= 100000)
-	supportedRelationKind = supportedRelationKind || relationKind ==
-							RELKIND_PARTITIONED_TABLE;
-#endif
+	char relationKind = get_rel_relkind(relationId);
+	bool supportedRelationKind = false;
 
-	/* Citus doesn't support bare inhereted tables (i.e., not a partition or partitioned table) */
-	supportedRelationKind = supportedRelationKind && !(IsChildTable(relation->rd_id) ||
-													   IsParentTable(relation->rd_id));
+	supportedRelationKind = RegularTable(relationId) ||
+							relationKind == RELKIND_FOREIGN_TABLE;
 
-	return supportedRelationKind;
+	/*
+	 * Citus doesn't support bare inherited tables (i.e., not a partition or
+	 * partitioned table)
+	 */
+	supportedRelationKind = supportedRelationKind && !(IsChildTable(relationId) ||
+													   IsParentTable(relationId));
+
+	if (!supportedRelationKind)
+	{
+		char *relationName = get_rel_name(relationId);
+
+		ereport(ERROR, (errcode(ERRCODE_WRONG_OBJECT_TYPE),
+						errmsg("%s is not a regular, foreign or partitioned table",
+							   relationName)));
+	}
 }
 
 
@@ -523,15 +524,12 @@ char *
 pg_get_tablecolumnoptionsdef_string(Oid tableRelationId)
 {
 	Relation relation = NULL;
-	char *relationName = NULL;
-	char relationKind = 0;
 	TupleDesc tupleDescriptor = NULL;
 	AttrNumber attributeIndex = 0;
 	List *columnOptionList = NIL;
 	ListCell *columnOptionCell = NULL;
 	bool firstOptionPrinted = false;
 	StringInfoData buffer = { NULL, 0, 0, 0 };
-	bool supportedRelationKind = false;
 
 	/*
 	 * Instead of retrieving values from system catalogs, we open the relation,
@@ -539,22 +537,8 @@ pg_get_tablecolumnoptionsdef_string(Oid tableRelationId)
 	 * This is primarily to maintain symmetry with pg_get_tableschemadef.
 	 */
 	relation = relation_open(tableRelationId, AccessShareLock);
-	relationName = generate_relation_name(tableRelationId, NIL);
 
-	relationKind = relation->rd_rel->relkind;
-	supportedRelationKind = (relationKind == RELKIND_RELATION || relationKind ==
-							 RELKIND_FOREIGN_TABLE);
-#if (PG_VERSION_NUM >= 100000)
-	supportedRelationKind = supportedRelationKind || relationKind ==
-							RELKIND_PARTITIONED_TABLE;
-#endif
-
-	if (!supportedRelationKind)
-	{
-		ereport(ERROR, (errcode(ERRCODE_WRONG_OBJECT_TYPE),
-						errmsg("%s is not a regular or foreign table or partitioned",
-							   relationName)));
-	}
+	EnsureRelationKindSupported(tableRelationId);
 
 	/*
 	 * Iterate over the table's columns. If a particular column is not dropped

src/backend/distributed/utils/colocation_utils.c

@@ -904,6 +904,15 @@ ColocatedTableId(Oid colocationId)
 	ScanKeyData scanKey[1];
 	int scanKeyCount = 1;
 
+	/*
+	 * We may have a distributed table whose colocation id is INVALID_COLOCATION_ID.
+	 * In this case, we do not want to send that table's id as colocated table id.
+	 */
+	if (colocationId == INVALID_COLOCATION_ID)
+	{
+		return colocatedTableId;
+	}
+
 	ScanKeyInit(&scanKey[0], Anum_pg_dist_partition_colocationid,
 				BTEqualStrategyNumber, F_INT4EQ, ObjectIdGetDatum(colocationId));
 

src/include/distributed/citus_ruleutils.h

@@ -34,6 +34,7 @@ extern char * pg_get_serverdef_string(Oid tableRelationId);
 extern char * pg_get_sequencedef_string(Oid sequenceRelid);
 extern Form_pg_sequence pg_get_sequencedef(Oid sequenceRelationId);
 extern char * pg_get_tableschemadef_string(Oid tableRelationId, bool forShardCreation);
+extern void EnsureRelationKindSupported(Oid relationId);
 extern char * pg_get_tablecolumnoptionsdef_string(Oid tableRelationId);
 extern void deparse_shard_index_statement(IndexStmt *origStmt, Oid distrelid,
 										  int64 shardid, StringInfo buffer);

src/include/distributed/master_metadata_utility.h

@@ -139,6 +139,9 @@ extern void UpdateShardPlacementState(uint64 placementId, char shardState);
 extern void DeleteShardPlacementRow(uint64 placementId);
 extern void UpdateColocationGroupReplicationFactor(uint32 colocationId,
 												   int replicationFactor);
+extern void CreateDistributedTable(Oid relationId, Var *distributionColumn,
+								   char distributionMethod, char *colocateWithTableName,
+								   bool viaDeprecatedAPI);
 extern void CreateTruncateTrigger(Oid relationId);
 
 /* Remaining metadata utility functions  */
@@ -147,6 +150,7 @@ extern void EnsureTablePermissions(Oid relationId, AclMode mode);
 extern void EnsureTableOwner(Oid relationId);
 extern void EnsureSuperUser(void);
 extern void EnsureReplicationSettings(Oid relationId, char replicationModel);
+extern bool RegularTable(Oid relationId);
 extern bool TableReferenced(Oid relationId);
 extern char * ConstructQualifiedShardName(ShardInterval *shardInterval);
 extern Datum StringToDatum(char *inputString, Oid dataType);

src/test/regress/expected/multi_create_shards.out

@@ -48,7 +48,7 @@ DETAIL:  Distributed relations must not specify the WITH (OIDS) option in their
 ALTER TABLE table_to_distribute SET WITHOUT OIDS;
 -- use an index instead of table name
 SELECT master_create_distributed_table('table_to_distribute_pkey', 'id', 'hash');
-ERROR:  table_to_distribute_pkey is not a regular or foreign table
+ERROR:  table_to_distribute_pkey is not a regular, foreign or partitioned table
 -- use a bad column name
 SELECT master_create_distributed_table('table_to_distribute', 'bad_column', 'hash');
 ERROR:  column "bad_column" of relation "table_to_distribute" does not exist

src/test/regress/expected/multi_generate_ddl_commands.out

@@ -146,7 +146,7 @@ NOTICE:  foreign-data wrapper "fake_fdw" does not have an extension defined
 -- propagating views is not supported
 CREATE VIEW local_view AS SELECT * FROM simple_table;
 SELECT table_ddl_command_array('local_view');
-ERROR:  public.local_view is not a regular, foreign, or partitioned table
+ERROR:  local_view is not a regular, foreign or partitioned table
 -- clean up
 DROP VIEW IF EXISTS local_view;
 DROP FOREIGN TABLE IF EXISTS foreign_table;

src/test/regress/expected/multi_partitioning_utils.out

@@ -415,8 +415,8 @@ SELECT table_inherited('date_partitioned_table');
 
 -- also these are not supported
 SELECT master_get_table_ddl_events('capitals');
-ERROR:  public.capitals is not a regular, foreign, or partitioned table
+ERROR:  capitals is not a regular, foreign or partitioned table
 SELECT master_get_table_ddl_events('cities');
-ERROR:  public.cities is not a regular, foreign, or partitioned table
+ERROR:  cities is not a regular, foreign or partitioned table
 -- dropping parents frop the partitions
 DROP TABLE date_partitioned_table, multi_column_partitioned, list_partitioned, partition_parent_schema.parent_table, cities, capitals;

src/test/regress/expected/multi_partitioning_utils_0.out

@@ -355,9 +355,9 @@ LINE 1: SELECT table_inherited('date_partitioned_table');
                                ^
 -- also these are not supported
 SELECT master_get_table_ddl_events('capitals');
-ERROR:  public.capitals is not a regular, foreign, or partitioned table
+ERROR:  capitals is not a regular, foreign or partitioned table
 SELECT master_get_table_ddl_events('cities');
-ERROR:  public.cities is not a regular, foreign, or partitioned table
+ERROR:  cities is not a regular, foreign or partitioned table
 -- dropping parents frop the partitions
 DROP TABLE date_partitioned_table, multi_column_partitioned, list_partitioned, partition_parent_schema.parent_table, cities, capitals;
 ERROR:  table "date_partitioned_table" does not exist