From 679af4a97528c140ac8e9cd85598f717b3f5cfb9 Mon Sep 17 00:00:00 2001 From: Daniele Varrazzo Date: Tue, 7 Jun 2011 01:20:25 +0100 Subject: [PATCH] Fixed copyfile refcount in copy_expert In case of early error, jumping to exit would have decref'd the borrowed reference to file. Issue spotted by Dave Malcolm, thanks! --- psycopg/cursor_type.c | 11 ++++++----- tests/test_copy.py | 9 +++++++++ 2 files changed, 15 insertions(+), 5 deletions(-) diff --git a/psycopg/cursor_type.c b/psycopg/cursor_type.c index f850ebe3..048ab6d8 100644 --- a/psycopg/cursor_type.c +++ b/psycopg/cursor_type.c @@ -1505,14 +1505,15 @@ psyco_curs_copy_expert(cursorObject *self, PyObject *args, PyObject *kwargs) 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 exit; } + if (pq_execute(self, Bytes_AS_STRING(sql), 0) == 1) { + res = Py_None; + Py_INCREF(res); + } - res = Py_None; - Py_INCREF(res); + self->copyfile = NULL; + Py_DECREF(file); exit: - self->copyfile = NULL; - Py_XDECREF(file); Py_XDECREF(sql); return res; diff --git a/tests/test_copy.py b/tests/test_copy.py index 9026abc5..7ec1b767 100755 --- a/tests/test_copy.py +++ b/tests/test_copy.py @@ -244,6 +244,15 @@ class CopyTests(unittest.TestCase): self.assertEqual(ntests, len(string.ascii_letters)) + def test_copy_expert_file_refcount(self): + class Whatever(object): + pass + + f = Whatever() + curs = self.conn.cursor() + self.assertRaises(TypeError, + curs.copy_expert, 'COPY tcopy (data) FROM STDIN', f) + decorate_all_tests(CopyTests, skip_if_green)