Copy operations correctly set the cursor.rowcount attribute

Fixes ticket #180.
This commit is contained in:
Daniele Varrazzo 2014-05-05 23:52:41 +01:00
parent a1344f30f3
commit e7fc781280
3 changed files with 49 additions and 8 deletions

2
NEWS
View File

@ -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).

View File

@ -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);
if (!rowcount || !rowcount[0])
curs->rowcount = -1;
else
curs->rowcount = atoi(rowcount);
_read_rowcount(curs);
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);

View File

@ -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)