1 files changed, 40 insertions, 0 deletions
diff --git a/psycopg/pqpath.c b/psycopg/pqpath.c
index 204a6b0..7e2f036 100644
--- a/psycopg/pqpath.c
+++ b/psycopg/pqpath.c
@@ -76,6 +76,35 @@ strip_severity(const char *msg)
         return msg;
 }
 
+/* Return a Python exception from a SQLSTATE from psycopg2.errors */
+BORROWED static PyObject *
+exception_from_module(const char *sqlstate)
+{
+    PyObject *rv = NULL;
+    PyObject *m = NULL;
+    PyObject *map = NULL;
+
+    if (!(m = PyImport_ImportModule("psycopg2.errors"))) { goto exit; }
+    if (!(map = PyObject_GetAttrString(m, "_by_sqlstate"))) { goto exit; }
+    if (!PyDict_Check(map)) {
+        Dprintf("'psycopg2.errors._by_sqlstate' is not a dict!");
+        goto exit;
+    }
+
+    /* get the sqlstate class (borrowed reference), or fail trying. */
+    rv = PyDict_GetItemString(map, sqlstate);
+
+exit:
+    /* We exit with a borrowed object, or a NULL but no error
+     * If an error did happen in this function, we don't want to clobber the
+     * database error. So better reporting it, albeit with the wrong class. */
+    PyErr_Clear();
+
+    Py_XDECREF(map);
+    Py_XDECREF(m);
+    return rv;
+}
+
 /* Returns the Python exception corresponding to an SQLSTATE error
    code.  A list of error codes can be found at:
 
@@ -83,6 +112,17 @@ strip_severity(const char *msg)
 BORROWED static PyObject *
 exception_from_sqlstate(const char *sqlstate)
 {
+    PyObject *exc;
+
+    /* First look up an exception of the proper class from the Python module */
+    exc = exception_from_module(sqlstate);
+    if (exc) {
+        return exc;
+    }
+    else {
+        PyErr_Clear();
+    }
+
     switch (sqlstate[0]) {
     case '0':
         switch (sqlstate[1]) {