Propagate read error messages in COPY FROM

Fix ticket psycopg#270.
wulczer · Feb 8, 2015 · 7ce7fef · 7ce7fef
1 parent d3c1ad5
commit 7ce7fef
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 3 deletions.
diff --git a/NEWS b/NEWS
@@ -21,6 +21,7 @@ What's new in psycopg 2.5.5
 
 - Named cursors used as context manager don't swallow the exception on exit
   (:ticket:`#262`).
+- Propagate read error messages in COPY FROM (:ticket:`#270`).
 - PostgreSQL time 24:00 is converted to Python 00:00 (:ticket:`#278`).
 
 

diff --git a/psycopg/pqpath.c b/psycopg/pqpath.c
@@ -1370,9 +1370,27 @@ _pq_copy_in_v3(cursorObject *curs)
         res = PQputCopyEnd(curs->conn->pgconn, NULL);
     else if (error == 2)
         res = PQputCopyEnd(curs->conn->pgconn, "error in PQputCopyData() call");
-    else
-        /* XXX would be nice to propagate the exception */
-        res = PQputCopyEnd(curs->conn->pgconn, "error in .read() call");
+    else {
+        char buf[1024];
+        strcpy(buf, "error in .read() call");
+        if (PyErr_Occurred()) {
+            PyObject *t, *ex, *tb;
+            PyErr_Fetch(&t, &ex, &tb);
+            if (ex) {
+                PyObject *str;
+                str = PyObject_Str(ex);
+                str = psycopg_ensure_bytes(str);
+                if (str) {
+                    PyOS_snprintf(buf, sizeof(buf),
+                        "error in .read() call: %s %s",
+                        ((PyTypeObject *)t)->tp_name, Bytes_AsString(str));
+                    Py_DECREF(str);
+                }
+            }
+            PyErr_Restore(t, ex, tb);
+        }
+        res = PQputCopyEnd(curs->conn->pgconn, buf);
+    }
 
     CLEARPGRES(curs->pgres);
 

diff --git a/tests/test_copy.py b/tests/test_copy.py
@@ -340,6 +340,34 @@ def test_copy_to_segfault(self):
         proc.communicate()
         self.assertEqual(0, proc.returncode)
 
+    def test_copy_from_propagate_error(self):
+        class BrokenRead(_base):
+            def read(self, size):
+                return 1/0
+
+            def readline(self):
+                return 1/0
+
+        curs = self.conn.cursor()
+        # It seems we cannot do this, but now at least we propagate the error
+        # self.assertRaises(ZeroDivisionError,
+        #     curs.copy_from, BrokenRead(), "tcopy")
+        try:
+            curs.copy_from(BrokenRead(), "tcopy")
+        except Exception, e:
+            self.assert_('ZeroDivisionError' in str(e))
+
+    def test_copy_to_propagate_error(self):
+        class BrokenWrite(_base):
+            def write(self, data):
+                return 1/0
+
+        curs = self.conn.cursor()
+        curs.execute("insert into tcopy values (10, 'hi')")
+        self.assertRaises(ZeroDivisionError,
+            curs.copy_to, BrokenWrite(), "tcopy")
+
+
 decorate_all_tests(CopyTests, skip_copy_if_green)