From c4412a46d3fe096337327da81f4cb7466e8445f8 Mon Sep 17 00:00:00 2001
From: Jeevan Chalke <jeevan.chalke@enterprisedb.com>
Date: Mon, 12 Oct 2020 20:02:17 +0530
Subject: [PATCH] Fix crash with COPY FROM and/or foreign partition routing
 operations.

COPY FROM and/or foreign partition routing code path in core assumes
that FDW has BeginForeignInsert() APIs present and thus later
executes ExecForeignInsert().  However, mongo_fdw does not support
routable foreign-table partitions and/or executing COPY FROM on
foreign tables and thus do not have BeginForeignInsert() API
implemented.  But as it has ExecForeignInsert() API, it gets called
for these operations and results in the server crash.  To fix this,
add the BeginForeignInsert() API that throws an error.  Also, add
EndForeignInsert() similar to the Begin API.

The issue was originally reported on MySQL's GitHub through #208.

FDW-226, Suraj Kharage, reviewed by Jeevan Chalke.
---
 mongo_fdw.c | 45 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 45 insertions(+)

diff --git a/mongo_fdw.c b/mongo_fdw.c
index 4427d15..40d9407 100644
--- a/mongo_fdw.c
+++ b/mongo_fdw.c
@@ -117,6 +117,12 @@ static void MongoExplainForeignModify(ModifyTableState *mtstate,
 static bool MongoAnalyzeForeignTable(Relation relation,
 									 AcquireSampleRowsFunc *func,
 									 BlockNumber *totalpages);
+#if PG_VERSION_NUM >= 110000
+static void MongoBeginForeignInsert(ModifyTableState *mtstate,
+									ResultRelInfo *resultRelInfo);
+static void MongoEndForeignInsert(EState *estate,
+								  ResultRelInfo *resultRelInfo);
+#endif
 
 /*
  * Helper functions
@@ -202,6 +208,12 @@ mongo_fdw_handler(PG_FUNCTION_ARGS)
 	/* Support for ANALYZE */
 	fdwRoutine->AnalyzeForeignTable = MongoAnalyzeForeignTable;
 
+#if PG_VERSION_NUM >= 110000
+	/* Partition routing and/or COPY from */
+	fdwRoutine->BeginForeignInsert = MongoBeginForeignInsert;
+	fdwRoutine->EndForeignInsert = MongoEndForeignInsert;
+#endif
+
 	PG_RETURN_POINTER(fdwRoutine);
 }
 
@@ -2186,3 +2198,36 @@ mongo_fdw_version(PG_FUNCTION_ARGS)
 {
 	PG_RETURN_INT32(CODE_VERSION);
 }
+
+#if PG_VERSION_NUM >= 110000
+/*
+ * MongoBeginForeignInsert
+ * 		Prepare for an insert operation triggered by partition routing
+ * 		or COPY FROM.
+ *
+ * This is not yet supported, so raise an error.
+ */
+static void
+MongoBeginForeignInsert(ModifyTableState *mtstate,
+						ResultRelInfo *resultRelInfo)
+{
+	ereport(ERROR,
+			(errcode(ERRCODE_FDW_UNABLE_TO_CREATE_EXECUTION),
+			 errmsg("COPY and foreign partition routing not supported in mongo_fdw")));
+}
+
+/*
+ * MongoEndForeignInsert
+ * 		BeginForeignInsert() is not yet implemented, hence we do not
+ * 		have anything to cleanup as of now. We throw an error here just
+ * 		to make sure when we do that we do not forget to cleanup
+ * 		resources.
+ */
+static void
+MongoEndForeignInsert(EState *estate, ResultRelInfo *resultRelInfo)
+{
+	ereport(ERROR,
+			(errcode(ERRCODE_FDW_UNABLE_TO_CREATE_EXECUTION),
+			 errmsg("COPY and foreign partition routing not supported in mongo_fdw")));
+}
+#endif