diff --git a/crates/stdlib/src/_sqlite3.rs b/crates/stdlib/src/_sqlite3.rs index e5953aef6ae..3eaa9533f71 100644 --- a/crates/stdlib/src/_sqlite3.rs +++ b/crates/stdlib/src/_sqlite3.rs @@ -49,6 +49,7 @@ mod _sqlite3 { atomic::{Ordering, PyAtomic, Radium}, hash::PyHash, lock::{PyMappedMutexGuard, PyMutex, PyMutexGuard}, + refcount::{try_defer_drop, with_deferred_drops}, static_cell, }; use rustpython_vm::{ @@ -515,7 +516,11 @@ mod _sqlite3 { let context = SqliteContext::from(context); let (cls, vm) = unsafe { (*context.user_data::()).retrieve() }; let args = unsafe { core::slice::from_raw_parts(argv, argc as usize) }; - let instance = context.aggregate_context::<*const PyObject>(); + let instance = context.aggregate_context(true); + if instance.is_null() { + unsafe { sqlite3_result_error_nomem(context.ctx) }; + return; + } if unsafe { (*instance).is_null() } { match cls.call((), vm) { Ok(obj) => unsafe { *instance = obj.into_raw().as_ptr() }, @@ -536,12 +541,22 @@ mod _sqlite3 { unsafe extern "C" fn finalize_callback(context: *mut sqlite3_context) { let context = SqliteContext::from(context); let (_, vm) = unsafe { (*context.user_data::()).retrieve() }; - let instance = context.aggregate_context::<*const PyObject>(); - let Some(instance) = (unsafe { (*instance).as_ref() }) else { + let instance = context.aggregate_context(false); + if instance.is_null() { + return; + } + // xFinal releases the reference acquired by xStep, even if the + // query or finalize() fails. Clear the slot before calling Python. + let object = unsafe { instance.replace(core::ptr::null()) }; + let Some(object) = NonNull::new(object.cast_mut()) else { return; }; + let instance = unsafe { PyObjectRef::from_raw(object) }; - Self::callback_result_from_method(context, instance, "finalize", vm); + Self::callback_result_from_method(context, &instance, "finalize", vm); + // __del__ and weakref callbacks must run after the caller releases + // its cursor, statement and connection locks. + try_defer_drop(move || drop(instance)); } unsafe extern "C" fn collation_callback( @@ -579,8 +594,10 @@ mod _sqlite3 { unsafe extern "C" fn value_callback(context: *mut sqlite3_context) { let context = SqliteContext::from(context); let (_, vm) = unsafe { (*context.user_data::()).retrieve() }; - let instance = context.aggregate_context::<*const PyObject>(); - let instance = unsafe { &**instance }; + let instance = context.aggregate_context(false); + let Some(instance) = (unsafe { instance.as_ref().and_then(|obj| obj.as_ref()) }) else { + return; + }; Self::callback_result_from_method(context, instance, "value", vm); } @@ -593,8 +610,10 @@ mod _sqlite3 { let context = SqliteContext::from(context); let (_, vm) = unsafe { (*context.user_data::()).retrieve() }; let args = unsafe { core::slice::from_raw_parts(argv, argc as usize) }; - let instance = context.aggregate_context::<*const PyObject>(); - let instance = unsafe { &**instance }; + let instance = context.aggregate_context(false); + let Some(instance) = (unsafe { instance.as_ref().and_then(|obj| obj.as_ref()) }) else { + return; + }; Self::call_method_with_args(context, instance, "inverse", args, vm); } @@ -995,44 +1014,50 @@ mod _sqlite3 { type Args = ConnectArgs; fn init(zelf: &Py, args: Self::Args, vm: &VirtualMachine) -> PyResult<()> { - let was_initialized = Radium::swap(&zelf.initialized, false, Ordering::AcqRel); + with_deferred_drops(|| { + let was_initialized = Radium::swap(&zelf.initialized, false, Ordering::AcqRel); - // Reset factories to their defaults, matching CPython's behavior. - zelf.reset_factories(vm); + // Reset factories to their defaults, matching CPython's behavior. + zelf.reset_factories(vm); - if was_initialized { - zelf.drop_db(vm)?; - } - - // Attempt to open the new database before mutating other state so failures leave - // the connection uninitialized (and subsequent operations raise ProgrammingError). - let db = Self::initialize_db(&args, vm)?; - - let ConnectArgs { - detect_types, - isolation_level, - check_same_thread, - autocommit, - .. - } = args; - - zelf.detect_types.store(detect_types, Ordering::Relaxed); - zelf.check_same_thread - .store(check_same_thread, Ordering::Relaxed); - *zelf.autocommit.lock() = autocommit; - *zelf.thread_ident.lock() = std::thread::current().id(); - let _ = unsafe { zelf.isolation_level.swap(isolation_level.0) }; + if was_initialized { + zelf.drop_db(vm)?; + } - let mut guard = zelf.db.lock(); - *guard = Some(db); - Radium::store(&zelf.initialized, true, Ordering::Release); - Ok(()) + // Attempt to open the new database before mutating other state so failures leave + // the connection uninitialized (and subsequent operations raise ProgrammingError). + let db = Self::initialize_db(&args, vm)?; + + let ConnectArgs { + detect_types, + isolation_level, + check_same_thread, + autocommit, + .. + } = args; + + zelf.detect_types.store(detect_types, Ordering::Relaxed); + zelf.check_same_thread + .store(check_same_thread, Ordering::Relaxed); + *zelf.autocommit.lock() = autocommit; + *zelf.thread_ident.lock() = std::thread::current().id(); + let _ = unsafe { zelf.isolation_level.swap(isolation_level.0) }; + + let mut guard = zelf.db.lock(); + *guard = Some(db); + Radium::store(&zelf.initialized, true, Ordering::Release); + Ok(()) + }) } } #[pyclass(with(Constructor, Callable, Initializer), flags(BASETYPE, HAS_WEAKREF))] impl Connection { fn drop_db(&self, vm: &VirtualMachine) -> PyResult<()> { + with_deferred_drops(|| Self::drop_db_inner(self, vm)) + } + + fn drop_db_inner(&self, vm: &VirtualMachine) -> PyResult<()> { let mut guard = self.db.lock(); let rollback_result = if let Some(db) = guard.as_ref() && *self.autocommit.lock() == AutocommitMode::Disabled @@ -1128,6 +1153,14 @@ mod _sqlite3 { zelf: PyRef, args: BlobOpenArgs, vm: &VirtualMachine, + ) -> PyResult> { + with_deferred_drops(|| Self::blobopen_inner(zelf, args, vm)) + } + + fn blobopen_inner( + zelf: PyRef, + args: BlobOpenArgs, + vm: &VirtualMachine, ) -> PyResult> { let table = args.table.to_cstring(vm)?; let column = args.column.to_cstring(vm)?; @@ -1170,6 +1203,10 @@ mod _sqlite3 { #[pymethod] fn commit(zelf: &Py, vm: &VirtualMachine) -> PyResult<()> { + with_deferred_drops(|| Self::commit_inner(zelf, vm)) + } + + fn commit_inner(zelf: &Py, vm: &VirtualMachine) -> PyResult<()> { let db = zelf.db_lock(vm)?; let mode = *zelf.autocommit.lock(); match mode { @@ -1184,6 +1221,10 @@ mod _sqlite3 { #[pymethod] fn rollback(zelf: &Py, vm: &VirtualMachine) -> PyResult<()> { + with_deferred_drops(|| Self::rollback_inner(zelf, vm)) + } + + fn rollback_inner(zelf: &Py, vm: &VirtualMachine) -> PyResult<()> { let db = zelf.db_lock(vm)?; let mode = *zelf.autocommit.lock(); match mode { @@ -1316,6 +1357,14 @@ mod _sqlite3 { zelf: &Py, args: CreateFunctionArgs, vm: &VirtualMachine, + ) -> PyResult<()> { + with_deferred_drops(|| Self::create_function_inner(zelf, args, vm)) + } + + fn create_function_inner( + zelf: &Py, + args: CreateFunctionArgs, + vm: &VirtualMachine, ) -> PyResult<()> { let name = args.name.to_cstring(vm)?; let flags = if args.deterministic { @@ -1357,6 +1406,14 @@ mod _sqlite3 { zelf: &Py, args: CreateAggregateArgs, vm: &VirtualMachine, + ) -> PyResult<()> { + with_deferred_drops(|| Self::create_aggregate_inner(zelf, args, vm)) + } + + fn create_aggregate_inner( + zelf: &Py, + args: CreateAggregateArgs, + vm: &VirtualMachine, ) -> PyResult<()> { let name = args.name.to_cstring(vm)?; let db = zelf.db_lock(vm)?; @@ -1394,6 +1451,15 @@ mod _sqlite3 { name: PyUtf8StrRef, callable: PyObjectRef, vm: &VirtualMachine, + ) -> PyResult<()> { + with_deferred_drops(|| Self::create_collation_inner(zelf, name, callable, vm)) + } + + fn create_collation_inner( + zelf: &Py, + name: PyUtf8StrRef, + callable: PyObjectRef, + vm: &VirtualMachine, ) -> PyResult<()> { let name = name.to_cstring(vm)?; let db = zelf.db_lock(vm)?; @@ -1440,6 +1506,18 @@ mod _sqlite3 { narg: c_int, aggregate_class: PyObjectRef, vm: &VirtualMachine, + ) -> PyResult<()> { + with_deferred_drops(|| { + Self::create_window_function_inner(zelf, name, narg, aggregate_class, vm) + }) + } + + fn create_window_function_inner( + zelf: &Py, + name: PyStrRef, + narg: c_int, + aggregate_class: PyObjectRef, + vm: &VirtualMachine, ) -> PyResult<()> { let name = name.to_cstring(vm)?; let db = zelf.db_lock(vm)?; @@ -1679,6 +1757,14 @@ mod _sqlite3 { } #[pygetset(setter)] fn set_autocommit(zelf: &Py, val: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { + with_deferred_drops(|| Self::set_autocommit_inner(zelf, val, vm)) + } + + fn set_autocommit_inner( + zelf: &Py, + val: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult<()> { let mode = AutocommitMode::try_from_borrowed_object(vm, &val)?; let db = zelf.db_lock(vm)?; *zelf.autocommit.lock() = mode; @@ -1912,6 +1998,15 @@ mod _sqlite3 { sql: PyUtf8StrRef, parameters: OptionalArg, vm: &VirtualMachine, + ) -> PyResult> { + with_deferred_drops(|| Self::execute_inner(zelf, sql, parameters, vm)) + } + + fn execute_inner( + zelf: PyRef, + sql: PyUtf8StrRef, + parameters: OptionalArg, + vm: &VirtualMachine, ) -> PyResult> { let mut inner = zelf.inner(vm)?; @@ -1994,6 +2089,15 @@ mod _sqlite3 { sql: PyUtf8StrRef, seq_of_params: ArgIterable, vm: &VirtualMachine, + ) -> PyResult> { + with_deferred_drops(|| Self::executemany_inner(zelf, sql, seq_of_params, vm)) + } + + fn executemany_inner( + zelf: PyRef, + sql: PyUtf8StrRef, + seq_of_params: ArgIterable, + vm: &VirtualMachine, ) -> PyResult> { let mut inner = zelf.inner(vm)?; @@ -2070,6 +2174,14 @@ mod _sqlite3 { zelf: PyRef, script: PyUtf8StrRef, vm: &VirtualMachine, + ) -> PyResult> { + with_deferred_drops(|| Self::executescript_inner(zelf, script, vm)) + } + + fn executescript_inner( + zelf: PyRef, + script: PyUtf8StrRef, + vm: &VirtualMachine, ) -> PyResult> { zelf.check_cursor_valid(vm)?; @@ -2146,6 +2258,10 @@ mod _sqlite3 { #[pymethod] fn close(zelf: &Py, vm: &VirtualMachine) -> PyResult<()> { + with_deferred_drops(|| Self::close_inner(zelf, vm)) + } + + fn close_inner(zelf: &Py, vm: &VirtualMachine) -> PyResult<()> { // Check if __init__ was called let mut guard = zelf.inner.lock(); @@ -2299,6 +2415,12 @@ mod _sqlite3 { impl SelfIter for Cursor {} impl IterNext for Cursor { fn next(zelf: &Py, vm: &VirtualMachine) -> PyResult { + with_deferred_drops(|| Self::next_inner(zelf, vm)) + } + } + + impl Cursor { + fn next_inner(zelf: &Py, vm: &VirtualMachine) -> PyResult { // Check if connection is closed first, and if so, clear statement to release file lock if zelf.connection.is_closed() { let mut guard = zelf.inner.lock(); @@ -2576,7 +2698,7 @@ mod _sqlite3 { impl Drop for BlobInner { fn drop(&mut self) { - unsafe { sqlite3_blob_close(self.blob.blob) }; + with_deferred_drops(|| unsafe { sqlite3_blob_close(self.blob.blob) }); } } @@ -2584,7 +2706,9 @@ mod _sqlite3 { impl Blob { #[pymethod] fn close(zelf: &Py) { - zelf.inner.lock().take(); + with_deferred_drops(|| { + zelf.inner.lock().take(); + }); } fn ensure_connection_open(&self, vm: &VirtualMachine) -> PyResult<()> { @@ -2742,6 +2866,10 @@ mod _sqlite3 { } fn subscript(&self, needle: &PyObject, vm: &VirtualMachine) -> PyResult { + with_deferred_drops(|| Self::subscript_inner(self, needle, vm)) + } + + fn subscript_inner(&self, needle: &PyObject, vm: &VirtualMachine) -> PyResult { self.ensure_connection_open(vm)?; let inner = self.inner(vm)?; if let Some(index) = needle.try_index_opt(vm) { @@ -2784,6 +2912,15 @@ mod _sqlite3 { needle: &PyObject, value: Option, vm: &VirtualMachine, + ) -> PyResult<()> { + with_deferred_drops(|| Self::ass_subscript_inner(self, needle, value, vm)) + } + + fn ass_subscript_inner( + &self, + needle: &PyObject, + value: Option, + vm: &VirtualMachine, ) -> PyResult<()> { self.ensure_connection_open(vm)?; let inner = self.inner(vm)?; @@ -3042,7 +3179,7 @@ mod _sqlite3 { impl Drop for Sqlite { fn drop(&mut self) { // Use sqlite3_close_v2 for safe closing even with active statements - unsafe { sqlite3_close_v2(self.raw.db) }; + with_deferred_drops(|| unsafe { sqlite3_close_v2(self.raw.db) }); } } @@ -3261,9 +3398,11 @@ mod _sqlite3 { impl Drop for SqliteStatement { fn drop(&mut self) { - unsafe { + // A cursor may be dropped without an explicit close(). Nesting + // preserves any outer scope that still owns SQLite locks. + with_deferred_drops(|| unsafe { sqlite3_finalize(self.raw.st); - } + }); } } @@ -3591,10 +3730,13 @@ mod _sqlite3 { unsafe { sqlite3_user_data(self.ctx).cast() } } - fn aggregate_context(self) -> *mut T { - unsafe { - sqlite3_aggregate_context(self.ctx, core::mem::size_of::() as c_int).cast() - } + fn aggregate_context(self, create: bool) -> *mut *const PyObject { + let size = if create { + core::mem::size_of::<*const PyObject>() as c_int + } else { + 0 + }; + unsafe { sqlite3_aggregate_context(self.ctx, size).cast() } } fn result_exception(self, vm: &VirtualMachine, exc: PyBaseExceptionRef, msg: &CStr) { diff --git a/extra_tests/snippets/stdlib_sqlite.py b/extra_tests/snippets/stdlib_sqlite.py index 83fc7db94f0..76c8b7adde1 100644 --- a/extra_tests/snippets/stdlib_sqlite.py +++ b/extra_tests/snippets/stdlib_sqlite.py @@ -70,3 +70,239 @@ def finalize(self): expected = b"t5i4 3l2b1" + data[10:] assert actual == expected, f"got {actual!r}, expected {expected!r}" blob.close() + + +# Aggregate instances belong to a query, including its error and cancellation paths. +import gc +import weakref + +from testutils import assert_raises + + +def check_aggregate_lifetime(failure): + refs = [] + + class TrackedAggregate: + def __init__(self): + refs.append(weakref.ref(self)) + + def step(self, value): + if failure == "step": + raise ValueError("step") + + def finalize(self): + if failure == "finalize": + raise ValueError("finalize") + return object() if failure == "result" else 1 + + cx.create_aggregate("tracked", 1, TrackedAggregate) + query = "SELECT tracked(key) FROM foo GROUP BY key % 2" + if failure: + with assert_raises(sqlite.OperationalError): + cx.execute(query).fetchall() + else: + assert cx.execute(query).fetchall() == [(1,), (1,)] + assert refs + gc.collect() + assert all(ref() is None for ref in refs), failure + + # xFinal may run without any preceding xStep. + refs.clear() + assert cx.execute("SELECT tracked(key) FROM foo WHERE 0").fetchall() == [(None,)] + assert not refs + + +for failure in (None, "step", "finalize", "result"): + check_aggregate_lifetime(failure) + + +# Window functions share aggregate finalization, including an unfinished cursor. +window_refs = [] + + +class TrackedWindow: + def __init__(self): + window_refs.append(weakref.ref(self)) + + def step(self, value): + pass + + def inverse(self, value): + pass + + def value(self): + return 1 + + def finalize(self): + return 1 + + +cx.create_window_function("tracked_window", 1, TrackedWindow) +window_query = """ + SELECT tracked_window(key) OVER ( + ORDER BY key ROWS BETWEEN 1 PRECEDING AND CURRENT ROW + ) FROM foo +""" +assert cx.execute(window_query).fetchall() == [(1,)] * 4 +assert window_refs +assert all(ref() is None for ref in window_refs) +window_refs.clear() +window_cursor = cx.execute(window_query) +assert window_cursor.fetchone() == (1,) +window_cursor.close() +gc.collect() +assert window_refs and all(ref() is None for ref in window_refs) + + +# Releasing the last reference can run Python that reads the cursor/connection. +def check_aggregate_destructor(operation): + connection = sqlite.connect(":memory:") + cursor = connection.cursor() + cursor_ref = weakref.ref(cursor) + released = [] + + class Aggregate: + def step(self, value): + pass + + def value(self): + return 1 + + finalize = value + + def __del__(self): + current = cursor_ref() + if current is not None: + try: + current.description + except sqlite.ProgrammingError: + pass # The cursor may already have been closed. + released.append(connection.total_changes) + + connection.create_aggregate("tracked", 1, Aggregate) + expected = 1 + if operation == "execute": + cursor.execute("SELECT tracked(1)") + elif operation == "script": + cursor.executescript("SELECT tracked(1);") + elif operation == "many": + connection.execute("CREATE TABLE results(value)") + cursor.executemany("INSERT INTO results SELECT tracked(?)", [(1,), (2,)]) + expected = 2 + else: + connection.create_window_function("tracked_window", 1, Aggregate) + cursor.execute(""" + SELECT tracked_window(x) OVER (ROWS UNBOUNDED PRECEDING) + FROM (SELECT 1 AS x UNION ALL SELECT 2 UNION ALL SELECT 3) + """) + assert not released + if operation == "fetch": + assert cursor.fetchall() == [(1,), (1,), (1,)] + elif operation == "close": + cursor.close() + elif operation == "replace": + cursor.execute("SELECT 2") + else: + del cursor + gc.collect() + assert len(released) == expected, (operation, released) + connection.close() + + +for operation in ("execute", "script", "many", "fetch", "close", "replace", "drop"): + check_aggregate_destructor(operation) + + +# A nested query must not release its aggregate while an outer callback holds locks. +def check_nested_aggregate_destructor(operation): + outer = sqlite.connect(":memory:", autocommit=False) + nested = sqlite.connect(":memory:") + released = [] + inspect = lambda: outer.total_changes + + class Aggregate: + def step(self, value): + pass + + def finalize(self): + return 1 + + def __del__(self): + try: + inspect() + except sqlite.ProgrammingError: + pass # Closing the outer connection may finish before cleanup. + released.append(True) + + nested.create_aggregate("tracked", 1, Aggregate) + outer.set_trace_callback(lambda sql: nested.execute("SELECT tracked(1)")) + if operation == "commit": + outer.commit() + elif operation == "rollback": + outer.rollback() + elif operation == "autocommit": + outer.autocommit = True + elif operation in ("blobopen", "blob_index"): + outer.execute("CREATE TABLE blobs(data BLOB)") + outer.execute("INSERT INTO blobs VALUES (zeroblob(1))") + released.clear() + + def progress(): + nested.execute("SELECT tracked(1)") + return 0 + + outer.set_progress_handler(progress, 1) + blob = outer.blobopen("blobs", "data", 1) + outer.set_progress_handler(None, 0) + if operation != "blobopen": + released.clear() + inspect = lambda: len(blob) + + class Index: + def __index__(self): + nested.execute("SELECT tracked(1)") + return 0 + + index = Index() + assert blob[index] == 0 + blob[index] = 1 + assert len(released) == 2 + blob.close() + elif operation == "replace_callback": + + class Callback: + def __call__(self): + return 1 + + def __del__(self): + nested.execute("SELECT tracked(1)") + + for register in ( + lambda cb: outer.create_function("hook", 0, cb), + lambda cb: outer.create_aggregate("hook", 0, cb), + lambda cb: outer.create_window_function("hook", 0, cb), + lambda cb: outer.create_collation("hook", cb), + ): + released.clear() + register(Callback()) + register(None) + assert released + else: + outer.close() + assert released, operation + if operation != "close": + outer.set_trace_callback(None) + outer.close() + nested.close() + + +for operation in ( + "commit", + "rollback", + "autocommit", + "close", + "blobopen", + "blob_index", + "replace_callback", +): + check_nested_aggregate_destructor(operation)