1 files changed, 69 insertions, 3 deletions

diff --git a/src/backend/tcop/cmdtag.c b/src/backend/tcop/cmdtag.c
index 262484f561f..2970f680ebf 100644
--- a/src/backend/tcop/cmdtag.c
+++ b/src/backend/tcop/cmdtag.c
@@ -15,18 +15,21 @@
 
 #include "miscadmin.h"
 #include "tcop/cmdtag.h"
+#include "utils/builtins.h"
 
 
 typedef struct CommandTagBehavior
 {
-	const char *name;
+	const char *name;			/* tag name, e.g. "SELECT" */
+	const uint8 namelen;		/* set to strlen(name) */
 	const bool	event_trigger_ok;
 	const bool	table_rewrite_ok;
-	const bool	display_rowcount;
+	const bool	display_rowcount;	/* should the number of rows affected be
+									 * shown in the command completion string */
 } CommandTagBehavior;
 
 #define PG_CMDTAG(tag, name, evtrgok, rwrok, rowcnt) \
-	{ name, evtrgok, rwrok, rowcnt },
+	{ name, (uint8) (sizeof(name) - 1), evtrgok, rwrok, rowcnt },
 
 static const CommandTagBehavior tag_behavior[COMMAND_TAG_NEXTTAG] = {
 #include "tcop/cmdtaglist.h"
@@ -47,6 +50,13 @@ GetCommandTagName(CommandTag commandTag)
 	return tag_behavior[commandTag].name;
 }
 
+const char *
+GetCommandTagNameAndLen(CommandTag commandTag, Size *len)
+{
+	*len = (Size) tag_behavior[commandTag].namelen;
+	return tag_behavior[commandTag].name;
+}
+
 bool
 command_tag_display_rowcount(CommandTag commandTag)
 {
@@ -96,3 +106,59 @@ GetCommandTagEnum(const char *commandname)
 	}
 	return CMDTAG_UNKNOWN;
 }
+
+/*
+ * BuildQueryCompletionString
+ *		Build a string containing the command tag name with the
+ *		QueryCompletion's nprocessed for command tags with display_rowcount
+ *		set.  Returns the strlen of the constructed string.
+ *
+ * The caller must ensure that buff is at least COMPLETION_TAG_BUFSIZE bytes.
+ *
+ * If nameonly is true, then the constructed string will contain only the tag
+ * name.
+ */
+Size
+BuildQueryCompletionString(char *buff, const QueryCompletion *qc,
+						   bool nameonly)
+{
+	CommandTag	tag = qc->commandTag;
+	Size		taglen;
+	const char *tagname = GetCommandTagNameAndLen(tag, &taglen);
+	char	   *bufp;
+
+	/*
+	 * We assume the tagname is plain ASCII and therefore requires no encoding
+	 * conversion.
+	 */
+	memcpy(buff, tagname, taglen);
+	bufp = buff + taglen;
+
+	/* ensure that the tagname isn't long enough to overrun the buffer */
+	Assert(taglen <= COMPLETION_TAG_BUFSIZE - MAXINT8LEN - 4);
+
+	/*
+	 * In PostgreSQL versions 11 and earlier, it was possible to create a
+	 * table WITH OIDS.  When inserting into such a table, INSERT used to
+	 * include the Oid of the inserted record in the completion tag.  To
+	 * maintain compatibility in the wire protocol, we now write a "0" (for
+	 * InvalidOid) in the location where we once wrote the new record's Oid.
+	 */
+	if (command_tag_display_rowcount(tag) && !nameonly)
+	{
+		if (tag == CMDTAG_INSERT)
+		{
+			*bufp++ = ' ';
+			*bufp++ = '0';
+		}
+		*bufp++ = ' ';
+		bufp += pg_ulltoa_n(qc->nprocessed, bufp);
+	}
+
+	/* and finally, NUL terminate the string */
+	*bufp = '\0';
+
+	Assert((bufp - buff) == strlen(buff));
+
+	return bufp - buff;
+}