Fix regression in allowed foreign keys on distributed tables (#6550)

DESCRIPTION: Fix regression in allowed foreign keys on distributed
tables

In commit eadc88a we changed how we skip foreign key validation. The
goal was to skip it in more cases. However, one change had the
unintended regression of introducing failures when trying to create
certain foreign keys. This reverts that part of the change.

The way of skipping validation of foreign keys that was introduced in
eadc88a was skipping validation during execution. The reason that
this caused this regression was because some foreign key validation
queries already fail during planning. In those cases it never gets to
the execution step where it would later be skipped.

Fixes #6543

(cherry picked from commit 7a7880aec9)

src/backend/distributed/commands/cascade_table_operation_for_connected_relations.c

@@ -623,18 +623,13 @@ ExecuteForeignKeyCreateCommand(const char *commandString, bool skip_validation)
 	 */
 	Assert(IsA(parseTree, AlterTableStmt));
 
-	bool oldSkipConstraintsValidationValue = SkipConstraintValidation;
-
 	if (skip_validation && IsA(parseTree, AlterTableStmt))
 	{
-		EnableSkippingConstraintValidation();
-
+		SkipForeignKeyValidationIfConstraintIsFkey((AlterTableStmt *) parseTree, true);
 		ereport(DEBUG4, (errmsg("skipping validation for foreign key create "
 								"command \"%s\"", commandString)));
 	}
 
 	ProcessUtilityParseTree(parseTree, commandString, PROCESS_UTILITY_QUERY,
 							NULL, None_Receiver, NULL);
-
-	SkipConstraintValidation = oldSkipConstraintsValidationValue;
 }

src/backend/distributed/commands/table.c

@@ -1878,7 +1878,8 @@ PreprocessAlterTableSchemaStmt(Node *node, const char *queryString,
  * ALTER TABLE ... ADD FOREIGN KEY command to skip the validation step.
  */
 void
-SkipForeignKeyValidationIfConstraintIsFkey(AlterTableStmt *alterTableStatement)
+SkipForeignKeyValidationIfConstraintIsFkey(AlterTableStmt *alterTableStatement,
+										   bool processLocalRelation)
 {
 	/* first check whether a distributed relation is affected */
 	if (alterTableStatement->relation == NULL)
@@ -1893,11 +1894,17 @@ SkipForeignKeyValidationIfConstraintIsFkey(AlterTableStmt *alterTableStatement)
 		return;
 	}
 
-	if (!IsCitusTable(leftRelationId))
+	if (!IsCitusTable(leftRelationId) && !processLocalRelation)
 	{
 		return;
 	}
 
+	/*
+	 * We check if there is a ADD FOREIGN CONSTRAINT command in sub commands
+	 * list. We set skip_validation to true to prevent PostgreSQL to verify
+	 * validity of the foreign constraint. Validity will be checked on the
+	 * shards anyway.
+	 */
 	AlterTableCmd *command = NULL;
 	foreach_ptr(command, alterTableStatement->cmds)
 	{
@@ -1909,9 +1916,8 @@ SkipForeignKeyValidationIfConstraintIsFkey(AlterTableStmt *alterTableStatement)
 			Constraint *constraint = (Constraint *) command->def;
 			if (constraint->contype == CONSTR_FOREIGN)
 			{
-				/* set the GUC skip_constraint_validation to on */
-				EnableSkippingConstraintValidation();
-				return;
+				/* foreign constraint validations will be done in shards. */
+				constraint->skip_validation = true;
 			}
 		}
 	}

src/backend/distributed/commands/utility_hook.c

@@ -594,7 +594,9 @@ ProcessUtilityInternal(PlannedStmt *pstmt,
 		 * Citus intervening. The only exception is partition column drop, in
 		 * which case we error out. Advanced Citus users use this to implement their
 		 * own DDL propagation. We also use it to avoid re-propagating DDL commands
-		 * when changing MX tables on workers.
+		 * when changing MX tables on workers. Below, we also make sure that DDL
+		 * commands don't run queries that might get intercepted by Citus and error
+		 * out during planning, specifically we skip validation in foreign keys.
 		 */
 
 		if (IsA(parsetree, AlterTableStmt))
@@ -613,7 +615,7 @@ ProcessUtilityInternal(PlannedStmt *pstmt,
 				 * Note validation is done on the shard level when DDL propagation
 				 * is enabled. The following eagerly executes some tasks on workers.
 				 */
-				SkipForeignKeyValidationIfConstraintIsFkey(alterTableStmt);
+				SkipForeignKeyValidationIfConstraintIsFkey(alterTableStmt, false);
 			}
 		}
 	}

src/include/distributed/commands.h

@@ -541,7 +541,8 @@ extern List * PreprocessAlterTableMoveAllStmt(Node *node, const char *queryStrin
 											  ProcessUtilityContext processUtilityContext);
 extern List * PreprocessAlterTableSchemaStmt(Node *node, const char *queryString,
 											 ProcessUtilityContext processUtilityContext);
-extern void SkipForeignKeyValidationIfConstraintIsFkey(AlterTableStmt *alterTableStmt);
+extern void SkipForeignKeyValidationIfConstraintIsFkey(AlterTableStmt *alterTableStmt,
+													   bool processLocalRelation);
 extern bool IsAlterTableRenameStmt(RenameStmt *renameStmt);
 extern void ErrorIfAlterDropsPartitionColumn(AlterTableStmt *alterTableStatement);
 extern void PostprocessAlterTableStmt(AlterTableStmt *pStmt);

src/test/regress/expected/issue_6543.out

@@ -0,0 +1,30 @@
+CREATE SCHEMA issue_6543;
+SET search_path TO issue_6543;
+SET citus.shard_count TO 4;
+SET citus.shard_replication_factor TO 1;
+SET citus.next_shard_id TO 67322500;
+CREATE TABLE event (
+  tenant_id varchar,
+  id bigint,
+  primary key (tenant_id, id)
+);
+CREATE TABLE page (
+  tenant_id varchar,
+  id int,
+  primary key (tenant_id, id)
+);
+SELECT create_distributed_table('event', 'tenant_id');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT create_distributed_table('page', 'tenant_id', colocate_with => 'event');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+alter table page add constraint fk21 foreign key (tenant_id, id) references event;
+SET client_min_messages TO WARNING;
+DROP SCHEMA issue_6543 CASCADE;

src/test/regress/multi_schedule

@@ -94,7 +94,7 @@ test: multi_dropped_column_aliases foreign_key_restriction_enforcement
 test: binary_protocol
 test: alter_table_set_access_method
 test: alter_distributed_table
-test: issue_5248 issue_5099
+test: issue_5248 issue_5099 issue_6543
 test: object_propagation_debug
 test: undistribute_table
 test: run_command_on_all_nodes

src/test/regress/sql/issue_6543.sql

@@ -0,0 +1,24 @@
+CREATE SCHEMA issue_6543;
+SET search_path TO issue_6543;
+SET citus.shard_count TO 4;
+SET citus.shard_replication_factor TO 1;
+SET citus.next_shard_id TO 67322500;
+
+CREATE TABLE event (
+  tenant_id varchar,
+  id bigint,
+  primary key (tenant_id, id)
+);
+
+CREATE TABLE page (
+  tenant_id varchar,
+  id int,
+  primary key (tenant_id, id)
+);
+
+SELECT create_distributed_table('event', 'tenant_id');
+SELECT create_distributed_table('page', 'tenant_id', colocate_with => 'event');
+alter table page add constraint fk21 foreign key (tenant_id, id) references event;
+
+SET client_min_messages TO WARNING;
+DROP SCHEMA issue_6543 CASCADE;