From a83630b966029b7a6ef7e63820fb7e4c08475e33 Mon Sep 17 00:00:00 2001
From: Masahiko Sawada <sawada.mshk@gmail.com>
Date: Wed, 4 Nov 2020 14:41:53 +0900
Subject: [PATCH v35 05/10] Add GetPrepareId API

Co-authored-by: Masahiko Sawada, Ashutosh Bapat
---
 src/backend/access/transam/fdwxact.c | 52 ++++++++++++++++++++++++----
 src/include/foreign/fdwapi.h         |  4 +++
 2 files changed, 49 insertions(+), 7 deletions(-)

diff --git a/src/backend/access/transam/fdwxact.c b/src/backend/access/transam/fdwxact.c
index 60d726af62..c8a75bcfbc 100644
--- a/src/backend/access/transam/fdwxact.c
+++ b/src/backend/access/transam/fdwxact.c
@@ -157,6 +157,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;
 
 /*
@@ -328,6 +329,7 @@ FdwXactRegisterXact(UserMapping *usermapping)
 	fdw_part->commit_foreign_xact_fn = routine->CommitForeignTransaction;
 	fdw_part->rollback_foreign_xact_fn = routine->RollbackForeignTransaction;
 	fdw_part->prepare_foreign_xact_fn = routine->PrepareForeignTransaction;
+	fdw_part->get_prepareid_fn = routine->GetPrepareId;
 
 	MemoryContextSwitchTo(old_ctx);
 	pfree(routine);
@@ -662,9 +664,10 @@ FdwXactPrepareForeignTransactions(TransactionId xid)
 }
 
 /*
- * Return a null-terminated foreign transaction identifier.  We generate an
- * unique identifier with in the form of
- * "fx_<random number>_<xid>_<umid> whose length is less than FDWXACT_ID_MAX_LEN.
+ * Return a null-terminated foreign transaction identifier.  If the given FDW
+ * supports getPrepareId callback we return the identifier returned from it.
+ * Otherwise we generate an unique identifier with in the form of
+ * "fx_<random number>_<xid>_<umid>" whose length is less than FDWXACT_ID_MAX_LEN.
  *
  * Returned string value is used to identify foreign transaction. The
  * identifier should not be same as any other concurrent prepared transaction
@@ -678,12 +681,47 @@ FdwXactPrepareForeignTransactions(TransactionId xid)
 static char *
 getFdwXactIdentifier(FdwXactParticipant *fdw_part, TransactionId xid)
 {
-	char		buf[FDWXACT_ID_MAX_LEN] = {0};
+	char *id;
+	int	id_len;
 
-	snprintf(buf, FDWXACT_ID_MAX_LEN, "fx_%ld_%u_%u", Abs(random()),
-			 xid, fdw_part->key);
+	/*
+	 * 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_%u", Abs(random()),
+				 xid, fdw_part->key);
+
+		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 e6e294e3ec..4bb6f3738b 100644
--- a/src/include/foreign/fdwapi.h
+++ b/src/include/foreign/fdwapi.h
@@ -179,9 +179,12 @@ typedef List *(*ReparameterizeForeignPathByChild_function) (PlannerInfo *root,
 															List *fdw_private,
 															RelOptInfo *child_rel);
 
+
 typedef void (*PrepareForeignTransaction_function) (FdwXactInfo *finfo);
 typedef void (*CommitForeignTransaction_function) (FdwXactInfo *finfo);
 typedef void (*RollbackForeignTransaction_function) (FdwXactInfo *finfo);
+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
@@ -266,6 +269,7 @@ typedef struct FdwRoutine
 	CommitForeignTransaction_function CommitForeignTransaction;
 	RollbackForeignTransaction_function RollbackForeignTransaction;
 	PrepareForeignTransaction_function PrepareForeignTransaction;
+	GetPrepareId_function GetPrepareId;
 } FdwRoutine;
 
 
-- 
2.27.0

