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

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

@@ -2255,7 +2255,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)
@@ -2270,11 +2271,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)
 	{
@@ -2286,9 +2293,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

@@ -608,7 +608,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))
@@ -627,7 +629,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

@@ -547,7 +547,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/expected/multi_alter_table_add_constraints.out

@@ -711,21 +711,6 @@ SET LOCAL application_name to 'citus_internal gpid=10000000001';
 SET citus.enable_ddl_propagation TO OFF;
 -- alter table triggers SELECT, and auto_explain catches that
 ALTER TABLE target_table ADD CONSTRAINT fkey_167 FOREIGN KEY (col_1) REFERENCES test_ref_table(key) ON DELETE CASCADE;
-LOG:  duration: xxxx ms  plan:
-{
-  "Query Text": "SELECT fk.\"col_1\" FROM ONLY \"test_auto_explain\".\"target_table\" fk LEFT OUTER JOIN ONLY \"test_auto_explain\".\"test_ref_table\" pk ON ( pk.\"key\" OPERATOR(pg_catalog.=) fk.\"col_1\") WHERE pk.\"key\" IS NULL AND (fk.\"col_1\" IS NOT NULL)",
-  "Plan": {
-    "Node Type": "Custom Scan",
-    "Custom Plan Provider": "Citus Adaptive",
-    "Parallel Aware": false,
-    "Startup Cost": 0.00,
-    "Total Cost": xxxx,
-    "Plan Rows": 100000,
-    "Plan Width": 4,
-    "Citus Explain Scan": "Explain for triggered constraint validation queries during ALTER TABLE commands are not supported by Citus"
-  }
-}
-CONTEXT:  SQL statement "SELECT fk."col_1" FROM ONLY "test_auto_explain"."target_table" fk LEFT OUTER JOIN ONLY "test_auto_explain"."test_ref_table" pk ON ( pk."key" OPERATOR(pg_catalog.=) fk."col_1") WHERE pk."key" IS NULL AND (fk."col_1" IS NOT NULL)"
 END;
 RESET citus.enable_ddl_propagation;
 SET client_min_messages to ERROR;

src/test/regress/multi_schedule

@@ -95,7 +95,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 issue_5763
+test: issue_5248 issue_5099 issue_5763 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;