From e5f01bddd8f4e53e3f6235d4f0ddc11aa97b973f Mon Sep 17 00:00:00 2001
From: Charles Leifer <coleifer@gmail.com>
Date: Fri, 15 Sep 2023 07:59:45 -0500
Subject: [PATCH] Add trace_v2 support.

---
 src/blob.c       |  2 +-
 src/connection.c | 67 ++++++++++++++++++++++++++++++++++++++++++++----
 2 files changed, 63 insertions(+), 6 deletions(-)

diff --git a/src/blob.c b/src/blob.c
index c3dabac..70b327a 100644
--- a/src/blob.c
+++ b/src/blob.c
@@ -474,7 +474,7 @@ static PyObject * pysqlite_blob_subscript(pysqlite_Blob *self, PyObject *item)
 
 static int pysqlite_blob_ass_subscript(pysqlite_Blob *self, PyObject *item, PyObject *value)
 {
-    int rc;
+    int rc = 0;
 
     if (!pysqlite_check_blob(self)) {
         return -1;
diff --git a/src/connection.c b/src/connection.c
index b6154c3..99c8e26 100644
--- a/src/connection.c
+++ b/src/connection.c
@@ -44,6 +44,10 @@
 #define HAVE_BACKUP_API
 #endif
 
+#if SQLITE_VERSION_NUMBER >= 3014002
+#define HAVE_TRACE_V2
+#endif
+
 #if SQLITE_VERSION_NUMBER >= 3025000
 #define HAVE_WINDOW_FUNCTION
 #endif
@@ -1192,16 +1196,60 @@ static int _progress_handler(void* user_arg)
     return rc;
 }
 
+#ifdef HAVE_TRACE_V2
+static int _trace_callback(unsigned int type, void *ctx, void *stmt, void *sql)
+{
+    if (type != SQLITE_TRACE_STMT) {
+        return 0;
+    }
+
+    PyGILState_STATE gilstate = PyGILState_Ensure();
+    PyObject *py_statement = NULL;
+    const char *expanded_sql = sqlite3_expanded_sql((sqlite3_stmt *)stmt);
+    if (expanded_sql == NULL) {
+        sqlite3 *db = sqlite3_db_handle((sqlite3_stmt *)stmt);
+        if (sqlite3_errcode(db) == SQLITE_NOMEM) {
+            (void)PyErr_NoMemory();
+            goto exit;
+        }
+        PyErr_SetString(pysqlite_DataError, "Expanded SQL string exceeds the maximum string length");
+        if (_pysqlite_enable_callback_tracebacks) {
+            PyErr_Print();
+        } else {
+            PyErr_Clear();
+        }
+
+        py_statement = PyUnicode_FromString((const char *)sql);
+    } else {
+        py_statement = PyUnicode_FromString(expanded_sql);
+        sqlite3_free((void *)expanded_sql);
+    }
+
+    if (py_statement) {
+        PyObject *ret = PyObject_CallFunctionObjArgs((PyObject*)ctx, py_statement, NULL);
+        Py_DECREF(py_statement);
+        Py_XDECREF(ret);
+    }
+
+    if (PyErr_Occurred()) {
+        if (_pysqlite_enable_callback_tracebacks) {
+            PyErr_Print();
+        } else {
+            PyErr_Clear();
+        }
+    }
+exit:
+    PyGILState_Release(gilstate);
+    return 0;
+}
+#else
 static void _trace_callback(void* user_arg, const char* statement_string)
 {
     PyObject *py_statement = NULL;
     PyObject *ret = NULL;
 
-    PyGILState_STATE gilstate;
-
-    gilstate = PyGILState_Ensure();
-    py_statement = PyUnicode_DecodeUTF8(statement_string,
-            strlen(statement_string), "replace");
+    PyGILState_STATE gilstate = PyGILState_Ensure();
+    py_statement = PyUnicode_FromString(statement_string);
     if (py_statement) {
         ret = PyObject_CallFunctionObjArgs((PyObject*)user_arg, py_statement, NULL);
         Py_DECREF(py_statement);
@@ -1219,6 +1267,7 @@ static void _trace_callback(void* user_arg, const char* statement_string)
 
     PyGILState_Release(gilstate);
 }
+#endif
 
 static PyObject* pysqlite_connection_set_authorizer(pysqlite_Connection* self, PyObject* args, PyObject* kwargs)
 {
@@ -1300,10 +1349,18 @@ static PyObject* pysqlite_connection_set_trace_callback(pysqlite_Connection* sel
 
     if (trace_callback == Py_None) {
         /* None clears the trace callback previously set */
+#ifdef HAVE_TRACE_V2
+        sqlite3_trace_v2(self->db, SQLITE_TRACE_STMT, NULL, (void*)0);
+#else
         sqlite3_trace(self->db, 0, (void*)0);
+#endif
         Py_XSETREF(self->function_pinboard_trace_callback, NULL);
     } else {
+#ifdef HAVE_TRACE_V2
+        sqlite3_trace_v2(self->db, SQLITE_TRACE_STMT, _trace_callback, trace_callback);
+#else
         sqlite3_trace(self->db, _trace_callback, trace_callback);
+#endif
         Py_INCREF(trace_callback);
         Py_XSETREF(self->function_pinboard_trace_callback, trace_callback);
     }