From 5c63b8da2b541f54857db9e894ee8250ca256f3e Mon Sep 17 00:00:00 2001
From: Jehan-Guillaume de Rorthais <jgdr@dalibo.com>
Date: Tue, 4 Feb 2020 18:05:17 +0100
Subject: [PATCH v2] Fix TRUNCATE on a partition to apply CASCADE to
 partitioned table.

Previously, when running TRUNCATE CASCADE on a child of a
partitioned table referenced by another partitioned table, the
truncate was not applied to referencing children. This leaves a
FK constraint violation in the referencing partitioned table.

This commit fixes that bug.

Reported-by: Christophe Courtois
Author: Jehan-Guillaume de Rorthais
---
 src/backend/catalog/heap.c             | 71 +++++++++++++++++++++++---
 src/test/regress/expected/truncate.out | 45 ++++++++++++++++
 src/test/regress/sql/truncate.sql      | 34 ++++++++++++
 3 files changed, 143 insertions(+), 7 deletions(-)

diff --git a/src/backend/catalog/heap.c b/src/backend/catalog/heap.c
index 046b3d37ce..4a1986f126 100644
--- a/src/backend/catalog/heap.c
+++ b/src/backend/catalog/heap.c
@@ -3396,30 +3396,45 @@ List *
 heap_truncate_find_FKs(List *relationIds)
 {
 	List	   *result = NIL;
+	List	   *parent_cons = NIL;
+	List	   *oids = list_copy(relationIds);
+	ListCell   *cell;
+	ScanKeyData key;
 	Relation	fkeyRel;
 	SysScanDesc fkeyScan;
 	HeapTuple	tuple;
 
+	fkeyRel = table_open(ConstraintRelationId, AccessShareLock);
+
 	/*
 	 * Must scan pg_constraint.  Right now, it is a seqscan because there is
 	 * no available index on confrelid.
 	 */
-	fkeyRel = table_open(ConstraintRelationId, AccessShareLock);
+	ScanKeyInit(&key,
+				Anum_pg_constraint_contype,
+				BTEqualStrategyNumber, F_CHAREQ,
+				CharGetDatum(CONSTRAINT_FOREIGN));
 
 	fkeyScan = systable_beginscan(fkeyRel, InvalidOid, false,
-								  NULL, 0, NULL);
+								  NULL, 1, &key);
 
 	while (HeapTupleIsValid(tuple = systable_getnext(fkeyScan)))
 	{
 		Form_pg_constraint con = (Form_pg_constraint) GETSTRUCT(tuple);
 
-		/* Not a foreign key */
-		if (con->contype != CONSTRAINT_FOREIGN)
+		/* Not referencing one of our list of tables */
+		if (!list_member_oid(oids, con->confrelid))
 			continue;
 
-		/* Not referencing one of our list of tables */
-		if (!list_member_oid(relationIds, con->confrelid))
-			continue;
+		/*
+		 * If this constraint has a parent constraint, keep track of that,
+		 * because we need to search possible cascaded constraints from that
+		 * one.  This will only happen once, because if the parent constraint
+		 * in turn has a parent constraint, the loop below will catch that.
+		 */
+		if (OidIsValid(con->conparentid) &&
+			!list_member_oid(parent_cons, con->conparentid))
+			parent_cons = lappend_oid(parent_cons, con->conparentid);
 
 		/* Add referencer to result, unless present in input list */
 		if (!list_member_oid(relationIds, con->conrelid))
@@ -3427,8 +3442,50 @@ heap_truncate_find_FKs(List *relationIds)
 	}
 
 	systable_endscan(fkeyScan);
+
+	/*
+	 * Process each parent constraint we found.  We do two things with them:
+	 * first, add its referenced relation to our output list.  Second, append
+	 * any possible parent constraint of the parent constraint to our own
+	 * list; that makes us process the whole partitioning hierarchy in a single
+	 * pass.
+	 */
+	foreach(cell, parent_cons)
+	{
+		Oid		parent = lfirst_oid(cell);
+
+		ScanKeyInit(&key,
+					Anum_pg_constraint_conparentid,
+					BTEqualStrategyNumber, F_OIDEQ,
+					ObjectIdGetDatum(parent));
+
+		fkeyScan = systable_beginscan(fkeyRel, ConstraintParentIndexId,
+									  true, NULL, 1, &key);
+
+		while (HeapTupleIsValid(tuple = systable_getnext(fkeyScan)))
+		{
+			Form_pg_constraint con = (Form_pg_constraint) GETSTRUCT(tuple);
+
+			/* If this constraint has a parent constraint, add that to the
+			 * list we're processing so that it's processed in this loop; this
+			 * way, we only have to do this the first time around.
+			 */
+			if (con->conparentid)
+				parent_cons = list_append_unique_oid(parent_cons,
+													 con->conparentid);
+
+			/* Add referencer to result list, unless already present. */
+			result = list_append_unique_oid(relationIds, con->conrelid);
+		}
+
+		systable_endscan(fkeyScan);
+	}
+
 	table_close(fkeyRel, AccessShareLock);
 
