From 56adc590fffbc76fa5e99aa64c657621a630cbe1 Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Fri, 6 Jun 2014 21:21:39 +0200
Subject: [PATCH] Fixed segfault if COPY statements are executed

Close ticket #219
---
 NEWS               |  2 ++
 psycopg/pqpath.c   | 16 +++++++++++++++-
 tests/test_copy.py | 40 +++++++++++++++++++++++++++++++++++++++-
 3 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/NEWS b/NEWS
index 51830ff1..63568de1 100644
--- a/NEWS
+++ b/NEWS
@@ -13,6 +13,8 @@ Bug fixes:
 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)