Skip to content

Commit 36ea504

Browse files
committed
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.
1 parent 6f07252 commit 36ea504

File tree

3 files changed

+106
-32
lines changed

3 files changed

+106
-32
lines changed

Modules/_sqlite/connection.c

Lines changed: 37 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@
3838
#include "pycore_pyerrors.h" // _PyErr_ChainExceptions1()
3939
#include "pycore_pylifecycle.h" // _Py_IsInterpreterFinalizing()
4040
#include "pycore_unicodeobject.h" // _PyUnicode_AsUTF8NoNUL
41+
#include "pycore_weakref.h"
4142

4243
#include <stdbool.h>
4344

@@ -283,10 +284,17 @@ pysqlite_connection_init_impl(pysqlite_Connection *self, PyObject *database,
283284
goto error;
284285
}
285286

286-
/* Create lists of weak references to blobs */
287+
/* Create lists of weak references to cursors and blobs */
288+
PyObject *cursors = PyList_New(0);
289+
if (cursors == NULL) {
290+
Py_DECREF(statement_cache);
291+
goto error;
292+
}
293+
287294
PyObject *blobs = PyList_New(0);
288295
if (blobs == NULL) {
289296
Py_DECREF(statement_cache);
297+
Py_DECREF(cursors);
290298
goto error;
291299
}
292300

