Handle SPIErrors raised directly in PL/Python code.

If a PL/Python function raises an SPIError (or one if its subclasses)
directly with python's raise statement, treat it the same as an SPIError
generated internally. In particular, if the user sets the sqlstate
attribute, preserve that.

Oskari Saarenmaa and Jan Urbański, reviewed by Karl O. Pinc.

4 files changed, 119 insertions, 5 deletions

diff --git a/src/pl/plpython/expected/plpython_error.out b/src/pl/plpython/expected/plpython_error.out
index e1ec9c2c132..be2ec9708ad 100644
--- a/src/pl/plpython/expected/plpython_error.out
+++ b/src/pl/plpython/expected/plpython_error.out
@@ -400,3 +400,29 @@ 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;
+/* setting a custom sqlstate should be handled
+ */
+CREATE FUNCTION plpy_raise_spiexception_override() RETURNS void AS $$
+exc = plpy.spiexceptions.DivisionByZero()
+exc.sqlstate = 'SILLY'
+raise exc
+$$ LANGUAGE plpythonu;
+DO $$
+BEGIN
+	SELECT plpy_raise_spiexception_override();
+EXCEPTION WHEN SQLSTATE 'SILLY' 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 6f74a5038f0..39c63c547a4 100644
--- a/src/pl/plpython/expected/plpython_error_0.out
+++ b/src/pl/plpython/expected/plpython_error_0.out
@@ -400,3 +400,29 @@ 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;
+/* setting a custom sqlstate should be handled
+ */
+CREATE FUNCTION plpy_raise_spiexception_override() RETURNS void AS $$
+exc = plpy.spiexceptions.DivisionByZero()
+exc.sqlstate = 'SILLY'
+raise exc
+$$ LANGUAGE plpythonu;
+DO $$
+BEGIN
+	SELECT plpy_raise_spiexception_override();
+EXCEPTION WHEN SQLSTATE 'SILLY' THEN
+	-- NOOP
+END
+$$ LANGUAGE plpgsql;
diff --git a/src/pl/plpython/plpy_elog.c b/src/pl/plpython/plpy_elog.c
index c375ac07fa8..70450d7d9e5 100644
--- a/src/pl/plpython/plpy_elog.c
+++ b/src/pl/plpython/plpy_elog.c
@@ -337,6 +337,31 @@ PLy_traceback(char **xmsg, char **tbmsg, int *tb_depth)
 }
 
 /*
+ * Extract error code from SPIError's sqlstate attribute.
+ */
+static void
+PLy_get_spi_sqlerrcode(PyObject *exc, int *sqlerrcode)
+{
+	PyObject	*sqlstate;
+	char		*buffer;
+
+	sqlstate = PyObject_GetAttrString(exc, "sqlstate");
+	if (sqlstate == NULL)
+		return;
+
+	buffer = PyString_AsString(sqlstate);
+	if (strlen(buffer) == 5 &&
+		strspn(buffer, "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ") == 5)
+	{
+		*sqlerrcode = MAKE_SQLSTATE(buffer[0], buffer[1], buffer[2],
+									buffer[3], buffer[4]);
+	}
+
+	Py_DECREF(sqlstate);
+}
+
+
+/*
  * Extract the error data from a SPIError
  */
 static void
@@ -345,13 +370,20 @@ PLy_get_spi_error_data(PyObject *exc, int *sqlerrcode, char **detail, char **hin
 	PyObject   *spidata = NULL;
 
 	spidata = PyObject_GetAttrString(exc, "spidata");
-	if (!spidata)
-		goto cleanup;
 
-	if (!PyArg_ParseTuple(spidata, "izzzi", sqlerrcode, detail, hint, query, position))
-		goto cleanup;
+	if (spidata != NULL)
+	{
+		PyArg_ParseTuple(spidata, "izzzi", sqlerrcode, detail, hint, query, position);
+	}
+	else
+	{
+		/*
+		 * If there's no spidata, at least set the sqlerrcode. This can happen
+		 * if someone explicitly raises a SPI exception from Python code.
+		 */
+		PLy_get_spi_sqlerrcode(exc, sqlerrcode);
+	}
 
-cleanup:
 	PyErr_Clear();
 	/* no elog here, we simply won't report the errhint, errposition etc */
 	Py_XDECREF(spidata);
diff --git a/src/pl/plpython/sql/plpython_error.sql b/src/pl/plpython/sql/plpython_error.sql
index 502bbec38f4..d0df7e607d3 100644
--- a/src/pl/plpython/sql/plpython_error.sql
+++ b/src/pl/plpython/sql/plpython_error.sql
@@ -298,3 +298,33 @@ 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;
+
+/* setting a custom sqlstate should be handled
+ */
+CREATE FUNCTION plpy_raise_spiexception_override() RETURNS void AS $$
+exc = plpy.spiexceptions.DivisionByZero()
+exc.sqlstate = 'SILLY'
+raise exc
+$$ LANGUAGE plpythonu;
+
+DO $$
+BEGIN
+	SELECT plpy_raise_spiexception_override();
+EXCEPTION WHEN SQLSTATE 'SILLY' THEN
+	-- NOOP
+END
+$$ LANGUAGE plpgsql;