From 37aaf1ba113461999c3cb8e142bc6d4590740c4f Mon Sep 17 00:00:00 2001
From: Jelte Fennema <jelte.fennema@microsoft.com>
Date: Fri, 16 Jun 2023 17:15:48 +0200
Subject: [PATCH v2] Support sending Close messages from libpq

This part of the protocol did not have a libpq implementation yet. That
is a shame because connection poolers can much more easily intercept
these Close messages than a DEALLOCATE query. Odyssey has prepared
statement support implemented using the Close message, and PgBouncer is
currently trying to do the same. But libpq based clients are not able
to use this feature of these connection poolers because they cannot send
the Close protocol message.
---
 doc/src/sgml/libpq.sgml                       | 120 +++++++++++++++-
 src/interfaces/libpq/exports.txt              |   4 +
 src/interfaces/libpq/fe-exec.c                | 133 ++++++++++++++++++
 src/interfaces/libpq/fe-protocol3.c           |  19 ++-
 src/interfaces/libpq/libpq-fe.h               |   6 +
 src/interfaces/libpq/libpq-int.h              |   2 +-
 .../modules/libpq_pipeline/libpq_pipeline.c   |  18 +++
 .../libpq_pipeline/traces/prepared.trace      |  16 +++
 8 files changed, 311 insertions(+), 7 deletions(-)

diff --git a/doc/src/sgml/libpq.sgml b/doc/src/sgml/libpq.sgml
index 2225e4e0ef3..293af6ae4f3 100644
--- a/doc/src/sgml/libpq.sgml
+++ b/doc/src/sgml/libpq.sgml
@@ -3360,6 +3360,66 @@ PGresult *PQdescribePortal(PGconn *conn, const char *portalName);
        </para>
       </listitem>
      </varlistentry>
+
+     <varlistentry id="libpq-PQclosePrepared">
+      <term><function>PQclosePrepared</function><indexterm><primary>PQclosePrepared</primary></indexterm></term>
+
+      <listitem>
+       <para>
+        Submits a request to obtain information about the specified
+        prepared statement, and waits for completion.
+<synopsis>
+PGresult *PQclosePrepared(PGconn *conn, const char *stmtName);
+</synopsis>
+       </para>
+
+       <para>
+        <xref linkend="libpq-PQclosePrepared"/> allows an application to close
+        a previously prepared statement.
+       </para>
+
+       <para>
+        <parameter>stmtName</parameter> can be <literal>""</literal> or <symbol>NULL</symbol> to reference
+        the unnamed statement, it is fine if no statement exists with this name.
+        On success, a <structname>PGresult</structname> with
+        status <literal>PGRES_COMMAND_OK</literal> is returned.
+       </para>
+      </listitem>
+     </varlistentry>
+
+     <varlistentry id="libpq-PQclosePortal">
+      <term><function>PQclosePortal</function><indexterm><primary>PQclosePortal</primary></indexterm></term>
+
+      <listitem>
+       <para>
+        Submits a request to close the the specified
+        portal, and waits for completion.
+<synopsis>
+PGresult *PQclosePortal(PGconn *conn, const char *portalName);
+</synopsis>
+       </para>
+
+       <para>
+        <xref linkend="libpq-PQclosePortal"/> allows an application to release
+        resources related to a portal previously created portal. If it was a
+        named portal the name can be reused after the close has completed.
+        (<application>libpq</application> does not provide any direct access to
+        portals, but you can use this function to inspect the properties
+        of a cursor created with a <command>DECLARE CURSOR</command> SQL command.)
+       </para>
+
+       <para>
+        <parameter>portalName</parameter> can be <literal>""</literal> or <symbol>NULL</symbol> to reference
+        the unnamed portal, it is fine if no portal exists with this name.
+        portal.  On success, a <structname>PGresult</structname> with status
+        <literal>PGRES_COMMAND_OK</literal> is returned.  The functions
+        <xref linkend="libpq-PQnfields"/>, <xref linkend="libpq-PQfname"/>,
+        <xref linkend="libpq-PQftype"/>, etc. can be applied to the
+        <structname>PGresult</structname> to obtain information about the result
+        columns (if any) of the portal.
+       </para>
+      </listitem>
+     </varlistentry>
     </variablelist>
    </para>
 
@@ -4851,15 +4911,19 @@ unsigned char *PQunescapeBytea(const unsigned char *from, size_t *to_length);
    <xref linkend="libpq-PQsendQueryParams"/>,
    <xref linkend="libpq-PQsendPrepare"/>,
    <xref linkend="libpq-PQsendQueryPrepared"/>,
