Use prepared statements for DirectModify.

We keep per-connection hashtable sql->prep statement name for checking whether
we have prepared it already or not. This is quick-and-dirty solution, because it
will lead to hashtable bloat if custom plans are used -- in this case we prepare
queries with substituted param values.

Author: arssher

contrib/postgres_fdw/connection.c

@@ -17,6 +17,7 @@
 #include "access/htup_details.h"
 #include "catalog/pg_user_mapping.h"
 #include "access/global_snapshot.h"
+#include "access/hash.h"
 #include "access/xact.h"
 #include "access/xtm.h"
 #include "access/transam.h"
@@ -33,41 +34,6 @@
 #include "utils/memutils.h"
 #include "utils/syscache.h"
 
-/*
- * Connection cache hash table entry
- *
- * The lookup key in this hash table is the user mapping OID. We use just one
- * connection per user mapping ID, which ensures that all the scans use the
- * same snapshot during a query.  Using the user mapping OID rather than
- * the foreign server OID + user OID avoids creating multiple connections when
- * the public user mapping applies to all user OIDs.
- *
- * The "conn" pointer can be NULL if we don't currently have a live connection.
- * When we do have a connection, xact_depth tracks the current depth of
- * transactions and subtransactions open on the remote side.  We need to issue
- * commands at the same nesting depth on the remote as we're executing at
- * ourselves, so that rolling back a subtransaction will kill the right
- * queries and not the wrong ones.
- */
-typedef Oid ConnCacheKey;
-
-struct ConnCacheEntry
-{
-	ConnCacheKey key;			/* hash key (must be first) */
-	PGconn	   *conn;			/* connection to foreign server, or NULL */
-	WaitEventSet *wait_set;		/* for data from server ready notifications */
-	/* Remaining fields are invalid when conn is NULL: */
-	int			xact_depth;		/* 0 = no xact open, 1 = main xact open, 2 =
-								 * one level of subxact open, etc */
-	bool		have_prep_stmt; /* have we prepared any stmts in this xact? */
-	bool		have_error;		/* have any subxacts aborted in this xact? */
-	bool		changing_xact_state;	/* xact state change in process */
-	bool		invalidated;	/* true if reconnect is pending */
-	uint32		server_hashvalue;	/* hash value of foreign server OID */
-	uint32		mapping_hashvalue;	/* hash value of user mapping OID */
-	bool		copy_from_started;	/* COPY FROM in progress on this conn */
-};
-
 /*
  * Connection cache (initialized on first use)
  */
@@ -117,6 +83,35 @@ static bool pgfdw_exec_cleanup_query(ConnCacheEntry *entry, const char *query,
 						 bool ignore_errors);
 static bool pgfdw_get_cleanup_result(ConnCacheEntry *entry, TimestampTz endtime,
 						 PGresult **result);
+static void cleanup_dm_prepared(ConnCacheEntry *entry);
+
+/* Adapted from string_hash */
+static uint32
+char_ptr_hash_fn(const void *key, Size keysize)
+{
+	char * const *keyptr = key;
+	return DatumGetUInt32(hash_any((const unsigned char *) (*keyptr), strlen(*keyptr)));
+}
+
+static int
+char_ptr_match_fn(const void *key1, const void *key2, Size keysize)
+{
+	char * const *keyptr1 = key1;
+	char * const *keyptr2 = key2;
+	return strcmp(*keyptr1, *keyptr2);
+}
+
+/* Allocate always from top-level, where hashtable lives */
+static void *
+char_ptr_keycopy_fn(void *dest, const void *src, Size keysize)
+{
+	char **destptr = dest;
+	char * const *srcptr = src;
+
+	*destptr = MemoryContextStrdup(CacheMemoryContext, *srcptr);
+	return NULL; /* not used */
+}
+
 /*
  * Get a ConnCacheEntry which can be used to execute queries on the remote PostgreSQL
  * server with the user's authorization.  A new connection is established
@@ -208,6 +203,7 @@ GetConnectionCopyFrom(UserMapping *user, bool will_prep_stmt,
 	if (entry->conn == NULL)
 	{
 		ForeignServer *server = GetForeignServer(user->serverid);
+		HASHCTL		ctl;
 
 		/* Reset all transient state fields, to be sure all are clean */
 		entry->xact_depth = 0;
