From b6e710b0fc022e009f51ca644cbcaa48f5207997 Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Tue, 7 Jun 2011 00:07:59 +0100
Subject: [PATCH] Fixed refcount bug in copy_to() and copy_expert() methods too

---
 psycopg/cursor_type.c | 31 ++++++++++++++-----------------
 scripts/ticket58.py   | 18 +++++++++++++++++-
 2 files changed, 31 insertions(+), 18 deletions(-)

diff --git a/psycopg/cursor_type.c b/psycopg/cursor_type.c
index 0aa88c6b..f850ebe3 100644
--- a/psycopg/cursor_type.c
+++ b/psycopg/cursor_type.c
@@ -1343,8 +1343,6 @@ static int
 _psyco_curs_has_write_check(PyObject* o, void* var)
 {
     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;
         return 1;
     }
@@ -1430,12 +1428,15 @@ psyco_curs_copy_to(cursorObject *self, PyObject *args, PyObject *kwargs)
     Dprintf("psyco_curs_copy_to: query = %s", query);
 
     self->copysize = 0;
+    Py_INCREF(file);
     self->copyfile = file;
 
     if (pq_execute(self, query, 0) == 1) {
         res = Py_None;
         Py_INCREF(Py_None);
     }
+
+    Py_DECREF(file);
     self->copyfile = NULL;
 
 exit:
@@ -1477,18 +1478,18 @@ psyco_curs_copy_expert(cursorObject *self, PyObject *args, PyObject *kwargs)
     EXC_IF_TPC_PREPARED(self->conn, copy_expert);
 
     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. */
-    
-    if (sql == NULL) { goto fail; }
+
+    if (sql == NULL) { goto exit; }
 
     /* This validation of file is rather weak, in that it doesn't enforce the
        assocation between "COPY FROM" -> "read" and "COPY TO" -> "write".
        However, the error handling in _pq_copy_[in|out] must be able to handle
        the case where the attempt to call file.read|write fails, so no harm
        done. */
-    
+
     if (   !PyObject_HasAttrString(file, "read")
         && !PyObject_HasAttrString(file, "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"
             " object for COPY FROM; a writeable file-like object for COPY TO."
           );
-        goto fail;
+        goto exit;
     }
 
     self->copysize = bufsize;
+    Py_INCREF(file);
     self->copyfile = file;
 
     /* 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;
     Py_INCREF(res);
-    goto cleanup;
- fail:
-    if (res != NULL) {
-        Py_DECREF(res);
-        res = NULL;
-    }
-    /* Fall through to cleanup */
- cleanup:
+
+exit:
     self->copyfile = NULL;
+    Py_XDECREF(file);
     Py_XDECREF(sql);
 
     return res;
diff --git a/scripts/ticket58.py b/scripts/ticket58.py
index 86cabac3..95520c1e 100644
--- a/scripts/ticket58.py
+++ b/scripts/ticket58.py
@@ -42,13 +42,29 @@ cur = conn.cursor()
 gc_thread.start()
 
 # Now do lots of "cursor.copy_from" calls:
+print "copy_from"
 for i in range(1000):
     f = StringIO("42\tfoo\n74\tbar\n")
     cur.copy_from(f, 'test', columns=('num', 'data'))
     # Assuming the other thread gets a chance to run during this call, expect a
     # 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.
-    
+
+# 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:
 done = 1