-   <xref linkend="libpq-PQsendDescribePrepared"/>, and
+   <xref linkend="libpq-PQsendDescribePrepared"/>,
    <xref linkend="libpq-PQsendDescribePortal"/>,
+   <xref linkend="libpq-PQsendClosePrepared"/>, and
+   <xref linkend="libpq-PQsendClosePortal"/>,
    which can be used with <xref linkend="libpq-PQgetResult"/> to duplicate
    the functionality of
    <xref linkend="libpq-PQexecParams"/>,
    <xref linkend="libpq-PQprepare"/>,
    <xref linkend="libpq-PQexecPrepared"/>,
-   <xref linkend="libpq-PQdescribePrepared"/>, and
+   <xref linkend="libpq-PQdescribePrepared"/>,
    <xref linkend="libpq-PQdescribePortal"/>
+   <xref linkend="libpq-PQclosePrepared"/>, and
+   <xref linkend="libpq-PQclosePortal"/>
    respectively.
 
    <variablelist>
@@ -5008,6 +5072,46 @@ int PQsendDescribePortal(PGconn *conn, const char *portalName);
      </listitem>
     </varlistentry>
 
+    <varlistentry id="libpq-PQsendClosePrepared">
+     <term><function>PQsendClosePrepared</function><indexterm><primary>PQsendClosePrepared</primary></indexterm></term>
+
+     <listitem>
+      <para>
+       Submits a request to close the specified prepared statement, without
+       waiting for completion.
+<synopsis>
+int PQsendClosePrepared(PGconn *conn, const char *stmtName);
+</synopsis>
+
+       This is an asynchronous version of <xref linkend="libpq-PQclosePrepared"/>:
+       it returns 1 if it was able to dispatch the request, and 0 if not.
+       After a successful call, call <xref linkend="libpq-PQgetResult"/> to
+       obtain the results.  The function's parameters are handled
+       identically to <xref linkend="libpq-PQclosePrepared"/>.
+      </para>
+     </listitem>
+    </varlistentry>
+
+    <varlistentry id="libpq-PQsendClosePortal">
+     <term><function>PQsendClosePortal</function><indexterm><primary>PQsendClosePortal</primary></indexterm></term>
+
+     <listitem>
+      <para>
+       Submits a request to close specified portal, without waiting for
+       completion.
+<synopsis>
+int PQsendClosePortal(PGconn *conn, const char *portalName);
+</synopsis>
+
+       This is an asynchronous version of <xref linkend="libpq-PQclosePortal"/>:
+       it returns 1 if it was able to dispatch the request, and 0 if not.
+       After a successful call, call <xref linkend="libpq-PQgetResult"/> to
+       obtain the results.  The function's parameters are handled
+       identically to <xref linkend="libpq-PQclosePortal"/>.
+      </para>
+     </listitem>
+    </varlistentry>
+
     <varlistentry id="libpq-PQgetResult">
      <term><function>PQgetResult</function><indexterm><primary>PQgetResult</primary></indexterm></term>
 
@@ -5019,7 +5123,9 @@ int PQsendDescribePortal(PGconn *conn, const char *portalName);
        <xref linkend="libpq-PQsendPrepare"/>,
        <xref linkend="libpq-PQsendQueryPrepared"/>,
        <xref linkend="libpq-PQsendDescribePrepared"/>,
-       <xref linkend="libpq-PQsendDescribePortal"/>, or
+       <xref linkend="libpq-PQsendDescribePortal"/>,
+       <xref linkend="libpq-PQsendClosePrepared"/>,
+       <xref linkend="libpq-PQsendClosePortal"/>, or
        <xref linkend="libpq-PQpipelineSync"/>
        call, and returns it.
        A null pointer is returned when the command is complete and there
@@ -5350,6 +5456,8 @@ int PQflush(PGconn *conn);
     <function>PQexecPrepared</function>,
     <function>PQdescribePrepared</function>,
     <function>PQdescribePortal</function>,
+    <function>PQclosePrepared</function>,
+    <function>PQclosePortal</function>,
     is an error condition.
     <function>PQsendQuery</function> is
     also disallowed, because it uses the simple query protocol.
@@ -5389,8 +5497,10 @@ int PQflush(PGconn *conn);
      establish a synchronization point in the pipeline,
      or when <xref linkend="libpq-PQflush"/> is called.
      The functions <xref linkend="libpq-PQsendPrepare"/>,
