port fdw patches

Author: kelvich

contrib/postgres_fdw/connection.c

@@ -17,6 +17,8 @@
 #include "access/htup_details.h"
 #include "catalog/pg_user_mapping.h"
 #include "access/xact.h"
+#include "access/xtm.h"
+#include "access/transam.h"
 #include "mb/pg_wchar.h"
 #include "miscadmin.h"
 #include "pgstat.h"
@@ -72,12 +74,18 @@ static unsigned int prep_stmt_number = 0;
 /* tracks whether any work is needed in callback functions */
 static bool xact_got_connection = false;
 
+typedef long long csn_t;
+static csn_t currentGlobalTransactionId = 0;
+static int	currentLocalTransactionId = 0;
+static PGconn* currentConnection = NULL;
+
 /* prototypes of private functions */
 static PGconn *connect_pg_server(ForeignServer *server, UserMapping *user);
 static void disconnect_pg_server(ConnCacheEntry *entry);
 static void check_conn_params(const char **keywords, const char **values);
 static void configure_remote_session(PGconn *conn);
 static void do_sql_command(PGconn *conn, const char *sql);
+static void do_sql_send_command(PGconn *conn, const char *sql);
 static void begin_remote_xact(ConnCacheEntry *entry);
 static void pgfdw_xact_callback(XactEvent event, void *arg);
 static void pgfdw_subxact_callback(SubXactEvent event,
@@ -403,6 +411,19 @@ do_sql_command(PGconn *conn, const char *sql)
 	PQclear(res);
 }
 
+static void
+do_sql_send_command(PGconn *conn, const char *sql)
+{
+	if (PQsendQuery(conn, sql) != PGRES_COMMAND_OK)
+	{
+		PGresult   *res = PQgetResult(conn);
+
+		elog(WARNING, "Failed to send command %s", sql);
+		pgfdw_report_error(ERROR, res, conn, true, sql);
+		PQclear(res);
+	}
+}
+
 /*
  * Start remote transaction or subtransaction, if needed.
  *
@@ -417,15 +438,26 @@ static void
 begin_remote_xact(ConnCacheEntry *entry)
 {
 	int			curlevel = GetCurrentTransactionNestLevel();
+	PGresult   *res;
+
 
 	/* Start main transaction if we haven't yet */
 	if (entry->xact_depth <= 0)
 	{
+		TransactionId gxid = GetTransactionManager()->GetGlobalTransactionId();
 		const char *sql;
 
 		elog(DEBUG3, "starting remote transaction on connection %p",
 			 entry->conn);
 
+		if (TransactionIdIsValid(gxid))
+		{
+			char	stmt[64];
+			snprintf(stmt, sizeof(stmt), "select public.dtm_join_transaction(%d)", gxid);
+			res = PQexec(entry->conn, stmt);
+			PQclear(res);
+		}
+
 		if (IsolationIsSerializable())
 			sql = "START TRANSACTION ISOLATION LEVEL SERIALIZABLE";
 		else
@@ -434,6 +466,41 @@ begin_remote_xact(ConnCacheEntry *entry)
 		do_sql_command(entry->conn, sql);
 		entry->xact_depth = 1;
 		entry->changing_xact_state = false;
+
+		if (UseTsDtmTransactions)
+		{
+			if (currentConnection == NULL)
+			{
+				currentConnection = entry->conn;
+			}
+			else if (entry->conn != currentConnection)
+			{
+				if (!currentGlobalTransactionId)
+				{
+					char	   *resp;
+					res = PQexec(currentConnection, psprintf("SELECT public.dtm_extend('%d.%d')",
+															 MyProcPid, ++currentLocalTransactionId));
+
+					if (PQresultStatus(res) != PGRES_TUPLES_OK)
+					{
+						pgfdw_report_error(ERROR, res, currentConnection, true, sql);
+					}
+					resp = PQgetvalue(res, 0, 0);
+					if (resp == NULL || (*resp) == '\0' || sscanf(resp, "%lld", &currentGlobalTransactionId) != 1)
+					{
+						pgfdw_report_error(ERROR, res, currentConnection, true, sql);
+					}
+					PQclear(res);
+				}
+				res = PQexec(entry->conn, psprintf("SELECT public.dtm_access(%llu, '%d.%d')", currentGlobalTransactionId, MyProcPid, currentLocalTransactionId));
+
+				if (PQresultStatus(res) != PGRES_TUPLES_OK)
+				{
+					pgfdw_report_error(ERROR, res, entry->conn, true, sql);
+				}
+				PQclear(res);
+			}
+		}
 	}
 
 	/*
@@ -643,6 +710,78 @@ pgfdw_report_error(int elevel, PGresult *res, PGconn *conn,
 		PQclear(res);
 }
 
+typedef bool (*DtmCommandResultHandler) (PGresult *result, void *arg);
+
+static bool
+RunDtmStatement(char const * sql, unsigned expectedStatus, DtmCommandResultHandler handler, void *arg)
+{
+	HASH_SEQ_STATUS scan;
+	ConnCacheEntry *entry;
+	bool		allOk = true;
+
+	hash_seq_init(&scan, ConnectionHash);
+	while ((entry = (ConnCacheEntry *) hash_seq_search(&scan)))
+	{
+		if (entry->xact_depth > 0)
+		{
+			do_sql_send_command(entry->conn, sql);
+		}
+	}
+
+	hash_seq_init(&scan, ConnectionHash);
+	while ((entry = (ConnCacheEntry *) hash_seq_search(&scan)))
+	{
+		if (entry->xact_depth > 0)
+		{
+			PGresult   *result = PQgetResult(entry->conn);
+
+			if (PQresultStatus(result) != expectedStatus || (handler && !handler(result, arg)))
+			{
+				elog(WARNING, "Failed command %s: status=%d, expected status=%d", sql, PQresultStatus(result), expectedStatus);
+				pgfdw_report_error(ERROR, result, entry->conn, true, sql);
+				allOk = false;
+			}
+			PQclear(result);
+			PQgetResult(entry->conn);	/* consume NULL result */
+		}
+	}
+	return allOk;
+}
+
+static bool
+RunDtmCommand(char const * sql)
+{
+	return RunDtmStatement(sql, PGRES_COMMAND_OK, NULL, NULL);
+}
+
+static bool
+RunDtmFunction(char const * sql)
+{
+	return RunDtmStatement(sql, PGRES_TUPLES_OK, NULL, NULL);
+}
+
+
+static bool
+DtmMaxCSN(PGresult *result, void *arg)
+{
+	char	   *resp = PQgetvalue(result, 0, 0);
+	csn_t	   *maxCSN = (csn_t *) arg;
+	csn_t		csn = 0;
+
+	if (resp == NULL || (*resp) == '\0' || sscanf(resp, "%lld", &csn) != 1)
+	{
+		return false;
+	}
+	else
+	{
+		if (*maxCSN < csn)
+		{
+			*maxCSN = csn;
+		}
+		return true;
+	}
+}
+
 /*
  * pgfdw_xact_callback --- cleanup at main-transaction end.
  */
