Correctly add exceptions to the plpy module for Python 3

The way the exception types where added to the module was wrong for
Python 3.  Exception classes were not actually available from plpy.
Fix that by factoring out code that is responsible for defining new
Python exceptions and make it work with Python 3.  New regression test
makes sure the plpy module has the expected contents.

Jan Urbanśki, slightly revised by me

3 files changed, 54 insertions, 11 deletions

diff --git a/src/pl/plpython/expected/plpython_test.out b/src/pl/plpython/expected/plpython_test.out
index a229b18f448..32938295726 100644
--- a/src/pl/plpython/expected/plpython_test.out
+++ b/src/pl/plpython/expected/plpython_test.out
@@ -35,6 +35,19 @@ select argument_test_one(users, fname, lname) from users where lname = 'doe' ord
  willem doe => {fname: willem, lname: doe, userid: 3, username: w_doe}
 (3 rows)
 
+-- check module contents
+CREATE FUNCTION module_contents() RETURNS text AS
+$$
+contents = list(filter(lambda x: not x.startswith("__"), dir(plpy)))
+contents.sort()
+return ", ".join(contents)
+$$ LANGUAGE plpythonu;
+select module_contents();
+                                      module_contents                                      
+-------------------------------------------------------------------------------------------
+ Error, Fatal, SPIError, debug, error, execute, fatal, info, log, notice, prepare, warning
+(1 row)
+
 CREATE FUNCTION elog_test() RETURNS void
 AS $$
 plpy.debug('debug')
diff --git a/src/pl/plpython/plpython.c b/src/pl/plpython/plpython.c
index 370f4f7ea28..18a523c8b3a 100644
--- a/src/pl/plpython/plpython.c
+++ b/src/pl/plpython/plpython.c
@@ -3268,11 +3268,37 @@ PLy_spi_execute_fetch_result(SPITupleTable *tuptable, int rows, int status)
  * language handler and interpreter initialization
  */
 
+/*
+ * Add exceptions to the plpy module
+ */
+static void
+PLy_add_exceptions(PyObject *plpy)
+{
+	PLy_exc_error = PyErr_NewException("plpy.Error", NULL, NULL);
+	PLy_exc_fatal = PyErr_NewException("plpy.Fatal", NULL, NULL);
+	PLy_exc_spi_error = PyErr_NewException("plpy.SPIError", NULL, NULL);
+
+	Py_INCREF(PLy_exc_error);
+	PyModule_AddObject(plpy, "Error", PLy_exc_error);
+	Py_INCREF(PLy_exc_fatal);
+	PyModule_AddObject(plpy, "Fatal", PLy_exc_fatal);
+	Py_INCREF(PLy_exc_spi_error);
+	PyModule_AddObject(plpy, "SPIError", PLy_exc_spi_error);
+}
+
 #if PY_MAJOR_VERSION >= 3
 static PyMODINIT_FUNC
 PyInit_plpy(void)
 {
-	return PyModule_Create(&PLy_module);
+	PyObject   *m;
+
+	m = PyModule_Create(&PLy_module);
+	if (m == NULL)
+		return NULL;
+
+	PLy_add_exceptions(m);
+
+	return m;
 }
 #endif
 
@@ -3363,8 +3389,7 @@ PLy_init_plpy(void)
 	PyObject   *main_mod,
 			   *main_dict,
 			   *plpy_mod;
-	PyObject   *plpy,
-			   *plpy_dict;
+	PyObject   *plpy;
 
 	/*
 	 * initialize plpy module
@@ -3376,20 +3401,14 @@ PLy_init_plpy(void)
 
 #if PY_MAJOR_VERSION >= 3
 	plpy = PyModule_Create(&PLy_module);
+	/* for Python 3 we initialized the exceptions in PyInit_plpy */
 #else
 	plpy = Py_InitModule("plpy", PLy_methods);
+	PLy_add_exceptions(plpy);
 #endif
-	plpy_dict = PyModule_GetDict(plpy);
 
 	/* PyDict_SetItemString(plpy, "PlanType", (PyObject *) &PLy_PlanType); */
 
-	PLy_exc_error = PyErr_NewException("plpy.Error", NULL, NULL);
-	PLy_exc_fatal = PyErr_NewException("plpy.Fatal", NULL, NULL);
-	PLy_exc_spi_error = PyErr_NewException("plpy.SPIError", NULL, NULL);
-	PyDict_SetItemString(plpy_dict, "Error", PLy_exc_error);
-	PyDict_SetItemString(plpy_dict, "Fatal", PLy_exc_fatal);
-	PyDict_SetItemString(plpy_dict, "SPIError", PLy_exc_spi_error);
-
 	/*
 	 * initialize main module, and add plpy
 	 */
diff --git a/src/pl/plpython/sql/plpython_test.sql b/src/pl/plpython/sql/plpython_test.sql
index 7cae124d98b..915189847a7 100644
--- a/src/pl/plpython/sql/plpython_test.sql
+++ b/src/pl/plpython/sql/plpython_test.sql
@@ -26,6 +26,17 @@ return words'
 select argument_test_one(users, fname, lname) from users where lname = 'doe' order by 1;
 
 
+-- check module contents
+CREATE FUNCTION module_contents() RETURNS text AS
+$$
+contents = list(filter(lambda x: not x.startswith("__"), dir(plpy)))
+contents.sort()
+return ", ".join(contents)
+$$ LANGUAGE plpythonu;
+
+select module_contents();
+
+
 CREATE FUNCTION elog_test() RETURNS void
 AS $$
 plpy.debug('debug')