Make sure to use correct execution mode for TRUNCATE

We used to set the execution mode in the truncate trigger. However,
when multiple tables are truncated with a single command, we could
set the execution mode very late. Instead, now set the execution mode
on the utility hook.

src/backend/distributed/executor/multi_utility.c

@@ -152,6 +152,7 @@ static void ErrorIfUnsupportedSeqStmt(CreateSeqStmt *createSeqStmt);
 static void ErrorIfDistributedAlterSeqOwnedBy(AlterSeqStmt *alterSeqStmt);
 static void ErrorIfUnsupportedTruncateStmt(TruncateStmt *truncateStatement);
 static void ProcessTruncateStatement(TruncateStmt *truncateStatement);
+static void ExecuteTruncateStmtSequentialIfNecessary(TruncateStmt *command);
 static void EnsurePartitionTableNotReplicatedForTruncate(TruncateStmt *truncateStatement);
 static void LockTruncatedRelationMetadataInWorkers(TruncateStmt *truncateStatement);
 static void AcquireDistributedLockOnRelations(List *relationIdList, LOCKMODE lockMode);
@@ -2940,10 +2941,55 @@ static void
 ProcessTruncateStatement(TruncateStmt *truncateStatement)
 {
 	EnsurePartitionTableNotReplicatedForTruncate(truncateStatement);
+	ExecuteTruncateStmtSequentialIfNecessary(truncateStatement);
 	LockTruncatedRelationMetadataInWorkers(truncateStatement);
 }
 
 
+/*
+ * ExecuteTruncateStmtSequentialIfNecessary decides if the TRUNCATE stmt needs
+ * to run sequential. If so, it calls SetLocalMultiShardModifyModeToSequential().
+ *
+ * If a reference table which has a foreign key from a distributed table is truncated
+ * we need to execute the command sequentially to avoid self-deadlock.
+ */
+static void
+ExecuteTruncateStmtSequentialIfNecessary(TruncateStmt *command)
+{
+	List *relationList = command->relations;
+	ListCell *relationCell = NULL;
+	bool failOK = false;
+
+	foreach(relationCell, relationList)
+	{
+		RangeVar *rangeVar = (RangeVar *) lfirst(relationCell);
+		Oid relationId = RangeVarGetRelid(rangeVar, NoLock, failOK);
+
+		if (IsDistributedTable(relationId) &&
+			PartitionMethod(relationId) == DISTRIBUTE_BY_NONE &&
+			TableReferenced(relationId))
+		{
+			char *relationName = get_rel_name(relationId);
+
+			ereport(DEBUG1, (errmsg("switching to sequential query execution mode"),
+							 errdetail(
+								 "Reference relation \"%s\" is modified, which might lead "
+								 "to data inconsistencies or distributed deadlocks via "
+								 "parallel accesses to hash distributed relations due to "
+								 "foreign keys. Any parallel modification to "
+								 "those hash distributed relations in the same "
+								 "transaction can only be executed in sequential query "
+								 "execution mode", relationName)));
+
+			SetLocalMultiShardModifyModeToSequential();
+
+			/* nothing to do more */
+			return;
+		}
+	}
+}
+
+
 /*
  * EnsurePartitionTableNotReplicatedForTruncate a simple wrapper around
  * EnsurePartitionTableNotReplicated for TRUNCATE command.

src/backend/distributed/master/master_modify_multiple_shards.c

@@ -61,7 +61,6 @@
 
 static List * ModifyMultipleShardsTaskList(Query *query, List *shardIntervalList, TaskType
 										   taskType);
-static bool ShouldExecuteTruncateStmtSequential(TruncateStmt *command);
 
 
 PG_FUNCTION_INFO_V1(master_modify_multiple_shards);
@@ -144,11 +143,6 @@ master_modify_multiple_shards(PG_FUNCTION_ARGS)
 		}
 
 		EnsureTablePermissions(relationId, ACL_TRUNCATE);
-
-		if (ShouldExecuteTruncateStmtSequential(truncateStatement))
-		{
-			SetLocalMultiShardModifyModeToSequential();
-		}
 	}
 	else
 	{
@@ -258,34 +252,3 @@ ModifyMultipleShardsTaskList(Query *query, List *shardIntervalList, TaskType tas
 
 	return taskList;
 }
-
-
-/*
- * ShouldExecuteTruncateStmtSequential decides if the TRUNCATE stmt needs
- * to run sequential. If so, it calls SetLocalMultiShardModifyModeToSequential().
- *
- * If a reference table which has a foreign key from a distributed table is truncated
- * we need to execute the command sequentially to avoid self-deadlock.
- */
-static bool
-ShouldExecuteTruncateStmtSequential(TruncateStmt *command)
-{
-	List *relationList = command->relations;
-	ListCell *relationCell = NULL;
-	bool failOK = false;
-
-	foreach(relationCell, relationList)
-	{
-		RangeVar *rangeVar = (RangeVar *) lfirst(relationCell);
-		Oid relationId = RangeVarGetRelid(rangeVar, NoLock, failOK);
-
-		if (IsDistributedTable(relationId) &&
-			PartitionMethod(relationId) == DISTRIBUTE_BY_NONE &&
-			TableReferenced(relationId))
-		{
-			return true;
-		}
-	}
-
-	return false;
-}

