1 files changed, 62 insertions, 3 deletions

diff --git a/Modules/_sqlite/connection.c b/Modules/_sqlite/connection.c
index e4969e3..76d635a 100644
--- a/Modules/_sqlite/connection.c
+++ b/Modules/_sqlite/connection.c
@@ -673,7 +673,6 @@ void _pysqlite_final_callback(sqlite3_context* context)
 {
     PyObject* function_result = NULL;
     PyObject** aggregate_instance;
-    PyObject* aggregate_class;
 
 #ifdef WITH_THREAD
     PyGILState_STATE threadstate;
@@ -681,8 +680,6 @@ void _pysqlite_final_callback(sqlite3_context* context)
     threadstate = PyGILState_Ensure();
 #endif
 
-    aggregate_class = (PyObject*)sqlite3_user_data(context);
-
     aggregate_instance = (PyObject**)sqlite3_aggregate_context(context, sizeof(PyObject*));
     if (!*aggregate_instance) {
         /* this branch is executed if there was an exception in the aggregate's
@@ -907,6 +904,38 @@ static int _progress_handler(void* user_arg)
     return rc;
 }
 
+static void _trace_callback(void* user_arg, const char* statement_string)
+{
+    PyObject *py_statement = NULL;
+    PyObject *ret = NULL;
+
+#ifdef WITH_THREAD
+    PyGILState_STATE gilstate;
+
+    gilstate = PyGILState_Ensure();
+#endif
+    py_statement = PyUnicode_DecodeUTF8(statement_string,
+            strlen(statement_string), "replace");
+    if (py_statement) {
+        ret = PyObject_CallFunctionObjArgs((PyObject*)user_arg, py_statement, NULL);
+        Py_DECREF(py_statement);
+    }
+
+    if (ret) {
+        Py_DECREF(ret);
+    } else {
+        if (_enable_callback_tracebacks) {
+            PyErr_Print();
+        } else {
+            PyErr_Clear();
+        }
+    }
+
+#ifdef WITH_THREAD
+    PyGILState_Release(gilstate);
+#endif
+}
+
 static PyObject* pysqlite_connection_set_authorizer(pysqlite_Connection* self, PyObject* args, PyObject* kwargs)
 {
     PyObject* authorizer_cb;
@@ -966,6 +995,34 @@ static PyObject* pysqlite_connection_set_progress_handler(pysqlite_Connection* s
     return Py_None;
 }
 
+static PyObject* pysqlite_connection_set_trace_callback(pysqlite_Connection* self, PyObject* args, PyObject* kwargs)
+{
+    PyObject* trace_callback;
+
+    static char *kwlist[] = { "trace_callback", NULL };
+
+    if (!pysqlite_check_thread(self) || !pysqlite_check_connection(self)) {
+        return NULL;
+    }
+
+    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O:set_trace_callback",
+                                      kwlist, &trace_callback)) {
+        return NULL;
+    }
+
+    if (trace_callback == Py_None) {
+        /* None clears the trace callback previously set */
+        sqlite3_trace(self->db, 0, (void*)0);
+    } else {
+        if (PyDict_SetItem(self->function_pinboard, trace_callback, Py_None) == -1)
+            return NULL;
+        sqlite3_trace(self->db, _trace_callback, trace_callback);
+    }
+
+    Py_INCREF(Py_None);
+    return Py_None;
+}
+
 #ifdef HAVE_LOAD_EXTENSION
 static PyObject* pysqlite_enable_load_extension(pysqlite_Connection* self, PyObject* args)
 {
@@ -1519,6 +1576,8 @@ static PyMethodDef connection_methods[] = {
     #endif
     {"set_progress_handler", (PyCFunction)pysqlite_connection_set_progress_handler, METH_VARARGS|METH_KEYWORDS,
         PyDoc_STR("Sets progress handler callback. Non-standard.")},
+    {"set_trace_callback", (PyCFunction)pysqlite_connection_set_trace_callback, METH_VARARGS|METH_KEYWORDS,
+        PyDoc_STR("Sets a trace callback called for each SQL statement (passed as unicode). Non-standard.")},
     {"execute", (PyCFunction)pysqlite_connection_execute, METH_VARARGS,
         PyDoc_STR("Executes a SQL statement. Non-standard.")},
     {"executemany", (PyCFunction)pysqlite_connection_executemany, METH_VARARGS,