From 7ae1815961c635fd1a6fe72acb89fdef741a45a8 Mon Sep 17 00:00:00 2001
From: Heikki Linnakangas <heikki.linnakangas@iki.fi>
Date: Wed, 3 Oct 2012 14:32:01 +0300
Subject: [PATCH] Return the number of rows processed when COPY is executed
 through SPI.

You can now get the number of rows processed by a COPY statement in a
PL/pgSQL function with "GET DIAGNOSTICS x = ROW_COUNT".

Pavel Stehule, reviewed by Amit Kapila, with some editing by me.
---
 doc/src/sgml/spi.sgml      |  5 ++++-
 src/backend/executor/spi.c | 24 +++++++++++++++---------
 2 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/doc/src/sgml/spi.sgml b/doc/src/sgml/spi.sgml
index 7162fdb7aa..3fe6be9bd1 100644
--- a/doc/src/sgml/spi.sgml
+++ b/doc/src/sgml/spi.sgml
@@ -377,7 +377,10 @@ SPI_execute("INSERT INTO foo SELECT * FROM bar", false, 5);
    global pointer <literal>SPITupleTable *SPI_tuptable</literal> to
    access the result rows.  Some utility commands (such as
    <command>EXPLAIN</>) also return row sets, and <literal>SPI_tuptable</>
-   will contain the result in these cases too.
+   will contain the result in these cases too. Some utility commands
+   (<command>COPY</>, <command>CREATE TABLE AS</>) don't return a row set, so
+   <literal>SPI_tuptable</> is NULL, but they still return the number of
+   rows processed in <varname>SPI_processed</>.
   </para>
 
   <para>
diff --git a/src/backend/executor/spi.c b/src/backend/executor/spi.c
index 70d6004905..bf8c4c7113 100644
--- a/src/backend/executor/spi.c
+++ b/src/backend/executor/spi.c
@@ -1922,25 +1922,31 @@ _SPI_execute_plan(SPIPlanPtr plan, ParamListInfo paramLI,
 					_SPI_current->processed = _SPI_current->tuptable->alloced -
 						_SPI_current->tuptable->free;
 
+				res = SPI_OK_UTILITY;
+
 				/*
-				 * CREATE TABLE AS is a messy special case for historical
-				 * reasons.  We must set _SPI_current->processed even though
-				 * the tuples weren't returned to the caller, and we must
-				 * return a special result code if the statement was spelled
-				 * SELECT INTO.
+				 * Some utility statements return a row count, even though the
+				 * tuples are not returned to the caller.
 				 */
 				if (IsA(stmt, CreateTableAsStmt))
 				{
 					Assert(strncmp(completionTag, "SELECT ", 7) == 0);
 					_SPI_current->processed = strtoul(completionTag + 7,
 													  NULL, 10);
+
+					/*
+					 * For historical reasons, if CREATE TABLE AS was spelled
+					 * as SELECT INTO, return a special return code.
+					 */
 					if (((CreateTableAsStmt *) stmt)->is_select_into)
 						res = SPI_OK_SELINTO;
-					else
-						res = SPI_OK_UTILITY;
 				}
-				else
-					res = SPI_OK_UTILITY;
+				else if (IsA(stmt, CopyStmt))
+				{
+					Assert(strncmp(completionTag, "COPY ", 5) == 0);
+					_SPI_current->processed = strtoul(completionTag + 5,
+													  NULL, 10);
+				}
 			}
 
 			/*