Fixed refcount bug in copy_to() and copy_expert() methods too

This commit is contained in:
Daniele Varrazzo 2011-06-07 00:07:59 +01:00
parent 1888bf41c0
commit b6e710b0fc
2 changed files with 31 additions and 18 deletions

View File

@ -1343,8 +1343,6 @@ static int
_psyco_curs_has_write_check(PyObject* o, void* var) _psyco_curs_has_write_check(PyObject* o, void* var)
{ {
if (PyObject_HasAttrString(o, "write")) { if (PyObject_HasAttrString(o, "write")) {
/* It's OK to store a borrowed reference, because it is only held for
* the duration of psyco_curs_copy_to. */
*((PyObject**)var) = o; *((PyObject**)var) = o;
return 1; return 1;
} }
@ -1430,12 +1428,15 @@ psyco_curs_copy_to(cursorObject *self, PyObject *args, PyObject *kwargs)
Dprintf("psyco_curs_copy_to: query = %s", query); Dprintf("psyco_curs_copy_to: query = %s", query);
self->copysize = 0; self->copysize = 0;
Py_INCREF(file);
self->copyfile = file; self->copyfile = file;
if (pq_execute(self, query, 0) == 1) { if (pq_execute(self, query, 0) == 1) {
res = Py_None; res = Py_None;
Py_INCREF(Py_None); Py_INCREF(Py_None);
} }
Py_DECREF(file);
self->copyfile = NULL; self->copyfile = NULL;
exit: exit:
@ -1478,10 +1479,10 @@ psyco_curs_copy_expert(cursorObject *self, PyObject *args, PyObject *kwargs)
sql = _psyco_curs_validate_sql_basic(self, sql); sql = _psyco_curs_validate_sql_basic(self, sql);
/* Any failure from here forward should 'goto fail' rather than /* Any failure from here forward should 'goto exit' rather than
'return NULL' directly. */ 'return NULL' directly. */
if (sql == NULL) { goto fail; } if (sql == NULL) { goto exit; }
/* This validation of file is rather weak, in that it doesn't enforce the /* This validation of file is rather weak, in that it doesn't enforce the
assocation between "COPY FROM" -> "read" and "COPY TO" -> "write". assocation between "COPY FROM" -> "read" and "COPY TO" -> "write".
@ -1496,26 +1497,22 @@ psyco_curs_copy_expert(cursorObject *self, PyObject *args, PyObject *kwargs)
PyErr_SetString(PyExc_TypeError, "file must be a readable file-like" PyErr_SetString(PyExc_TypeError, "file must be a readable file-like"
" object for COPY FROM; a writeable file-like object for COPY TO." " object for COPY FROM; a writeable file-like object for COPY TO."
); );
goto fail; goto exit;
} }
self->copysize = bufsize; self->copysize = bufsize;
Py_INCREF(file);
self->copyfile = file; self->copyfile = file;
/* At this point, the SQL statement must be str, not unicode */ /* At this point, the SQL statement must be str, not unicode */
if (pq_execute(self, Bytes_AS_STRING(sql), 0) != 1) { goto fail; } if (pq_execute(self, Bytes_AS_STRING(sql), 0) != 1) { goto exit; }
res = Py_None; res = Py_None;
Py_INCREF(res); Py_INCREF(res);
goto cleanup;
fail: exit:
if (res != NULL) {
Py_DECREF(res);
res = NULL;
}
/* Fall through to cleanup */
cleanup:
self->copyfile = NULL; self->copyfile = NULL;
Py_XDECREF(file);
Py_XDECREF(sql); Py_XDECREF(sql);
return res; return res;

View File

@ -42,6 +42,7 @@ cur = conn.cursor()
gc_thread.start() gc_thread.start()
# Now do lots of "cursor.copy_from" calls: # Now do lots of "cursor.copy_from" calls:
print "copy_from"
for i in range(1000): for i in range(1000):
f = StringIO("42\tfoo\n74\tbar\n") f = StringIO("42\tfoo\n74\tbar\n")
cur.copy_from(f, 'test', columns=('num', 'data')) cur.copy_from(f, 'test', columns=('num', 'data'))
@ -49,6 +50,21 @@ for i in range(1000):
# build of python (with assertions enabled) to bail out here with: # build of python (with assertions enabled) to bail out here with:
# python: Modules/gcmodule.c:277: visit_decref: Assertion `gc->gc.gc_refs != 0' failed. # python: Modules/gcmodule.c:277: visit_decref: Assertion `gc->gc.gc_refs != 0' failed.
# Also exercise the copy_to code path
print "copy_to"
cur.execute("truncate test")
f = StringIO("42\tfoo\n74\tbar\n")
cur.copy_from(f, 'test', columns=('num', 'data'))
for i in range(1000):
f = StringIO()
cur.copy_to(f, 'test', columns=('num', 'data'))
# And copy_expert too
print "copy_expert"
cur.execute("truncate test")
for i in range(1000):
f = StringIO("42\tfoo\n74\tbar\n")
cur.copy_expert("copy test to stdout", f)
# Terminate the GC thread's loop: # Terminate the GC thread's loop:
done = 1 done = 1