Copy operations correctly set the cursor.rowcount attribute

Fixes ticket #180.

NEWS

@@ -3,6 +3,8 @@ What's new in psycopg 2.4.7
 
   - Work around `pip issue #1630 <https://github.com/pypa/pip/issues/1630>`__
     making installation via ``pip -e git+url`` impossible (:ticket:`#18`).
+  - Copy operations correctly set the `cursor.rowcount` attribute
+    (:ticket:`#180`).
   - It is now possible to call `get_transaction_status()` on closed
     connections.
   - Properly cleanup memory of broken connections (ticket #142).

psycopg/pqpath.c

@@ -1231,6 +1231,20 @@ exit:
     return rv;
 }
 
+void
+_read_rowcount(cursorObject *curs)
+{
+    const char *rowcount;
+
+    rowcount = PQcmdTuples(curs->pgres);
+    Dprintf("_read_rowcount: PQcmdTuples = %s", rowcount);
+    if (!rowcount || !rowcount[0]) {
+        curs->rowcount = -1;
+    } else {
+        curs->rowcount = atoi(rowcount);
+    }
+}
+
 static int
 _pq_copy_in_v3(cursorObject *curs)
 {
@@ -1348,6 +1362,7 @@ _pq_copy_in_v3(cursorObject *curs)
 
             if (NULL == curs->pgres)
                 break;
+            _read_rowcount(curs);
             if (PQresultStatus(curs->pgres) == PGRES_FATAL_ERROR)
                 pq_raise(curs->conn, curs, NULL);
             IFCLEARPGRES(curs->pgres);
@@ -1424,6 +1439,7 @@ _pq_copy_out_v3(cursorObject *curs)
 
         if (NULL == curs->pgres)
             break;
+        _read_rowcount(curs);
         if (PQresultStatus(curs->pgres) == PGRES_FATAL_ERROR)
             pq_raise(curs->conn, curs, NULL);
         IFCLEARPGRES(curs->pgres);
@@ -1439,7 +1455,6 @@ int
 pq_fetch(cursorObject *curs, int no_result)
 {
     int pgstatus, ex = -1;
-    const char *rowcount;
 
     /* even if we fail, we remove any information about the previous query */
     curs_reset(curs);
@@ -1471,11 +1486,7 @@ pq_fetch(cursorObject *curs, int no_result)
 
     case PGRES_COMMAND_OK:
         Dprintf("pq_fetch: command returned OK (no tuples)");
-        rowcount = PQcmdTuples(curs->pgres);
+        _read_rowcount(curs);
-        if (!rowcount || !rowcount[0])
-          curs->rowcount = -1;
-        else
-          curs->rowcount = atoi(rowcount);
         curs->lastoid = PQoidValue(curs->pgres);
         CLEARPGRES(curs->pgres);
         ex = 1;
@@ -1483,8 +1494,8 @@ pq_fetch(cursorObject *curs, int no_result)
 
     case PGRES_COPY_OUT:
         Dprintf("pq_fetch: data from a COPY TO (no tuples)");
-        ex = _pq_copy_out_v3(curs);
         curs->rowcount = -1;
+        ex = _pq_copy_out_v3(curs);
         /* error caught by out glorious notice handler */
         if (PyErr_Occurred()) ex = -1;
         IFCLEARPGRES(curs->pgres);
@@ -1492,8 +1503,8 @@ pq_fetch(cursorObject *curs, int no_result)
 
     case PGRES_COPY_IN:
         Dprintf("pq_fetch: data from a COPY FROM (no tuples)");
-        ex = _pq_copy_in_v3(curs);
         curs->rowcount = -1;
+        ex = _pq_copy_in_v3(curs);
         /* error caught by out glorious notice handler */
         if (PyErr_Occurred()) ex = -1;
         IFCLEARPGRES(curs->pgres);

tests/test_copy.py

@@ -285,6 +285,34 @@ class CopyTests(unittest.TestCase):
         curs.execute("select count(*) from manycols;")
         self.assertEqual(curs.fetchone()[0], 2)
 
+    def test_copy_rowcount(self):
+        curs = self.conn.cursor()
+
+        curs.copy_from(StringIO('aaa\nbbb\nccc\n'), 'tcopy', columns=['data'])
+        self.assertEqual(curs.rowcount, 3)
+
+        curs.copy_expert(
+            "copy tcopy (data) from stdin",
+            StringIO('ddd\neee\n'))
+        self.assertEqual(curs.rowcount, 2)
+
+        curs.copy_to(StringIO(), "tcopy")
+        self.assertEqual(curs.rowcount, 5)
+
+        curs.execute("insert into tcopy (data) values ('fff')")
+        curs.copy_expert("copy tcopy to stdout", StringIO())
+        self.assertEqual(curs.rowcount, 6)
+
+    def test_copy_rowcount_error(self):
+        curs = self.conn.cursor()
+
+        curs.execute("insert into tcopy (data) values ('fff')")
+        self.assertEqual(curs.rowcount, 1)
+
+        self.assertRaises(psycopg2.DataError,
+            curs.copy_from, StringIO('aaa\nbbb\nccc\n'), 'tcopy')
+        self.assertEqual(curs.rowcount, -1)
+
 
 decorate_all_tests(CopyTests, skip_if_green)