-     <xref linkend="libpq-PQsendDescribePrepared"/>, and
-     <xref linkend="libpq-PQsendDescribePortal"/> also work in pipeline mode.
+     <xref linkend="libpq-PQsendDescribePrepared"/>,
+     <xref linkend="libpq-PQsendDescribePortal"/>,
+     <xref linkend="libpq-PQsendClosePrepared"/>, and
+     <xref linkend="libpq-PQsendClosePortal"/> also work in pipeline mode.
      Result processing is described below.
     </para>
 
diff --git a/src/interfaces/libpq/exports.txt b/src/interfaces/libpq/exports.txt
index 7ded77aff37..850734ac96c 100644
--- a/src/interfaces/libpq/exports.txt
+++ b/src/interfaces/libpq/exports.txt
@@ -187,3 +187,7 @@ PQsetTraceFlags           184
 PQmblenBounded            185
 PQsendFlushRequest        186
 PQconnectionUsedGSSAPI    187
+PQclosePrepared           188
+PQclosePortal             189
+PQsendClosePrepared       190
+PQsendClosePortal         191
diff --git a/src/interfaces/libpq/fe-exec.c b/src/interfaces/libpq/fe-exec.c
index 14d706efd57..ce70d32185e 100644
--- a/src/interfaces/libpq/fe-exec.c
+++ b/src/interfaces/libpq/fe-exec.c
@@ -79,6 +79,8 @@ static bool PQexecStart(PGconn *conn);
 static PGresult *PQexecFinish(PGconn *conn);
 static int	PQsendDescribe(PGconn *conn, char desc_type,
 						   const char *desc_target);
+static int	PQsendClose(PGconn *conn, char desc_type,
+						const char *desc_target);
 static int	check_field_number(const PGresult *res, int field_num);
 static void pqPipelineProcessQueue(PGconn *conn);
 static int	pqPipelineFlush(PGconn *conn);
@@ -2534,6 +2536,137 @@ sendFailed:
 	return 0;
 }
 
+/*
+ * PQclosePrepared
+ *	  Obtain information about a previously prepared statement
+ *
+ * If the query was not even sent, return NULL; conn->errorMessage is set to
+ * a relevant message.
+ * If the query was sent, a new PGresult is returned (which could indicate
+ * either success or failure).  On success, the PGresult contains status
+ * PGRES_COMMAND_OK, and its parameter and column-heading fields close
+ * the statement's inputs and outputs respectively.
+ * The user is responsible for freeing the PGresult via PQclear()
+ * when done with it.
+ */
+PGresult *
+PQclosePrepared(PGconn *conn, const char *stmt)
+{
+	if (!PQexecStart(conn))
+		return NULL;
+	if (!PQsendClose(conn, 'S', stmt))
+		return NULL;
+	return PQexecFinish(conn);
+}
+
+/*
+ * PQclosePortal
+ *	  Obtain information about a previously created portal
+ *
+ * This is much like PQclosePrepared, except that no parameter info is
+ * returned.  Note that at the moment, libpq doesn't really expose portals
+ * to the client; but this can be used with a portal created by a SQL
+ * DECLARE CURSOR command.
+ */
+PGresult *
+PQclosePortal(PGconn *conn, const char *portal)
+{
+	if (!PQexecStart(conn))
+		return NULL;
+	if (!PQsendClose(conn, 'P', portal))
+		return NULL;
+	return PQexecFinish(conn);
+}
+
+/*
+ * PQsendClosePrepared
+ *	 Submit a Close Statement command, but don't wait for it to finish
+ *
+ * Returns: 1 if successfully submitted
+ *			0 if error (conn->errorMessage is set)
+ */
+int
+PQsendClosePrepared(PGconn *conn, const char *stmt)
+{
+	return PQsendClose(conn, 'S', stmt);
+}
+
+/*
+ * PQsendClosePortal
+ *	 Submit a Close Portal command, but don't wait for it to finish
+ *
+ * Returns: 1 if successfully submitted
+ *			0 if error (conn->errorMessage is set)
+ */
+int
+PQsendClosePortal(PGconn *conn, const char *portal)
+{
+	return PQsendClose(conn, 'P', portal);
+}
+
+/*
+ * PQsendClose
+ *	 Common code to send a Close command
+ *
+ * Available options for close_type are
+ *	 'S' to close a prepared statement; or
+ *	 'P' to close a portal.
+ * Returns 1 on success and 0 on failure.
+ */
+static int
+PQsendClose(PGconn *conn, char close_type, const char *close_target)
+{
+	PGcmdQueueEntry *entry = NULL;
+
+	/* Treat null close_target as empty string */
+	if (!close_target)
+		close_target = "";
+
+	if (!PQsendQueryStart(conn, true))
+		return 0;
+
+	entry = pqAllocCmdQueueEntry(conn);
+	if (entry == NULL)
+		return 0;				/* error msg already set */
+
+	/* construct the Close message */
+	if (pqPutMsgStart('C', conn) < 0 ||
+		pqPutc(close_type, conn) < 0 ||
+		pqPuts(close_target, conn) < 0 ||
+		pqPutMsgEnd(conn) < 0)
+		goto sendFailed;
+
+	/* construct the Sync message */
+	if (conn->pipelineStatus == PQ_PIPELINE_OFF)
+	{
+		if (pqPutMsgStart('S', conn) < 0 ||
+			pqPutMsgEnd(conn) < 0)
+			goto sendFailed;
+	}
+
+	/* remember we are doing a Close */
+	entry->queryclass = PGQUERY_CLOSE;
+
+	/*
+	 * Give the data a push (in pipeline mode, only if we're past the size
+	 * threshold).  In nonblock mode, don't complain if we're unable to send
+	 * it all; PQgetResult() will do any additional flushing needed.
+	 */
+	if (pqPipelineFlush(conn) < 0)
+		goto sendFailed;
+
+	/* OK, it's launched! */
+	pqAppendCmdQueueEntry(conn, entry);
+
+	return 1;
+
+sendFailed:
+	pqRecycleCmdQueueEntry(conn, entry);
+	/* error message should be set up already */
+	return 0;
+}
+
+
 /*
  * PQnotifies
  *	  returns a PGnotify* structure of the latest async notification
diff --git a/src/interfaces/libpq/fe-protocol3.c b/src/interfaces/libpq/fe-protocol3.c
index 32b66d561cb..7bc6355d17f 100644
--- a/src/interfaces/libpq/fe-protocol3.c
+++ b/src/interfaces/libpq/fe-protocol3.c
@@ -278,8 +278,25 @@ pqParseInput3(PGconn *conn)
 					}
 					break;
 				case '2':		/* Bind Complete */