src/test/regress/expected/foreign_key_restriction_enforcement.out

@@ -457,6 +457,8 @@ BEGIN;
 (1 row)
 
 	TRUNCATE referece_table CASCADE;
+DEBUG:  switching to sequential query execution mode
+DETAIL:  Reference relation "referece_table" is modified, which might lead to data inconsistencies or distributed deadlocks via parallel accesses to hash distributed relations due to foreign keys. Any parallel modification to those hash distributed relations in the same transaction can only be executed in sequential query execution mode
 NOTICE:  truncate cascades to table "on_update_fkey_table"
 ERROR:  cannot execute DDL on reference relation "referece_table" because there was a parallel SELECT access to distributed relation "on_update_fkey_table" in the same transaction
 HINT:  Try re-running the transaction with "SET LOCAL citus.multi_shard_modify_mode TO 'sequential';"
@@ -470,6 +472,8 @@ BEGIN;
 (1 row)
 
 	TRUNCATE referece_table CASCADE;
+DEBUG:  switching to sequential query execution mode
+DETAIL:  Reference relation "referece_table" is modified, which might lead to data inconsistencies or distributed deadlocks via parallel accesses to hash distributed relations due to foreign keys. Any parallel modification to those hash distributed relations in the same transaction can only be executed in sequential query execution mode
 NOTICE:  truncate cascades to table "on_update_fkey_table"
 DEBUG:  truncate cascades to table "on_update_fkey_table_2380002"
 DETAIL:  NOTICE from localhost:57638

src/test/regress/expected/foreign_key_restriction_enforcement_0.out

@@ -457,6 +457,8 @@ BEGIN;
 (1 row)
 
 	TRUNCATE referece_table CASCADE;
+DEBUG:  switching to sequential query execution mode
+DETAIL:  Reference relation "referece_table" is modified, which might lead to data inconsistencies or distributed deadlocks via parallel accesses to hash distributed relations due to foreign keys. Any parallel modification to those hash distributed relations in the same transaction can only be executed in sequential query execution mode
 NOTICE:  truncate cascades to table "on_update_fkey_table"
 ERROR:  cannot execute DDL on reference relation "referece_table" because there was a parallel SELECT access to distributed relation "on_update_fkey_table" in the same transaction
 HINT:  Try re-running the transaction with "SET LOCAL citus.multi_shard_modify_mode TO 'sequential';"
@@ -470,6 +472,8 @@ BEGIN;
 (1 row)
 
 	TRUNCATE referece_table CASCADE;
