>From 633fe3cf5872f57fcc33c5462048c0cfea81d907 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jan=20Urba=C5=84ski?= <wulczer@wulczer.org>
Date: Wed, 31 Oct 2012 10:27:59 +0100
Subject: [PATCH] Set SQL error code for SPI errors raised from userland
 Python code.

If a PL/Python function was terminated by a SPIError raised from user
code, the Postgres-level error code was always set to internal_error.

This makes it pay attention to the exception's sqlstate attribute and
set the error code accordingly.
---
 src/pl/plpython/expected/plpython_error.out   |   26 ++++++++++++++++
 src/pl/plpython/expected/plpython_error_0.out |   26 ++++++++++++++++
 src/pl/plpython/plpy_elog.c                   |   40 +++++++++++++++++++++----
 src/pl/plpython/sql/plpython_error.sql        |   30 +++++++++++++++++++
 4 files changed, 117 insertions(+), 5 deletions(-)

diff --git a/src/pl/plpython/expected/plpython_error.out b/src/pl/plpython/expected/plpython_error.out
index e1ec9c2..be2ec97 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 6f74a50..39c63c5 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 c375ac0..e3044e7 100644
--- a/src/pl/plpython/plpy_elog.c
+++ b/src/pl/plpython/plpy_elog.c
@@ -336,6 +336,29 @@ PLy_traceback(char **xmsg, char **tbmsg, int *tb_depth)
 	Py_DECREF(e);
 }
 
+
+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 +368,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 502bbec..d0df7e6 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;
-- 
1.7.10.4

