From b1bb4eaf19701a2c49c041bf9ef09a356506648e Mon Sep 17 00:00:00 2001
From: Masahiko Sawada <sawada.mshk@gmail.com>
Date: Mon, 21 Sep 2020 17:00:21 +0900
Subject: [PATCH v26 05/11] postgres_fdw supports prepare API and support
 commit/rollback prepared.

Co-authored-by: Masahiko Sawada, Ashutosh Bapat
---
 contrib/postgres_fdw/connection.c   | 185 ++++++++++++++++++++++------
 contrib/postgres_fdw/postgres_fdw.c |   1 +
 contrib/postgres_fdw/postgres_fdw.h |   1 +
 3 files changed, 151 insertions(+), 36 deletions(-)

diff --git a/contrib/postgres_fdw/connection.c b/contrib/postgres_fdw/connection.c
index 10a2815c64..fa8aa6d5df 100644
--- a/contrib/postgres_fdw/connection.c
+++ b/contrib/postgres_fdw/connection.c
@@ -75,7 +75,7 @@ static unsigned int prep_stmt_number = 0;
 static bool xact_got_connection = false;
 
 /* prototypes of private functions */
-static PGconn *connect_pg_server(ForeignServer *server, UserMapping *user);
+static void connect_pg_server(ConnCacheEntry *entry, UserMapping *user);
 static void disconnect_pg_server(ConnCacheEntry *entry);
 static void check_conn_params(const char **keywords, const char **values, UserMapping *user);
 static void configure_remote_session(PGconn *conn);