@@ -226,6 +222,19 @@ GetConnectionCopyFrom(UserMapping *user, bool will_prep_stmt,
 		/* Now try to make the connection */
 		connect_pg_server(entry, server, user);
 
+		/* Create hash table of prepared statemetns for DirectModify */
+		ctl.keysize = sizeof(char *);
+		ctl.entrysize = sizeof(DirectModifyPrepStmtHashEnt);
+		ctl.hash = char_ptr_hash_fn;
+		ctl.match = char_ptr_match_fn;
+		ctl.keycopy = char_ptr_keycopy_fn;
+		ctl.hcxt = CacheMemoryContext;
+
+		entry->dm_prepared = hash_create("DirectModify prepared stmts",
+										 16, &ctl,
+										 HASH_ELEM | HASH_FUNCTION | HASH_COMPARE |
+										 HASH_KEYCOPY | HASH_CONTEXT);
+
 		elog(DEBUG3, "new postgres_fdw connection %p for server \"%s\" (user mapping oid %u, userid %u)",
 			 entry->conn, server->servername, user->umid, user->userid);
 	}
@@ -364,6 +373,8 @@ disconnect_pg_server(ConnCacheEntry *entry)
 {
 	if (entry->conn != NULL)
 	{
+		cleanup_dm_prepared(entry);
+		hash_destroy(entry->dm_prepared);
 		Assert(entry->wait_set);
 		FreeWaitEventSet(entry->wait_set);
 		entry->wait_set = NULL;
@@ -1021,6 +1032,9 @@ pgfdw_xact_callback(XactEvent event, void *arg)
 						entry->have_error = false;
 					}
 
+					/* We have deallocated all prepared statements */
+					cleanup_dm_prepared(entry);
+
 					/* Disarm changing_xact_state if it all worked. */
 					entry->changing_xact_state = abort_cleanup_failure;
 					break;
@@ -1118,11 +1132,27 @@ deallocate_prepared_stmts(ConnCacheEntry *entry)
 	{
 		res = PQexec(entry->conn, "DEALLOCATE ALL");
 		PQclear(res);
+		cleanup_dm_prepared(entry);
 	}
 	entry->have_prep_stmt = false;
 	entry->have_error = false;
 }
 
+static void cleanup_dm_prepared(ConnCacheEntry *entry)
+{
+	HASH_SEQ_STATUS scan;
+	DirectModifyPrepStmtHashEnt *prep_stmt_entry;
+
+	hash_seq_init(&scan, entry->dm_prepared);
+	while ((prep_stmt_entry = (DirectModifyPrepStmtHashEnt *) hash_seq_search(&scan)))
+	{
+		/* save the key to free it */
+		char *sql = prep_stmt_entry->sql;
+		hash_search(entry->dm_prepared, &prep_stmt_entry->sql, HASH_REMOVE, NULL);
+		pfree(sql);
+	}
+}
+
 /*
  * pgfdw_subxact_callback --- cleanup at subtransaction end.
  */

contrib/postgres_fdw/postgres_fdw.c

@@ -2337,7 +2337,7 @@ postgresBeginDirectModify(ForeignScanState *node, int eflags)
 	 * Get connection to the foreign server.  Connection manager will
 	 * establish new connection if necessary.
 	 */
-	dmstate->conn_entry = GetConnection(user, false);
+	dmstate->conn_entry = GetConnection(user, true);
 
 	/* Initialize state variable */
 	dmstate->num_tuples = -1;	/* -1 means not set yet */
@@ -3323,6 +3323,8 @@ execute_dml_stmt(ForeignScanState *node)
 	ExprContext *econtext = node->ss.ps.ps_ExprContext;
 	int			numParams = dmstate->numParams;
 	const char **values = dmstate->param_values;
+	DirectModifyPrepStmtHashEnt *prep_stmt_entry;
+	bool found;
 
 	/*
 	 * Construct array of query parameter values in text format.
@@ -3334,22 +3336,59 @@ execute_dml_stmt(ForeignScanState *node)
 							 dmstate->param_exprs,
 							 values);
 
+	/*
+	 * Prepare the statement, if we have never seen it before.
+	 */
+	prep_stmt_entry = hash_search(entry->dm_prepared, &dmstate->query,
+								  HASH_FIND, &found);
+	if (!found)
+	{
+		char p_name[NAMEDATALEN];
+		long num = hash_get_num_entries(entry->dm_prepared) + 1;
+		PGresult   *res;
+
+		snprintf(p_name, NAMEDATALEN, "postgres_fdw:%d:%ld", MyProcPid, num);
+
+		if (!PQsendPrepare(conn,
+						   p_name,
+						   dmstate->query,
+						   0,
+						   NULL))
+			pgfdw_report_error(ERROR, NULL, conn, false, dmstate->query);
+
+		/*
+		 * Get the result, and check for success.
+		 *
+		 * We don't use a PG_TRY block here, so be careful not to throw error
+		 * without releasing the PGresult.
+		 */
+		res = pgfdw_get_result(entry, dmstate->query);
+		if (PQresultStatus(res) != PGRES_COMMAND_OK)
+			pgfdw_report_error(ERROR, res, conn, true, dmstate->query);
+		PQclear(res);
+
+		/* Now, when it is successfully prepared, add it to the hashtable */
+		prep_stmt_entry = hash_search(entry->dm_prepared, &dmstate->query,
+									  HASH_ENTER, &found);
+		strlcpy(prep_stmt_entry->prep_name, p_name, NAMEDATALEN);
+	}
+
 	/*
 	 * Notice that we pass NULL for paramTypes, thus forcing the remote server
 	 * to infer types for all parameters.  Since we explicitly cast every
 	 * parameter (see deparse.c), the "inference" is trivial and will produce
 	 * the desired result.  This allows us to avoid assuming that the remote
 	 * server has the same OIDs we do for the parameters' types.
 	 */
