diff -r 44749e501982 -r 3df81c81c481 Lib/sqlite3/test/hooks.py --- a/Lib/sqlite3/test/hooks.py Thu Mar 24 22:38:25 2011 +0200 +++ b/Lib/sqlite3/test/hooks.py Sat Mar 26 22:43:26 2011 +0100 @@ -175,10 +175,56 @@ con.execute("select 1 union select 2 union select 3").fetchall() self.assertEqual(action, 0, "progress handler was not cleared") +class TraceCallbackTests(unittest.TestCase): + def CheckTraceCallbackUsed(self): + """ + Test that the trace callback is invoked once it is set. + """ + con = sqlite.connect(":memory:") + traced_statements = [] + def trace(statement): + traced_statements.append(statement) + con.set_trace_callback(trace) + con.execute("create table foo(a, b)") + self.assertTrue(traced_statements) + self.assertTrue([x for x in traced_statements if x.find("create table foo") != -1]) + + def CheckClearTraceCallback(self): + """ + Test that setting the trace callback to None clears the previously set callback. + """ + con = sqlite.connect(":memory:") + traced_statements = [] + def trace(statement): + traced_statements.append(statement) + con.set_trace_callback(trace) + con.set_trace_callback(None) + con.execute("create table foo(a, b)") + self.assertFalse(traced_statements, "trace callback was not cleared") + + def CheckUnicodeContent(self): + """ + Test that the statement can contain unicode literals. + """ + unicode_value = u'\xf6\xe4\xfc\xd6\xc4\xdc\xdf\u20ac' + con = sqlite.connect(":memory:") + traced_statements = [] + def trace(statement): + traced_statements.append(statement) + con.set_trace_callback(trace) + con.execute("create table foo(x)") + con.execute("insert into foo(x) values (?)", (unicode_value,)) + con.commit() + self.assertTrue([x for x in traced_statements if x.find(unicode_value) != -1], + "Unicode data garbled in trace callback") + + + def suite(): collation_suite = unittest.makeSuite(CollationTests, "Check") progress_suite = unittest.makeSuite(ProgressTests, "Check") - return unittest.TestSuite((collation_suite, progress_suite)) + trace_suite = unittest.makeSuite(TraceCallbackTests, "Check") + return unittest.TestSuite((collation_suite, progress_suite, trace_suite)) def test(): runner = unittest.TextTestRunner() diff -r 44749e501982 -r 3df81c81c481 Modules/_sqlite/connection.c --- a/Modules/_sqlite/connection.c Thu Mar 24 22:38:25 2011 +0200 +++ b/Modules/_sqlite/connection.c Sat Mar 26 22:43:26 2011 +0100 @@ -976,6 +976,38 @@ 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; @@ -1035,6 +1067,34 @@ 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 { + sqlite3_trace(self->db, _trace_callback, trace_callback); + if (PyDict_SetItem(self->function_pinboard, trace_callback, Py_None) == -1) + return NULL; + } + + Py_INCREF(Py_None); + return Py_None; +} + #ifdef HAVE_LOAD_EXTENSION static PyObject* pysqlite_enable_load_extension(pysqlite_Connection* self, PyObject* args) { @@ -1579,6 +1639,8 @@ #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,