@@ -95,6 +95,8 @@ static bool pgfdw_get_cleanup_result(PGconn *conn, TimestampTz endtime,
 static bool UserMappingPasswordRequired(UserMapping *user);
 static void pgfdw_cleanup_after_transaction(ConnCacheEntry *entry);
 static ConnCacheEntry *GetConnectionCacheEntry(Oid umid);
+static void pgfdw_end_prepared_xact(ConnCacheEntry *entry, char *fdwxact_id,
+									bool is_commit);
 
 /*
  * Get a PGconn which can be used to execute queries on the remote PostgreSQL
@@ -132,35 +134,8 @@ GetConnection(UserMapping *user, bool will_prep_stmt)
 	 * require some overhead.  Broken connection will be detected when the
 	 * connection is actually used.
 	 */
-
-	/*
-	 * If cache entry doesn't have a connection, we have to establish a new
-	 * connection.  (If connect_pg_server throws an error, the cache entry
-	 * will remain in a valid empty state, ie conn == NULL.)
-	 */
 	if (entry->conn == NULL)
-	{
-		ForeignServer *server = GetForeignServer(user->serverid);
-
-		/* Reset all transient state fields, to be sure all are clean */
-		entry->xact_depth = 0;
-		entry->have_prep_stmt = false;
-		entry->have_error = false;
-		entry->changing_xact_state = false;
-		entry->invalidated = false;
-		entry->server_hashvalue =
-			GetSysCacheHashValue1(FOREIGNSERVEROID,
-								  ObjectIdGetDatum(server->serverid));
-		entry->mapping_hashvalue =
-			GetSysCacheHashValue1(USERMAPPINGOID,
-								  ObjectIdGetDatum(user->umid));
-
-		/* Now try to make the connection */
-		entry->conn = connect_pg_server(server, user);
-
-		elog(DEBUG3, "new postgres_fdw connection %p for server \"%s\" (user mapping oid %u, userid %u)",
-			 entry->conn, server->servername, user->umid, user->userid);
-	}
+		connect_pg_server(entry, user);
 
 	/*
 	 * Start a new transaction or subtransaction if needed.
@@ -229,11 +204,29 @@ GetConnectionCacheEntry(Oid umid)
 
 /*
  * Connect to remote server using specified server and user mapping properties.
+
+ * If cache entry doesn't have a connection, we have to establish a new
+ * connection.  (If connect_pg_server throws an error, the cache entry
+ * will remain in a valid empty state, ie conn == NULL.)
  */
-static PGconn *
-connect_pg_server(ForeignServer *server, UserMapping *user)
+static void
+connect_pg_server(ConnCacheEntry *entry, UserMapping *user)
 {
 	PGconn	   *volatile conn = NULL;
+	ForeignServer *server = GetForeignServer(user->serverid);
+
+	/* Reset all transient state fields, to be sure all are clean */
+	entry->xact_depth = 0;
+	entry->have_prep_stmt = false;
+	entry->have_error = false;
+	entry->changing_xact_state = false;
+	entry->invalidated = false;
+	entry->server_hashvalue =
+		GetSysCacheHashValue1(FOREIGNSERVEROID,
+							  ObjectIdGetDatum(server->serverid));
+	entry->mapping_hashvalue =
+		GetSysCacheHashValue1(USERMAPPINGOID,
+							  ObjectIdGetDatum(user->umid));
 
 	/*
 	 * Use PG_TRY block to ensure closing connection on error.
@@ -344,7 +337,9 @@ connect_pg_server(ForeignServer *server, UserMapping *user)
 	}
 	PG_END_TRY();
 
-	return conn;
+	entry->conn = conn;
+	elog(DEBUG3, "new postgres_fdw connection %p for server \"%s\" (user mapping oid %u, userid %u)",
+		 entry->conn, server->servername, user->umid, user->userid);
 }
 
 /*
@@ -1084,12 +1079,26 @@ void
 postgresCommitForeignTransaction(FdwXactRslvState *frstate)
 {
 	ConnCacheEntry *entry;
+	bool			is_onephase = (frstate->flags & FDWXACT_FLAG_ONEPHASE) != 0;
 	PGresult   *res;
 
-	Assert((frstate->flags & FDWXACT_FLAG_ONEPHASE) != 0);
-
 	entry = GetConnectionCacheEntry(frstate->usermapping->umid);
 
+	if (!is_onephase)
+	{
+		/*
+		 * In two-phase commit case, the foreign transaction has been prepared
+		 * and closed, so we might not have a connection to it. We get a
+		 * connection if necessary.
+		 */
+		if (!entry->conn)
+			connect_pg_server(entry, frstate->usermapping);
+
+		/* COMMIT PREPARED the transaction and cleanup */
+		pgfdw_end_prepared_xact(entry, frstate->fdwxact_id, true);
+		return;
+	}
+
 	Assert(entry->conn);
 
 	/*
@@ -1135,16 +1144,31 @@ void
 postgresRollbackForeignTransaction(FdwXactRslvState *frstate)
 {
 	ConnCacheEntry *entry = NULL;
+	bool			is_onephase = (frstate->flags & FDWXACT_FLAG_ONEPHASE) != 0;
 	bool abort_cleanup_failure = false;
 
-	Assert((frstate->flags & FDWXACT_FLAG_ONEPHASE) != 0);
-
 	/*
 	 * In simple rollback case, we must have a connection to the foreign server
 	 * because the foreign transaction is not closed yet. We get the connection
 	 * entry from the cache.
 	 */
 	entry = GetConnectionCacheEntry(frstate->usermapping->umid);
+
+	if (!is_onephase)
+	{
+		/*
+		 * In two-phase commit case, the foreign transaction has been prepared
+		 * and closed, so we might not have a connection to it. We get a
+		 * connection if necessary.
+		 */
+		if (!entry->conn)
+			connect_pg_server(entry, frstate->usermapping);
+
+		/* ROLLBACK PREPARED the transaction and cleanup */
+		pgfdw_end_prepared_xact(entry, frstate->fdwxact_id, false);
+		return;
+	}
+
 	Assert(entry);
 
 	/*
@@ -1221,6 +1245,46 @@ postgresRollbackForeignTransaction(FdwXactRslvState *frstate)
 	return;
 }
 
+/*
+ * Prepare a transaction on foreign server.
+ */
+void
+postgresPrepareForeignTransaction(FdwXactRslvState *frstate)
+{
+	ConnCacheEntry *entry = NULL;
+	PGresult	*res;
+	StringInfo	command;
+
+	/* The transaction should have started already get the cache entry */
+	entry = GetConnectionCacheEntry(frstate->usermapping->umid);
+	Assert(entry->conn);
+
+	pgfdw_reject_incomplete_xact_state_change(entry);
+
+	command = makeStringInfo();
+	appendStringInfo(command, "PREPARE TRANSACTION '%s'", frstate->fdwxact_id);
+
+	/* Do prepare foreign transaction */
+	entry->changing_xact_state = true;
+	res = pgfdw_exec_query(entry->conn, command->data);
+	entry->changing_xact_state = false;
+
+	if (PQresultStatus(res) != PGRES_COMMAND_OK)
+		ereport(ERROR, (errmsg("could not prepare transaction on server %s with ID %s",
+							   frstate->server->servername, frstate->fdwxact_id)));
+
+	elog(DEBUG1, "prepared foreign transaction on server %s with ID %s",
+		 frstate->server->servername, frstate->fdwxact_id);
+
+	if (entry->have_prep_stmt && entry->have_error)
+	{
+		res = PQexec(entry->conn, "DEALLOCATE ALL");
+		PQclear(res);
+	}
+
+	pgfdw_cleanup_after_transaction(entry);
+}
+
 /* Cleanup at main-transaction end */
 static void
 pgfdw_cleanup_after_transaction(ConnCacheEntry *entry)
@@ -1247,3 +1311,52 @@ pgfdw_cleanup_after_transaction(ConnCacheEntry *entry)
 	/* Also reset cursor numbering for next transaction */
 	cursor_number = 0;
 }
