From c984484d6520645c62a3d5f53f638ca5b67c263c Mon Sep 17 00:00:00 2001 From: Emilien Macchi Date: Thu, 3 Sep 2026 21:53:47 -0400 Subject: [PATCH] fix(server): preserve in-memory SQLite across reconnects Keep a dedicated SQLite connection alive so pool connection replacement retains the shared in-memory schema and objects. Closes #3173 Signed-off-by: Emilien Macchi --- architecture/gateway.md | 4 ++ .../src/persistence/sqlite.rs | 48 +++++++++++++++-- .../openshell-server/src/persistence/tests.rs | 52 +++++++++++++++++++ 3 files changed, 100 insertions(+), 4 deletions(-) diff --git a/architecture/gateway.md b/architecture/gateway.md index ba325ccc2c..769f57f6a0 100644 --- a/architecture/gateway.md +++ b/architecture/gateway.md @@ -336,6 +336,10 @@ This keeps the gateway data model portable across storage backends and leaves room for future stores that can provide the same object, label, version, and scope semantics. +For in-memory SQLite, the adapter retains a dedicated keepalive connection for +the store lifetime. Operational connection replacement therefore preserves the +shared in-memory schema and objects instead of creating an empty database. + The SQLite adapter tightens the on-disk database file to mode `0o600` on every connect so that provider API keys, SSH session tokens, and sandbox metadata are not readable by other local users on shared hosts. The same restriction is diff --git a/crates/openshell-server/src/persistence/sqlite.rs b/crates/openshell-server/src/persistence/sqlite.rs index c945c417f9..3e96040e34 100644 --- a/crates/openshell-server/src/persistence/sqlite.rs +++ b/crates/openshell-server/src/persistence/sqlite.rs @@ -14,35 +14,57 @@ use openshell_core::SetResourceVersion; use openshell_core::paths::set_file_owner_only; use openshell_core::proto::Sandbox; use prost::Message; -use sqlx::sqlite::{SqliteConnectOptions, SqlitePoolOptions}; +use sqlx::sqlite::{SqliteConnectOptions, SqliteConnection, SqlitePoolOptions}; use sqlx::{Connection, QueryBuilder, Row, Sqlite, SqlitePool}; use std::path::{Path, PathBuf}; use std::str::FromStr; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; +use tokio::sync::Mutex; static SQLITE_MIGRATOR: sqlx::migrate::Migrator = sqlx::migrate!("./migrations/sqlite"); +static IN_MEMORY_DB_SEQUENCE: AtomicU64 = AtomicU64::new(0); use super::{DELETE_MANY_BATCH_SIZE, DRAFT_CHUNK_OBJECT_TYPE, POLICY_OBJECT_TYPE}; #[derive(Debug, Clone)] pub struct SqliteStore { pool: SqlitePool, + #[cfg_attr(not(any(test, feature = "test-support")), allow(dead_code))] + in_memory_keepalive: Option>>>, +} + +#[cfg(test)] +pub(super) async fn replace_pool_connection(store: &SqliteStore) -> PersistenceResult<()> { + let connection = store.pool.acquire().await.map_err(|e| map_db_error(&e))?; + connection.close().await.map_err(|e| map_db_error(&e)) } impl SqliteStore { /// Closes the connection pool. #[cfg(test)] pub(crate) async fn close_for_test(&self) { - self.pool.close().await; + self.close().await; } pub async fn connect(url: &str) -> PersistenceResult { let is_in_memory = url.contains(":memory:") || url.contains("mode=memory"); let max_connections = if is_in_memory { 1 } else { 5 }; - let options = SqliteConnectOptions::from_str(url) + let mut options = SqliteConnectOptions::from_str(url) .map_err(|e| map_db_error(&e))? .create_if_missing(true); + if is_in_memory { + if options.get_filename().as_os_str().is_empty() + || options.get_filename() == Path::new(":memory:") + { + let sequence = IN_MEMORY_DB_SEQUENCE.fetch_add(1, Ordering::Relaxed); + options = options.filename(format!("file:openshell-in-memory-{sequence}")); + } + options = options.shared_cache(true); + } + let mut pool_options = SqlitePoolOptions::new() .max_connections(max_connections) .min_connections(max_connections); @@ -55,6 +77,15 @@ impl SqliteStore { // so we can restrict the permissions after the database is connected. let db_path = (!is_in_memory).then(|| options.get_filename().to_path_buf()); + let in_memory_keepalive = if is_in_memory { + let connection = SqliteConnection::connect_with(&options) + .await + .map_err(|e| map_db_error(&e))?; + Some(Arc::new(Mutex::new(Some(connection)))) + } else { + None + }; + let pool = pool_options .connect_with(options) .await @@ -65,7 +96,10 @@ impl SqliteStore { restrict_db_file_permissions(&path)?; } - Ok(Self { pool }) + Ok(Self { + pool, + in_memory_keepalive, + }) } pub async fn migrate(&self) -> PersistenceResult<()> { @@ -88,6 +122,12 @@ impl SqliteStore { #[cfg(any(test, feature = "test-support"))] pub async fn close(&self) { self.pool.close().await; + if let Some(keepalive) = &self.in_memory_keepalive { + let connection = keepalive.lock().await.take(); + if let Some(connection) = connection { + let _ = connection.close().await; + } + } } pub async fn put( diff --git a/crates/openshell-server/src/persistence/tests.rs b/crates/openshell-server/src/persistence/tests.rs index 7882c9246c..37fc469d66 100644 --- a/crates/openshell-server/src/persistence/tests.rs +++ b/crates/openshell-server/src/persistence/tests.rs @@ -132,6 +132,58 @@ async fn sqlite_connect_runs_embedded_migrations() { assert!(records.is_empty()); } +#[tokio::test] +async fn sqlite_in_memory_store_survives_pool_connection_replacement() { + for url in ["sqlite::memory:", "sqlite://?mode=memory"] { + let store = super::sqlite::SqliteStore::connect(url) + .await + .expect("connect to in-memory SQLite"); + store.migrate().await.expect("migrate in-memory SQLite"); + store + .put( + "sandbox", + "before-replacement", + "before-replacement", + "default", + b"before", + None, + ) + .await + .expect("write before connection replacement"); + + super::sqlite::replace_pool_connection(&store) + .await + .expect("replace operational pool connection"); + + let preserved = store + .get("sandbox", "before-replacement") + .await + .expect("schema survives connection replacement") + .expect("existing object survives connection replacement"); + assert_eq!(preserved.payload, b"before", "database URL: {url}"); + + store + .put( + "sandbox", + "after-replacement", + "after-replacement", + "default", + b"after", + None, + ) + .await + .expect("write after connection replacement"); + assert!( + store + .get("sandbox", "after-replacement") + .await + .expect("read after connection replacement") + .is_some(), + "database URL: {url}" + ); + } +} + #[cfg(unix)] #[tokio::test] async fn sqlite_connect_restricts_db_file_permissions() {