@@ -299,11 +307,11 @@ pysqlite_connection_init_impl(pysqlite_Connection *self, PyObject *database,
299307
self->check_same_thread = check_same_thread;
300308
self->thread_ident = PyThread_get_thread_ident();
301309
self->statement_cache = statement_cache;
310+
self->cursors = cursors;
302311
self->blobs = blobs;
312+
self->close_attempted_in_callback = 0;
303313
self->row_factory = Py_NewRef(Py_None);
304314
self->text_factory = Py_NewRef(&PyUnicode_Type);
305-
self->in_callback = 0;
306-
self->close_attempted_in_callback = 0;
307315
self->trace_ctx = NULL;
308316
self->progress_ctx = NULL;
309317
self->authorizer_ctx = NULL;
@@ -383,6 +391,7 @@ connection_traverse(PyObject *op, visitproc visit, void *arg)
383391
pysqlite_Connection *self = _pysqlite_Connection_CAST(op);
384392
Py_VISIT(Py_TYPE(self));
385393
Py_VISIT(self->statement_cache);
394+
Py_VISIT(self->cursors);
386395
Py_VISIT(self->blobs);
387396
Py_VISIT(self->row_factory);
388397
Py_VISIT(self->text_factory);
@@ -407,6 +416,7 @@ connection_clear(PyObject *op)
407416
{
408417
pysqlite_Connection *self = _pysqlite_Connection_CAST(op);
409418
Py_CLEAR(self->statement_cache);
419+
Py_CLEAR(self->cursors);
410420
Py_CLEAR(self->blobs);
411421
Py_CLEAR(self->row_factory);
412422
Py_CLEAR(self->text_factory);
@@ -657,14 +667,30 @@ pysqlite_connection_close_impl(pysqlite_Connection *self)
657667
return NULL;
658668
}
659669

660-
if (self->in_callback > 0) {
661-
self->close_attempted_in_callback = 1;
662-
PyTypeObject *tp = Py_TYPE(self);
663-
pysqlite_state *state = pysqlite_get_state_by_type(tp);
664-
PyErr_SetString(state->ProgrammingError,
665-
"Cannot close the database connection "
666-
"from within a callback function.");
667-
return NULL;
670+
/* Check if any cursor is locked (actively executing a query);
671+
* closing during a callback is illegal per the SQLite C API docs. */
672+
assert(PyList_CheckExact(self->cursors));
673+
Py_ssize_t n = PyList_GET_SIZE(self->cursors);
674+
for (Py_ssize_t i = 0; i < n; i++) {
675+
PyObject *weakref = PyList_GET_ITEM(self->cursors, i);
676+
if (_PyWeakref_IsDead(weakref)) {
677+
continue;
678+
}
679+
PyObject *obj;
680+
if (!PyWeakref_GetRef(weakref, &obj)) {
681+
continue;
682+
}
683+
int locked = ((pysqlite_Cursor *)obj)->locked;
684+
Py_DECREF(obj);
685+
if (locked) {
686+
self->close_attempted_in_callback = 1;
687+
PyTypeObject *tp = Py_TYPE(self);
688+
pysqlite_state *state = pysqlite_get_state_by_type(tp);
689+
PyErr_SetString(state->ProgrammingError,
690+
"Cannot close the database connection "
691+
"from within a callback function.");
692+
return NULL;
693+
}
668694
}
669695

670696
pysqlite_close_all_blobs(self);

Modules/_sqlite/connection.h

Lines changed: 5 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -65,21 +65,18 @@ typedef struct
6565

6666
int initialized;
6767

68-
/* set to 1 while a SQLite callback (UDF, aggregate, progress handler,
69-
* etc.) is executing; used to prevent closing the connection from
70-
* within a callback, which is illegal per the SQLite C API docs */
71-
int in_callback;
72-
73-
/* set to 1 when close() is attempted during a callback; checked after
74-
* stmt_step() returns to raise the appropriate ProgrammingError */
68+
/* set to 1 when close() is attempted while a cursor is locked (actively
69+
* executing); checked after stmt_step() returns to raise the appropriate
70+
* ProgrammingError */
7571
int close_attempted_in_callback;
7672

7773
/* thread identification of the thread the connection was created in */
7874
unsigned long thread_ident;
7975

8076
PyObject *statement_cache;
8177

82-
/* Lists of weak references to blobs used within this connection */
78+
/* Lists of weak references to cursors and blobs used within this connection */
79+
PyObject *cursors;
8380
PyObject *blobs;
8481

8582
PyObject* row_factory;

Modules/_sqlite/cursor.c

Lines changed: 64 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,31 @@ check_cursor_locked(pysqlite_Cursor *cur)
5858
return 1;
5959
}
6060

61+
/*
62+
* Check if any cursor other than 'exclude' is locked on this connection.
63+
* Used to determine if we're in a nested callback scenario.
64+
*/
65+
static int
66+
any_other_cursor_locked(pysqlite_Connection *conn, pysqlite_Cursor *exclude)
67+
{
68+
assert(PyList_CheckExact(conn->cursors));
69+
Py_ssize_t n = PyList_GET_SIZE(conn->cursors);
70+
for (Py_ssize_t i = 0; i < n; i++) {
71+
PyObject *weakref = PyList_GET_ITEM(conn->cursors, i);
72+
PyObject *obj;
73+
if (!PyWeakref_GetRef(weakref, &obj)) {
74+
continue;
75+
}
76+
pysqlite_Cursor *cursor = (pysqlite_Cursor *)obj;
77+
int locked = cursor->locked;
78+
Py_DECREF(obj);
79+
if (locked && cursor != exclude) {
80+
return 1;
81+
}
82+
}
83+
return 0;
84+
}
85+
6186
static pysqlite_state *
6287
get_module_state_by_cursor(pysqlite_Cursor *cursor)
6388
{
@@ -99,6 +124,28 @@ class _sqlite3.Cursor "pysqlite_Cursor *" "clinic_state()->CursorType"
99124
[clinic start generated code]*/
100125
/*[clinic end generated code: output=da39a3ee5e6b4b0d input=3c5b8115c5cf30f1]*/
101126

127+
/*
128+
* Registers a cursor with the connection.
129+
*
130+
* 0 => error; 1 => ok
131+
*/
132+
static int
133+
register_cursor(pysqlite_Connection *connection, PyObject *cursor)
134+
{
135+
PyObject *weakref = PyWeakref_NewRef(cursor, NULL);
136+
if (weakref == NULL) {
137+
return 0;
138+
}
139+
140+
if (PyList_Append(connection->cursors, weakref) < 0) {
141+
Py_DECREF(weakref);
142+
return 0;
143+
}
144+
145+
Py_DECREF(weakref);
146+
return 1;
147+
}
148+
102149
/*[clinic input]
103150
_sqlite3.Cursor.__init__ as pysqlite_cursor_init
104151
@@ -138,6 +185,10 @@ pysqlite_cursor_init_impl(pysqlite_Cursor *self,
138185
return -1;
139186
}
140187

188+
if (!register_cursor(connection, (PyObject *)self)) {
189+
return -1;
190+
}
191+
141192
self->initialized = 1;
142193

143194
return 0;
@@ -905,13 +956,12 @@ _pysqlite_query_execute(pysqlite_Cursor* self, int multiple, PyObject* operation
905956
goto error;
906957
}
907958

908-
self->connection->in_callback++;
909959
rc = stmt_step(self->statement->st);
910-
self->connection->in_callback--;
911-
if (self->connection->close_attempted_in_callback
912-
&& self->connection->in_callback == 0)
913-
{
914-
self->connection->close_attempted_in_callback = 0;
960+
if (self->connection->close_attempted_in_callback) {
961+
/* Only clear the flag if no other cursor is locked (outermost) */
962+
if (!any_other_cursor_locked(self->connection, self)) {
963+
self->connection->close_attempted_in_callback = 0;
964+
}
915965
PyErr_Clear();
916966
PyErr_SetString(state->ProgrammingError,
917967
"Cannot close the database connection "
@@ -1163,18 +1213,19 @@ pysqlite_cursor_iternext(PyObject *op)
11631213

11641214
self->locked = 1; // GH-80254: Prevent recursive use of cursors.
11651215
PyObject *row = _pysqlite_fetch_one_row(self);
1166-
self->locked = 0;
11671216
if (row == NULL) {
1217+
self->locked = 0;
11681218
return NULL;
11691219
}
1170-
self->connection->in_callback++;
11711220
int rc = stmt_step(stmt);
1172-
self->connection->in_callback--;
1173-
if (self->connection->close_attempted_in_callback
1174-
&& self->connection->in_callback == 0)
1175-
{
1176-
self->connection->close_attempted_in_callback = 0;
1221+
self->locked = 0;
1222+
if (self->connection->close_attempted_in_callback) {
1223+
/* Only clear the flag if no other cursor is locked (outermost) */
1224+
if (!any_other_cursor_locked(self->connection, self)) {
1225+
self->connection->close_attempted_in_callback = 0;
1226+
}
11771227
Py_DECREF(row);
1228+
stmt_reset(self->statement);
11781229
Py_CLEAR(self->statement);
11791230
PyErr_Clear();
11801231
PyErr_SetString(self->connection->state->ProgrammingError,

0 commit comments

Comments
 (0)