LRU caching for SQLite

sqlx-core/src/common/statement_cache.rs

@@ -3,11 +3,11 @@ use lru_cache::LruCache;
 /// A cache for prepared statements. When full, the least recently used
 /// statement gets removed.
 #[derive(Debug)]
-pub struct StatementCache {
-    inner: LruCache<String, u32>,
+pub struct StatementCache<T> {
+    inner: LruCache<String, T>,
 }
 
-impl StatementCache {
+impl<T> StatementCache<T> {
     /// Create a new cache with the given capacity.
     pub fn new(capacity: usize) -> Self {
         Self {
@@ -17,19 +17,19 @@ impl StatementCache {
 
     /// Returns a mutable reference to the value corresponding to the given key
     /// in the cache, if any.
-    pub fn get_mut(&mut self, k: &str) -> Option<&mut u32> {
+    pub fn get_mut(&mut self, k: &str) -> Option<&mut T> {
         self.inner.get_mut(k)
     }
 
     /// Inserts a new statement to the cache, returning the least recently used
     /// statement id if the cache is full, or if inserting with an existing key,
     /// the replaced existing statement.
-    pub fn insert(&mut self, k: &str, v: u32) -> Option<u32> {
+    pub fn insert(&mut self, k: &str, v: T) -> Option<T> {
         let mut lru_item = None;
 
         if self.inner.capacity() == self.len() && !self.inner.contains_key(k) {
             lru_item = self.remove_lru();
-        } else if self.inner.contains_key(k) {
+        } else if self.contains_key(k) {
             lru_item = self.inner.remove(k);
         }
 
@@ -44,12 +44,17 @@ impl StatementCache {
     }
 
     /// Removes the least recently used item from the cache.
-    pub fn remove_lru(&mut self) -> Option<u32> {
+    pub fn remove_lru(&mut self) -> Option<T> {
         self.inner.remove_lru().map(|(_, v)| v)
     }
 
     /// Clear all cached statements from the cache.
     pub fn clear(&mut self) {
         self.inner.clear();
     }
+
+    /// True if cache has a value for the given key.
+    pub fn contains_key(&mut self, k: &str) -> bool {
+        self.inner.contains_key(k)
+    }
 }

sqlx-core/src/mysql/connection/mod.rs

@@ -37,7 +37,7 @@ pub struct MySqlConnection {
     pub(crate) stream: MySqlStream,
 
     // cache by query string to the statement id
-    cache_statement: StatementCache,
+    cache_statement: StatementCache<u32>,
 
     // working memory for the active row's column information
     // this allows us to re-use these allocations unless the user is persisting the

sqlx-core/src/postgres/connection/mod.rs

@@ -48,7 +48,7 @@ pub struct PgConnection {
     next_statement_id: u32,
 
     // cache statement by query string to the id and columns
-    cache_statement: StatementCache,
+    cache_statement: StatementCache<u32>,
 
     // cache user-defined types by id <-> info
     cache_type_info: HashMap<u32, PgTypeInfo>,

sqlx-core/src/sqlite/connection/establish.rs

@@ -1,7 +1,6 @@
 use std::io;
 use std::ptr::{null, null_mut};
 
-use hashbrown::HashMap;
 use libsqlite3_sys::{
     sqlite3_busy_timeout, sqlite3_extended_result_codes, sqlite3_open_v2, SQLITE_OK,
     SQLITE_OPEN_CREATE, SQLITE_OPEN_MEMORY, SQLITE_OPEN_NOMUTEX, SQLITE_OPEN_PRIVATECACHE,
@@ -12,7 +11,10 @@ use sqlx_rt::blocking;
 use crate::error::Error;
 use crate::sqlite::connection::handle::ConnectionHandle;
 use crate::sqlite::statement::StatementWorker;
-use crate::sqlite::{SqliteConnectOptions, SqliteConnection, SqliteError};
+use crate::{
+    common::StatementCache,
+    sqlite::{SqliteConnectOptions, SqliteConnection, SqliteError},
+};
 
 pub(super) async fn establish(options: &SqliteConnectOptions) -> Result<SqliteConnection, Error> {
     let mut filename = options
@@ -87,7 +89,7 @@ pub(super) async fn establish(options: &SqliteConnectOptions) -> Result<SqliteCo
     Ok(SqliteConnection {
         handle,
         worker: StatementWorker::new(),
-        statements: HashMap::new(),
+        statements: StatementCache::new(options.statement_cache_size),
         statement: None,
         scratch_row_column_names: Default::default(),
     })

sqlx-core/src/sqlite/connection/executor.rs

@@ -6,6 +6,7 @@ use futures_core::stream::BoxStream;
 use futures_util::TryStreamExt;
 use hashbrown::HashMap;
 
+use crate::common::StatementCache;
 use crate::describe::{Column, Describe};
 use crate::error::Error;
 use crate::executor::{Execute, Executor};
@@ -16,7 +17,7 @@ use crate::sqlite::{Sqlite, SqliteArguments, SqliteConnection, SqliteRow};
 
 fn prepare<'a>(
     conn: &mut ConnectionHandle,
-    statements: &'a mut HashMap<String, SqliteStatement>,
+    statements: &'a mut StatementCache<SqliteStatement>,
     statement: &'a mut Option<SqliteStatement>,
     query: &str,
     persistent: bool,
@@ -28,7 +29,10 @@ fn prepare<'a>(
 
     if !statements.contains_key(query) {
         let statement = SqliteStatement::prepare(conn, query, false)?;
-        statements.insert(query.to_owned(), statement);
+
+        if let Some(mut statement) = statements.insert(query, statement) {
+            statement.reset();
+        }
     }
 
     let statement = statements.get_mut(query).unwrap();

sqlx-core/src/sqlite/connection/mod.rs

@@ -6,6 +6,8 @@ use futures_util::future;
 use hashbrown::HashMap;
 use libsqlite3_sys::sqlite3;
 
+use crate::caching_connection::CachingConnection;
+use crate::common::StatementCache;
 use crate::connection::{Connect, Connection};
 use crate::error::Error;
 use crate::ext::ustr::UStr;
@@ -25,7 +27,7 @@ pub struct SqliteConnection {
     pub(crate) worker: StatementWorker,
 
     // cache of semi-persistent statements
-    pub(crate) statements: HashMap<String, SqliteStatement>,
+    pub(crate) statements: StatementCache<SqliteStatement>,
 
     // most recent non-persistent statement
     pub(crate) statement: Option<SqliteStatement>,
@@ -47,6 +49,19 @@ impl Debug for SqliteConnection {
     }
 }
 
+impl CachingConnection for SqliteConnection {
+    fn cached_statements_count(&self) -> usize {
+        self.statements.len()
+    }
+
+    fn clear_cached_statements(&mut self) -> BoxFuture<'_, Result<(), Error>> {
+        Box::pin(async move {
+            self.statements.clear();
+            Ok(())
+        })
+    }
+}
+
 impl Connection for SqliteConnection {
     type Database = Sqlite;
 

sqlx-core/src/sqlite/options.rs

@@ -1,5 +1,5 @@
 use std::path::PathBuf;
-use std::str::FromStr;
+use std::{io, str::FromStr};
 
 use crate::error::BoxDynError;
 
@@ -10,6 +10,7 @@ use crate::error::BoxDynError;
 pub struct SqliteConnectOptions {
     pub(crate) filename: PathBuf,
     pub(crate) in_memory: bool,
+    pub(crate) statement_cache_size: usize,
 }
 
 impl Default for SqliteConnectOptions {
@@ -23,6 +24,7 @@ impl SqliteConnectOptions {
         Self {
             filename: PathBuf::from(":memory:"),
             in_memory: false,
+            statement_cache_size: 100,
         }
     }
 }
@@ -34,17 +36,34 @@ impl FromStr for SqliteConnectOptions {
         let mut options = Self {
             filename: PathBuf::new(),
             in_memory: false,
+            statement_cache_size: 100,
         };
 
         // remove scheme
         s = s
             .trim_start_matches("sqlite://")
             .trim_start_matches("sqlite:");
 
-        if s == ":memory:" {
-            options.in_memory = true;
-        } else {
-            options.filename = s.parse()?;
+        let mut splitted = s.split("?");
+
+        match splitted.next() {
+            Some(":memory:") => options.in_memory = true,
+            Some(s) => options.filename = s.parse()?,
+            None => unreachable!(),
+        }
+
+        match splitted.next().map(|s| s.split("=")) {
+            Some(mut splitted) => {
+                if splitted.next() == Some("statement-cache-size") {
+                    options.statement_cache_size = splitted
+                        .next()
+                        .ok_or_else(|| {
+                            io::Error::new(io::ErrorKind::InvalidInput, "Invalid connection string")
+                        })?
+                        .parse()?
+                }
+            }
+            _ => (),
         }
 
         Ok(options)

tests/sqlite/sqlite.rs

@@ -1,6 +1,7 @@
 use futures::TryStreamExt;
 use sqlx::{
-    query, sqlite::Sqlite, Connect, Connection, Executor, Row, SqliteConnection, SqlitePool,
+    query, sqlite::Sqlite, CachingConnection, Connect, Connection, Executor, Row, SqliteConnection,
+    SqlitePool,
 };
 use sqlx_test::new;
 
@@ -269,3 +270,25 @@ SELECT id, text FROM _sqlx_test;
 
     Ok(())
 }
+
+#[sqlx_macros::test]
+async fn it_caches_statements() -> anyhow::Result<()> {
+    let mut conn = new::<Sqlite>().await?;
+
+    for i in 0..2 {
+        let row = sqlx::query("SELECT ? AS val")
+            .bind(i)
+            .fetch_one(&mut conn)
+            .await?;
+
+        let val: i32 = row.get("val");
+
+        assert_eq!(i, val);
+    }
+
+    assert_eq!(1, conn.cached_statements_count());
+    conn.clear_cached_statements().await?;
+    assert_eq!(0, conn.cached_statements_count());
+
+    Ok(())
+}