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.

diff --git a/doc/src/sgml/spi.sgml b/doc/src/sgml/spi.sgml
index 7162fdb7aa343ccb4490715ea4f4065f8148b43b..3fe6be9bd1d3cf2f3a899d3617c84d682dcdacd6 100644 (file)
--- 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 70d600490507f6797a5d4079f5cb51ad19d28018..bf8c4c711369a4637d936e1e2b7d88b22083891c 100644 (file)
--- 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);
+               }
            }
 
            /*