From ed63df0fdeac43443b8e15709c19fe6c1f38a1f8 Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@paquier.xyz>
Date: Mon, 20 Jan 2020 22:39:34 +0900
Subject: [PATCH 1/2] Add FDW callback for support of TRUNCATE

---
 src/include/foreign/fdwapi.h               |   7 ++
 src/backend/commands/tablecmds.c           | 110 ++++++++++++++++++++-
 src/test/regress/expected/foreign_data.out |   8 +-
 doc/src/sgml/fdwhandler.sgml               |  36 +++++++
 4 files changed, 153 insertions(+), 8 deletions(-)

diff --git a/src/include/foreign/fdwapi.h b/src/include/foreign/fdwapi.h
index 95556dfb15..0a9f36735e 100644
--- a/src/include/foreign/fdwapi.h
+++ b/src/include/foreign/fdwapi.h
@@ -151,6 +151,10 @@ typedef bool (*AnalyzeForeignTable_function) (Relation relation,
 typedef List *(*ImportForeignSchema_function) (ImportForeignSchemaStmt *stmt,
 											   Oid serverOid);
 
+typedef void (*ExecForeignTruncate_function) (List *frels_list,
+											  DropBehavior behavior,
+											  bool restart_seqs);
+
 typedef Size (*EstimateDSMForeignScan_function) (ForeignScanState *node,
 												 ParallelContext *pcxt);
 typedef void (*InitializeDSMForeignScan_function) (ForeignScanState *node,
@@ -236,6 +240,9 @@ typedef struct FdwRoutine
 	/* Support functions for IMPORT FOREIGN SCHEMA */
 	ImportForeignSchema_function ImportForeignSchema;
 
+	/* Support functions for TRUNCATE */
+	ExecForeignTruncate_function ExecForeignTruncate;
+
 	/* Support functions for parallelism under Gather node */
 	IsForeignScanParallelSafe_function IsForeignScanParallelSafe;
 	EstimateDSMForeignScan_function EstimateDSMForeignScan;
diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
index 30b72b6297..f83f88a82f 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -60,6 +60,7 @@
 #include "commands/user.h"
 #include "executor/executor.h"
 #include "foreign/foreign.h"
+#include "foreign/fdwapi.h"
 #include "miscadmin.h"
 #include "nodes/makefuncs.h"
 #include "nodes/nodeFuncs.h"
@@ -295,6 +296,20 @@ struct DropRelationCallbackState
 #define		ATT_FOREIGN_TABLE		0x0020
 #define		ATT_PARTITIONED_INDEX	0x0040
 
+/*
+ * ForeignTruncateInfo
+ *
+ * Information related to truncation of foreign tables.  This
+ * is a single entry part of a hash table using the server OID
+ * as key, with a list of all foreign tables involved.
+ */
+typedef struct
+{
+	Oid		server_oid;
+	List   *frels_list;
+} ForeignTruncateInfo;
+
+
 /*
  * Partition tables are expected to be dropped when the parent partitioned
  * table gets dropped. Hence for partitioning we use AUTO dependency.
@@ -1647,6 +1662,7 @@ ExecuteTruncateGuts(List *explicit_rels, List *relids, List *relids_logged,
 {
 	List	   *rels;
 	List	   *seq_relids = NIL;
+	HTAB	   *ft_htab = NULL;
 	EState	   *estate;
 	ResultRelInfo *resultRelInfos;
 	ResultRelInfo *resultRelInfo;
@@ -1792,6 +1808,57 @@ ExecuteTruncateGuts(List *explicit_rels, List *relids, List *relids_logged,
 		if (rel->rd_rel->relkind == RELKIND_PARTITIONED_TABLE)
 			continue;
 
+		/*
+		 * If truncating a foreign table, the foreign data wrapper
+		 * callback for TRUNCATE is called once for each server
+		 * with a list of all the relations to process linked to this
+		 * server.  The list of relations for each server is saved as
+		 * a single entry of a hash table that uses the server OID as
+		 * lookup key.  Once the full set of lists is built, all the
+		 * entries of the hash table are scanned, and the list of relations
+		 * associated to the server is passed down to the TRUNCATE callback
+		 * of its foreign data wrapper.
+		 */
+		if (rel->rd_rel->relkind == RELKIND_FOREIGN_TABLE)
+		{
+			Oid		frel_oid = RelationGetRelid(rel);
+			Oid		server_oid = GetForeignServerIdByRelId(frel_oid);
+			bool	found;
+			ForeignTruncateInfo *ft_info;
+
+			/* if the hash table does not exist yet, initialize it */
+			if (!ft_htab)
+			{
+				HASHCTL		hctl;
+
+				memset(&hctl, 0, sizeof(HASHCTL));
+				hctl.keysize = sizeof(Oid);
+				hctl.entrysize = sizeof(ForeignTruncateInfo);
+				hctl.hcxt = CurrentMemoryContext;
+
+				ft_htab = hash_create("TRUNCATE for Foreign Tables",
+									  32,	/* start small and extend */
+									  &hctl,
+									  HASH_ELEM | HASH_BLOBS | HASH_CONTEXT);
+			}
+
+			/*
+			 * Look after the entry of the server in the hash table,
+			 * and initialize it if the entry does not exist yet.
+			 */
+			ft_info = hash_search(ft_htab, &server_oid, HASH_ENTER, &found);
+			if (!found)
+			{
+				ft_info->server_oid = server_oid;
+				ft_info->frels_list = NIL;
+
+			}
+
+			/* save the relation in the list */
+			ft_info->frels_list = lappend(ft_info->frels_list, rel);
+			continue;
+		}
+
 		/*
 		 * Normally, we need a transaction-safe truncation here.  However, if
 		 * the table was either created in the current (sub)transaction or has
@@ -1852,6 +1919,29 @@ ExecuteTruncateGuts(List *explicit_rels, List *relids, List *relids_logged,
 		pgstat_count_truncate(rel);
 	}
 
+	/*
+	 * Now go through the hash table, and process each entry associated
+	 * to the servers involved in the TRUNCATE.
+	 */
+	if (ft_htab)
+	{
+		ForeignTruncateInfo *ft_info;
+		HASH_SEQ_STATUS	seq;
+
+		hash_seq_init(&seq, ft_htab);
+
+		while ((ft_info = hash_seq_search(&seq)) != NULL)
+		{
+			FdwRoutine *routine = GetFdwRoutineByServerId(ft_info->server_oid);
+
+			/* truncate_check_rel() has checked that already */
+			Assert(routine->ExecForeignTruncate != NULL);
+			routine->ExecForeignTruncate(ft_info->frels_list,
+										 behavior,
+										 restart_seqs);
+		}
+	}
+
 	/*
 	 * Restart owned sequences if we were asked to.
 	 */
@@ -1939,12 +2029,24 @@ truncate_check_rel(Oid relid, Form_pg_class reltuple)
 	char	   *relname = NameStr(reltuple->relname);
 
 	/*
-	 * Only allow truncate on regular tables and partitioned tables (although,
+	 * Only allow truncate on regular tables, foreign tables using foreign
+	 * data wrappers supporting TRUNCATE and partitioned tables (although,
 	 * the latter are only being included here for the following checks; no
-	 * physical truncation will occur in their case.)
+	 * physical truncation will occur in their case.).
 	 */
-	if (reltuple->relkind != RELKIND_RELATION &&
-		reltuple->relkind != RELKIND_PARTITIONED_TABLE)
+	if (reltuple->relkind == RELKIND_FOREIGN_TABLE)
+	{
+		Oid			server_id = GetForeignServerIdByRelId(relid);
+		FdwRoutine *fdwroutine = GetFdwRoutineByServerId(server_id);
+
+		if (!fdwroutine->ExecForeignTruncate)
+			ereport(ERROR,
+					(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+					 errmsg("cannot truncate foreign table \"%s\"",
+							relname)));
+	}
+	else if (reltuple->relkind != RELKIND_RELATION &&
+			 reltuple->relkind != RELKIND_PARTITIONED_TABLE)
 		ereport(ERROR,
 				(errcode(ERRCODE_WRONG_OBJECT_TYPE),
 				 errmsg("\"%s\" is not a table", relname)));
diff --git a/src/test/regress/expected/foreign_data.out b/src/test/regress/expected/foreign_data.out
index b9e25820bc..e2c0bcea51 100644
--- a/src/test/regress/expected/foreign_data.out
+++ b/src/test/regress/expected/foreign_data.out
@@ -1807,9 +1807,9 @@ Inherits: fd_pt1
 
 -- TRUNCATE doesn't work on foreign tables, either directly or recursively
 TRUNCATE ft2;  -- ERROR
-ERROR:  "ft2" is not a table
+ERROR:  foreign-data wrapper "dummy" has no handler
 TRUNCATE fd_pt1;  -- ERROR
-ERROR:  "ft2" is not a table
+ERROR:  foreign-data wrapper "dummy" has no handler
 DROP TABLE fd_pt1 CASCADE;
 NOTICE:  drop cascades to foreign table ft2
 -- IMPORT FOREIGN SCHEMA
@@ -2032,9 +2032,9 @@ ALTER FOREIGN TABLE fd_pt2_1 ADD CONSTRAINT fd_pt2chk1 CHECK (c1 > 0);
 ALTER TABLE fd_pt2 ATTACH PARTITION fd_pt2_1 FOR VALUES IN (1);
 -- TRUNCATE doesn't work on foreign tables, either directly or recursively
 TRUNCATE fd_pt2_1;  -- ERROR
-ERROR:  "fd_pt2_1" is not a table
+ERROR:  foreign-data wrapper "dummy" has no handler
 TRUNCATE fd_pt2;  -- ERROR
-ERROR:  "fd_pt2_1" is not a table
+ERROR:  foreign-data wrapper "dummy" has no handler
 DROP FOREIGN TABLE fd_pt2_1;
 DROP TABLE fd_pt2;
 -- foreign table cannot be part of partition tree made of temporary
diff --git a/doc/src/sgml/fdwhandler.sgml b/doc/src/sgml/fdwhandler.sgml
index 6587678af2..f2416c9074 100644
--- a/doc/src/sgml/fdwhandler.sgml
+++ b/doc/src/sgml/fdwhandler.sgml
@@ -968,6 +968,42 @@ EndDirectModify(ForeignScanState *node);
 
    </sect2>
 
+   <sect2 id="fdw-callbacks-truncate">
+    <title>FDW Routines for Truncate</title>
+<programlisting>
+void
+ExecForeignTruncate(List *frels_list,
+                    DropBehavior behavior, bool restart_seqs);
+</programlisting>
+    <para>
+     Truncate a set of foreign tables defined by
+     <literal>frels_list</literal> belonging to the same foreign server.
+     This optional function is called during execution of
+     <command>TRUNCATE</command> for each foreign server being involved
+     in one <command>TRUNCATE</command> command (note that invocations
+     are not per foreign table).
+    </para>
+    
+    <para>
+     If the <function>ExecForeignTruncate</function> pointer is set to
+     <literal>NULL</literal>, attempts to truncate the foreign table will
+     fail with an error message.
+    </para>
+    
+    <para>
+     <literal>behavior</literal> defines how foreign tables should
+     be truncated, using as possible values <literal>DROP_RESTRICT</literal>
+     and <literal>DROP_CASCADE</literal> (to map with the equivalents of
+     <command>TRUNCATE</command>).
+    </para>
+    
+    <para>
+     <literal>restart_seqs</literal> is set to <literal>true</literal>
+     if <literal>RESTART IDENTITY</literal> was supplied in the
+     <command>TRUNCATE</command>.
+    </para>
+   </sect2>
+
    <sect2 id="fdw-callbacks-row-locking">
     <title>FDW Routines for Row Locking</title>
 
-- 
2.25.0