@@ -652,10 +791,55 @@ pgfdw_xact_callback(XactEvent event, void *arg)
 	HASH_SEQ_STATUS scan;
 	ConnCacheEntry *entry;
 
+	/* Do nothing for this events */
+	switch (event)
+	{
+		case XACT_EVENT_START:
+		case XACT_EVENT_COMMIT_PREPARED:
+		case XACT_EVENT_ABORT_PREPARED:
+			return;
+		default:
+			break;
+	}
+
 	/* Quick exit if no connections were touched in this transaction. */
 	if (!xact_got_connection)
 		return;
 
+	if (currentGlobalTransactionId != 0)
+	{
+		switch (event)
+		{
+			case XACT_EVENT_PARALLEL_PRE_COMMIT:
+			case XACT_EVENT_PRE_COMMIT:
+				{
+					csn_t		maxCSN = 0;
+
+					if (!RunDtmCommand(psprintf("PREPARE TRANSACTION '%d.%d'",
+									MyProcPid, currentLocalTransactionId)) ||
+						!RunDtmFunction(psprintf("SELECT public.dtm_begin_prepare('%d.%d')",
+									MyProcPid, currentLocalTransactionId)) ||
+						!RunDtmStatement(psprintf("SELECT public.dtm_prepare('%d.%d',0)",
+												  MyProcPid, currentLocalTransactionId), PGRES_TUPLES_OK, DtmMaxCSN, &maxCSN) ||
+						!RunDtmFunction(psprintf("SELECT public.dtm_end_prepare('%d.%d',%lld)",
+							MyProcPid, currentLocalTransactionId, maxCSN)) ||
+						!RunDtmCommand(psprintf("COMMIT PREPARED '%d.%d'",
+									  MyProcPid, currentLocalTransactionId)))
+					{
+						RunDtmCommand(psprintf("ROLLBACK PREPARED '%d.%d'",
+									  MyProcPid, currentLocalTransactionId));
+						ereport(ERROR,
+								(errcode(ERRCODE_TRANSACTION_ROLLBACK),
+								 errmsg("transaction was aborted at one of the shards")));
+						break;
+					}
+					return;
+				}
+			default:
+				break;
+		}
+	}
+
 	/*
 	 * Scan all connection cache entries to find open remote transactions, and
 	 * close them.
@@ -689,10 +873,34 @@ pgfdw_xact_callback(XactEvent event, void *arg)
 					pgfdw_reject_incomplete_xact_state_change(entry);
 
 					/* Commit all remote transactions during pre-commit */
