mirror of
https://github.com/psycopg/psycopg2.git
synced 2024-11-22 00:46:33 +03:00
Copy operations correctly set the cursor.rowcount attribute
Fixes ticket #180.
This commit is contained in:
parent
a1344f30f3
commit
e7fc781280
2
NEWS
2
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).
|
||||
|
|
|
@ -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);
|
||||
|
|
|
@ -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)
|
||||
|
||||
|
|
Loading…
Reference in New Issue
Block a user