From 66fdaeaf0990e7d8b0658dae31cf577907e62113 Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Sun, 8 Feb 2015 01:42:21 +0000
Subject: [PATCH] Propagate read error messages in COPY FROM

Fix ticket #270.
---
 NEWS               |  1 +
 psycopg/pqpath.c   | 24 +++++++++++++++++++++---
 tests/test_copy.py | 28 ++++++++++++++++++++++++++++
 3 files changed, 50 insertions(+), 3 deletions(-)

diff --git a/NEWS b/NEWS
index 5b5f55cc..fcc1fd90 100644
--- a/NEWS
+++ b/NEWS
@@ -6,6 +6,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
index 3512f639..2de27905 100644
--- a/psycopg/pqpath.c
+++ b/psycopg/pqpath.c
@@ -1376,9 +1376,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
index 5b28b20d..32134215 100755
--- a/tests/test_copy.py
+++ b/tests/test_copy.py
@@ -340,6 +340,34 @@ conn.close()
         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)