+
+/*
+ * Commit or rollback prepared transaction on the foreign server.
+ */
+static void
+pgfdw_end_prepared_xact(ConnCacheEntry *entry, char *fdwxact_id, bool is_commit)
+{
+	StringInfo	command;
+	PGresult	*res;
+
+	command = makeStringInfo();
+	appendStringInfo(command, "%s PREPARED '%s'",
+					 is_commit ? "COMMIT" : "ROLLBACK",
+					 fdwxact_id);
+
+	res = pgfdw_exec_query(entry->conn, command->data);
+
+	if (PQresultStatus(res) != PGRES_COMMAND_OK)
+	{
+		int		sqlstate;
+		char	*diag_sqlstate = PQresultErrorField(res, PG_DIAG_SQLSTATE);
+
+		if (diag_sqlstate)
+		{
+			sqlstate = MAKE_SQLSTATE(diag_sqlstate[0],
+									 diag_sqlstate[1],
+									 diag_sqlstate[2],
+									 diag_sqlstate[3],
+									 diag_sqlstate[4]);
+		}
+		else
+			sqlstate = ERRCODE_CONNECTION_FAILURE;
+
+		/*
+		 * As core global transaction manager states, it's possible that the
+		 * given foreign transaction doesn't exist on the foreign server. So
+		 * we should accept an UNDEFINED_OBJECT error.
+		 */
+		if (sqlstate != ERRCODE_UNDEFINED_OBJECT)
+			pgfdw_report_error(ERROR, res, entry->conn, false, command->data);
+	}
+
+	elog(DEBUG1, "%s prepared foreign transaction with ID %s",
+		 is_commit ? "commit" : "rollback",
+		 fdwxact_id);
+
+	/* Cleanup transaction status */
+	pgfdw_cleanup_after_transaction(entry);
+}
diff --git a/contrib/postgres_fdw/postgres_fdw.c b/contrib/postgres_fdw/postgres_fdw.c
index 8ed6577d35..caccad9846 100644
--- a/contrib/postgres_fdw/postgres_fdw.c
+++ b/contrib/postgres_fdw/postgres_fdw.c
@@ -561,6 +561,7 @@ postgres_fdw_handler(PG_FUNCTION_ARGS)
 	/* Support functions for foreign transactions */
 	routine->CommitForeignTransaction = postgresCommitForeignTransaction;
 	routine->RollbackForeignTransaction = postgresRollbackForeignTransaction;
+	routine->PrepareForeignTransaction = postgresPrepareForeignTransaction;
 
 	PG_RETURN_POINTER(routine);
 }
diff --git a/contrib/postgres_fdw/postgres_fdw.h b/contrib/postgres_fdw/postgres_fdw.h
index e3b2897495..659222b97a 100644
--- a/contrib/postgres_fdw/postgres_fdw.h
+++ b/contrib/postgres_fdw/postgres_fdw.h
@@ -140,6 +140,7 @@ extern void pgfdw_report_error(int elevel, PGresult *res, PGconn *conn,
 							   bool clear, const char *sql);
 extern void postgresCommitForeignTransaction(FdwXactRslvState *frstate);
 extern void postgresRollbackForeignTransaction(FdwXactRslvState *frstate);
+extern void postgresPrepareForeignTransaction(FdwXactRslvState *frstate);
 
 /* in option.c */
 extern int	ExtractConnectionOptions(List *defelems,
-- 
2.23.0

