gh-145040: Use cursor->locked to prevent closing sqlite3 connection in callback

Replace the in_callback counter with cursor->locked checks.
Restore the cursor weakref list (partially rolling back 74c1f41)
so that close() can iterate cursors and detect locked ones.

Author: raminfp

Modules/_sqlite/connection.c

@@ -38,6 +38,7 @@
 #include "pycore_pyerrors.h"      // _PyErr_ChainExceptions1()
 #include "pycore_pylifecycle.h"   // _Py_IsInterpreterFinalizing()
 #include "pycore_unicodeobject.h" // _PyUnicode_AsUTF8NoNUL
+#include "pycore_weakref.h"
 
 #include <stdbool.h>
 
@@ -283,10 +284,17 @@ pysqlite_connection_init_impl(pysqlite_Connection *self, PyObject *database,
         goto error;
     }
 
-    /* Create lists of weak references to blobs */
+    /* Create lists of weak references to cursors and blobs */
+    PyObject *cursors = PyList_New(0);
+    if (cursors == NULL) {
+        Py_DECREF(statement_cache);
+        goto error;
+    }
+
     PyObject *blobs = PyList_New(0);
     if (blobs == NULL) {
         Py_DECREF(statement_cache);
+        Py_DECREF(cursors);
         goto error;
     }
 
