From ad18b8b4eab13d87002b1d331779fe639e604a5d Mon Sep 17 00:00:00 2001 From: 64673a <96230516+64673a@users.noreply.github.com> Date: Thu, 24 Sep 2026 16:27:03 +0800 Subject: [PATCH] fix(sqlite): run remote batch saves in a worker transaction --- crates/dbx-core/src/query/mod.rs | 107 +++++++++++++++++++++ crates/dbx-drivers/src/db/sqlite_worker.rs | 72 ++++++++++---- crates/dbx-sqlite-worker/src/protocol.rs | 21 +++- 3 files changed, 176 insertions(+), 24 deletions(-) diff --git a/crates/dbx-core/src/query/mod.rs b/crates/dbx-core/src/query/mod.rs index e538db1add..1f21dba027 100644 --- a/crates/dbx-core/src/query/mod.rs +++ b/crates/dbx-core/src/query/mod.rs @@ -5306,6 +5306,9 @@ async fn exec_tx_sqlite_inner( ) -> Result { let statements = statements.to_vec(); let query_timeout = budget.query_timeout; + if let Some(worker) = pool.worker() { + return exec_tx_sqlite_worker_inner(worker, statements, start, query_timeout).await; + } tokio::task::spawn_blocking(move || { pool.with_connection(|conn| { conn.execute_batch("BEGIN").map_err(|e| format!("Failed to begin transaction: {}", e))?; @@ -5476,6 +5479,55 @@ async fn exec_tx_sqlite_inner( .map_err(|e| e.to_string())? } +async fn exec_tx_sqlite_worker_inner( + worker: Arc, + statements: Vec, + start: std::time::Instant, + query_timeout: Option, +) -> Result { + // Detached like the local spawn_blocking path, so a dropped caller cannot leave the + // worker connection inside an open transaction. One session keeps other requests out. + tokio::spawn(async move { + let mut session = worker.session().await; + session.query("BEGIN", None).await.map_err(|e| format!("Failed to begin transaction: {e}"))?; + // ponytail: the worker cannot interrupt a running statement, so the budget is only + // checked between statements; bounding one slow statement needs a worker interrupt op. + let within_budget = || match query_timeout { + Some(timeout) if start.elapsed() >= timeout => { + Err(format!("Query timed out after {} seconds", timeout.as_secs())) + } + _ => Ok(()), + }; + let outcome: Result = async { + let mut total_affected = 0; + for (i, sql) in statements.iter().enumerate() { + within_budget()?; + total_affected += session + .query(sql, None) + .await + .map_err(|e| { + query_error_with_omitted_sql_context(&format!("Statement {} failed: {}", i + 1, e), sql) + })? + .affected_rows; + } + within_budget()?; + let committed = session.query("COMMIT", None).await.map_err(|e| format!("COMMIT failed: {e}"))?; + Ok(db::QueryResult { + affected_rows: total_affected, + execution_time_ms: start.elapsed().as_millis(), + ..committed + }) + } + .await; + if outcome.is_err() { + let _ = session.query("ROLLBACK", None).await; + } + outcome + }) + .await + .map_err(|e| e.to_string())? +} + async fn exec_tx_explicit_inner( state: &AppState, pool_key: &str, @@ -7304,6 +7356,61 @@ mod tests { assert!(error.contains("Statement 1 failed") && error.contains("interrupted_at"), "unexpected error: {error}"); } + #[tokio::test] + async fn sqlite_ssh_worker_transaction_rolls_back_on_failure_and_commits_on_success() { + let remote = Arc::new(Mutex::new(rusqlite::Connection::open_in_memory().expect("open remote SQLite"))); + remote.lock().unwrap().execute_batch("CREATE TABLE t (id INTEGER PRIMARY KEY)").expect("create table"); + let (client_stream, worker_stream) = tokio::io::duplex(64 * 1024); + tokio::spawn(fake_sqlite_ssh_worker(worker_stream, remote.clone())); + let pool = db::sqlite::SqliteHandle::from_worker(Arc::new( + db::sqlite_worker::SqliteWorkerClient::from_test_stream(client_stream), + )); + let budget = DbOperationBudget::with_defaults(); + let row_count = + || remote.lock().unwrap().query_row("SELECT COUNT(*) FROM t", [], |row| row.get::<_, i64>(0)).unwrap(); + + let error = exec_tx_sqlite_inner( + pool.clone(), + &["INSERT INTO t VALUES (1)".to_string(), "INSERT INTO missing VALUES (2)".to_string()], + std::time::Instant::now(), + &budget, + ) + .await + .expect_err("second statement fails"); + assert!(error.contains("Statement 2 failed") && error.contains("no such table: missing"), "{error}"); + assert_eq!(row_count(), 0); + + let result = exec_tx_sqlite_inner( + pool, + &["INSERT INTO t VALUES (1)".to_string(), "INSERT INTO t VALUES (2)".to_string()], + std::time::Instant::now(), + &budget, + ) + .await + .expect("transaction commits"); + assert_eq!(result.affected_rows, 2); + assert_eq!(row_count(), 2); + } + + /// Answers SQLite worker JSONL requests from a real SQLite connection. + async fn fake_sqlite_ssh_worker(stream: tokio::io::DuplexStream, conn: Arc>) { + use tokio::io::{AsyncBufReadExt, AsyncWriteExt}; + let (reader, mut writer) = tokio::io::split(stream); + let mut lines = tokio::io::BufReader::new(reader).lines(); + while let Ok(Some(line)) = lines.next_line().await { + let request: serde_json::Value = serde_json::from_str(&line).expect("worker request"); + let mut response = { + let conn = conn.lock().unwrap(); + match conn.execute_batch(request["sql"].as_str().unwrap_or_default()) { + Ok(()) => serde_json::json!({ "affected_rows": conn.changes() }), + Err(error) => serde_json::json!({ "error": error.to_string() }), + } + }; + response["id"] = request["id"].clone(); + writer.write_all(format!("{response}\n").as_bytes()).await.expect("worker response"); + } + } + #[tokio::test] async fn sqlite_external_interrupt_without_budget_does_not_panic() { // Regression: with query_timeout = None (MCP policy can set it to 0 = diff --git a/crates/dbx-drivers/src/db/sqlite_worker.rs b/crates/dbx-drivers/src/db/sqlite_worker.rs index deb6247081..9ab352033c 100644 --- a/crates/dbx-drivers/src/db/sqlite_worker.rs +++ b/crates/dbx-drivers/src/db/sqlite_worker.rs @@ -102,8 +102,15 @@ enum WorkerIo { Closed, } -impl SqliteWorkerClient { - pub async fn query(&self, sql: &str, max_rows: Option) -> Result { +/// Holds the worker I/O lock so consecutive requests, such as the statements of +/// one transaction, cannot interleave with other callers on the worker's single connection. +pub struct SqliteWorkerSession<'a> { + client: &'a SqliteWorkerClient, + io: tokio::sync::MutexGuard<'a, WorkerIo>, +} + +impl SqliteWorkerSession<'_> { + pub async fn query(&mut self, sql: &str, max_rows: Option) -> Result { match self.roundtrip(WorkerOp::Query { sql: sql.to_string(), max_rows }).await? { WorkerBody::Ok { columns, column_types, rows, affected_rows, truncated, .. } => Ok(QueryResult { columns: columns.unwrap_or_default(), @@ -125,6 +132,47 @@ impl SqliteWorkerClient { } } + async fn roundtrip(&mut self, op: WorkerOp) -> Result { + let id = self.client.next_id.fetch_add(1, Ordering::SeqCst); + let mut encoded = serde_json::to_vec(&WorkerRequest { id, op }).map_err(|e| e.to_string())?; + encoded.push(b'\n'); + match &mut *self.io { + WorkerIo::Process { stdin, stdout, .. } => { + stdin.write_all(&encoded).await.map_err(|e| e.to_string())?; + stdin.flush().await.map_err(|e| e.to_string())?; + let mut line = String::new(); + stdout.read_line(&mut line).await.map_err(|e| e.to_string())?; + parse_response(id, &line) + } + WorkerIo::Closed => Err("SQLite worker session is closed".to_string()), + WorkerIo::Ssh { stream, .. } => { + stream.write_all(&encoded).await.map_err(|e| e.to_string())?; + stream.flush().await.map_err(|e| e.to_string())?; + parse_response(id, &read_jsonl_line(stream).await?) + } + } + } +} + +impl SqliteWorkerClient { + #[cfg(any(test, feature = "test-support"))] + pub fn from_test_stream(stream: impl AsyncRead + AsyncWrite + Send + Unpin + 'static) -> Self { + Self { + io: AsyncMutex::new(WorkerIo::Ssh { stream: BufReader::new(Box::pin(stream)) }), + next_id: AtomicU64::new(1), + ssh_session: None, + remove_remote_path: None, + } + } + + pub async fn session(&self) -> SqliteWorkerSession<'_> { + SqliteWorkerSession { client: self, io: self.io.lock().await } + } + + pub async fn query(&self, sql: &str, max_rows: Option) -> Result { + self.session().await.query(sql, max_rows).await + } + pub async fn backup(&self, dest: &str) -> Result<(), String> { match self.roundtrip(WorkerOp::Backup { dest: dest.to_string() }).await? { WorkerBody::Ok { .. } => Ok(()), @@ -208,25 +256,7 @@ impl SqliteWorkerClient { } async fn roundtrip(&self, op: WorkerOp) -> Result { - let id = self.next_id.fetch_add(1, Ordering::SeqCst); - let mut encoded = serde_json::to_vec(&WorkerRequest { id, op }).map_err(|e| e.to_string())?; - encoded.push(b'\n'); - let mut io = self.io.lock().await; - match &mut *io { - WorkerIo::Process { stdin, stdout, .. } => { - stdin.write_all(&encoded).await.map_err(|e| e.to_string())?; - stdin.flush().await.map_err(|e| e.to_string())?; - let mut line = String::new(); - stdout.read_line(&mut line).await.map_err(|e| e.to_string())?; - parse_response(id, &line) - } - WorkerIo::Closed => Err("SQLite worker session is closed".to_string()), - WorkerIo::Ssh { stream, .. } => { - stream.write_all(&encoded).await.map_err(|e| e.to_string())?; - stream.flush().await.map_err(|e| e.to_string())?; - parse_response(id, &read_jsonl_line(stream).await?) - } - } + self.session().await.roundtrip(op).await } } diff --git a/crates/dbx-sqlite-worker/src/protocol.rs b/crates/dbx-sqlite-worker/src/protocol.rs index a61a194822..27a07551a7 100644 --- a/crates/dbx-sqlite-worker/src/protocol.rs +++ b/crates/dbx-sqlite-worker/src/protocol.rs @@ -28,6 +28,11 @@ pub struct WorkerResponse { #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] #[serde(untagged)] pub enum WorkerBody { + // Untagged variants are tried in order. `Ok` accepts any object because all of + // its fields are optional, so `Err` must come first or error payloads decode as success. + Err { + error: String, + }, Ok { #[serde(default, skip_serializing_if = "Option::is_none")] columns: Option>, @@ -42,9 +47,6 @@ pub enum WorkerBody { #[serde(default, skip_serializing_if = "Option::is_none")] pong: Option, }, - Err { - error: String, - }, } impl WorkerBody { @@ -97,4 +99,17 @@ mod tests { let decoded: WorkerRequest = serde_json::from_str(&encoded).unwrap(); assert_eq!(decoded, request); } + + #[test] + fn response_roundtrip_keeps_errors_distinct_from_success() { + for body in [ + WorkerBody::err("no such table: t"), + WorkerBody::ok(), + WorkerBody::query(vec!["id".into()], Vec::new(), vec![vec![serde_json::json!(1)]], 0, false), + ] { + let encoded = serde_json::to_string(&WorkerResponse { id: 7, body: body.clone() }).unwrap(); + let decoded: WorkerResponse = serde_json::from_str(&encoded).unwrap(); + assert_eq!(decoded.body, body, "{encoded}"); + } + } }