From 9cd9209aa366b826aee1ea2713e1e3122e2b19ab Mon Sep 17 00:00:00 2001
From: Andrew Whitehead <cywolf@gmail.com>
Date: Tue, 11 Aug 2020 19:25:13 -0700
Subject: [PATCH] add support for pooling in-memory sqlite DB, enabling shared
 cache

---
 sqlx-core/src/sqlite/connection/establish.rs | 10 ++++-
 sqlx-core/src/sqlite/options/mod.rs          |  2 +
 sqlx-core/src/sqlite/options/parse.rs        | 44 +++++++++++++++++++-
 3 files changed, 53 insertions(+), 3 deletions(-)

diff --git a/sqlx-core/src/sqlite/connection/establish.rs b/sqlx-core/src/sqlite/connection/establish.rs
index 954cac4f..3311019c 100644
--- a/sqlx-core/src/sqlite/connection/establish.rs
+++ b/sqlx-core/src/sqlite/connection/establish.rs
@@ -8,7 +8,7 @@ use crate::{
 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,
-    SQLITE_OPEN_READONLY, SQLITE_OPEN_READWRITE,
+    SQLITE_OPEN_READONLY, SQLITE_OPEN_READWRITE, SQLITE_OPEN_SHAREDCACHE,
 };
 use sqlx_rt::blocking;
 use std::io;
@@ -35,7 +35,7 @@ pub(crate) async fn establish(options: &SqliteConnectOptions) -> Result<SqliteCo
     // [SQLITE_OPEN_NOMUTEX] will instruct [sqlite3_open_v2] to return an error if it
     // cannot satisfy our wish for a thread-safe, lock-free connection object
 
-    let mut flags = SQLITE_OPEN_NOMUTEX | SQLITE_OPEN_PRIVATECACHE;
+    let mut flags = SQLITE_OPEN_NOMUTEX;
 
     flags |= if options.read_only {
         SQLITE_OPEN_READONLY
@@ -49,6 +49,12 @@ pub(crate) async fn establish(options: &SqliteConnectOptions) -> Result<SqliteCo
         flags |= SQLITE_OPEN_MEMORY;
     }
 
+    flags |= if options.shared_cache {
+        SQLITE_OPEN_SHAREDCACHE
+    } else {
+        SQLITE_OPEN_PRIVATECACHE
+    };
+
     let busy_timeout = options.busy_timeout;
 
     let handle = blocking!({
diff --git a/sqlx-core/src/sqlite/options/mod.rs b/sqlx-core/src/sqlite/options/mod.rs
index fce09967..824e0bcb 100644
--- a/sqlx-core/src/sqlite/options/mod.rs
+++ b/sqlx-core/src/sqlite/options/mod.rs
@@ -48,6 +48,7 @@ pub struct SqliteConnectOptions {
     pub(crate) create_if_missing: bool,
     pub(crate) journal_mode: SqliteJournalMode,
     pub(crate) foreign_keys: bool,
+    pub(crate) shared_cache: bool,
     pub(crate) statement_cache_capacity: usize,
     pub(crate) busy_timeout: Duration,
 }
@@ -66,6 +67,7 @@ impl SqliteConnectOptions {
             read_only: false,
             create_if_missing: false,
             foreign_keys: true,
+            shared_cache: false,
             statement_cache_capacity: 100,
             journal_mode: SqliteJournalMode::Wal,
             busy_timeout: Duration::from_secs(5),
diff --git a/sqlx-core/src/sqlite/options/parse.rs b/sqlx-core/src/sqlite/options/parse.rs
index 119d7c3b..7c21adf4 100644
--- a/sqlx-core/src/sqlite/options/parse.rs
+++ b/sqlx-core/src/sqlite/options/parse.rs
@@ -2,11 +2,14 @@ use crate::error::Error;
 use crate::sqlite::SqliteConnectOptions;
 use percent_encoding::percent_decode_str;
 use std::borrow::Cow;
-use std::path::Path;
+use std::path::{Path, PathBuf};
 use std::str::FromStr;
+use std::sync::atomic::{AtomicUsize, Ordering};
 
 // https://www.sqlite.org/uri.html
 
+static IN_MEMORY_DB_SEQ: AtomicUsize = AtomicUsize::new(0);
+
 impl FromStr for SqliteConnectOptions {
     type Err = Error;
 
@@ -24,6 +27,9 @@ impl FromStr for SqliteConnectOptions {
 
         if database == ":memory:" {
             options.in_memory = true;
+            options.shared_cache = true;
+            let seqno = IN_MEMORY_DB_SEQ.fetch_add(1, Ordering::Relaxed);
+            options.filename = Cow::Owned(PathBuf::from(format!("file:sqlx-in-memory-{}", seqno)));
         } else {
             // % decode to allow for `?` or `#` in the filename
             options.filename = Cow::Owned(
@@ -58,6 +64,7 @@ impl FromStr for SqliteConnectOptions {
 
                             "memory" => {
                                 options.in_memory = true;
+                                options.shared_cache = true;
                             }
 
                             _ => {
@@ -68,6 +75,25 @@ impl FromStr for SqliteConnectOptions {
                         }
                     }
 
+                    // The cache query parameter specifies the cache behaviour across multiple
+                    // connections to the same database within the process. A shared cache is
+                    // essential for persisting data across connections to an in-memory database.
+                    "cache" => match &*value {
+                        "private" => {
+                            options.shared_cache = false;
+                        }
+
+                        "shared" => {
+                            options.shared_cache = true;
+                        }
+
+                        _ => {
+                            return Err(Error::Configuration(
+                                format!("unknown value {:?} for `cache`", value).into(),
+                            ));
+                        }
+                    },
+
                     _ => {
                         return Err(Error::Configuration(
                             format!(
@@ -89,12 +115,19 @@ impl FromStr for SqliteConnectOptions {
 fn test_parse_in_memory() -> Result<(), Error> {
     let options: SqliteConnectOptions = "sqlite::memory:".parse()?;
     assert!(options.in_memory);
+    assert!(options.shared_cache);
 
     let options: SqliteConnectOptions = "sqlite://?mode=memory".parse()?;
     assert!(options.in_memory);
+    assert!(options.shared_cache);
 
     let options: SqliteConnectOptions = "sqlite://:memory:".parse()?;
     assert!(options.in_memory);
+    assert!(options.shared_cache);
+
+    let options: SqliteConnectOptions = "sqlite://?mode=memory&cache=private".parse()?;
+    assert!(options.in_memory);
+    assert!(!options.shared_cache);
 
     Ok(())
 }
@@ -107,3 +140,12 @@ fn test_parse_read_only() -> Result<(), Error> {
 
     Ok(())
 }
+
+#[test]
+fn test_parse_shared_in_memory() -> Result<(), Error> {
+    let options: SqliteConnectOptions = "sqlite://a.db?cache=shared".parse()?;
+    assert!(options.shared_cache);
+    assert_eq!(&*options.filename.to_string_lossy(), "a.db");
+
+    Ok(())
+}