Return the number of rows processed when COPY is executed through SPI.

hlinnaka · hlinnaka · commit 7ae1815961c6 · 2012-10-03T14:38:22.000+03:00
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.
diff --git 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
@@ -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);
+				}
 			}
 
 			/*
