gh-145040: Prevent closing sqlite3 connection from within a callback

Instead of detecting a closed connection after the damage has been done,
prevent Connection.close() from succeeding while a SQLite callback is
executing. This aligns with the SQLite C API docs, which state that
applications must not close the database connection from within a
callback.

Add an in_callback counter to the connection object, incremented before
stmt_step() and decremented after. If close() is called while the
counter is positive, ProgrammingError is raised and the database
connection remains open. A counter (rather than a boolean flag) is used
to correctly handle nested callbacks.

Also convert test docstrings to comments per reviewer feedback, and add
a test for the nested callback scenario.

Author: raminfp

Lib/test/test_sqlite3/test_userfunctions.py

@@ -588,7 +588,7 @@ def value(self): return 1 << 65
                                self.cur.execute, self.query % "err_val_ret")
 
     def test_close_conn_in_window_func_value(self):
-        """gh-145040: closing connection in window function value() callback."""
+        # gh-145040: closing connection in window function value() callback.
         con = sqlite.connect(":memory:", autocommit=True)
         con.execute("CREATE TABLE t(x INTEGER)")
         con.executemany("INSERT INTO t VALUES(?)",
@@ -751,7 +751,7 @@ def test_agg_keyword_args(self):
             self.con.create_aggregate("test", 1, aggregate_class=AggrText)
 
     def test_aggr_close_conn_in_step(self):
-        """Connection.close() in an aggregate step callback must not crash."""
+        # Connection.close() in an aggregate step callback must not crash.
         con = sqlite.connect(":memory:", autocommit=True)
         cur = con.cursor()
         cur.execute("CREATE TABLE t(x INTEGER)")
@@ -774,8 +774,31 @@ def finalize(self):
         with self.assertRaisesRegex(sqlite.ProgrammingError, msg):
             con.execute("SELECT agg_close(x) FROM t")
 
+    def test_close_conn_in_nested_callback(self):
+        # gh-145040: close() must be prevented even in nested callbacks.
+        con = sqlite.connect(":memory:", autocommit=True)
+        con.execute("CREATE TABLE t(x INTEGER)")
+        for i in range(5):
+            con.execute("INSERT INTO t VALUES(?)", (i,))
+
+        def outer_func(x):
+            con.close()
+            return x
+
+        def inner_func(x):
+            return x * 10
+
+        con.create_function("outer_func", 1, outer_func)
+        con.create_function("inner_func", 1, inner_func)
+        msg = "from within a callback"
+        with self.assertRaisesRegex(sqlite.ProgrammingError, msg):
+            con.execute("SELECT outer_func(inner_func(x)) FROM t")
+        # Connection must still be usable after the failed close attempt.
+        self.assertEqual(con.execute("SELECT 1").fetchone(), (1,))
+        con.close()
+
     def test_close_conn_in_udf_during_executemany(self):
-        """gh-145040: closing connection in UDF during executemany."""
+        # gh-145040: closing connection in UDF during executemany.
         con = sqlite.connect(":memory:", autocommit=True)
         con.execute("CREATE TABLE t(x)")
 
@@ -790,7 +813,7 @@ def close_conn(x):
                             [(i,) for i in range(10)])
 
     def test_close_conn_in_progress_handler_during_iternext(self):
-        """gh-145040: closing connection in progress handler during iteration."""
+        # gh-145040: closing connection in progress handler during iteration.
         con = sqlite.connect(":memory:", autocommit=True)
         con.execute("CREATE TABLE t(x)")
         con.executemany("INSERT INTO t VALUES(?)",
@@ -817,7 +840,7 @@ def close_progress():
             gc_collect()
 
     def test_close_conn_in_collation_callback(self):
-        """gh-145040: closing connection in collation callback."""
+        # gh-145040: closing connection in collation callback.
         con = sqlite.connect(":memory:", autocommit=True)
         con.execute("CREATE TABLE t(x TEXT)")
         con.executemany("INSERT INTO t VALUES(?)",

Modules/_sqlite/connection.c

@@ -302,6 +302,8 @@ pysqlite_connection_init_impl(pysqlite_Connection *self, PyObject *database,
     self->blobs = blobs;
     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;
@@ -655,6 +657,16 @@ 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;
+    }
+
     pysqlite_close_all_blobs(self);
     Py_CLEAR(self->statement_cache);
     if (connection_close(self) < 0) {

Modules/_sqlite/connection.h

@@ -65,6 +65,15 @@ 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 */
+    int close_attempted_in_callback;
+
     /* thread identification of the thread the connection was created in */
     unsigned long thread_ident;
 

Modules/_sqlite/cursor.c

@@ -905,8 +905,12 @@ _pysqlite_query_execute(pysqlite_Cursor* self, int multiple, PyObject* operation
             goto error;
         }
 
+        self->connection->in_callback++;
         rc = stmt_step(self->statement->st);
-        if (self->connection->db == NULL) {
+        self->connection->in_callback--;
+        if (self->connection->close_attempted_in_callback) {
+            self->connection->close_attempted_in_callback = 0;
+            PyErr_Clear();
             PyErr_SetString(state->ProgrammingError,
                             "Cannot close the database connection "
                             "from within a callback function.");
@@ -1161,10 +1165,14 @@ pysqlite_cursor_iternext(PyObject *op)
     if (row == NULL) {
         return NULL;
     }
+    self->connection->in_callback++;
     int rc = stmt_step(stmt);
-    if (self->connection->db == NULL) {
+    self->connection->in_callback--;
+    if (self->connection->close_attempted_in_callback) {
+        self->connection->close_attempted_in_callback = 0;
         Py_DECREF(row);
         Py_CLEAR(self->statement);
+        PyErr_Clear();
         PyErr_SetString(self->connection->state->ProgrammingError,
                         "Cannot close the database connection "
                         "from within a callback function.");