From 5f0d996720bd20d5614c168cfd2e124717986965 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 v36 7/9] Add GetPrepareId API

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

diff --git a/src/backend/access/transam/fdwxact.c b/src/backend/access/transam/fdwxact.c
index bc4540d3ae..df9ff87ae3 100644
--- a/src/backend/access/transam/fdwxact.c
+++ b/src/backend/access/transam/fdwxact.c
@@ -181,6 +181,7 @@ typedef struct FdwXactEntry
 	CommitForeignTransaction_function commit_foreign_xact_fn;
 	RollbackForeignTransaction_function rollback_foreign_xact_fn;
 	PrepareForeignTransaction_function prepare_foreign_xact_fn;
+	GetPrepareId_function get_prepareid_fn;
 } FdwXactEntry;
 
 /*
@@ -385,6 +386,7 @@ FdwXactRegisterXact(UserMapping *usermapping, bool modified)
 	fdwent->commit_foreign_xact_fn = routine->CommitForeignTransaction;
 	fdwent->rollback_foreign_xact_fn = routine->RollbackForeignTransaction;
 	fdwent->prepare_foreign_xact_fn = routine->PrepareForeignTransaction;
+	fdwent->get_prepareid_fn = routine->GetPrepareId;
 
 	MemoryContextSwitchTo(old_ctx);
 
@@ -871,9 +873,10 @@ FdwXactPrepareForeignTransactions(TransactionId xid, bool prepare_all)
 }
 
 /*
- * 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
@@ -887,12 +890,47 @@ FdwXactPrepareForeignTransactions(TransactionId xid, bool prepare_all)
 static char *
 getFdwXactIdentifier(FdwXactEntry *fdwent, 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, fdwent->umid);
+	/*
+	 * If FDW doesn't provide the callback function, generate an unique
+	 * identifier.
+	 */
+	if (!fdwent->get_prepareid_fn)
+	{
+		char		buf[FDWXACT_ID_MAX_LEN] = {0};
+
+		snprintf(buf, FDWXACT_ID_MAX_LEN, "fx_%ld_%u_%u", Abs(random()),
+				 xid, fdwent->umid);
+
+		return pstrdup(buf);
+	}
+
+	/* Get an unique identifier from callback function */
+	id = fdwent->get_prepareid_fn(xid, fdwent->server->serverid,
+								  fdwent->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 55298f8dae..0f6d0543c7 100644
--- a/src/include/foreign/fdwapi.h
+++ b/src/include/foreign/fdwapi.h
@@ -195,6 +195,8 @@ typedef void (*ForeignAsyncNotify_function) (AsyncRequest *areq);
 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);
 
 
 /*
@@ -289,6 +291,7 @@ typedef struct FdwRoutine
 	CommitForeignTransaction_function CommitForeignTransaction;
 	RollbackForeignTransaction_function RollbackForeignTransaction;
 	PrepareForeignTransaction_function PrepareForeignTransaction;
+	GetPrepareId_function GetPrepareId;
 } FdwRoutine;
 
 
-- 
2.24.3 (Apple Git-128)