diff --git a/sqlx-sqlite/src/connection/mod.rs b/sqlx-sqlite/src/connection/mod.rs index b94ad91c4d..0a0645af09 100644 --- a/sqlx-sqlite/src/connection/mod.rs +++ b/sqlx-sqlite/src/connection/mod.rs @@ -12,8 +12,14 @@ use futures_core::future::BoxFuture; use futures_intrusive::sync::MutexGuard; use futures_util::future; use libsqlite3_sys::{ - sqlite3, sqlite3_commit_hook, sqlite3_get_autocommit, sqlite3_progress_handler, - sqlite3_rollback_hook, sqlite3_update_hook, SQLITE_DELETE, SQLITE_INSERT, SQLITE_UPDATE, + sqlite3, sqlite3_commit_hook, sqlite3_db_release_memory, sqlite3_db_status, + sqlite3_get_autocommit, sqlite3_progress_handler, sqlite3_rollback_hook, sqlite3_update_hook, + SQLITE_DBSTATUS_CACHE_HIT, SQLITE_DBSTATUS_CACHE_MISS, SQLITE_DBSTATUS_CACHE_USED, + SQLITE_DBSTATUS_CACHE_USED_SHARED, SQLITE_DBSTATUS_CACHE_WRITE, SQLITE_DBSTATUS_DEFERRED_FKS, + SQLITE_DBSTATUS_LOOKASIDE_HIT, SQLITE_DBSTATUS_LOOKASIDE_MISS_FULL, + SQLITE_DBSTATUS_LOOKASIDE_MISS_SIZE, SQLITE_DBSTATUS_LOOKASIDE_USED, + SQLITE_DBSTATUS_SCHEMA_USED, SQLITE_DBSTATUS_STMT_USED, SQLITE_DELETE, SQLITE_INSERT, + SQLITE_OK, SQLITE_UPDATE, }; #[cfg(feature = "preupdate-hook")] pub use preupdate_hook::*; @@ -77,6 +83,54 @@ pub enum SqliteOperation { Unknown(i32), } +/// Database status parameters for the sqlite3_db_status function. +#[derive(Debug, PartialEq, Eq, Clone, Copy)] +pub enum SqliteDatabaseStatus { + /// Current number of bytes used by lookaside allocations + LookasideUsed, + /// Current number of bytes of pager cache used + CacheUsed, + /// Current number of bytes used by the schema + SchemaUsed, + /// Current number of bytes used by prepared statements + StmtUsed, + /// Number of lookaside malloc hits + LookasideHit, + /// Number of lookaside malloc misses due to size + LookasideMissSize, + /// Number of lookaside malloc misses due to full buffer + LookasideMissFull, + /// Number of pager cache hits + CacheHit, + /// Number of pager cache misses + CacheMiss, + /// Number of dirty cache pages written + CacheWrite, + /// Number of foreign key constraint violations detected + DeferredFks, + /// Maximum cache used in shared cache mode + CacheUsedShared, +} + +impl From for i32 { + fn from(status: SqliteDatabaseStatus) -> Self { + match status { + SqliteDatabaseStatus::LookasideUsed => SQLITE_DBSTATUS_LOOKASIDE_USED, + SqliteDatabaseStatus::CacheUsed => SQLITE_DBSTATUS_CACHE_USED, + SqliteDatabaseStatus::SchemaUsed => SQLITE_DBSTATUS_SCHEMA_USED, + SqliteDatabaseStatus::StmtUsed => SQLITE_DBSTATUS_STMT_USED, + SqliteDatabaseStatus::LookasideHit => SQLITE_DBSTATUS_LOOKASIDE_HIT, + SqliteDatabaseStatus::LookasideMissSize => SQLITE_DBSTATUS_LOOKASIDE_MISS_SIZE, + SqliteDatabaseStatus::LookasideMissFull => SQLITE_DBSTATUS_LOOKASIDE_MISS_FULL, + SqliteDatabaseStatus::CacheHit => SQLITE_DBSTATUS_CACHE_HIT, + SqliteDatabaseStatus::CacheMiss => SQLITE_DBSTATUS_CACHE_MISS, + SqliteDatabaseStatus::CacheWrite => SQLITE_DBSTATUS_CACHE_WRITE, + SqliteDatabaseStatus::DeferredFks => SQLITE_DBSTATUS_DEFERRED_FKS, + SqliteDatabaseStatus::CacheUsedShared => SQLITE_DBSTATUS_CACHE_USED_SHARED, + } + } +} + impl From for SqliteOperation { fn from(value: i32) -> Self { match value { @@ -557,6 +611,52 @@ impl LockedSqliteHandle<'_> { let ret = unsafe { sqlite3_get_autocommit(self.as_raw_handle().as_ptr()) }; ret == 0 } + + /// Retrieves statistics about a database connection. + /// + /// This function is used to retrieve runtime status information about a single database connection. + /// The `status` parameter determines which statistic to retrieve. + /// + /// Returns a tuple containing `(current_value, highest_value_since_reset)`. + /// If `reset` is true, the highest value is reset to the current value after retrieval. + /// + /// See: https://www.sqlite.org/c3ref/db_status.html + pub fn db_status( + &mut self, + status: SqliteDatabaseStatus, + reset: bool, + ) -> Result<(i32, i32), Error> { + let mut current = 0i32; + let mut highest = 0i32; + + let result = unsafe { + sqlite3_db_status( + self.as_raw_handle().as_ptr(), + status.into(), + &mut current, + &mut highest, + if reset { 1 } else { 0 }, + ) + }; + + if result == SQLITE_OK { + Ok((current, highest)) + } else { + Err(self.guard.handle.expect_error().into()) + } + } + + /// Attempts to free as much heap memory as possible from the database connection. + /// + /// This function causes SQLite to release some memory used by the database connection, + /// such as memory used to cache prepared statements. + /// + /// Returns the number of bytes of memory released. + /// + /// See: https://www.sqlite.org/c3ref/db_release_memory.html + pub fn db_release_memory(&mut self) -> i32 { + unsafe { sqlite3_db_release_memory(self.as_raw_handle().as_ptr()) } + } } impl Drop for ConnectionState { diff --git a/sqlx-sqlite/src/lib.rs b/sqlx-sqlite/src/lib.rs index e4a122b6bd..b176c45e53 100644 --- a/sqlx-sqlite/src/lib.rs +++ b/sqlx-sqlite/src/lib.rs @@ -49,7 +49,9 @@ pub use column::SqliteColumn; pub use connection::serialize::SqliteOwnedBuf; #[cfg(feature = "preupdate-hook")] pub use connection::PreupdateHookResult; -pub use connection::{LockedSqliteHandle, SqliteConnection, SqliteOperation, UpdateHookResult}; +pub use connection::{ + LockedSqliteHandle, SqliteConnection, SqliteDatabaseStatus, SqliteOperation, UpdateHookResult, +}; pub use database::Sqlite; pub use error::SqliteError; pub use options::{ @@ -124,6 +126,63 @@ impl_encode_for_option!(Sqlite); #[doc(hidden)] pub static CREATE_DB_WAL: AtomicBool = AtomicBool::new(true); +/// Sets the soft heap limit for SQLite. +/// +/// This function sets a soft limit on the amount of heap memory that can be allocated by SQLite +/// across all database connections within a single process. +/// +/// - `Some(limit)` sets the heap limit to the specified number of bytes +/// - `None` disables the heap limit +/// +/// Returns the previous heap limit in bytes, or `None` if no limit was previously set. +/// +/// # Errors +/// +/// Returns an error if the limit exceeds `i64::MAX`. +/// +/// See: https://www.sqlite.org/c3ref/hard_heap_limit64.html +pub fn set_soft_heap_limit( + limit: Option, +) -> Result, Error> { + use libsqlite3_sys::sqlite3_soft_heap_limit64; + + let limit_value = match limit { + Some(n) => { + let value = n.get(); + if value > i64::MAX as u64 { + return Err(Error::Configuration("heap limit exceeds i64::MAX".into())); + } + value as i64 + } + None => 0, + }; + + let previous = unsafe { sqlite3_soft_heap_limit64(limit_value) }; + + Ok(if previous > 0 { + Some(std::num::NonZeroU64::new(previous as u64).unwrap()) + } else { + None + }) +} + +/// Gets the current soft heap limit for SQLite. +/// +/// Returns the current heap limit in bytes, or `None` if no limit is set. +/// +/// See: https://www.sqlite.org/c3ref/hard_heap_limit64.html +pub fn soft_heap_limit() -> Option { + use libsqlite3_sys::sqlite3_soft_heap_limit64; + + let current = unsafe { sqlite3_soft_heap_limit64(-1) }; + + if current > 0 { + Some(std::num::NonZeroU64::new(current as u64).unwrap()) + } else { + None + } +} + /// UNSTABLE: for use by `sqlite-macros-core` only. #[doc(hidden)] pub fn describe_blocking(query: &str, database_url: &str) -> Result, Error> { diff --git a/tests/sqlite/sqlite.rs b/tests/sqlite/sqlite.rs index 4d24b07412..f74155df92 100644 --- a/tests/sqlite/sqlite.rs +++ b/tests/sqlite/sqlite.rs @@ -1,7 +1,9 @@ use futures::TryStreamExt; use rand::{Rng, SeedableRng}; use rand_xoshiro::Xoshiro256PlusPlus; -use sqlx::sqlite::{SqliteConnectOptions, SqliteOperation, SqlitePoolOptions}; +use sqlx::sqlite::{ + SqliteConnectOptions, SqliteDatabaseStatus, SqliteOperation, SqlitePoolOptions, +}; use sqlx::{ query, sqlite::Sqlite, sqlite::SqliteRow, Column, ConnectOptions, Connection, Executor, Row, SqliteConnection, SqlitePool, Statement, TypeInfo, @@ -1367,3 +1369,306 @@ enum SqliteTransactionState { Read, Write, } + +#[sqlx_macros::test] +async fn test_soft_heap_limit() -> anyhow::Result<()> { + let mut conn = new::().await?; + let mut handle = conn.lock_handle().await?; + + // Test setting and getting heap limit + let old_limit = handle.soft_heap_limit(1024 * 1024); // 1MB + let current_limit = handle.soft_heap_limit(2 * 1024 * 1024); // 2MB + assert_eq!(current_limit, 1024 * 1024); + + // Disable heap limit + let disabled_limit = handle.soft_heap_limit(0); + assert_eq!(disabled_limit, 2 * 1024 * 1024); + + // Restore old limit if it existed + if old_limit > 0 { + handle.soft_heap_limit(old_limit); + } + + Ok(()) +} + +#[sqlx_macros::test] +async fn test_db_status() -> anyhow::Result<()> { + let mut conn = new::().await?; + let mut handle = conn.lock_handle().await?; + + // Test various status options + let (current, highest) = handle.db_status(SqliteDatabaseStatus::CacheUsed, false)?; + assert!(current >= 0); + assert!(highest >= 0); // Highest can be 0 if no tracking occurred yet + + let (current, highest) = handle.db_status(SqliteDatabaseStatus::SchemaUsed, false)?; + assert!(current >= 0); + assert!(highest >= 0); // Highest can be 0 if no tracking occurred yet + + let (current, highest) = handle.db_status(SqliteDatabaseStatus::StmtUsed, false)?; + assert!(current >= 0); + assert!(highest >= 0); // Highest can be 0 if no tracking occurred yet + + // Test reset functionality + let (current1, highest1) = handle.db_status(SqliteDatabaseStatus::CacheUsed, false)?; + let (current2, highest2) = handle.db_status(SqliteDatabaseStatus::CacheUsed, true)?; + + // After reset, the highest value should be equal to current + assert_eq!(current1, current2); + assert_eq!(highest1, highest2); + + // Verify that all status types can be queried without error + let status_types = [ + SqliteDatabaseStatus::LookasideUsed, + SqliteDatabaseStatus::CacheUsed, + SqliteDatabaseStatus::SchemaUsed, + SqliteDatabaseStatus::StmtUsed, + SqliteDatabaseStatus::LookasideHit, + SqliteDatabaseStatus::LookasideMissSize, + SqliteDatabaseStatus::LookasideMissFull, + SqliteDatabaseStatus::CacheHit, + SqliteDatabaseStatus::CacheMiss, + SqliteDatabaseStatus::CacheWrite, + SqliteDatabaseStatus::DeferredFks, + SqliteDatabaseStatus::CacheUsedShared, + ]; + + for status in status_types { + let result = handle.db_status(status, false); + assert!(result.is_ok(), "Failed to query status: {:?}", status); + } + + Ok(()) +} + +#[sqlx_macros::test] +async fn test_db_status_with_usage() -> anyhow::Result<()> { + let mut conn = new::().await?; + + // Create some schema to generate usage + conn.execute("CREATE TEMPORARY TABLE test_table (id INTEGER, data TEXT)") + .await?; + + // Insert some data to trigger cache usage + for i in 0..10 { + sqlx::query("INSERT INTO test_table (id, data) VALUES (?, ?)") + .bind(i) + .bind(format!("test_data_{}", i)) + .execute(&mut conn) + .await?; + } + + let mut handle = conn.lock_handle().await?; + + // Check that schema usage is non-zero + let (current, highest) = handle.db_status(SqliteDatabaseStatus::SchemaUsed, false)?; + assert!(current > 0, "Schema should be using some memory"); + assert!(highest >= 0); // Highest can be 0 if no tracking occurred yet + + // Check cache usage after some operations + let (_current, highest) = handle.db_status(SqliteDatabaseStatus::CacheUsed, false)?; + assert!(highest >= 0); // Highest can be 0 if no tracking occurred yet + + Ok(()) +} + +#[sqlx_macros::test] +async fn test_db_release_memory() -> anyhow::Result<()> { + let mut conn = new::().await?; + + // Create some prepared statements to use memory + conn.execute("CREATE TEMPORARY TABLE test_table (id INTEGER, data TEXT)") + .await?; + + for i in 0..5 { + sqlx::query("INSERT INTO test_table (id, data) VALUES (?, ?)") + .bind(i) + .bind(format!("test_data_{}", i)) + .execute(&mut conn) + .await?; + } + + let mut handle = conn.lock_handle().await?; + + // Get initial memory usage + let (_initial_cache, _) = handle.db_status(SqliteDatabaseStatus::CacheUsed, false)?; + + // Release memory + let released = handle.db_release_memory(); + + // released should be non-negative + assert!(released >= 0); + + // Memory usage should potentially be lower (though this isn't guaranteed) + let (after_release_cache, _) = handle.db_status(SqliteDatabaseStatus::CacheUsed, false)?; + assert!(after_release_cache >= 0); + + Ok(()) +} + +#[sqlx_macros::test] +async fn test_memory_functions_edge_cases() -> anyhow::Result<()> { + let mut conn = new::().await?; + let mut handle = conn.lock_handle().await?; + + // Test soft_heap_limit with edge values + let original_limit = handle.soft_heap_limit(-1); // Should be ignored/treated as 0 + assert!(original_limit >= 0); + + // Test very large limit + let large_limit = handle.soft_heap_limit(i64::MAX); + assert!(large_limit >= 0); + + // Test zero (disable limit) + let zero_result = handle.soft_heap_limit(0); + assert_eq!(zero_result, i64::MAX); + + // Test negative values (should be ignored) + let negative_result = handle.soft_heap_limit(-1000); + assert!(negative_result >= 0); + + // Restore original limit + handle.soft_heap_limit(original_limit); + + Ok(()) +} + +#[sqlx_macros::test] +async fn test_db_status_reset_functionality() -> anyhow::Result<()> { + let mut conn = new::().await?; + let mut handle = conn.lock_handle().await?; + + // Test reset functionality for various status types + let status_types_to_test = [ + SqliteDatabaseStatus::CacheUsed, + SqliteDatabaseStatus::SchemaUsed, + SqliteDatabaseStatus::StmtUsed, + SqliteDatabaseStatus::LookasideUsed, + ]; + + for status in status_types_to_test { + // Get status without reset + let (current1, highest1) = handle.db_status(status, false)?; + + // Get status with reset - highest should now equal current + let (current2, highest2) = handle.db_status(status, true)?; + + // Current values should be the same + assert_eq!(current1, current2); + + // All values should be non-negative (reset behavior may vary) + assert!(current1 >= 0); + assert!(current2 >= 0); + assert!(highest1 >= 0); + assert!(highest2 >= 0); + } + + Ok(()) +} + +#[sqlx_macros::test] +async fn test_db_release_memory_multiple_calls() -> anyhow::Result<()> { + let mut conn = new::().await?; + let mut handle = conn.lock_handle().await?; + + // Test multiple consecutive calls to db_release_memory + for _ in 0..5 { + let released = handle.db_release_memory(); + // Should always return a non-negative value, even if no memory was released + assert!(released >= 0); + } + + Ok(()) +} + +#[sqlx_macros::test] +async fn test_memory_functions_on_fresh_connection() -> anyhow::Result<()> { + // Test functions on a completely fresh connection with no activity + let mut conn = new::().await?; + let mut handle = conn.lock_handle().await?; + + // Test db_status on fresh connection + let (current, highest) = handle.db_status(SqliteDatabaseStatus::CacheUsed, false)?; + assert!(current >= 0); + assert!(highest >= 0); // Highest can be 0 if no tracking occurred yet + + // Test db_release_memory on fresh connection + let released = handle.db_release_memory(); + assert!(released >= 0); + + // Test soft_heap_limit on fresh connection + let old_limit = handle.soft_heap_limit(1024); + let restored = handle.soft_heap_limit(old_limit); + assert_eq!(restored, 1024); + + Ok(()) +} + +#[sqlx_macros::test] +async fn test_db_status_all_types_no_error() -> anyhow::Result<()> { + let mut conn = new::().await?; + let mut handle = conn.lock_handle().await?; + + // Ensure every single status type can be queried without error + // This is the most important test - ensuring none of the status types cause panics or errors + + // Test with reset = false + let result1 = handle.db_status(SqliteDatabaseStatus::LookasideUsed, false); + assert!(result1.is_ok(), "LookasideUsed failed: {:?}", result1.err()); + + let result2 = handle.db_status(SqliteDatabaseStatus::CacheUsed, false); + assert!(result2.is_ok(), "CacheUsed failed: {:?}", result2.err()); + + let result3 = handle.db_status(SqliteDatabaseStatus::SchemaUsed, false); + assert!(result3.is_ok(), "SchemaUsed failed: {:?}", result3.err()); + + let result4 = handle.db_status(SqliteDatabaseStatus::StmtUsed, false); + assert!(result4.is_ok(), "StmtUsed failed: {:?}", result4.err()); + + let result5 = handle.db_status(SqliteDatabaseStatus::LookasideHit, false); + assert!(result5.is_ok(), "LookasideHit failed: {:?}", result5.err()); + + let result6 = handle.db_status(SqliteDatabaseStatus::LookasideMissSize, false); + assert!( + result6.is_ok(), + "LookasideMissSize failed: {:?}", + result6.err() + ); + + let result7 = handle.db_status(SqliteDatabaseStatus::LookasideMissFull, false); + assert!( + result7.is_ok(), + "LookasideMissFull failed: {:?}", + result7.err() + ); + + let result8 = handle.db_status(SqliteDatabaseStatus::CacheHit, false); + assert!(result8.is_ok(), "CacheHit failed: {:?}", result8.err()); + + let result9 = handle.db_status(SqliteDatabaseStatus::CacheMiss, false); + assert!(result9.is_ok(), "CacheMiss failed: {:?}", result9.err()); + + let result10 = handle.db_status(SqliteDatabaseStatus::CacheWrite, false); + assert!(result10.is_ok(), "CacheWrite failed: {:?}", result10.err()); + + let result11 = handle.db_status(SqliteDatabaseStatus::DeferredFks, false); + assert!(result11.is_ok(), "DeferredFks failed: {:?}", result11.err()); + + let result12 = handle.db_status(SqliteDatabaseStatus::CacheUsedShared, false); + assert!( + result12.is_ok(), + "CacheUsedShared failed: {:?}", + result12.err() + ); + + // Test with reset = true + let result13 = handle.db_status(SqliteDatabaseStatus::CacheUsed, true); + assert!( + result13.is_ok(), + "CacheUsed with reset failed: {:?}", + result13.err() + ); + + Ok(()) +}