From b711634f5d109b51e7deca7905bd297f35bc04ce Mon Sep 17 00:00:00 2001
From: Masahiko Sawada <sawada.mshk@gmail.com>
Date: Mon, 21 Sep 2020 17:09:41 +0900
Subject: [PATCH v26 06/11] Add GetPrepareID API.

Co-authored-by: Masahiko Sawada, Ashutosh Bapat
---
 src/backend/access/fdwxact/fdwxact.c | 50 +++++++++++++++++++++++-----
 src/include/foreign/fdwapi.h         |  3 ++
 2 files changed, 45 insertions(+), 8 deletions(-)

diff --git a/src/backend/access/fdwxact/fdwxact.c b/src/backend/access/fdwxact/fdwxact.c
index bf5daedaa9..c6f6a92752 100644
--- a/src/backend/access/fdwxact/fdwxact.c
+++ b/src/backend/access/fdwxact/fdwxact.c
@@ -136,6 +136,7 @@ typedef struct FdwXactParticipant
 	CommitForeignTransaction_function commit_foreign_xact_fn;
 	RollbackForeignTransaction_function rollback_foreign_xact_fn;
 	PrepareForeignTransaction_function prepare_foreign_xact_fn;
+	GetPrepareId_function get_prepareid_fn;
 } FdwXactParticipant;
 
 /*
@@ -408,9 +409,10 @@ FdwXactPrepareForeignTransactions(void)
 }
 
 /*
- * Return a null-terminated foreign transaction identifier.  We generate an
- * unique identifier with in the form of
- * "fx_<random number>_<xid>_<serverid>_<userid> whose length is
+ * Return a null-terminated foreign transaction identifier.  If the given
+ * foreign server's FDW provides getPrepareId callback we return the identifier
+ * returned from it. Otherwise we generate an unique identifier with in the
+ * form of "fx_<random number>_<xid>_<serverid>_<userid> whose length is
  * less than FDWXACT_ID_MAX_LEN.
  *
  * Returned string value is used to identify foreign transaction. The
@@ -425,13 +427,45 @@ FdwXactPrepareForeignTransactions(void)
 static char *
 get_fdwxact_identifier(FdwXactParticipant *fdw_part, TransactionId xid)
 {
-	char		buf[FDWXACT_ID_MAX_LEN] = {0};
+	/*
+	 * If FDW doesn't provide the callback function, generate an unique
+	 * identifier.
+	 */
+	if (!fdw_part->get_prepareid_fn)
+	{
+		char		buf[FDWXACT_ID_MAX_LEN] = {0};
 
-	snprintf(buf, FDWXACT_ID_MAX_LEN, "fx_%ld_%u_%d_%d",
-			 Abs(random()), xid, fdw_part->server->serverid,
-			 fdw_part->usermapping->userid);
+		snprintf(buf, FDWXACT_ID_MAX_LEN, "fx_%ld_%u_%d_%d",
+				 Abs(random()), xid, fdw_part->server->serverid,
+				 fdw_part->usermapping->userid);
+
+		return pstrdup(buf);
+	}
+
+	/* Get an unique identifier from callback function */
+	id = fdw_part->get_prepareid_fn(xid, fdw_part->server->serverid,
+									fdw_part->usermapping->userid,
+									&id_len);
+
+	if (id == NULL)
+		ereport(ERROR,
+				(errcode(ERRCODE_UNDEFINED_OBJECT),
+				 (errmsg("foreign transaction identifier is not provided"))));
+
+	/* Check length of foreign transaction identifier */
+	if (id_len > FDWXACT_ID_MAX_LEN)
+	{
+		id[FDWXACT_ID_MAX_LEN] = '\0';
+		ereport(ERROR,
+				(errcode(ERRCODE_NAME_TOO_LONG),
+				 errmsg("foreign transaction identifier \"%s\" is too long",
+						id),
+				 errdetail("Foreign transaction identifier must be less than %d characters.",
+						   FDWXACT_ID_MAX_LEN)));
+	}
 
-	return pstrdup(buf);
+	id[id_len] = '\0';
+	return pstrdup(id);
 }
 
 /*
diff --git a/src/include/foreign/fdwapi.h b/src/include/foreign/fdwapi.h
index 89cec9aa96..91db4f5bfc 100644
--- a/src/include/foreign/fdwapi.h
+++ b/src/include/foreign/fdwapi.h
@@ -174,6 +174,8 @@ typedef List *(*ReparameterizeForeignPathByChild_function) (PlannerInfo *root,
 typedef void (*PrepareForeignTransaction_function) (FdwXactRslvState *frstate);
 typedef void (*CommitForeignTransaction_function) (FdwXactRslvState *frstate);
 typedef void (*RollbackForeignTransaction_function) (FdwXactRslvState *frstate);
+typedef char *(*GetPrepareId_function) (TransactionId xid, Oid serverid,
+										Oid userid, int *prep_id_len);
 
 /*
  * FdwRoutine is the struct returned by a foreign-data wrapper's handler
@@ -256,6 +258,7 @@ typedef struct FdwRoutine
 	CommitForeignTransaction_function CommitForeignTransaction;
 	RollbackForeignTransaction_function RollbackForeignTransaction;
 	PrepareForeignTransaction_function PrepareForeignTransaction;
+	GetPrepareId_function GetPrepareId;
 } FdwRoutine;
 
 
-- 
2.23.0