+	list_free(parent_cons);
+	list_free(oids);
+
 	/* Now sort and de-duplicate the result list */
 	list_sort(result, list_oid_cmp);
 	list_deduplicate_oid(result);
diff --git a/src/test/regress/expected/truncate.out b/src/test/regress/expected/truncate.out
index cc68274dca..b07c429158 100644
--- a/src/test/regress/expected/truncate.out
+++ b/src/test/regress/expected/truncate.out
@@ -542,3 +542,48 @@ SELECT * FROM tp_chk_data();
 
 DROP TABLE truncprim, truncpart;
 DROP FUNCTION tp_ins_data(), tp_chk_data();
+-- test cascade when referencing a partitioned table
+CREATE TABLE trunc_a (a INT PRIMARY KEY) PARTITION BY RANGE (a);
+CREATE TABLE trunc_a1 PARTITION OF trunc_a FOR VALUES FROM (0) TO (10);
+CREATE TABLE trunc_a2 PARTITION OF trunc_a FOR VALUES FROM (10) TO (20);
+CREATE TABLE trunc_a3 PARTITION OF trunc_a FOR VALUES FROM (20) TO (30);
+INSERT INTO trunc_a VALUES (0), (5), (10), (15), (20), (25);
+-- truncate a partition cascading to a table
+CREATE TABLE ref_b (
+    b INT PRIMARY KEY,
+    a INT REFERENCES trunc_a(a) ON DELETE CASCADE
+);
+INSERT INTO ref_b VALUES (10, 0), (50, 5), (100, 10), (150, 15);
+TRUNCATE TABLE trunc_a1 CASCADE;
+NOTICE:  truncate cascades to table "ref_b"
+SELECT a FROM ref_b;
+ a 
+---
+(0 rows)
+
+DROP TABLE ref_b;
+-- truncate a partition cascading to a partitioned table
+CREATE TABLE ref_c (
+    c INT PRIMARY KEY,
+    a INT REFERENCES trunc_a(a) ON DELETE CASCADE
+) PARTITION BY RANGE (c);
+CREATE TABLE ref_c1 PARTITION OF ref_c FOR VALUES FROM (100) TO (200);
+CREATE TABLE ref_c2 PARTITION OF ref_c FOR VALUES FROM (200) TO (300);
+INSERT INTO ref_c VALUES (100, 10), (150, 15), (200, 20), (250, 25);
+TRUNCATE TABLE trunc_a2 CASCADE;
+NOTICE:  truncate cascades to table "ref_c"
+NOTICE:  truncate cascades to table "ref_c1"
+NOTICE:  truncate cascades to table "ref_c2"
+SELECT a FROM ref_c;
+ a 
+---
+(0 rows)
+
+SELECT a FROM trunc_a ORDER BY a;
+ a  
+----
+ 20
+ 25
+(2 rows)
+
+DROP TABLE trunc_a, ref_c;
diff --git a/src/test/regress/sql/truncate.sql b/src/test/regress/sql/truncate.sql
index 28395e82bf..b00311a72f 100644
--- a/src/test/regress/sql/truncate.sql
+++ b/src/test/regress/sql/truncate.sql
@@ -289,3 +289,37 @@ TRUNCATE TABLE truncpart;
 SELECT * FROM tp_chk_data();
 DROP TABLE truncprim, truncpart;
 DROP FUNCTION tp_ins_data(), tp_chk_data();
+
+-- test cascade when referencing a partitioned table
+CREATE TABLE trunc_a (a INT PRIMARY KEY) PARTITION BY RANGE (a);
+CREATE TABLE trunc_a1 PARTITION OF trunc_a FOR VALUES FROM (0) TO (10);
+CREATE TABLE trunc_a2 PARTITION OF trunc_a FOR VALUES FROM (10) TO (20);
+CREATE TABLE trunc_a3 PARTITION OF trunc_a FOR VALUES FROM (20) TO (30);
+INSERT INTO trunc_a VALUES (0), (5), (10), (15), (20), (25);
+
+-- truncate a partition cascading to a table
+CREATE TABLE ref_b (
+    b INT PRIMARY KEY,
+    a INT REFERENCES trunc_a(a) ON DELETE CASCADE
+);
+INSERT INTO ref_b VALUES (10, 0), (50, 5), (100, 10), (150, 15);
+
+TRUNCATE TABLE trunc_a1 CASCADE;
+SELECT a FROM ref_b;
+
+DROP TABLE ref_b;
+
+-- truncate a partition cascading to a partitioned table
+CREATE TABLE ref_c (
+    c INT PRIMARY KEY,
+    a INT REFERENCES trunc_a(a) ON DELETE CASCADE
+) PARTITION BY RANGE (c);
+CREATE TABLE ref_c1 PARTITION OF ref_c FOR VALUES FROM (100) TO (200);
+CREATE TABLE ref_c2 PARTITION OF ref_c FOR VALUES FROM (200) TO (300);
+INSERT INTO ref_c VALUES (100, 10), (150, 15), (200, 20), (250, 25);
+
+TRUNCATE TABLE trunc_a2 CASCADE;
+SELECT a FROM ref_c;
+SELECT a FROM trunc_a ORDER BY a;
+
+DROP TABLE trunc_a, ref_c;
-- 
2.20.1