+DEBUG:  switching to sequential query execution mode
+DETAIL:  Reference relation "referece_table" is modified, which might lead to data inconsistencies or distributed deadlocks via parallel accesses to hash distributed relations due to foreign keys. Any parallel modification to those hash distributed relations in the same transaction can only be executed in sequential query execution mode
 NOTICE:  truncate cascades to table "on_update_fkey_table"
 DEBUG:  truncate cascades to table "on_update_fkey_table_2380002"
 DETAIL:  NOTICE from localhost:57638

src/test/regress/expected/foreign_key_to_reference_table.out

@@ -1885,6 +1885,27 @@ CONTEXT:  while executing command on localhost:57638
 INSERT INTO referenced_table VALUES(5,5);
 -- should succeed since both of the foreign constraints are positive
 INSERT INTO referencing_table VALUES (0, 5);
+-- TRUNCATE should work in any way
+TRUNCATE referencing_table, referenced_table;
+TRUNCATE referenced_table, referencing_table;
+BEGIN;
+  TRUNCATE referencing_table, referenced_table;
+  ALTER TABLE referencing_table ADD COLUMN x INT;
+  SELECT * FROM referencing_table;
+ id | value_1 | x 
+----+---------+---
+(0 rows)
+
+ROLLBACK;
+BEGIN;
+  TRUNCATE referenced_table, referencing_table;
+  ALTER TABLE referencing_table ADD COLUMN x INT;
+  SELECT * FROM referencing_table;
+ id | value_1 | x 
+----+---------+---
+(0 rows)
+
+ROLLBACK;
 DROP TABLE referenced_table CASCADE;
 NOTICE:  drop cascades to constraint fkey_to_ref on table referencing_table_4
 DROP TABLE referencing_table;

src/test/regress/expected/foreign_key_to_reference_table_0.out

@@ -1904,6 +1904,27 @@ INSERT INTO referencing_table VALUES (0, 5);
 ERROR:  relation "referencing_table" does not exist
 LINE 1: INSERT INTO referencing_table VALUES (0, 5);
                     ^
+-- TRUNCATE should work in any way
+TRUNCATE referencing_table, referenced_table;
+ERROR:  relation "referencing_table" does not exist
+TRUNCATE referenced_table, referencing_table;
+ERROR:  relation "referencing_table" does not exist
+BEGIN;
+  TRUNCATE referencing_table, referenced_table;
+ERROR:  relation "referencing_table" does not exist
+  ALTER TABLE referencing_table ADD COLUMN x INT;
+ERROR:  current transaction is aborted, commands ignored until end of transaction block
+  SELECT * FROM referencing_table;
+ERROR:  current transaction is aborted, commands ignored until end of transaction block
+ROLLBACK;
+BEGIN;
+  TRUNCATE referenced_table, referencing_table;
+ERROR:  relation "referencing_table" does not exist
+  ALTER TABLE referencing_table ADD COLUMN x INT;
+ERROR:  current transaction is aborted, commands ignored until end of transaction block
+  SELECT * FROM referencing_table;
+ERROR:  current transaction is aborted, commands ignored until end of transaction block
+ROLLBACK;
 DROP TABLE referenced_table CASCADE;
 DROP TABLE referencing_table;
 ERROR:  table "referencing_table" does not exist

src/test/regress/sql/foreign_key_to_reference_table.sql

@@ -949,6 +949,22 @@ INSERT INTO referenced_table VALUES(5,5);
 -- should succeed since both of the foreign constraints are positive
 INSERT INTO referencing_table VALUES (0, 5);
 
+-- TRUNCATE should work in any way
+TRUNCATE referencing_table, referenced_table;
+TRUNCATE referenced_table, referencing_table;
+
+BEGIN;
+  TRUNCATE referencing_table, referenced_table;
+  ALTER TABLE referencing_table ADD COLUMN x INT;
+  SELECT * FROM referencing_table;
+ROLLBACK;
+
+BEGIN;
+  TRUNCATE referenced_table, referencing_table;
+  ALTER TABLE referencing_table ADD COLUMN x INT;
+  SELECT * FROM referencing_table;
+ROLLBACK;
+
 DROP TABLE referenced_table CASCADE;
 DROP TABLE referencing_table;