From 43daba38e7563ffc0369ed07211464aff597e21c Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Sat, 14 Jan 2012 17:34:09 +0000
Subject: [PATCH] Make Error and subclasses picklable

Useful for multiprocessing interaction.
Closes ticket #90.
---
 NEWS                    |  2 ++
 psycopg/psycopgmodule.c | 77 +++++++++++++++++++++++++++++++++++++++--
 tests/test_module.py    | 35 ++++++++++++++++++-
 3 files changed, 111 insertions(+), 3 deletions(-)

diff --git a/NEWS b/NEWS
index 21319d7b..24811fcd 100644
--- a/NEWS
+++ b/NEWS
@@ -5,6 +5,8 @@ What's new in psycopg 2.4.5
     (ticket #84).
   - Use lo_creat() instead of lo_create() when possible for better
     interaction with pgpool-II (ticket #88).
+  - Error and its subclasses are picklable, useful for multiprocessing
+    interaction (ticket #90).
 
 
 What's new in psycopg 2.4.4
diff --git a/psycopg/psycopgmodule.c b/psycopg/psycopgmodule.c
index 3b2b0609..ea04e658 100644
--- a/psycopg/psycopgmodule.c
+++ b/psycopg/psycopgmodule.c
@@ -380,7 +380,59 @@ static struct {
     {NULL}  /* Sentinel */
 };
 
-static void
+
+/* Error.__reduce_ex__
+ *
+ * The method is required to make exceptions picklable: set the cursor
+ * attribute to None. Only working from Py 2.5: previous versions
+ * would require implementing __getstate__, and as of 2012 it's a little
+ * bit too late to care. */
+static PyObject *
+psyco_error_reduce_ex(PyObject *self, PyObject *args)
+{
+    PyObject *proto = NULL;
+    PyObject *super = NULL;
+    PyObject *tuple = NULL;
+    PyObject *dict = NULL;
+    PyObject *rv = NULL;
+
+    /* tuple = Exception.__reduce_ex__(self, proto) */
+    if (!PyArg_ParseTuple(args, "O", &proto)) {
+        goto error;
+    }
+    if (!(super = PyObject_GetAttrString(PyExc_Exception, "__reduce_ex__"))) {
+        goto error;
+    }
+    if (!(tuple = PyObject_CallFunctionObjArgs(super, self, proto, NULL))) {
+        goto error;
+    }
+
+    /* tuple[2]['cursor'] = None
+     *
+     * If these checks fail, we can still return a valid object. Pickle
+     * will likely fail downstream, but there's nothing else we can do here */
+    if (!PyTuple_Check(tuple)) { goto exit; }
+    if (3 > PyTuple_GET_SIZE(tuple)) { goto exit; }
+    dict = PyTuple_GET_ITEM(tuple, 2);      /* borrowed */
+    if (!PyDict_Check(dict)) { goto exit; }
+
+    /* Modify the tuple inplace and return it */
+    if (0 != PyDict_SetItemString(dict, "cursor", Py_None)) {
+        goto error;
+    }
+
+exit:
+    rv = tuple;
+    tuple = NULL;
+
+error:
+    Py_XDECREF(tuple);
+    Py_XDECREF(super);
+
+    return rv;
+}
+
+static int
 psyco_errors_init(void)
 {
     /* the names of the exceptions here reflect the oranization of the
@@ -391,6 +443,11 @@ psyco_errors_init(void)
     PyObject *dict;
     PyObject *base;
     PyObject *str;
+    PyObject *descr;
+    int rv = -1;
+
+    static PyMethodDef psyco_error_reduce_ex_def =
+        {"__reduce_ex__", psyco_error_reduce_ex, METH_VARARGS, "pickle helper"};
 
     for (i=0; exctable[i].name; i++) {
         dict = PyDict_New();
@@ -420,6 +477,22 @@ psyco_errors_init(void)
     PyObject_SetAttrString(Error, "pgerror", Py_None);
     PyObject_SetAttrString(Error, "pgcode", Py_None);
     PyObject_SetAttrString(Error, "cursor", Py_None);
+
+    /* install __reduce_ex__ on Error to make all the subclasses picklable */
+    if (!(descr = PyDescr_NewMethod((PyTypeObject *)Error,
+            &psyco_error_reduce_ex_def))) {
+        goto exit;
+    }
+    if (0 != PyObject_SetAttrString(Error,
+            psyco_error_reduce_ex_def.ml_name, descr)) {
+        goto exit;
+    }
+    Py_DECREF(descr);
+
+    rv = 0;
+
+exit:
+    return rv;
 }
 
 void
@@ -869,7 +942,7 @@ INIT_MODULE(_psycopg)(void)
     psyco_adapters_init(dict);
 
     /* create a standard set of exceptions and add them to the module's dict */
-    psyco_errors_init();
+    if (0 != psyco_errors_init()) { goto exit; }
     psyco_errors_fill(dict);
 
     /* Solve win32 build issue about non-constant initializer element */
diff --git a/tests/test_module.py b/tests/test_module.py
index 9c130f3f..49889693 100755
--- a/tests/test_module.py
+++ b/tests/test_module.py
@@ -22,7 +22,8 @@
 # FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public
 # License for more details.
 
-from testutils import unittest
+from testutils import unittest, skip_before_python
+from testconfig import dsn
 
 import psycopg2
 
@@ -127,6 +128,38 @@ class ConnectTestCase(unittest.TestCase):
         self.assertEqual(self.args[0], r"dbname='\\every thing\''")
 
 
+class ExceptionsTestCase(unittest.TestCase):
+    def setUp(self):
+        self.conn = psycopg2.connect(dsn)
+
+    def tearDown(self):
+        self.conn.close()
+
+    def test_attributes(self):
+        cur = self.conn.cursor()
+        try: cur.execute("select * from nonexist")
+        except psycopg2.Error, e: pass
+
+        self.assertEqual(e.pgcode, '42P01')
+        self.assert_(e.pgerror)
+        self.assert_(e.cursor is cur)
+
+    @skip_before_python(2, 5)
+    def test_pickle(self):
+        import pickle
+        cur = self.conn.cursor()
+        try:
+            cur.execute("select * from nonexist")
+        except psycopg2.Error, e:
+            pass
+
+        e1 = pickle.loads(pickle.dumps(e))
+
+        self.assertEqual(e.pgerror, e1.pgerror)
+        self.assertEqual(e.pgcode, e1.pgcode)
+        self.assert_(e1.cursor is None)
+
+
 def test_suite():
     return unittest.TestLoader().loadTestsFromName(__name__)