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__)