Disable some udfs for tenant tables (#6965)

Disables following UDFs for tenant tables:
* update_distributed_table_colocation // i) table_name cannot be a
tenant table ii) colocate_with cannot be a tenant table
* undistribute_table
* alter_distributed_table // i) table_name cannot be a tenant table ii)
colocate_with cannot be a tenant table

Also, see:
https://gist.github.com/onurctirtir/4c20217200f29b1b1fdaf187d1ecb4f3?permalink_comment_id=4587463#gistcomment-4587463

src/backend/distributed/commands/alter_table.c

@@ -196,6 +196,7 @@ static void EnsureTableNotReferencing(Oid relationId, char conversionType);
 static void EnsureTableNotReferenced(Oid relationId, char conversionType);
 static void EnsureTableNotForeign(Oid relationId);
 static void EnsureTableNotPartition(Oid relationId);
+static void ErrorIfColocateWithTenantTable(char *colocateWith);
 static TableConversionState * CreateTableConversion(TableConversionParameters *params);
 static void CreateDistributedTableLike(TableConversionState *con);
 static void CreateCitusTableLike(TableConversionState *con);
@@ -380,6 +381,8 @@ UndistributeTable(TableConversionParameters *params)
 							   "because the table is not distributed")));
 	}
 
+	ErrorIfTenantTable(params->relationId, "undistribute_table");
+
 	if (!params->cascadeViaForeignKeys)
 	{
 		EnsureTableNotReferencing(params->relationId, UNDISTRIBUTE_TABLE);
@@ -435,6 +438,9 @@ AlterDistributedTable(TableConversionParameters *params)
 							   "is not distributed")));
 	}
 
+	ErrorIfTenantTable(params->relationId, "alter_distributed_table");
+	ErrorIfColocateWithTenantTable(params->colocateWith);
+
 	EnsureTableNotForeign(params->relationId);
 	EnsureTableNotPartition(params->relationId);
 	EnsureHashDistributedTable(params->relationId);
@@ -1180,6 +1186,24 @@ EnsureTableNotPartition(Oid relationId)
 }
 
 
+/*
+ * ErrorIfColocateWithTenantTable errors out if given colocateWith text refers to
+ * a tenant table.
+ */
+void
+ErrorIfColocateWithTenantTable(char *colocateWith)
+{
+	if (colocateWith != NULL &&
+		!IsColocateWithDefault(colocateWith) &&
+		!IsColocateWithNone(colocateWith))
+	{
+		text *colocateWithTableNameText = cstring_to_text(colocateWith);
+		Oid colocateWithTableId = ResolveRelationId(colocateWithTableNameText, false);
+		ErrorIfTenantTable(colocateWithTableId, "colocate_with");
+	}
+}
+
+
 TableConversionState *
 CreateTableConversion(TableConversionParameters *params)
 {

src/backend/distributed/commands/schema_based_sharding.c

@@ -242,3 +242,19 @@ 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",
+							   generate_qualified_relation_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,21 @@ 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:  tenant_2.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:  tenant_2.test_table is not allowed for colocate_with because it is a tenant table
+-- verify we don't allow undistribute_table for tenant tables
+SELECT undistribute_table('tenant_2.test_table');
+ERROR:  tenant_2.test_table is not allowed for undistribute_table because it is a tenant table
+-- verify we don't allow alter_distributed_table for tenant tables
+SELECT alter_distributed_table('tenant_2.test_table', colocate_with => 'none');
+ERROR:  tenant_2.test_table is not allowed for alter_distributed_table because it is a tenant table
+-- verify we also don't allow colocate_with a tenant table
+SELECT alter_distributed_table('regular_schema.test_table', colocate_with => 'tenant_2.test_table');
+ERROR:  tenant_2.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 +256,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,17 @@ 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');
+-- verify we don't allow undistribute_table for tenant tables
+SELECT undistribute_table('tenant_2.test_table');
+-- verify we don't allow alter_distributed_table for tenant tables
+SELECT alter_distributed_table('tenant_2.test_table', colocate_with => 'none');
+-- verify we also don't allow colocate_with a tenant table
+SELECT alter_distributed_table('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');