psycopg/psycopg2
1
1
Fork 0
mirror of https://github.com/psycopg/psycopg2.git synced 2024-11-22 17:06:33 +03:00
Code Issues Packages Projects Releases Wiki Activity

Copy operations correctly set the cursor.rowcount attribute

Browse Source 
Fixes ticket #180.
This commit is contained in:
Daniele Varrazzo 2014-05-05 23:52:41 +01:00
parent abd975ae40
commit 81b5f1fb26
3 changed files with 49 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
NEWS
View File

@ -15,6 +15,8 @@ What's new in psycopg 2.5.3
- 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.
- Fixed unsafe access to object names causing assertion failures in
Python 3 debug builds (:ticket:`#188`).

27
psycopg/pqpath.c
View File

@ -1264,6 +1264,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)
{
@ -1381,6 +1395,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);
CLEARPGRES(curs->pgres);
@ -1457,6 +1472,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);
CLEARPGRES(curs->pgres);
@ -1472,7 +1488,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);
@ -1504,11 +1519,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;
@ -1516,8 +1527,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;
CLEARPGRES(curs->pgres);
@ -1525,8 +1536,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;
CLEARPGRES(curs->pgres);

28
tests/test_copy.py
View File

@ -272,6 +272,34 @@ class CopyTests(ConnectingTestCase):
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_copy_if_green)