Add support for schema-based-sharding via a GUC

src/backend/distributed/commands/citus_add_local_table_to_metadata.c

@@ -18,6 +18,7 @@
  */
 
 #include "postgres.h"
+#include "miscadmin.h"
 
 #include "access/genam.h"
 #include "access/htup_details.h"
@@ -1500,3 +1501,38 @@ FinalizeCitusLocalTableCreation(Oid relationId)
 		InvalidateForeignKeyGraph();
 	}
 }
+
+
+/*
+ * ShouldAddNewTableToMetadata takes a relationId and returns true if we need to add a
+ * newly created table to metadata, false otherwise.
+ * For partitions and temporary tables, ShouldAddNewTableToMetadata returns false.
+ * For other tables created, returns true, if we are on a coordinator that is added
+ * as worker, and ofcourse, if the GUC use_citus_managed_tables is set to on.
+ */
+bool
+ShouldAddNewTableToMetadata(Oid relationId)
+{
+	if (get_rel_persistence(relationId) == RELPERSISTENCE_TEMP ||
+		PartitionTableNoLock(relationId))
+	{
+		/*
+		 * Shouldn't add table to metadata if it's a temp table, or a partition.
+		 * Creating partitions of a table that is added to metadata is already handled.
+		 */
+		return false;
+	}
+
+	if (AddAllLocalTablesToMetadata && !IsBinaryUpgrade &&
+		IsCoordinator() && CoordinatorAddedAsWorkerNode())
+	{
+		/*
+		 * We have verified that the GUC is set to true, and we are not upgrading,
+		 * and we are on the coordinator that is added as worker node.
+		 * So return true here, to add this newly created table to metadata.
+		 */
+		return true;
+	}
+
+	return false;
+}

src/backend/distributed/commands/create_distributed_table.c

@@ -111,8 +111,8 @@ typedef struct
 {
 	int shardCount;
 	bool shardCountIsStrict;
-	char *colocateWithTableName;
 	char *distributionColumnName;
+	ColocationParam colocationParam;
 } DistributedTableParams;
 
 
@@ -296,7 +296,11 @@ create_distributed_table(PG_FUNCTION_ARGS)
 								   "when the distribution column is null ")));
 		}
 
-		CreateSingleShardTable(relationId, colocateWithTableName);
+		ColocationParam colocationParam = {
+			.colocationParamType = COLOCATE_WITH_TABLE_LIKE_OPT,
+			.colocateWithTableName = colocateWithTableName,
+		};
+		CreateSingleShardTable(relationId, colocationParam);
 	}
 
 	PG_RETURN_VOID();
