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.

Author: hlinnaka

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;

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;

src/pl/plpython/plpy_elog.c

@@ -336,6 +336,31 @@ PLy_traceback(char **xmsg, char **tbmsg, int *tb_depth)
 	Py_DECREF(e);
 }
 
+/*
+ * 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
  */
@@ -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);

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;