-					entry->changing_xact_state = true;
-					do_sql_command(entry->conn, "COMMIT TRANSACTION");
-					entry->changing_xact_state = false;
+					do_sql_send_command(entry->conn, "COMMIT TRANSACTION");
+					continue;
+
+				case XACT_EVENT_PRE_PREPARE:
+
+					/*
+					 * We disallow remote transactions that modified anything,
+					 * since it's not very reasonable to hold them open until
+					 * the prepared transaction is committed.  For the moment,
+					 * throw error unconditionally; later we might allow
+					 * read-only cases.  Note that the error will cause us to
+					 * come right back here with event == XACT_EVENT_ABORT, so
+					 * we'll clean up the connection state at that point.
+					 */
+					ereport(ERROR,
+							(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+							 errmsg("cannot prepare a transaction that modified remote tables")));
+					break;
 
+				case XACT_EVENT_PARALLEL_COMMIT:
+				case XACT_EVENT_COMMIT:
+				case XACT_EVENT_PREPARE:
+					if (!currentGlobalTransactionId)
+					{
+						entry->changing_xact_state = true;
+						do_sql_command(entry->conn, "COMMIT TRANSACTION");
+						entry->changing_xact_state = false;
+					}
 					/*
 					 * If there were any errors in subtransactions, and we
 					 * made prepared statements, do a DEALLOCATE ALL to make
@@ -716,27 +924,7 @@ pgfdw_xact_callback(XactEvent event, void *arg)
 					entry->have_prep_stmt = false;
 					entry->have_error = false;
 					break;
-				case XACT_EVENT_PRE_PREPARE:
 
-					/*
-					 * We disallow remote transactions that modified anything,
-					 * since it's not very reasonable to hold them open until
-					 * the prepared transaction is committed.  For the moment,
-					 * throw error unconditionally; later we might allow
-					 * read-only cases.  Note that the error will cause us to
-					 * come right back here with event == XACT_EVENT_ABORT, so
-					 * we'll clean up the connection state at that point.
-					 */
-					ereport(ERROR,
-							(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-							 errmsg("cannot prepare a transaction that modified remote tables")));
-					break;
-				case XACT_EVENT_PARALLEL_COMMIT:
-				case XACT_EVENT_COMMIT:
-				case XACT_EVENT_PREPARE:
-					/* Pre-commit should have closed the open transaction */
-					elog(ERROR, "missed cleaning up connection during pre-commit");
-					break;
 				case XACT_EVENT_PARALLEL_ABORT:
 				case XACT_EVENT_ABORT:
 
@@ -800,6 +988,11 @@ pgfdw_xact_callback(XactEvent event, void *arg)
 					/* Disarm changing_xact_state if it all worked. */
 					entry->changing_xact_state = abort_cleanup_failure;
 					break;
+
+				case XACT_EVENT_START:
+				case XACT_EVENT_COMMIT_PREPARED:
+				case XACT_EVENT_ABORT_PREPARED:
+					break;
 			}
 		}
 
@@ -818,16 +1011,22 @@ pgfdw_xact_callback(XactEvent event, void *arg)
 			disconnect_pg_server(entry);
 		}
 	}
+	if (event != XACT_EVENT_PARALLEL_PRE_COMMIT && event != XACT_EVENT_PRE_COMMIT)
+	{
+		/*
+		 * Regardless of the event type, we can now mark ourselves as out of
+		 * the transaction.  (Note: if we are here during PRE_COMMIT or
+		 * PRE_PREPARE, this saves a useless scan of the hashtable during
+		 * COMMIT or PREPARE.)
+		 */
+		xact_got_connection = false;
 
-	/*
-	 * Regardless of the event type, we can now mark ourselves as out of the
-	 * transaction.  (Note: if we are here during PRE_COMMIT or PRE_PREPARE,
-	 * this saves a useless scan of the hashtable during COMMIT or PREPARE.)
-	 */
-	xact_got_connection = false;
+		/* Also reset cursor numbering for next transaction */
+		cursor_number = 0;
 
-	/* Also reset cursor numbering for next transaction */
-	cursor_number = 0;
+		currentGlobalTransactionId = 0;
+		currentConnection = NULL;
+	}
 }
 
 /*

contrib/postgres_fdw/postgres_fdw--1.0.sql

@@ -16,3 +16,8 @@ LANGUAGE C STRICT;
 CREATE FOREIGN DATA WRAPPER postgres_fdw
   HANDLER postgres_fdw_handler
   VALIDATOR postgres_fdw_validator;
+
+CREATE FUNCTION postgres_fdw_exec(relid oid, sql cstring)
+RETURNS void
+AS 'MODULE_PATHNAME'
+LANGUAGE C STRICT;