+					/* Nothing to do for this message type */
+					break;
 				case '3':		/* Close Complete */
-					/* Nothing to do for these message types */
+					/* If we're doing PQsendClose, we're done; else ignore */
+					if (conn->cmd_queue_head &&
+						conn->cmd_queue_head->queryclass == PGQUERY_CLOSE)
+					{
+						if (!pgHavePendingResult(conn))
+						{
+							conn->result = PQmakeEmptyPGresult(conn,
+															   PGRES_COMMAND_OK);
+							if (!conn->result)
+							{
+								libpq_append_conn_error(conn, "out of memory");
+								pqSaveErrorResult(conn);
+							}
+						}
+						conn->asyncStatus = PGASYNC_READY;
+					}
 					break;
 				case 'S':		/* parameter status */
 					if (getParameterStatus(conn))
diff --git a/src/interfaces/libpq/libpq-fe.h b/src/interfaces/libpq/libpq-fe.h
index 7476dbe0e90..97762d56f5d 100644
--- a/src/interfaces/libpq/libpq-fe.h
+++ b/src/interfaces/libpq/libpq-fe.h
@@ -548,6 +548,12 @@ extern PGresult *PQdescribePortal(PGconn *conn, const char *portal);
 extern int	PQsendDescribePrepared(PGconn *conn, const char *stmt);
 extern int	PQsendDescribePortal(PGconn *conn, const char *portal);
 
+/* Close prepared statements and portals */
+extern PGresult *PQclosePrepared(PGconn *conn, const char *stmt);
+extern PGresult *PQclosePortal(PGconn *conn, const char *portal);
+extern int	PQsendClosePrepared(PGconn *conn, const char *stmt);
+extern int	PQsendClosePortal(PGconn *conn, const char *portal);
+
 /* Delete a PGresult */
 extern void PQclear(PGresult *res);
 
diff --git a/src/interfaces/libpq/libpq-int.h b/src/interfaces/libpq/libpq-int.h
index 0045f83cbfd..02681ae7a16 100644
--- a/src/interfaces/libpq/libpq-int.h
+++ b/src/interfaces/libpq/libpq-int.h
@@ -324,7 +324,7 @@ typedef enum
 	PGQUERY_PREPARE,			/* Parse only (PQprepare) */
 	PGQUERY_DESCRIBE,			/* Describe Statement or Portal */
 	PGQUERY_SYNC,				/* Sync (at end of a pipeline) */