-	if (!PQsendQueryParams(conn, dmstate->query, numParams,
-						   NULL, values, NULL, NULL, 0))
+
+	if (!PQsendQueryPrepared(conn,
+							 prep_stmt_entry->prep_name,
+							 numParams,
+							 values,
+							 NULL,
+							 NULL,
+							 0))
 		pgfdw_report_error(ERROR, NULL, conn, false, dmstate->query);
-	// }
-	// else
-	// {
-	// 	if (!PQsendQuery(dmstate->conn, dmstate->query))
-	// 		pgfdw_report_error(ERROR, NULL, dmstate->conn, false, dmstate->query);
-	// }
 
 	/*
 	 * Get the result, and check for success.

contrib/postgres_fdw/postgres_fdw.h

@@ -17,16 +17,11 @@
 #include "commands/copy.h"
 #include "lib/stringinfo.h"
 #include "nodes/relation.h"
+#include "storage/latch.h"
 #include "utils/relcache.h"
 
 #include "libpq-fe.h"
 
-/*
- * Encapsulates connection to foreign server. Contents should be unknown
- * outside connection.c
- */
-typedef struct ConnCacheEntry ConnCacheEntry;
-
 /*
  * FDW-specific planner information kept in RelOptInfo.fdw_private for a
  * postgres_fdw foreign table.  For a baserel, this struct is created by
@@ -117,6 +112,50 @@ typedef struct PgFdwRelationInfo
 	int			relation_index;
 } PgFdwRelationInfo;
 
+/*
+ * Connection cache hash table entry
+ *
+ * The lookup key in this hash table is the user mapping OID. We use just one
+ * connection per user mapping ID, which ensures that all the scans use the
+ * same snapshot during a query.  Using the user mapping OID rather than
+ * the foreign server OID + user OID avoids creating multiple connections when
+ * the public user mapping applies to all user OIDs.
+ *
+ * The "conn" pointer can be NULL if we don't currently have a live connection.
+ * When we do have a connection, xact_depth tracks the current depth of
+ * transactions and subtransactions open on the remote side.  We need to issue
+ * commands at the same nesting depth on the remote as we're executing at
+ * ourselves, so that rolling back a subtransaction will kill the right
+ * queries and not the wrong ones.
+ */
+typedef Oid ConnCacheKey;
+
+typedef struct ConnCacheEntry ConnCacheEntry;
+struct ConnCacheEntry
+{
+	ConnCacheKey key;			/* hash key (must be first) */
+	PGconn	   *conn;			/* connection to foreign server, or NULL */
+	WaitEventSet *wait_set;		/* for data from server ready notifications */
+	/* Remaining fields are invalid when conn is NULL: */
+	int			xact_depth;		/* 0 = no xact open, 1 = main xact open, 2 =
+								 * one level of subxact open, etc */
+	bool		have_prep_stmt; /* have we prepared any stmts in this xact? */
+	bool		have_error;		/* have any subxacts aborted in this xact? */
+	bool		changing_xact_state;	/* xact state change in process */
+	bool		invalidated;	/* true if reconnect is pending */
+	uint32		server_hashvalue;	/* hash value of foreign server OID */
+	uint32		mapping_hashvalue;	/* hash value of user mapping OID */
+	bool		copy_from_started;	/* COPY FROM in progress on this conn */
+	HTAB		*dm_prepared; /* prepared statements for DirectModify */
+};
+
+/* sql -> prepared statement hashtable */
+typedef struct DirectModifyPrepStmtHashEnt
+{
+	char  *sql; /* SQL of the statement, the key; arbitrary size */
+	char  prep_name[NAMEDATALEN]; /* name of prepared statement */
+} DirectModifyPrepStmtHashEnt;
+
 /* in postgres_fdw.c */
 extern int	set_transmission_modes(void);
 extern void reset_transmission_modes(int nestlevel);