Make the COPY command return a command tag that includes the number of

rows copied.  Backend side of Volkan Yazici's recent patch, with
corrections and documentation.

5 files changed, 50 insertions, 5 deletions

diff --git a/doc/src/sgml/protocol.sgml b/doc/src/sgml/protocol.sgml
index 0a4d191f42..9b9a9b7a6a 100644
--- a/doc/src/sgml/protocol.sgml
+++ b/doc/src/sgml/protocol.sgml
@@ -2069,7 +2069,7 @@ CommandComplete (B)
         String
 </term>
 <listitem>
-<para>
+       <para>
         The command tag.  This is usually a single
         word that identifies which SQL command was completed.
        </para>
@@ -2109,7 +2109,16 @@ CommandComplete (B)
         <literal>FETCH <replaceable>rows</replaceable></literal> where
         <replaceable>rows</replaceable> is the number of rows that
         have been retrieved from the cursor.
-</para>
+       </para>
+
+       <para>
+        For a <command>COPY</command> command, the tag is
+        <literal>COPY <replaceable>rows</replaceable></literal> where
+        <replaceable>rows</replaceable> is the number of rows copied.
+        (Note: the row count appears only in
+        <productname>PostgreSQL</productname> 8.2 and later.)
+       </para>
+
 </listitem>
 </varlistentry>
 </variablelist>
diff --git a/doc/src/sgml/ref/copy.sgml b/doc/src/sgml/ref/copy.sgml
index 36558ac21f..a101cd7249 100644
--- a/doc/src/sgml/ref/copy.sgml
+++ b/doc/src/sgml/ref/copy.sgml
@@ -254,6 +254,20 @@ COPY <replaceable class="parameter">tablename</replaceable> [ ( <replaceable cla
  </refsect1>
 
  <refsect1>
+  <title>Outputs</title>
+
+  <para>
+   On successful completion, a <command>COPY</> command returns a command
+   tag of the form
+<screen>
+COPY <replaceable class="parameter">count</replaceable>
+</screen>
+   The <replaceable class="parameter">count</replaceable> is the number
+   of rows inserted into or copied from the table.
+  </para>
+ </refsect1>
+
+ <refsect1>
   <title>Notes</title>
 
    <para>
diff --git a/src/backend/commands/copy.c b/src/backend/commands/copy.c
index 0af1adc916..970581c7b0 100644
--- a/src/backend/commands/copy.c
+++ b/src/backend/commands/copy.c
@@ -102,6 +102,7 @@ typedef struct CopyStateData
 	int			client_encoding;	/* remote side's character encoding */
 	bool		need_transcoding;		/* client encoding diff from server? */
 	bool		encoding_embeds_ascii;	/* ASCII can be non-first byte? */
+	uint64		processed;		/* # of tuples processed */
 
 	/* parameters from the COPY command */
 	Relation	rel;			/* relation to copy to or from */
@@ -710,7 +711,7 @@ CopyLoadRawBuf(CopyState cstate)
  * Do not allow the copy if user doesn't have proper permission to access
  * the table.
  */
-void
+uint64
 DoCopy(const CopyStmt *stmt)
 {
 	CopyState	cstate;
@@ -724,6 +725,7 @@ DoCopy(const CopyStmt *stmt)
 	AclMode		required_access = (is_from ? ACL_INSERT : ACL_SELECT);
 	AclResult	aclresult;
 	ListCell   *option;
+	uint64		processed;
 
 	/* Allocate workspace and zero all fields */
 	cstate = (CopyStateData *) palloc0(sizeof(CopyStateData));
@@ -1019,6 +1021,7 @@ DoCopy(const CopyStmt *stmt)
 	cstate->line_buf_converted = false;
 	cstate->raw_buf = (char *) palloc(RAW_BUF_SIZE + 1);
 	cstate->raw_buf_index = cstate->raw_buf_len = 0;
+	cstate->processed = 0;
 
 	/* Set up encoding conversion info */
 	cstate->client_encoding = pg_get_client_encoding();
@@ -1161,10 +1164,14 @@ DoCopy(const CopyStmt *stmt)
 	heap_close(cstate->rel, (is_from ? NoLock : AccessShareLock));
 
 	/* Clean up storage (probably not really necessary) */
+	processed = cstate->processed;
+
 	pfree(cstate->attribute_buf.data);
 	pfree(cstate->line_buf.data);
 	pfree(cstate->raw_buf);
 	pfree(cstate);
+
+	return processed;
 }
 
 
@@ -1401,6 +1408,8 @@ CopyTo(CopyState cstate)
 		CopySendEndOfRow(cstate);
 
 		MemoryContextSwitchTo(oldcontext);
+		
+		cstate->processed++;
 	}
 
 	heap_endscan(scandesc);
@@ -2002,6 +2011,13 @@ CopyFrom(CopyState cstate)
 
 			/* AFTER ROW INSERT Triggers */
 			ExecARInsertTriggers(estate, resultRelInfo, tuple);
+
+			/*
+			 * We count only tuples not suppressed by a BEFORE INSERT trigger;
+			 * this is the same definition used by execMain.c for counting
+			 * tuples inserted by an INSERT command.
+			 */
+			cstate->processed++;
 		}
 	}
 
diff --git a/src/backend/tcop/utility.c b/src/backend/tcop/utility.c
index 4b797b535e..73d56be104 100644
--- a/src/backend/tcop/utility.c
+++ b/src/backend/tcop/utility.c
@@ -636,7 +636,13 @@ ProcessUtility(Node *parsetree,
 			break;
 
 		case T_CopyStmt:
-			DoCopy((CopyStmt *) parsetree);
+			{
+				uint64	processed = DoCopy((CopyStmt *) parsetree);
+
+				if (completionTag)
+					snprintf(completionTag, COMPLETION_TAG_BUFSIZE,
+							 "COPY " UINT64_FORMAT, processed);
+			}
 			break;
 
 		case T_PrepareStmt:
diff --git a/src/include/commands/copy.h b/src/include/commands/copy.h
index 65a37c5cdb..76780b8028 100644
--- a/src/include/commands/copy.h
+++ b/src/include/commands/copy.h
@@ -17,6 +17,6 @@
 #include "nodes/parsenodes.h"
 
 
-extern void DoCopy(const CopyStmt *stmt);
+extern uint64 DoCopy(const CopyStmt *stmt);
 
 #endif   /* COPY_H */