>From 8dcda806e6dd529c41196a5ae32c5f14aab3f6fe Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jan=20Urba=C5=84ski?= <wulczer@wulczer.org>
Date: Sun, 2 Jan 2011 11:21:30 +0100
Subject: [PATCH 3/7] 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 to make
sure the plpy module has the expected contents.
---
 src/pl/plpython/expected/plpython_test.out |   13 ++++++
 src/pl/plpython/plpython.c                 |   63 +++++++++++++++++++++++-----
 src/pl/plpython/sql/plpython_test.sql      |    9 ++++
 3 files changed, 74 insertions(+), 11 deletions(-)

diff --git a/src/pl/plpython/expected/plpython_test.out b/src/pl/plpython/expected/plpython_test.out
index df463c6..583a149 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 f205768..be39799 100644
--- a/src/pl/plpython/plpython.c
+++ b/src/pl/plpython/plpython.c
@@ -3226,6 +3226,46 @@ PLy_spi_execute_fetch_result(SPITupleTable *tuptable, int rows, int status)
 	return (PyObject *) result;
 }
 
+/*
+ * Exception support
+ */
+
+/*
+ * Add an exception to the module. The first parameter is either the actual
+ * module (for Python 3) or the module dictionary
+ */
+static void PLy_add_exception(PyObject *mod, const char *name, PyObject *exc)
+{
+#if PY_MAJOR_VERSION >= 3
+	Py_INCREF(exc);
+	PyModule_AddObject(mod, name, exc);
+#else
+	PyDict_SetItemString(mod, name, exc);
+#endif
+}
+
+/* Add exceptions to the plpy module */
+static void
+PLy_initialize_exceptions(PyObject *plpy)
+{
+	PyObject	*mod;
+
+#if PY_MAJOR_VERSION < 3
+	/* For Python <3 we add the exceptions to the module dictionary */
+	mod = PyModule_GetDict(plpy);
+#else
+	/* In Python 3 you add them directly into the module */
+	mod = plpy;
+#endif
+
+	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);
+
+	PLy_add_exception(mod, "Error", PLy_exc_error);
+	PLy_add_exception(mod, "Fatal", PLy_exc_fatal);
+	PLy_add_exception(mod, "SPIError", PLy_exc_spi_error);
+}
 
 /*
  * language handler and interpreter initialization
@@ -3235,7 +3275,15 @@ PLy_spi_execute_fetch_result(SPITupleTable *tuptable, int rows, int status)
 static PyMODINIT_FUNC
 PyInit_plpy(void)
 {
-	return PyModule_Create(&PLy_module);
+	PyObject	*m;
+
+	m = PyModule_Create(&PLy_module);
+	if (m == NULL)
+		return NULL;
+
+	PLy_initialize_exceptions(m);
+
+	return m;
 }
 #endif
 
@@ -3326,8 +3374,7 @@ PLy_init_plpy(void)
 	PyObject   *main_mod,
 			   *main_dict,
 			   *plpy_mod;
-	PyObject   *plpy,
-			   *plpy_dict;
+	PyObject   *plpy;
 
 	/*
 	 * initialize plpy module
@@ -3341,18 +3388,12 @@ PLy_init_plpy(void)
 	plpy = PyModule_Create(&PLy_module);
 #else
 	plpy = Py_InitModule("plpy", PLy_methods);
+	/* for Python 3 we initialized the exceptions in PyInit_plpy */
+	PLy_initialize_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 7cae124..298411a 100644
--- a/src/pl/plpython/sql/plpython_test.sql
+++ b/src/pl/plpython/sql/plpython_test.sql
@@ -25,6 +25,15 @@ 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 $$
-- 
1.7.2.3

