Disable update_distributed_table_colocation for tenant tables

src/backend/distributed/commands/schema_based_sharding.c

@@ -242,3 +242,18 @@ citus_internal_unregister_tenant_schema_globally(PG_FUNCTION_ARGS)
 
 	PG_RETURN_VOID();
 }
+
+
+/*
+ * ErrorIfTenantTable errors out with the given operation name,
+ * if the given relation is a tenant table.
+ */
+void
+ErrorIfTenantTable(Oid relationId, char *operationName)
+{
+	if (IsTenantSchema(get_rel_namespace(relationId)))
+	{
+		ereport(ERROR, (errmsg("%s is not allowed for %s because it is a tenant table",
+			get_rel_name(relationId), operationName)));
+	}
+}

src/backend/distributed/utils/colocation_utils.c

@@ -116,6 +116,7 @@ update_distributed_table_colocation(PG_FUNCTION_ARGS)
 	text *colocateWithTableNameText = PG_GETARG_TEXT_P(1);
 
 	EnsureTableOwner(targetRelationId);
+	ErrorIfTenantTable(targetRelationId, "update_distributed_table_colocation");
 
 	char *colocateWithTableName = text_to_cstring(colocateWithTableNameText);
 	if (IsColocateWithNone(colocateWithTableName))
@@ -126,6 +127,7 @@ update_distributed_table_colocation(PG_FUNCTION_ARGS)
 	else
 	{
 		Oid colocateWithTableId = ResolveRelationId(colocateWithTableNameText, false);
+		ErrorIfTenantTable(colocateWithTableId, "colocate_with");
 		EnsureTableOwner(colocateWithTableId);
 		MarkTablesColocated(colocateWithTableId, targetRelationId);
 	}

src/include/distributed/commands.h

@@ -789,5 +789,6 @@ extern bool IsTenantSchema(Oid schemaId);
 extern void ErrorIfIllegalPartitioningInTenantSchema(Oid parentRelationId,
 													 Oid partitionRelationId);
 extern void CreateTenantSchemaTable(Oid relationId);
+extern void ErrorIfTenantTable(Oid relationId, char *operationName);
 
 #endif /*CITUS_COMMANDS_H */

src/test/regress/expected/schema_based_sharding.out

@@ -34,6 +34,12 @@ SELECT 1 FROM citus_remove_node('localhost', :worker_2_port);
 (1 row)
 
 CREATE TABLE regular_schema.test_table(a int, b text);
+SELECT create_distributed_table('regular_schema.test_table', 'a');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
 SET citus.enable_schema_based_sharding TO ON;
 -- show that regular_schema doesn't show up in pg_dist_tenant_schema
 SELECT COUNT(*)=0 FROM pg_dist_tenant_schema WHERE schemaid::regnamespace::text = 'regular_schema';
@@ -68,6 +74,12 @@ SELECT create_reference_table('tenant_2.test_table');
 ERROR:  table "test_table" is already distributed
 SELECT citus_add_local_table_to_metadata('tenant_2.test_table');
 ERROR:  table "test_table" is already distributed
+-- verify we don't allow update_distributed_table_colocation for tenant tables
+SELECT update_distributed_table_colocation('tenant_2.test_table', colocate_with => 'none');
+ERROR:  test_table is not allowed for update_distributed_table_colocation because it is a tenant table
+-- verify we also don't allow colocate_with a tenant table
+SELECT update_distributed_table_colocation('regular_schema.test_table', colocate_with => 'tenant_2.test_table');
+ERROR:  test_table is not allowed for colocate_with because it is a tenant table
 -- (on coordinator) verify that colocation id is set for empty tenants too
 SELECT colocationid > 0 FROM pg_dist_tenant_schema
 WHERE schemaid::regnamespace::text IN ('tenant_1', 'tenant_3');
@@ -235,8 +247,8 @@ SELECT EXISTS(
 (1 row)
 
 INSERT INTO tenant_4.another_partitioned_table VALUES (1, 'a');
-ERROR:  insert or update on table "another_partitioned_table_child_1920008" violates foreign key constraint "another_partitioned_table_a_fkey_1920007"
-DETAIL:  Key (a)=(1) is not present in table "partitioned_table_1920005".
+ERROR:  insert or update on table "another_partitioned_table_child_1920040" violates foreign key constraint "another_partitioned_table_a_fkey_1920039"
+DETAIL:  Key (a)=(1) is not present in table "partitioned_table_1920037".
 CONTEXT:  while executing command on localhost:xxxxx
 INSERT INTO tenant_4.partitioned_table VALUES (1, 'a');
 INSERT INTO tenant_4.another_partitioned_table VALUES (1, 'a');

src/test/regress/sql/schema_based_sharding.sql

@@ -25,7 +25,7 @@ SELECT citus_internal_unregister_tenant_schema_globally('regular_schema'::regnam
 SELECT 1 FROM citus_remove_node('localhost', :worker_2_port);
 
 CREATE TABLE regular_schema.test_table(a int, b text);
-
+SELECT create_distributed_table('regular_schema.test_table', 'a');
 SET citus.enable_schema_based_sharding TO ON;
 
 -- show that regular_schema doesn't show up in pg_dist_tenant_schema
@@ -55,6 +55,11 @@ SELECT create_distributed_table('tenant_2.test_table', 'a');
 SELECT create_reference_table('tenant_2.test_table');
 SELECT citus_add_local_table_to_metadata('tenant_2.test_table');
 
+-- verify we don't allow update_distributed_table_colocation for tenant tables
+SELECT update_distributed_table_colocation('tenant_2.test_table', colocate_with => 'none');
+-- verify we also don't allow colocate_with a tenant table
+SELECT update_distributed_table_colocation('regular_schema.test_table', colocate_with => 'tenant_2.test_table');
+
 -- (on coordinator) verify that colocation id is set for empty tenants too
 SELECT colocationid > 0 FROM pg_dist_tenant_schema
 WHERE schemaid::regnamespace::text IN ('tenant_1', 'tenant_3');