diff --git a/NEWS b/NEWS index 0dd423fb..21f6d201 100644 --- a/NEWS +++ b/NEWS @@ -4,6 +4,8 @@ Current release What's new in psycopg 2.5.4 ^^^^^^^^^^^^^^^^^^^^^^^^^^^ +- Fixed segfault if COPY statements are executed instead of using the + proper methods (:ticket:`#219`). - Don't ignore silently the `cursor.callproc` argument without a length. diff --git a/psycopg/pqpath.c b/psycopg/pqpath.c index 05d22ee9..42460f5b 100644 --- a/psycopg/pqpath.c +++ b/psycopg/pqpath.c @@ -1288,6 +1288,13 @@ _pq_copy_in_v3(cursorObject *curs) Py_ssize_t length = 0; int res, error = 0; + if (!curs->copyfile) { + PyErr_SetString(ProgrammingError, + "can't execute COPY FROM: use the copy_from() method instead"); + error = 1; + goto exit; + } + if (!(func = PyObject_GetAttrString(curs->copyfile, "read"))) { Dprintf("_pq_copy_in_v3: can't get o.read"); error = 1; @@ -1411,7 +1418,8 @@ exit: static int _pq_copy_out_v3(cursorObject *curs) { - PyObject *tmp = NULL, *func; + PyObject *tmp = NULL; + PyObject *func = NULL; PyObject *obj = NULL; int ret = -1; int is_text; @@ -1419,6 +1427,12 @@ _pq_copy_out_v3(cursorObject *curs) char *buffer; Py_ssize_t len; + if (!curs->copyfile) { + PyErr_SetString(ProgrammingError, + "can't execute COPY TO: use the copy_to() method instead"); + goto exit; + } + if (!(func = PyObject_GetAttrString(curs->copyfile, "write"))) { Dprintf("_pq_copy_out_v3: can't get o.write"); goto exit; diff --git a/tests/test_copy.py b/tests/test_copy.py index 43fa5973..5b28b20d 100755 --- a/tests/test_copy.py +++ b/tests/test_copy.py @@ -28,10 +28,13 @@ from testutils import unittest, ConnectingTestCase, decorate_all_tests from testutils import skip_if_no_iobase, skip_before_postgres from cStringIO import StringIO from itertools import cycle, izip +from subprocess import Popen, PIPE import psycopg2 import psycopg2.extensions -from testutils import skip_copy_if_green +from testutils import skip_copy_if_green, script_to_py3 +from testconfig import dsn + if sys.version_info[0] < 3: _base = object @@ -301,6 +304,41 @@ class CopyTests(ConnectingTestCase): curs.copy_from, StringIO('aaa\nbbb\nccc\n'), 'tcopy') self.assertEqual(curs.rowcount, -1) + def test_copy_from_segfault(self): + # issue #219 + script = ("""\ +import psycopg2 +conn = psycopg2.connect(%(dsn)r) +curs = conn.cursor() +curs.execute("create table copy_segf (id int)") +try: + curs.execute("copy copy_segf from stdin") +except psycopg2.ProgrammingError: + pass +conn.close() +""" % { 'dsn': dsn,}) + + proc = Popen([sys.executable, '-c', script_to_py3(script)]) + proc.communicate() + self.assertEqual(0, proc.returncode) + + def test_copy_to_segfault(self): + # issue #219 + script = ("""\ +import psycopg2 +conn = psycopg2.connect(%(dsn)r) +curs = conn.cursor() +curs.execute("create table copy_segf (id int)") +try: + curs.execute("copy copy_segf to stdout") +except psycopg2.ProgrammingError: + pass +conn.close() +""" % { 'dsn': dsn,}) + + proc = Popen([sys.executable, '-c', script_to_py3(script)], stdout=PIPE) + proc.communicate() + self.assertEqual(0, proc.returncode) decorate_all_tests(CopyTests, skip_copy_if_green)