-	PGQUERY_CLOSE
+	PGQUERY_CLOSE				/* Close Statoment or Portal */
 } PGQueryClass;
 
 /*
diff --git a/src/test/modules/libpq_pipeline/libpq_pipeline.c b/src/test/modules/libpq_pipeline/libpq_pipeline.c
index f5b4d4d1ff2..a3067228dad 100644
--- a/src/test/modules/libpq_pipeline/libpq_pipeline.c
+++ b/src/test/modules/libpq_pipeline/libpq_pipeline.c
@@ -950,6 +950,15 @@ test_prepared(PGconn *conn)
 	if (PQexitPipelineMode(conn) != 1)
 		pg_fatal("could not exit pipeline mode: %s", PQerrorMessage(conn));
 
+	res = PQclosePrepared(conn, "select_one");
+	if (PQresultStatus(res) != PGRES_COMMAND_OK)
+		pg_fatal("expected COMMAND_OK, got %s", PQresStatus(PQresultStatus(res)));
+
+	/* Now that it's closed we should get an error when describing */
+	res = PQdescribePrepared(conn, "select_one");
+	if (PQresultStatus(res) != PGRES_FATAL_ERROR)
+		pg_fatal("expected FATAL_ERROR, got %s", PQresStatus(PQresultStatus(res)));
+
 	PQexec(conn, "BEGIN");
 	PQexec(conn, "DECLARE cursor_one CURSOR FOR SELECT 1");
 	PQenterPipelineMode(conn);
@@ -978,6 +987,15 @@ test_prepared(PGconn *conn)
 	if (PQexitPipelineMode(conn) != 1)
 		pg_fatal("could not exit pipeline mode: %s", PQerrorMessage(conn));
 
+	res = PQclosePortal(conn, "cursor_one");
+	if (PQresultStatus(res) != PGRES_COMMAND_OK)
+		pg_fatal("expected COMMAND_OK, got %s", PQresStatus(PQresultStatus(res)));
+
+	/* Now that it's closed we should get an error when describing */
+	res = PQdescribePortal(conn, "cursor_one");
+	if (PQresultStatus(res) != PGRES_FATAL_ERROR)
+		pg_fatal("expected COMMAND_OK, got %s", PQresStatus(PQresultStatus(res)));
+
 	fprintf(stderr, "ok\n");
 }
 
diff --git a/src/test/modules/libpq_pipeline/traces/prepared.trace b/src/test/modules/libpq_pipeline/traces/prepared.trace
index 1a7de5c3e65..9c95838393a 100644
--- a/src/test/modules/libpq_pipeline/traces/prepared.trace
+++ b/src/test/modules/libpq_pipeline/traces/prepared.trace
@@ -5,6 +5,14 @@ B	4	ParseComplete
 B	10	ParameterDescription	 1 NNNN
 B	113	RowDescription	 4 "?column?" NNNN 0 NNNN 4 -1 0 "?column?" NNNN 0 NNNN 65535 -1 0 "numeric" NNNN 0 NNNN 65535 -1 0 "interval" NNNN 0 NNNN 16 -1 0
 B	5	ReadyForQuery	 I
+F	16	Close	 S "select_one"
+F	4	Sync
+B	4	CloseComplete
+B	5	ReadyForQuery	 I
+F	16	Describe	 S "select_one"
+F	4	Sync
+B	NN	ErrorResponse	 S "ERROR" V "ERROR" C "26000" M "prepared statement "select_one" does not exist" F "SSSS" L "SSSS" R "SSSS" \x00
+B	5	ReadyForQuery	 I
 F	10	Query	 "BEGIN"
 B	10	CommandComplete	 "BEGIN"
 B	5	ReadyForQuery	 T
@@ -15,4 +23,12 @@ F	16	Describe	 P "cursor_one"
 F	4	Sync
 B	33	RowDescription	 1 "?column?" NNNN 0 NNNN 4 -1 0
 B	5	ReadyForQuery	 T
+F	16	Close	 P "cursor_one"
+F	4	Sync
+B	4	CloseComplete
+B	5	ReadyForQuery	 T
+F	16	Describe	 P "cursor_one"
+F	4	Sync
+B	NN	ErrorResponse	 S "ERROR" V "ERROR" C "34000" M "portal "cursor_one" does not exist" F "SSSS" L "SSSS" R "SSSS" \x00
+B	5	ReadyForQuery	 E
 F	4	Terminate

base-commit: c0d951262c80f42b3bfe037f940e103a24da84f4
-- 
2.34.1