@@ -1006,7 +1010,10 @@ CreateDistributedTable(Oid relationId, char *distributionColumnName,
 	}
 
 	DistributedTableParams distributedTableParams = {
-		.colocateWithTableName = colocateWithTableName,
+		.colocationParam = {
+			.colocateWithTableName = colocateWithTableName,
+			.colocationParamType = COLOCATE_WITH_TABLE_LIKE_OPT
+		},
 		.shardCount = shardCount,
 		.shardCountIsStrict = shardCountIsStrict,
 		.distributionColumnName = distributionColumnName
@@ -1031,10 +1038,10 @@ CreateReferenceTable(Oid relationId)
  * single shard distributed table that doesn't have a shard key.
  */
 void
-CreateSingleShardTable(Oid relationId, char *colocateWithTableName)
+CreateSingleShardTable(Oid relationId, ColocationParam colocationParam)
 {
 	DistributedTableParams distributedTableParams = {
-		.colocateWithTableName = colocateWithTableName,
+		.colocationParam = colocationParam,
 		.shardCount = 1,
 		.shardCountIsStrict = true,
 		.distributionColumnName = NULL
@@ -1155,9 +1162,23 @@ CreateCitusTable(Oid relationId, CitusTableType tableType,
 	 * ColocationIdForNewTable assumes caller acquires lock on relationId. In our case,
 	 * our caller already acquired lock on relationId.
 	 */
-	uint32 colocationId = ColocationIdForNewTable(relationId, tableType,
-												  distributedTableParams,
-												  distributionColumn);
+	uint32 colocationId = INVALID_COLOCATION_ID;
+	if (distributedTableParams &&
+		distributedTableParams->colocationParam.colocationParamType ==
+		COLOCATE_WITH_COLOCATION_ID)
+	{
+		colocationId = distributedTableParams->colocationParam.colocationId;
+	}
+	else
+	{
+		/*
+		 * ColocationIdForNewTable assumes caller acquires lock on relationId. In our case,
+		 * our caller already acquired lock on relationId.
+		 */
+		colocationId = ColocationIdForNewTable(relationId, tableType,
+											   distributedTableParams,
+											   distributionColumn);
+	}
 
 	EnsureRelationCanBeDistributed(relationId, distributionColumn,
 								   citusTableParams.distributionMethod,
@@ -1257,7 +1278,10 @@ CreateCitusTable(Oid relationId, CitusTableType tableType,
 			MemoryContextReset(citusPartitionContext);
 
 			DistributedTableParams childDistributedTableParams = {
-				.colocateWithTableName = parentRelationName,
+				.colocationParam = {
+					.colocationParamType = COLOCATE_WITH_TABLE_LIKE_OPT,
+					.colocateWithTableName = parentRelationName,
+				},
 				.shardCount = distributedTableParams->shardCount,
 				.shardCountIsStrict = false,
 				.distributionColumnName = distributedTableParams->distributionColumnName,
@@ -1308,30 +1332,39 @@ DecideCitusTableParams(CitusTableType tableType,
 	{
 		case HASH_DISTRIBUTED:
 		{
+			Assert(distributedTableParams->colocationParam.colocationParamType ==
+				   COLOCATE_WITH_TABLE_LIKE_OPT);
+
 			citusTableParams.distributionMethod = DISTRIBUTE_BY_HASH;
 			citusTableParams.replicationModel =
 				DecideDistTableReplicationModel(DISTRIBUTE_BY_HASH,
-												distributedTableParams->
+												distributedTableParams->colocationParam.
 												colocateWithTableName);
 			break;
 		}
 
 		case APPEND_DISTRIBUTED:
 		{
+			Assert(distributedTableParams->colocationParam.colocationParamType ==
+				   COLOCATE_WITH_TABLE_LIKE_OPT);
+
 			citusTableParams.distributionMethod = DISTRIBUTE_BY_APPEND;
 			citusTableParams.replicationModel =
 				DecideDistTableReplicationModel(APPEND_DISTRIBUTED,
-												distributedTableParams->
+												distributedTableParams->colocationParam.
 												colocateWithTableName);
 			break;
 		}
 
 		case RANGE_DISTRIBUTED:
 		{
+			Assert(distributedTableParams->colocationParam.colocationParamType ==
+				   COLOCATE_WITH_TABLE_LIKE_OPT);
+
 			citusTableParams.distributionMethod = DISTRIBUTE_BY_RANGE;
 			citusTableParams.replicationModel =
 				DecideDistTableReplicationModel(RANGE_DISTRIBUTED,
-												distributedTableParams->
+												distributedTableParams->colocationParam.
 												colocateWithTableName);
 			break;
 		}
@@ -1768,7 +1801,11 @@ ColocationIdForNewTable(Oid relationId, CitusTableType tableType,
 
 	if (tableType == APPEND_DISTRIBUTED || tableType == RANGE_DISTRIBUTED)
 	{
-		if (!IsColocateWithDefault(distributedTableParams->colocateWithTableName))
+		Assert(distributedTableParams->colocationParam.colocationParamType ==
+			   COLOCATE_WITH_TABLE_LIKE_OPT);
+		char *colocateWithTableName =
+			distributedTableParams->colocationParam.colocateWithTableName;
+		if (!IsColocateWithDefault(colocateWithTableName))
 		{
 			ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
 							errmsg("cannot distribute relation"),
@@ -1795,8 +1832,13 @@ ColocationIdForNewTable(Oid relationId, CitusTableType tableType,
 		Oid distributionColumnCollation =
 			distributionColumn ? get_typcollation(distributionColumnType) : InvalidOid;
 
+		Assert(distributedTableParams->colocationParam.colocationParamType ==
+			   COLOCATE_WITH_TABLE_LIKE_OPT);
+		char *colocateWithTableName =
+			distributedTableParams->colocationParam.colocateWithTableName;
+
 		/* get an advisory lock to serialize concurrent default group creations */
-		if (IsColocateWithDefault(distributedTableParams->colocateWithTableName))
+		if (IsColocateWithDefault(colocateWithTableName))
 		{
 			AcquireColocationDefaultLock();
 		}
@@ -1808,10 +1850,9 @@ ColocationIdForNewTable(Oid relationId, CitusTableType tableType,
 													distributedTableParams->shardCount,
 													distributedTableParams->
 													shardCountIsStrict,
-													distributedTableParams->
 													colocateWithTableName);
 
-		if (IsColocateWithDefault(distributedTableParams->colocateWithTableName) &&
+		if (IsColocateWithDefault(colocateWithTableName) &&
 			(colocationId != INVALID_COLOCATION_ID))
 		{
 			/*
@@ -1824,7 +1865,7 @@ ColocationIdForNewTable(Oid relationId, CitusTableType tableType,
 
 		if (colocationId == INVALID_COLOCATION_ID)
 		{
-			if (IsColocateWithDefault(distributedTableParams->colocateWithTableName))
+			if (IsColocateWithDefault(colocateWithTableName))
 			{
 				/*
 				 * Generate a new colocation ID and insert a pg_dist_colocation
@@ -1835,7 +1876,7 @@ ColocationIdForNewTable(Oid relationId, CitusTableType tableType,
 													 distributionColumnType,
 													 distributionColumnCollation);
 			}
-			else if (IsColocateWithNone(distributedTableParams->colocateWithTableName))
+			else if (IsColocateWithNone(colocateWithTableName))
 			{
 				/*
 				 * Generate a new colocation ID and insert a pg_dist_colocation

src/backend/distributed/commands/distribute_object_ops.c

@@ -294,8 +294,8 @@ static DistributeObjectOps Any_CreateForeignServer = {
 static DistributeObjectOps Any_CreateSchema = {
 	.deparse = DeparseCreateSchemaStmt,
 	.qualify = NULL,
-	.preprocess = PreprocessCreateSchemaStmt,
-	.postprocess = NULL,
+	.preprocess = NULL,
+	.postprocess = PostprocessCreateSchemaStmt,
 	.operationType = DIST_OPS_CREATE,
 	.address = CreateSchemaStmtObjectAddress,
 	.markDistributed = true,

src/backend/distributed/commands/drop_distributed_table.c

@@ -90,7 +90,27 @@ master_remove_partition_metadata(PG_FUNCTION_ARGS)
 
 	DeletePartitionRow(relationId);
 
-	DeleteColocationGroupIfNoTablesBelong(colocationId);
+	/*
+	 * We want to keep using the same colocation group for the tenant even if
+	 * all the tables that belong to it are dropped and new tables are created
+	 * for the tenant etc. For this reason, if a colocation group belongs to a
+	 * tenant schema, we don't delete the colocation group even if there are no
+	 * tables that belong to it.
+	 *
+	 * We do the same if system catalog cannot find the schema of the table
+	 * because this means that the whole schema is dropped.
+	 *
+	 * In that case, we want to delete the colocation group regardless of
+	 * whether the schema is a tenant schema or not. Even more, calling
+	 * IsTenantSchema() with InvalidOid would cause an error, hence we check
+	 * whether the schema is valid beforehand.
+	 */
+	bool missingOk = true;
+	Oid schemaId = get_namespace_oid(schemaName, missingOk);
+	if (!OidIsValid(schemaId) || !IsTenantSchema(schemaId))
+	{
+		DeleteColocationGroupIfNoTablesBelong(colocationId);
+	}
 
 	PG_RETURN_VOID();
 }

src/backend/distributed/commands/schema.c

@@ -19,6 +19,7 @@
 #include "catalog/namespace.h"
 #include "catalog/pg_class.h"
 #include "catalog/pg_namespace.h"
+#include "distributed/colocation_utils.h"
 #include "distributed/commands.h"
 #include <distributed/connection_management.h>
 #include "distributed/commands/utility_hook.h"
@@ -33,6 +34,7 @@
 #include "distributed/resource_lock.h"
 #include <distributed/remote_commands.h>
 #include <distributed/remote_commands.h>
+#include "distributed/tenant_schema_metadata.h"
 #include "distributed/version_compat.h"
 #include "nodes/parsenodes.h"
 #include "utils/fmgroids.h"
@@ -48,13 +50,14 @@ static List * GetGrantCommandsFromCreateSchemaStmt(Node *node);
 
 
 /*
- * PreprocessCreateSchemaStmt is called during the planning phase for
+ * PostprocessCreateSchemaStmt is called during the planning phase for
  * CREATE SCHEMA ..
  */
 List *
-PreprocessCreateSchemaStmt(Node *node, const char *queryString,
-						   ProcessUtilityContext processUtilityContext)
+PostprocessCreateSchemaStmt(Node *node, const char *queryString)
 {
+	CreateSchemaStmt *createSchemaStmt = castNode(CreateSchemaStmt, node);
+
 	if (!ShouldPropagateCreateSchemaStmt())
 	{
 		return NIL;
@@ -74,6 +77,28 @@ PreprocessCreateSchemaStmt(Node *node, const char *queryString,
 
 	commands = list_concat(commands, GetGrantCommandsFromCreateSchemaStmt(node));
 
+	if (ShouldUseSchemaBasedSharding(createSchemaStmt->schemaname))
+	{
+		bool missingOk = false;
+		Oid schemaId = get_namespace_oid(createSchemaStmt->schemaname, missingOk);
+
+		/*
+		 * Register the tenant schema on the coordinator and save the command
+		 * to register it on the workers.
+		 */
+		int shardCount = 1;
+		int replicationFactor = 1;
+		Oid distributionColumnType = InvalidOid;
+		Oid distributionColumnCollation = InvalidOid;
+		uint32 colocationId = CreateColocationGroup(
+			shardCount, replicationFactor, distributionColumnType,
+			distributionColumnCollation);
+
+		InsertTenantSchemaLocally(schemaId, colocationId);
+
+		commands = lappend(commands, TenantSchemaInsertCommand(schemaId, colocationId));
+	}
+
 	commands = lappend(commands, ENABLE_DDL_PROPAGATION);
 
 	return NodeDDLTaskList(NON_COORDINATOR_NODES, commands);