[PATCH] PL/Python: Add spidata to all spiexceptions

PL/Python maps Python SPIError exceptions with 'spidata' attribute into SQL
errors.  PL/Python also creates classes in plpy.spiexceptions for all known
errors but does not initialize their spidata, so when a PL/Python function
raises such an exception it is not recognized properly and is always
reported as an internal error.

This allows PL/Python code to raise exceptions that PL/pgSQL can catch and
which are correctly reported in logs instead of always showing up as XX000.
---
 src/pl/plpython/expected/plpython_error.out   | 12 ++++++++++++
 src/pl/plpython/expected/plpython_error_0.out | 12 ++++++++++++
 src/pl/plpython/plpy_plpymodule.c             |  9 +++++++++
 src/pl/plpython/sql/plpython_error.sql        | 14 ++++++++++++++
 4 files changed, 47 insertions(+)

diff --git a/src/pl/plpython/expected/plpython_error.out b/src/pl/plpython/expected/plpython_error.out
index e1ec9c2..c1c36d9 100644
--- a/src/pl/plpython/expected/plpython_error.out
+++ b/src/pl/plpython/expected/plpython_error.out
@@ -400,3 +400,15 @@ CONTEXT:  Traceback (most recent call last):
   PL/Python function "manual_subxact_prepared", line 4, in <module>
     plpy.execute(save)
 PL/Python function "manual_subxact_prepared"
+/* raising plpy.spiexception.* from python code should preserve sqlstate
+ */
+CREATE FUNCTION plpy_raise_spiexception() RETURNS void AS $$
+raise plpy.spiexceptions.DivisionByZero()
+$$ LANGUAGE plpythonu;
+DO $$
+BEGIN
+	SELECT plpy_raise_spiexception();
+EXCEPTION WHEN division_by_zero THEN
+	-- NOOP
+END
+$$ LANGUAGE plpgsql;
diff --git a/src/pl/plpython/expected/plpython_error_0.out b/src/pl/plpython/expected/plpython_error_0.out
index 6f74a50..61d95e3 100644
--- a/src/pl/plpython/expected/plpython_error_0.out
+++ b/src/pl/plpython/expected/plpython_error_0.out
@@ -400,3 +400,15 @@ CONTEXT:  Traceback (most recent call last):
   PL/Python function "manual_subxact_prepared", line 4, in <module>
     plpy.execute(save)
 PL/Python function "manual_subxact_prepared"
+/* raising plpy.spiexception.* from python code should preserve sqlstate
+ */
+CREATE FUNCTION plpy_raise_spiexception() RETURNS void AS $$
+raise plpy.spiexceptions.DivisionByZero()
+$$ LANGUAGE plpythonu;
+DO $$
+BEGIN
+	SELECT plpy_raise_spiexception();
+EXCEPTION WHEN division_by_zero THEN
+	-- NOOP
+END
+$$ LANGUAGE plpgsql;
diff --git a/src/pl/plpython/plpy_plpymodule.c b/src/pl/plpython/plpy_plpymodule.c
index 37ea2a4..4213e83 100644
--- a/src/pl/plpython/plpy_plpymodule.c
+++ b/src/pl/plpython/plpy_plpymodule.c
@@ -247,6 +247,7 @@ PLy_generate_spi_exceptions(PyObject *mod, PyObject *base)
 		PyObject   *exc;
 		PLyExceptionEntry *entry;
 		PyObject   *sqlstate;
+		PyObject   *spidata;
 		PyObject   *dict = PyDict_New();
 
 		if (dict == NULL)
@@ -258,6 +259,14 @@ PLy_generate_spi_exceptions(PyObject *mod, PyObject *base)
 
 		PyDict_SetItemString(dict, "sqlstate", sqlstate);
 		Py_DECREF(sqlstate);
+
+		spidata = Py_BuildValue("izzzi", exception_map[i].sqlstate,
+								NULL, NULL, NULL, 0);
+		if (spidata == NULL)
+			PLy_elog(ERROR, "could not generate SPI exceptions");
+		PyDict_SetItemString(dict, "spidata", spidata);
+		Py_DECREF(spidata);
+
 		exc = PyErr_NewException(exception_map[i].name, base, dict);
 		PyModule_AddObject(mod, exception_map[i].classname, exc);
 		entry = hash_search(PLy_spi_exceptions, &exception_map[i].sqlstate,
diff --git a/src/pl/plpython/sql/plpython_error.sql b/src/pl/plpython/sql/plpython_error.sql
index 502bbec..ec93144 100644
--- a/src/pl/plpython/sql/plpython_error.sql
+++ b/src/pl/plpython/sql/plpython_error.sql
@@ -298,3 +298,17 @@ plpy.execute(rollback)
 $$ LANGUAGE plpythonu;
 
 SELECT manual_subxact_prepared();
+
+/* raising plpy.spiexception.* from python code should preserve sqlstate
+ */
+CREATE FUNCTION plpy_raise_spiexception() RETURNS void AS $$
+raise plpy.spiexceptions.DivisionByZero()
+$$ LANGUAGE plpythonu;
+
+DO $$
+BEGIN
+	SELECT plpy_raise_spiexception();
+EXCEPTION WHEN division_by_zero THEN
+	-- NOOP
+END
+$$ LANGUAGE plpgsql;
-- 
1.7.12.1