@@ -299,11 +307,11 @@ pysqlite_connection_init_impl(pysqlite_Connection *self, PyObject *database,
     self->check_same_thread = check_same_thread;
     self->thread_ident = PyThread_get_thread_ident();
     self->statement_cache = statement_cache;
+    self->cursors = cursors;
     self->blobs = blobs;
+    self->close_attempted_in_callback = 0;
     self->row_factory = Py_NewRef(Py_None);
     self->text_factory = Py_NewRef(&PyUnicode_Type);
-    self->in_callback = 0;
-    self->close_attempted_in_callback = 0;
     self->trace_ctx = NULL;
     self->progress_ctx = NULL;
     self->authorizer_ctx = NULL;
@@ -383,6 +391,7 @@ connection_traverse(PyObject *op, visitproc visit, void *arg)
     pysqlite_Connection *self = _pysqlite_Connection_CAST(op);
     Py_VISIT(Py_TYPE(self));
     Py_VISIT(self->statement_cache);
+    Py_VISIT(self->cursors);
     Py_VISIT(self->blobs);
     Py_VISIT(self->row_factory);
     Py_VISIT(self->text_factory);
@@ -407,6 +416,7 @@ connection_clear(PyObject *op)
 {
     pysqlite_Connection *self = _pysqlite_Connection_CAST(op);
     Py_CLEAR(self->statement_cache);
+    Py_CLEAR(self->cursors);
     Py_CLEAR(self->blobs);
     Py_CLEAR(self->row_factory);
     Py_CLEAR(self->text_factory);
@@ -657,14 +667,30 @@ pysqlite_connection_close_impl(pysqlite_Connection *self)
         return NULL;
     }
 
-    if (self->in_callback > 0) {
-        self->close_attempted_in_callback = 1;
-        PyTypeObject *tp = Py_TYPE(self);
-        pysqlite_state *state = pysqlite_get_state_by_type(tp);
-        PyErr_SetString(state->ProgrammingError,
-                        "Cannot close the database connection "
-                        "from within a callback function.");
-        return NULL;
+    /* Check if any cursor is locked (actively executing a query);
+     * closing during a callback is illegal per the SQLite C API docs. */
+    assert(PyList_CheckExact(self->cursors));
+    Py_ssize_t n = PyList_GET_SIZE(self->cursors);
+    for (Py_ssize_t i = 0; i < n; i++) {
+        PyObject *weakref = PyList_GET_ITEM(self->cursors, i);
+        if (_PyWeakref_IsDead(weakref)) {
+            continue;
+        }
+        PyObject *obj;
+        if (!PyWeakref_GetRef(weakref, &obj)) {
+            continue;
+        }
+        int locked = ((pysqlite_Cursor *)obj)->locked;
+        Py_DECREF(obj);
+        if (locked) {
+            self->close_attempted_in_callback = 1;
+            PyTypeObject *tp = Py_TYPE(self);
+            pysqlite_state *state = pysqlite_get_state_by_type(tp);
+            PyErr_SetString(state->ProgrammingError,
+                            "Cannot close the database connection "
+                            "from within a callback function.");
+            return NULL;
+        }
     }
 
     pysqlite_close_all_blobs(self);

Modules/_sqlite/connection.h

@@ -65,21 +65,18 @@ typedef struct
 
     int initialized;
 
-    /* set to 1 while a SQLite callback (UDF, aggregate, progress handler,
-     * etc.) is executing; used to prevent closing the connection from
-     * within a callback, which is illegal per the SQLite C API docs */
-    int in_callback;
-
-    /* set to 1 when close() is attempted during a callback; checked after
-     * stmt_step() returns to raise the appropriate ProgrammingError */
+    /* set to 1 when close() is attempted while a cursor is locked (actively
+     * executing); checked after stmt_step() returns to raise the appropriate
+     * ProgrammingError */
     int close_attempted_in_callback;
 
     /* thread identification of the thread the connection was created in */
     unsigned long thread_ident;
 
     PyObject *statement_cache;
 
-    /* Lists of weak references to blobs used within this connection */
+    /* Lists of weak references to cursors and blobs used within this connection */
+    PyObject *cursors;
     PyObject *blobs;
 
     PyObject* row_factory;

Modules/_sqlite/cursor.c

@@ -58,6 +58,31 @@ check_cursor_locked(pysqlite_Cursor *cur)
     return 1;
 }
 
+/*
+ * Check if any cursor other than 'exclude' is locked on this connection.
+ * Used to determine if we're in a nested callback scenario.
+ */
+static int
+any_other_cursor_locked(pysqlite_Connection *conn, pysqlite_Cursor *exclude)
+{
+    assert(PyList_CheckExact(conn->cursors));
+    Py_ssize_t n = PyList_GET_SIZE(conn->cursors);
+    for (Py_ssize_t i = 0; i < n; i++) {
+        PyObject *weakref = PyList_GET_ITEM(conn->cursors, i);
+        PyObject *obj;
+        if (!PyWeakref_GetRef(weakref, &obj)) {
+            continue;
+        }
+        pysqlite_Cursor *cursor = (pysqlite_Cursor *)obj;
+        int locked = cursor->locked;
+        Py_DECREF(obj);
+        if (locked && cursor != exclude) {
+            return 1;
+        }
+    }
+    return 0;
+}
+
 static pysqlite_state *
 get_module_state_by_cursor(pysqlite_Cursor *cursor)
 {
@@ -99,6 +124,28 @@ class _sqlite3.Cursor "pysqlite_Cursor *" "clinic_state()->CursorType"
 [clinic start generated code]*/
 /*[clinic end generated code: output=da39a3ee5e6b4b0d input=3c5b8115c5cf30f1]*/
 
+/*
+ * Registers a cursor with the connection.
+ *
+ * 0 => error; 1 => ok
+ */
+static int
+register_cursor(pysqlite_Connection *connection, PyObject *cursor)
+{
+    PyObject *weakref = PyWeakref_NewRef(cursor, NULL);
+    if (weakref == NULL) {
+        return 0;
+    }
+
+    if (PyList_Append(connection->cursors, weakref) < 0) {
+        Py_DECREF(weakref);
+        return 0;
+    }
+
+    Py_DECREF(weakref);
+    return 1;
+}
+
 /*[clinic input]
 _sqlite3.Cursor.__init__ as pysqlite_cursor_init
 
@@ -138,6 +185,10 @@ pysqlite_cursor_init_impl(pysqlite_Cursor *self,
         return -1;
     }
 
+    if (!register_cursor(connection, (PyObject *)self)) {
+        return -1;
+    }
+
     self->initialized = 1;
 
     return 0;
@@ -905,13 +956,12 @@ _pysqlite_query_execute(pysqlite_Cursor* self, int multiple, PyObject* operation
             goto error;
         }
 
-        self->connection->in_callback++;
         rc = stmt_step(self->statement->st);
-        self->connection->in_callback--;
-        if (self->connection->close_attempted_in_callback
-            && self->connection->in_callback == 0)
-        {
-            self->connection->close_attempted_in_callback = 0;
+        if (self->connection->close_attempted_in_callback) {
+            /* Only clear the flag if no other cursor is locked (outermost) */
+            if (!any_other_cursor_locked(self->connection, self)) {
+                self->connection->close_attempted_in_callback = 0;
+            }
             PyErr_Clear();
             PyErr_SetString(state->ProgrammingError,
                             "Cannot close the database connection "
@@ -1163,18 +1213,19 @@ pysqlite_cursor_iternext(PyObject *op)
 
     self->locked = 1;  // GH-80254: Prevent recursive use of cursors.
     PyObject *row = _pysqlite_fetch_one_row(self);
-    self->locked = 0;
     if (row == NULL) {
+        self->locked = 0;
         return NULL;
     }
-    self->connection->in_callback++;
     int rc = stmt_step(stmt);
-    self->connection->in_callback--;
-    if (self->connection->close_attempted_in_callback
-        && self->connection->in_callback == 0)
-    {
-        self->connection->close_attempted_in_callback = 0;
+    self->locked = 0;
+    if (self->connection->close_attempted_in_callback) {
+        /* Only clear the flag if no other cursor is locked (outermost) */
+        if (!any_other_cursor_locked(self->connection, self)) {
+            self->connection->close_attempted_in_callback = 0;
+        }
         Py_DECREF(row);
+        stmt_reset(self->statement);
         Py_CLEAR(self->statement);
         PyErr_Clear();
         PyErr_SetString(self->connection->state->ProgrammingError,