diff --git a/crates-cli/yaak-cli/src/commands/environment.rs b/crates-cli/yaak-cli/src/commands/environment.rs index 2c060a92..856417a7 100644 --- a/crates-cli/yaak-cli/src/commands/environment.rs +++ b/crates-cli/yaak-cli/src/commands/environment.rs @@ -49,8 +49,11 @@ fn schema(pretty: bool) -> CommandResult { fn list(ctx: &CliContext, workspace_id: Option<&str>) -> CommandResult { let workspace_id = resolve_workspace_id(ctx, workspace_id, "environment list")?; let environments = ctx - .db() - .list_environments_ensure_base(&workspace_id) + .query_manager() + .with_tx(|tx| { + tx.ensure_base_environment(&workspace_id)?; + tx.list_environments(&workspace_id) + }) .map_err(|e| format!("Failed to list environments: {e}"))?; if environments.is_empty() { @@ -111,8 +114,8 @@ fn create( } let created = ctx - .db() - .upsert_environment(&environment, &UpdateSource::Sync) + .query_manager() + .with_tx(|tx| tx.upsert_environment(&environment, &UpdateSource::Sync)) .map_err(|e| format!("Failed to create environment: {e}"))?; println!("Created environment: {}", created.id); @@ -133,8 +136,8 @@ fn create( }; let created = ctx - .db() - .upsert_environment(&environment, &UpdateSource::Sync) + .query_manager() + .with_tx(|tx| tx.upsert_environment(&environment, &UpdateSource::Sync)) .map_err(|e| format!("Failed to create environment: {e}"))?; println!("Created environment: {}", created.id); @@ -152,8 +155,8 @@ fn update(ctx: &CliContext, json: Option, json_input: Option) -> let updated = apply_merge_patch(&existing, &patch, &id, "environment update")?; let saved = ctx - .db() - .upsert_environment(&updated, &UpdateSource::Sync) + .query_manager() + .with_tx(|tx| tx.upsert_environment(&updated, &UpdateSource::Sync)) .map_err(|e| format!("Failed to update environment: {e}"))?; println!("Updated environment: {}", saved.id); @@ -167,8 +170,8 @@ fn delete(ctx: &CliContext, environment_id: &str, yes: bool) -> CommandResult { } let deleted = ctx - .db() - .delete_environment_by_id(environment_id, &UpdateSource::Sync) + .query_manager() + .with_tx(|tx| tx.delete_environment_by_id(environment_id, &UpdateSource::Sync)) .map_err(|e| format!("Failed to delete environment: {e}"))?; println!("Deleted environment: {}", deleted.id); diff --git a/crates-cli/yaak-cli/src/commands/folder.rs b/crates-cli/yaak-cli/src/commands/folder.rs index 71eab72c..7b90815d 100644 --- a/crates-cli/yaak-cli/src/commands/folder.rs +++ b/crates-cli/yaak-cli/src/commands/folder.rs @@ -102,8 +102,8 @@ fn create( )?; let created = ctx - .db() - .upsert_folder(&folder, &UpdateSource::Sync) + .query_manager() + .with_tx(|tx| tx.upsert_folder(&folder, &UpdateSource::Sync)) .map_err(|e| format!("Failed to create folder: {e}"))?; println!("Created folder: {}", created.id); @@ -118,8 +118,8 @@ fn create( let folder = Folder { workspace_id, name, ..Default::default() }; let created = ctx - .db() - .upsert_folder(&folder, &UpdateSource::Sync) + .query_manager() + .with_tx(|tx| tx.upsert_folder(&folder, &UpdateSource::Sync)) .map_err(|e| format!("Failed to create folder: {e}"))?; println!("Created folder: {}", created.id); @@ -135,8 +135,8 @@ fn update(ctx: &CliContext, json: Option, json_input: Option) -> let updated = apply_merge_patch(&existing, &patch, &id, "folder update")?; let saved = ctx - .db() - .upsert_folder(&updated, &UpdateSource::Sync) + .query_manager() + .with_tx(|tx| tx.upsert_folder(&updated, &UpdateSource::Sync)) .map_err(|e| format!("Failed to update folder: {e}"))?; println!("Updated folder: {}", saved.id); @@ -150,8 +150,8 @@ fn delete(ctx: &CliContext, folder_id: &str, yes: bool) -> CommandResult { } let deleted = ctx - .db() - .delete_folder_by_id(folder_id, &UpdateSource::Sync) + .query_manager() + .with_tx(|tx| tx.delete_folder_by_id(folder_id, &UpdateSource::Sync)) .map_err(|e| format!("Failed to delete folder: {e}"))?; println!("Deleted folder: {}", deleted.id); diff --git a/crates-cli/yaak-cli/src/commands/plugin.rs b/crates-cli/yaak-cli/src/commands/plugin.rs index 1e224f02..23e38454 100644 --- a/crates-cli/yaak-cli/src/commands/plugin.rs +++ b/crates-cli/yaak-cli/src/commands/plugin.rs @@ -348,17 +348,19 @@ async fn install_from_directory(context: &CliContext, source: &str) -> CommandRe ui::info(&format!("Installing plugin from directory {}", plugin_dir.display())); let plugin = context - .db() - .upsert_plugin( - &Plugin { - directory: plugin_dir_str, - url: None, - enabled: true, - source: PluginSource::Filesystem, - ..Default::default() - }, - &UpdateSource::Background, - ) + .query_manager() + .with_tx(|tx| { + tx.upsert_plugin( + &Plugin { + directory: plugin_dir_str, + url: None, + enabled: true, + source: PluginSource::Filesystem, + ..Default::default() + }, + &UpdateSource::Background, + ) + }) .map_err(|err| format!("Failed to save plugin in database: {err}"))?; let plugin_context = PluginContext::new(Some("cli".to_string()), None); diff --git a/crates-cli/yaak-cli/src/commands/request.rs b/crates-cli/yaak-cli/src/commands/request.rs index ec37789a..39070a73 100644 --- a/crates-cli/yaak-cli/src/commands/request.rs +++ b/crates-cli/yaak-cli/src/commands/request.rs @@ -424,8 +424,8 @@ fn create( )?; let created = ctx - .db() - .upsert_http_request(&request, &UpdateSource::Sync) + .query_manager() + .with_tx(|tx| tx.upsert_http_request(&request, &UpdateSource::Sync)) .map_err(|e| format!("Failed to create request: {e}"))?; println!("Created request: {}", created.id); @@ -443,8 +443,8 @@ fn create( } let created = ctx - .db() - .upsert_http_request(&request, &UpdateSource::Sync) + .query_manager() + .with_tx(|tx| tx.upsert_http_request(&request, &UpdateSource::Sync)) .map_err(|e| format!("Failed to create request: {e}"))?; println!("Created request: {}", created.id); @@ -463,8 +463,8 @@ fn update(ctx: &CliContext, json: Option, json_input: Option) -> let updated = apply_merge_patch(&existing, &patch, &id, "request update")?; let saved = ctx - .db() - .upsert_http_request(&updated, &UpdateSource::Sync) + .query_manager() + .with_tx(|tx| tx.upsert_http_request(&updated, &UpdateSource::Sync)) .map_err(|e| format!("Failed to update request: {e}"))?; println!("Updated request: {}", saved.id); @@ -487,8 +487,8 @@ fn delete(ctx: &CliContext, request_id: &str, yes: bool) -> CommandResult { } let deleted = ctx - .db() - .delete_http_request_by_id(request_id, &UpdateSource::Sync) + .query_manager() + .with_tx(|tx| tx.delete_http_request_by_id(request_id, &UpdateSource::Sync)) .map_err(|e| format!("Failed to delete request: {e}"))?; println!("Deleted request: {}", deleted.id); Ok(()) diff --git a/crates-cli/yaak-cli/src/commands/response.rs b/crates-cli/yaak-cli/src/commands/response.rs index 3a8b3ba8..57177a53 100644 --- a/crates-cli/yaak-cli/src/commands/response.rs +++ b/crates-cli/yaak-cli/src/commands/response.rs @@ -118,8 +118,8 @@ fn delete(ctx: &CliContext, id: &str, yes: bool) -> CommandResult { return Ok(()); } let count = ctx - .db() - .delete_all_http_responses_for_request(id, &UpdateSource::Sync) + .query_manager() + .with_tx(|tx| tx.delete_all_http_responses_for_request(id, &UpdateSource::Sync)) .map_err(|e| format!("Failed to delete responses: {e}"))?; println!("Deleted {count} responses for request {id}"); return Ok(()); @@ -131,8 +131,10 @@ fn delete(ctx: &CliContext, id: &str, yes: bool) -> CommandResult { return Ok(()); } let count = ctx - .db() - .delete_all_http_responses_for_workspace(&workspace_id, &UpdateSource::Sync) + .query_manager() + .with_tx(|tx| { + tx.delete_all_http_responses_for_workspace(&workspace_id, &UpdateSource::Sync) + }) .map_err(|e| format!("Failed to delete responses: {e}"))?; println!("Deleted {count} responses for workspace {workspace_id}"); Ok(()) diff --git a/crates-cli/yaak-cli/src/commands/workspace.rs b/crates-cli/yaak-cli/src/commands/workspace.rs index 9a05219b..cc1fc08c 100644 --- a/crates-cli/yaak-cli/src/commands/workspace.rs +++ b/crates-cli/yaak-cli/src/commands/workspace.rs @@ -84,8 +84,8 @@ fn create( .map_err(|e| format!("Failed to parse workspace create JSON: {e}"))?; let created = ctx - .db() - .upsert_workspace(&workspace, &UpdateSource::Sync) + .query_manager() + .with_tx(|tx| tx.upsert_workspace(&workspace, &UpdateSource::Sync)) .map_err(|e| format!("Failed to create workspace: {e}"))?; println!("Created workspace: {}", created.id); return Ok(()); @@ -97,8 +97,8 @@ fn create( let workspace = Workspace { name, ..Default::default() }; let created = ctx - .db() - .upsert_workspace(&workspace, &UpdateSource::Sync) + .query_manager() + .with_tx(|tx| tx.upsert_workspace(&workspace, &UpdateSource::Sync)) .map_err(|e| format!("Failed to create workspace: {e}"))?; println!("Created workspace: {}", created.id); Ok(()) @@ -115,8 +115,8 @@ fn update(ctx: &CliContext, json: Option, json_input: Option) -> let updated = apply_merge_patch(&existing, &patch, &id, "workspace update")?; let saved = ctx - .db() - .upsert_workspace(&updated, &UpdateSource::Sync) + .query_manager() + .with_tx(|tx| tx.upsert_workspace(&updated, &UpdateSource::Sync)) .map_err(|e| format!("Failed to update workspace: {e}"))?; println!("Updated workspace: {}", saved.id); @@ -130,8 +130,10 @@ fn delete(ctx: &CliContext, workspace_id: &str, yes: bool) -> CommandResult { } let deleted = ctx - .db() - .delete_workspace_by_id(workspace_id, &UpdateSource::Sync, ctx.blob_manager()) + .query_manager() + .with_tx(|tx| { + tx.delete_workspace_by_id(workspace_id, &UpdateSource::Sync, ctx.blob_manager()) + }) .map_err(|e| format!("Failed to delete workspace: {e}"))?; println!("Deleted workspace: {}", deleted.id); Ok(()) diff --git a/crates-cli/yaak-cli/src/context.rs b/crates-cli/yaak-cli/src/context.rs index be30439f..9822e0cb 100644 --- a/crates-cli/yaak-cli/src/context.rs +++ b/crates-cli/yaak-cli/src/context.rs @@ -51,11 +51,9 @@ impl CliContext { }; // Guest: the desktop may have this DB open, so only what's safe beside a live session - let _ = yaak_lifecycle::on_launch( - &yaak_lifecycle::Host::guest(), - &query_manager.connect(), - &blob_manager, - ); + let _ = query_manager.with_tx(|tx| { + yaak_lifecycle::on_launch(&yaak_lifecycle::Host::guest(), tx, &blob_manager) + }); let encryption_manager = Arc::new(EncryptionManager::new(query_manager.clone(), app_id)); diff --git a/crates-cli/yaak-cli/tests/common/mod.rs b/crates-cli/yaak-cli/tests/common/mod.rs index 2a874939..d500eca2 100644 --- a/crates-cli/yaak-cli/tests/common/mod.rs +++ b/crates-cli/yaak-cli/tests/common/mod.rs @@ -40,8 +40,7 @@ pub fn seed_workspace(data_dir: &Path, workspace_id: &str) { }; query_manager(data_dir) - .connect() - .upsert_workspace(&workspace, &UpdateSource::Sync) + .with_tx(|tx| tx.upsert_workspace(&workspace, &UpdateSource::Sync)) .expect("Failed to seed workspace"); } @@ -56,8 +55,7 @@ pub fn seed_request(data_dir: &Path, workspace_id: &str, request_id: &str) { }; query_manager(data_dir) - .connect() - .upsert_http_request(&request, &UpdateSource::Sync) + .with_tx(|tx| tx.upsert_http_request(&request, &UpdateSource::Sync)) .expect("Failed to seed request"); } @@ -70,8 +68,7 @@ pub fn seed_folder(data_dir: &Path, workspace_id: &str, folder_id: &str) { }; query_manager(data_dir) - .connect() - .upsert_folder(&folder, &UpdateSource::Sync) + .with_tx(|tx| tx.upsert_folder(&folder, &UpdateSource::Sync)) .expect("Failed to seed folder"); } @@ -85,8 +82,7 @@ pub fn seed_grpc_request(data_dir: &Path, workspace_id: &str, request_id: &str) }; query_manager(data_dir) - .connect() - .upsert_grpc_request(&request, &UpdateSource::Sync) + .with_tx(|tx| tx.upsert_grpc_request(&request, &UpdateSource::Sync)) .expect("Failed to seed gRPC request"); } @@ -100,7 +96,6 @@ pub fn seed_websocket_request(data_dir: &Path, workspace_id: &str, request_id: & }; query_manager(data_dir) - .connect() - .upsert_websocket_request(&request, &UpdateSource::Sync) + .with_tx(|tx| tx.upsert_websocket_request(&request, &UpdateSource::Sync)) .expect("Failed to seed WebSocket request"); } diff --git a/crates-cli/yaak-cli/tests/import_export_commands.rs b/crates-cli/yaak-cli/tests/import_export_commands.rs index 5f9d44a2..922e04d5 100644 --- a/crates-cli/yaak-cli/tests/import_export_commands.rs +++ b/crates-cli/yaak-cli/tests/import_export_commands.rs @@ -161,8 +161,7 @@ fn import_postman_environment_uses_workspace_id() { let query_manager = query_manager(data_dir); let db = query_manager.connect(); - let environments = - db.list_environments_ensure_base(&workspace_id).expect("list imported environments"); + let environments = db.list_environments(&workspace_id).expect("list imported environments"); let imported_environment = environments.iter().find(|e| e.name == "Local").expect("postman environment imported"); @@ -299,7 +298,10 @@ fn re_import_leaves_deleted_resources_alone() { .into_iter() .find(|r| r.name == "Request B") .expect("request B imported"); - db.delete_http_request_by_id(&request_b.id, &UpdateSource::Sync).expect("delete request B"); + drop(db); + query_manager + .with_tx(|tx| tx.delete_http_request_by_id(&request_b.id, &UpdateSource::Sync)) + .expect("delete request B"); workspace_id }; diff --git a/crates-cli/yaak-cli/tests/send_commands.rs b/crates-cli/yaak-cli/tests/send_commands.rs index 6ed03382..2ff09b4b 100644 --- a/crates-cli/yaak-cli/tests/send_commands.rs +++ b/crates-cli/yaak-cli/tests/send_commands.rs @@ -25,8 +25,7 @@ fn top_level_send_folder_sends_http_requests_and_prints_summary() { ..Default::default() }; query_manager(data_dir) - .connect() - .upsert_http_request(&request, &UpdateSource::Sync) + .with_tx(|tx| tx.upsert_http_request(&request, &UpdateSource::Sync)) .expect("Failed to seed folder request"); cli_cmd(data_dir) diff --git a/crates-tauri/yaak-app-client/src/lib.rs b/crates-tauri/yaak-app-client/src/lib.rs index 2003cdbb..4c29fd0d 100644 --- a/crates-tauri/yaak-app-client/src/lib.rs +++ b/crates-tauri/yaak-app-client/src/lib.rs @@ -330,18 +330,20 @@ async fn cmd_grpc_go( let settings = app_handle.db().get_settings(); let client_cert = find_client_certificate(&request.url, &settings.client_certificates); - let conn = app_handle.db().upsert_grpc_connection( - &GrpcConnection { - workspace_id: request.workspace_id.clone(), - request_id: request.id.clone(), - status: -1, - elapsed: 0, - state: GrpcConnectionState::Initialized, - url: request.url.clone(), - ..Default::default() - }, - &UpdateSource::from_window_label(window.label()), - )?; + let conn = app_handle.with_tx(|tx| { + tx.upsert_grpc_connection( + &GrpcConnection { + workspace_id: request.workspace_id.clone(), + request_id: request.id.clone(), + status: -1, + elapsed: 0, + state: GrpcConnectionState::Initialized, + url: request.url.clone(), + ..Default::default() + }, + &UpdateSource::from_window_label(window.label()), + ) + })?; let conn_id = conn.id.clone(); @@ -386,15 +388,17 @@ async fn cmd_grpc_go( let connection = match connection { Ok(c) => c, Err(err) => { - app_handle.db().upsert_grpc_connection( - &GrpcConnection { - elapsed: start.elapsed().as_millis() as i32, - error: Some(err.to_string()), - state: GrpcConnectionState::Closed, - ..conn.clone() - }, - &UpdateSource::from_window_label(window.label()), - )?; + app_handle.with_tx(|tx| { + tx.upsert_grpc_connection( + &GrpcConnection { + elapsed: start.elapsed().as_millis() as i32, + error: Some(err.to_string()), + state: GrpcConnectionState::Closed, + ..conn.clone() + }, + &UpdateSource::from_window_label(window.label()), + ) + })?; return Ok(conn_id); } }; @@ -495,15 +499,17 @@ async fn cmd_grpc_go( .await?; let msg = strip_json_comments(&msg); - app_handle.db().upsert_grpc_event( - &GrpcEvent { - content: format!("Connecting to {}", req.url), - event_type: GrpcEventType::ConnectionStart, - metadata: metadata.clone(), - ..base_event.clone() - }, - &UpdateSource::from_window_label(window.label()), - )?; + app_handle.with_tx(|tx| { + tx.upsert_grpc_event( + &GrpcEvent { + content: format!("Connecting to {}", req.url), + event_type: GrpcEventType::ConnectionStart, + metadata: metadata.clone(), + ..base_event.clone() + }, + &UpdateSource::from_window_label(window.label()), + ) + })?; async move { // Create callback for streaming methods that handles both success and error @@ -513,24 +519,28 @@ async fn cmd_grpc_go( let window_label = window.label().to_string(); move |result: std::result::Result| match result { Ok(msg) => { - let _ = app_handle.db().upsert_grpc_event( - &GrpcEvent { - content: msg, - event_type: GrpcEventType::ClientMessage, - ..base_event.clone() - }, - &UpdateSource::from_window_label(&window_label), - ); + let _ = app_handle.with_tx(|tx| { + tx.upsert_grpc_event( + &GrpcEvent { + content: msg, + event_type: GrpcEventType::ClientMessage, + ..base_event.clone() + }, + &UpdateSource::from_window_label(&window_label), + ) + }); } Err(error) => { - let _ = app_handle.db().upsert_grpc_event( - &GrpcEvent { - content: format!("Failed to send message: {}", error), - event_type: GrpcEventType::Error, - ..base_event.clone() - }, - &UpdateSource::from_window_label(&window_label), - ); + let _ = app_handle.with_tx(|tx| { + tx.upsert_grpc_event( + &GrpcEvent { + content: format!("Failed to send message: {}", error), + event_type: GrpcEventType::Error, + ..base_event.clone() + }, + &UpdateSource::from_window_label(&window_label), + ) + }); } } }; @@ -583,36 +593,38 @@ async fn cmd_grpc_go( if !method_desc.is_client_streaming() { app_handle - .db() - .upsert_grpc_event( - &GrpcEvent { - event_type: GrpcEventType::ClientMessage, - content: msg, - ..base_event.clone() - }, - &UpdateSource::from_window_label(window.label()), - ) + .with_tx(|tx| { + tx.upsert_grpc_event( + &GrpcEvent { + event_type: GrpcEventType::ClientMessage, + content: msg, + ..base_event.clone() + }, + &UpdateSource::from_window_label(window.label()), + ) + }) .unwrap(); } match maybe_msg { Some(Ok(msg)) => { app_handle - .db() - .upsert_grpc_event( - &GrpcEvent { - metadata: metadata_to_map(msg.metadata().clone()), - content: if msg.metadata().len() == 0 { - "Received response" - } else { - "Received response with metadata" - } - .to_string(), - event_type: GrpcEventType::Info, - ..base_event.clone() - }, - &UpdateSource::from_window_label(window.label()), - ) + .with_tx(|tx| { + tx.upsert_grpc_event( + &GrpcEvent { + metadata: metadata_to_map(msg.metadata().clone()), + content: if msg.metadata().len() == 0 { + "Received response" + } else { + "Received response with metadata" + } + .to_string(), + event_type: GrpcEventType::Info, + ..base_event.clone() + }, + &UpdateSource::from_window_label(window.label()), + ) + }) .unwrap(); let response_message = msg.into_inner(); let content = match connection @@ -622,83 +634,88 @@ async fn cmd_grpc_go( Ok(content) => content, Err(err) => { app_handle - .db() - .upsert_grpc_event( - &GrpcEvent { - content: "Failed to read response".to_string(), - error: Some(err.to_string()), - status: Some(Code::Internal as i32), - event_type: GrpcEventType::ConnectionEnd, - ..base_event.clone() - }, - &UpdateSource::from_window_label(window.label()), - ) + .with_tx(|tx| { + tx.upsert_grpc_event( + &GrpcEvent { + content: "Failed to read response".to_string(), + error: Some(err.to_string()), + status: Some(Code::Internal as i32), + event_type: GrpcEventType::ConnectionEnd, + ..base_event.clone() + }, + &UpdateSource::from_window_label(window.label()), + ) + }) .unwrap(); return; } }; app_handle - .db() - .upsert_grpc_event( - &GrpcEvent { - content, - event_type: GrpcEventType::ServerMessage, - ..base_event.clone() - }, - &UpdateSource::from_window_label(window.label()), - ) + .with_tx(|tx| { + tx.upsert_grpc_event( + &GrpcEvent { + content, + event_type: GrpcEventType::ServerMessage, + ..base_event.clone() + }, + &UpdateSource::from_window_label(window.label()), + ) + }) .unwrap(); app_handle - .db() - .upsert_grpc_event( - &GrpcEvent { - content: "Connection complete".to_string(), - event_type: GrpcEventType::ConnectionEnd, - status: Some(Code::Ok as i32), - ..base_event.clone() - }, - &UpdateSource::from_window_label(window.label()), - ) + .with_tx(|tx| { + tx.upsert_grpc_event( + &GrpcEvent { + content: "Connection complete".to_string(), + event_type: GrpcEventType::ConnectionEnd, + status: Some(Code::Ok as i32), + ..base_event.clone() + }, + &UpdateSource::from_window_label(window.label()), + ) + }) .unwrap(); } Some(Err(yaak_grpc::error::Error::GrpcStreamError(e))) => { app_handle - .db() - .upsert_grpc_event( - &(match e.status { - Some(s) => GrpcEvent { - error: Some(s.message().to_string()), - status: Some(s.code() as i32), - content: "Request failed".to_string(), - metadata: metadata_to_map(s.metadata().clone()), - event_type: GrpcEventType::ConnectionEnd, - ..base_event.clone() - }, - None => GrpcEvent { - error: Some(e.message), + .with_tx(|tx| { + tx.upsert_grpc_event( + &(match e.status { + Some(s) => GrpcEvent { + error: Some(s.message().to_string()), + status: Some(s.code() as i32), + content: "Request failed".to_string(), + metadata: metadata_to_map(s.metadata().clone()), + event_type: GrpcEventType::ConnectionEnd, + ..base_event.clone() + }, + None => GrpcEvent { + error: Some(e.message), + status: Some(Code::Unknown as i32), + content: "Request failed".to_string(), + event_type: GrpcEventType::ConnectionEnd, + ..base_event.clone() + }, + }), + &UpdateSource::from_window_label(window.label()), + ) + }) + .unwrap(); + } + Some(Err(e)) => { + app_handle + .with_tx(|tx| { + tx.upsert_grpc_event( + &GrpcEvent { + error: Some(e.to_string()), status: Some(Code::Unknown as i32), content: "Request failed".to_string(), event_type: GrpcEventType::ConnectionEnd, ..base_event.clone() }, - }), - &UpdateSource::from_window_label(window.label()), - ) - .unwrap(); - } - Some(Err(e)) => { - app_handle - .db() - .upsert_grpc_event( - &GrpcEvent { - error: Some(e.to_string()), - status: Some(Code::Unknown as i32), - content: "Request failed".to_string(), - event_type: GrpcEventType::ConnectionEnd, - ..base_event.clone() - }, - &UpdateSource::from_window_label(window.label()), - ) + &UpdateSource::from_window_label(window.label()), + ) + }) .unwrap(); } None => { @@ -709,64 +726,67 @@ async fn cmd_grpc_go( let mut stream = match maybe_stream { Some(Ok(stream)) => { app_handle - .db() - .upsert_grpc_event( - &GrpcEvent { - metadata: metadata_to_map(stream.metadata().clone()), - content: if stream.metadata().len() == 0 { - "Received response" - } else { - "Received response with metadata" - } - .to_string(), - event_type: GrpcEventType::Info, - ..base_event.clone() - }, - &UpdateSource::from_window_label(window.label()), - ) + .with_tx(|tx| { + tx.upsert_grpc_event( + &GrpcEvent { + metadata: metadata_to_map(stream.metadata().clone()), + content: if stream.metadata().len() == 0 { + "Received response" + } else { + "Received response with metadata" + } + .to_string(), + event_type: GrpcEventType::Info, + ..base_event.clone() + }, + &UpdateSource::from_window_label(window.label()), + ) + }) .unwrap(); stream.into_inner() } Some(Err(yaak_grpc::error::Error::GrpcStreamError(e))) => { warn!("GRPC stream error {e:?}"); app_handle - .db() - .upsert_grpc_event( - &(match e.status { - Some(s) => GrpcEvent { - error: Some(s.message().to_string()), - status: Some(s.code() as i32), - content: "Stream failed".to_string(), - metadata: metadata_to_map(s.metadata().clone()), - event_type: GrpcEventType::ConnectionEnd, - ..base_event.clone() - }, - None => GrpcEvent { - error: Some(e.message), - status: Some(Code::Unknown as i32), - content: "Stream failed".to_string(), - event_type: GrpcEventType::ConnectionEnd, - ..base_event.clone() - }, - }), - &UpdateSource::from_window_label(window.label()), - ) + .with_tx(|tx| { + tx.upsert_grpc_event( + &(match e.status { + Some(s) => GrpcEvent { + error: Some(s.message().to_string()), + status: Some(s.code() as i32), + content: "Stream failed".to_string(), + metadata: metadata_to_map(s.metadata().clone()), + event_type: GrpcEventType::ConnectionEnd, + ..base_event.clone() + }, + None => GrpcEvent { + error: Some(e.message), + status: Some(Code::Unknown as i32), + content: "Stream failed".to_string(), + event_type: GrpcEventType::ConnectionEnd, + ..base_event.clone() + }, + }), + &UpdateSource::from_window_label(window.label()), + ) + }) .unwrap(); return; } Some(Err(e)) => { app_handle - .db() - .upsert_grpc_event( - &GrpcEvent { - error: Some(e.to_string()), - status: Some(Code::Unknown as i32), - content: "Stream failed".to_string(), - event_type: GrpcEventType::ConnectionEnd, - ..base_event.clone() - }, - &UpdateSource::from_window_label(window.label()), - ) + .with_tx(|tx| { + tx.upsert_grpc_event( + &GrpcEvent { + error: Some(e.to_string()), + status: Some(Code::Unknown as i32), + content: "Stream failed".to_string(), + event_type: GrpcEventType::ConnectionEnd, + ..base_event.clone() + }, + &UpdateSource::from_window_label(window.label()), + ) + }) .unwrap(); return; } @@ -783,65 +803,69 @@ async fn cmd_grpc_go( Ok(message) => message, Err(err) => { app_handle - .db() - .upsert_grpc_event( - &GrpcEvent { - content: "Failed to read response".to_string(), - error: Some(err.to_string()), - status: Some(Code::Internal as i32), - event_type: GrpcEventType::ConnectionEnd, - ..base_event.clone() - }, - &UpdateSource::from_window_label(window.label()), - ) + .with_tx(|tx| { + tx.upsert_grpc_event( + &GrpcEvent { + content: "Failed to read response".to_string(), + error: Some(err.to_string()), + status: Some(Code::Internal as i32), + event_type: GrpcEventType::ConnectionEnd, + ..base_event.clone() + }, + &UpdateSource::from_window_label(window.label()), + ) + }) .unwrap(); break; } }; app_handle - .db() - .upsert_grpc_event( - &GrpcEvent { - content: message, - event_type: GrpcEventType::ServerMessage, - ..base_event.clone() - }, - &UpdateSource::from_window_label(window.label()), - ) + .with_tx(|tx| { + tx.upsert_grpc_event( + &GrpcEvent { + content: message, + event_type: GrpcEventType::ServerMessage, + ..base_event.clone() + }, + &UpdateSource::from_window_label(window.label()), + ) + }) .unwrap(); } Ok(None) => { let trailers = stream.trailers().await.unwrap_or_default().unwrap_or_default(); app_handle - .db() - .upsert_grpc_event( - &GrpcEvent { - content: "Connection complete".to_string(), - status: Some(Code::Ok as i32), - metadata: metadata_to_map(trailers), - event_type: GrpcEventType::ConnectionEnd, - ..base_event.clone() - }, - &UpdateSource::from_window_label(window.label()), - ) + .with_tx(|tx| { + tx.upsert_grpc_event( + &GrpcEvent { + content: "Connection complete".to_string(), + status: Some(Code::Ok as i32), + metadata: metadata_to_map(trailers), + event_type: GrpcEventType::ConnectionEnd, + ..base_event.clone() + }, + &UpdateSource::from_window_label(window.label()), + ) + }) .unwrap(); break; } Err(status) => { app_handle - .db() - .upsert_grpc_event( - &GrpcEvent { - content: "Stream failed".to_string(), - error: Some(status.message().to_string()), - status: Some(status.code() as i32), - metadata: metadata_to_map(status.metadata().clone()), - event_type: GrpcEventType::ConnectionEnd, - ..base_event.clone() - }, - &UpdateSource::from_window_label(window.label()), - ) + .with_tx(|tx| { + tx.upsert_grpc_event( + &GrpcEvent { + content: "Stream failed".to_string(), + error: Some(status.message().to_string()), + status: Some(status.code() as i32), + metadata: metadata_to_map(status.metadata().clone()), + event_type: GrpcEventType::ConnectionEnd, + ..base_event.clone() + }, + &UpdateSource::from_window_label(window.label()), + ) + }) .unwrap(); break; } @@ -874,21 +898,24 @@ async fn cmd_grpc_go( }).unwrap(); }, _ = cancelled_rx.changed() => { - w.db().upsert_grpc_event( - &GrpcEvent { - content: "Cancelled".to_string(), - event_type: GrpcEventType::ConnectionEnd, - status: Some(Code::Cancelled as i32), - ..base_msg.clone() - }, - &UpdateSource::from_window_label(window.label()), - ).unwrap(); + w.with_tx(|tx| { + tx.upsert_grpc_event( + &GrpcEvent { + content: "Cancelled".to_string(), + event_type: GrpcEventType::ConnectionEnd, + status: Some(Code::Cancelled as i32), + ..base_msg.clone() + }, + &UpdateSource::from_window_label(window.label()), + ) + }) + .unwrap(); w.with_tx(|c| { c.upsert_grpc_connection( &GrpcConnection{ - elapsed: start.elapsed().as_millis() as i32, - status: Code::Cancelled as i32, - state: GrpcConnectionState::Closed, + elapsed: start.elapsed().as_millis() as i32, + status: Code::Cancelled as i32, + state: GrpcConnectionState::Closed, ..c.get_grpc_connection( &conn_id).unwrap().clone() }, &UpdateSource::from_window_label(window.label()), @@ -1066,15 +1093,17 @@ async fn cmd_send_http_request( let request = app_handle.db().get_http_request(&request_id)?; let blobs = app_handle.blob_manager(); - let response = app_handle.db().upsert_http_response( - &HttpResponse { - request_id: request.id.clone(), - workspace_id: request.workspace_id.clone(), - ..Default::default() - }, - &UpdateSource::from_window_label(window.label()), - &blobs, - )?; + let response = app_handle.with_tx(|tx| { + tx.upsert_http_response( + &HttpResponse { + request_id: request.id.clone(), + workspace_id: request.workspace_id.clone(), + ..Default::default() + }, + &UpdateSource::from_window_label(window.label()), + &blobs, + ) + })?; let (cancel_tx, mut cancel_rx) = tokio::sync::watch::channel(false); app_handle.listen_any(format!("cancel_http_response_{}", response.id), move |_event| { @@ -1112,15 +1141,17 @@ async fn cmd_send_http_request( Ok(sent) => sent.response, Err(e) => { let resp = app_handle.db().get_http_response(&response.id)?; - app_handle.db().upsert_http_response( - &HttpResponse { - state: HttpResponseState::Closed, - error: Some(e.to_string()), - ..resp - }, - &UpdateSource::from_window_label(window.label()), - &blobs, - )? + app_handle.with_tx(|tx| { + tx.upsert_http_response( + &HttpResponse { + state: HttpResponseState::Closed, + error: Some(e.to_string()), + ..resp + }, + &UpdateSource::from_window_label(window.label()), + &blobs, + ) + })? } }; @@ -1254,8 +1285,8 @@ pub fn run() { .setup(|app| { let lifecycle_host = yaak_lifecycle::Host::owner() .with_responses_dir(app.path().app_data_dir()?.join("responses")); - if let Err(e) = - yaak_lifecycle::on_launch(&lifecycle_host, &app.db(), &app.blob_manager()) + if let Err(e) = app + .with_tx(|tx| yaak_lifecycle::on_launch(&lifecycle_host, tx, &app.blob_manager())) { error!("on_launch hook failed: {e:?}"); } diff --git a/crates-tauri/yaak-app-client/src/models_ext.rs b/crates-tauri/yaak-app-client/src/models_ext.rs index 835547dd..5584d14e 100644 --- a/crates-tauri/yaak-app-client/src/models_ext.rs +++ b/crates-tauri/yaak-app-client/src/models_ext.rs @@ -10,7 +10,7 @@ use tauri::plugin::TauriPlugin; use tauri::{Emitter, Manager, Runtime, State}; use tauri_plugin_dialog::{DialogExt, MessageDialogKind}; use yaak_models::blob_manager::BlobManager; -use yaak_models::client_db::ClientDb; +use yaak_models::client_db::{ClientDb, WriteDb}; use yaak_models::error::Result; use yaak_models::query_manager::QueryManager; use yaak_models::util::{ModelPayload, UpdateSource}; @@ -95,7 +95,7 @@ pub trait QueryManagerExt<'a, R> { fn db(&'a self) -> ClientDb<'a>; fn with_tx(&'a self, func: F) -> Result where - F: FnOnce(&ClientDb) -> Result; + F: FnOnce(&WriteDb) -> Result; } impl<'a, R: Runtime, M: Manager> QueryManagerExt<'a, R> for M { @@ -110,7 +110,7 @@ impl<'a, R: Runtime, M: Manager> QueryManagerExt<'a, R> for M { fn with_tx(&'a self, func: F) -> Result where - F: FnOnce(&ClientDb) -> Result, + F: FnOnce(&WriteDb) -> Result, { let qm = self.state::(); qm.inner().with_tx(func) diff --git a/crates-tauri/yaak-app-client/src/notifications.rs b/crates-tauri/yaak-app-client/src/notifications.rs index c653c55e..829afe42 100644 --- a/crates-tauri/yaak-app-client/src/notifications.rs +++ b/crates-tauri/yaak-app-client/src/notifications.rs @@ -55,12 +55,15 @@ impl YaakNotifier { seen.push(id.to_string()); debug!("Marked notification as seen {}", id); let seen_json = serde_json::to_string(&seen)?; - window.db().set_key_value_raw( - KV_NAMESPACE, - KV_KEY, - seen_json.as_str(), - &UpdateSource::from_window_label(window.label()), - ); + window.with_tx(|tx| { + tx.set_key_value_raw( + KV_NAMESPACE, + KV_KEY, + seen_json.as_str(), + &UpdateSource::from_window_label(window.label()), + ); + Ok(()) + })?; Ok(()) } diff --git a/crates-tauri/yaak-app-client/src/plugin_events.rs b/crates-tauri/yaak-app-client/src/plugin_events.rs index 9071313c..3f891922 100644 --- a/crates-tauri/yaak-app-client/src/plugin_events.rs +++ b/crates-tauri/yaak-app-client/src/plugin_events.rs @@ -111,7 +111,7 @@ async fn handle_host_plugin_request( } let new_plugin = Plugin { updated_at: Utc::now().naive_utc(), ..plugin }; - app_handle.db().upsert_plugin(&new_plugin, &UpdateSource::Plugin)?; + app_handle.with_tx(|tx| tx.upsert_plugin(&new_plugin, &UpdateSource::Plugin))?; } if !req.silent { @@ -294,15 +294,17 @@ async fn handle_host_plugin_request( HttpResponse::default() } else { let blobs = window.blob_manager(); - window.db().upsert_http_response( - &HttpResponse { - request_id: http_request.id.clone(), - workspace_id: http_request.workspace_id.clone(), - ..Default::default() - }, - &UpdateSource::from_window_label(window.label()), - &blobs, - )? + window.with_tx(|tx| { + tx.upsert_http_response( + &HttpResponse { + request_id: http_request.id.clone(), + workspace_id: http_request.workspace_id.clone(), + ..Default::default() + }, + &UpdateSource::from_window_label(window.label()), + &blobs, + ) + })? }; let http_response = send_http_request_with_context( diff --git a/crates-tauri/yaak-app-client/src/plugins_ext.rs b/crates-tauri/yaak-app-client/src/plugins_ext.rs index e1ca87ce..8e26c647 100644 --- a/crates-tauri/yaak-app-client/src/plugins_ext.rs +++ b/crates-tauri/yaak-app-client/src/plugins_ext.rs @@ -202,16 +202,18 @@ pub async fn cmd_plugins_install_from_directory( // Resolve the manager before writing the row so startup's plugin snapshot // can't include it and boot it a second time let plugin_manager = Arc::new(plugin_manager(&window).await?); - let plugin = window.db().upsert_plugin( - &Plugin { - directory: directory.into(), - url: None, - enabled: true, - source: PluginSource::Filesystem, - ..Default::default() - }, - &UpdateSource::from_window_label(window.label()), - )?; + let plugin = window.with_tx(|tx| { + tx.upsert_plugin( + &Plugin { + directory: directory.into(), + url: None, + enabled: true, + source: PluginSource::Filesystem, + ..Default::default() + }, + &UpdateSource::from_window_label(window.label()), + ) + })?; plugin_manager.add_plugin(&window.plugin_context(), &plugin).await?; diff --git a/crates-tauri/yaak-app-client/src/sync_ext.rs b/crates-tauri/yaak-app-client/src/sync_ext.rs index 223573ce..ab26fccc 100644 --- a/crates-tauri/yaak-app-client/src/sync_ext.rs +++ b/crates-tauri/yaak-app-client/src/sync_ext.rs @@ -12,8 +12,8 @@ use tokio::sync::watch; use yaak_rpc_schema::WatchResult; use yaak_sync::error::Error::InvalidSyncDirectory; use yaak_sync::sync::{ - FsCandidate, SyncOp, apply_sync_ops, apply_sync_state_ops, compute_sync_ops, get_db_candidates, - get_fs_candidates, + FsCandidate, SyncOp, apply_db_sync_ops, apply_fs_sync_ops, apply_sync_state_ops, + compute_sync_ops, get_db_candidates, get_fs_candidates, }; use yaak_sync::watch::{WatchEvent, watch_directory}; @@ -49,11 +49,14 @@ pub(crate) async fn cmd_sync_apply( sync_dir: &Path, workspace_id: &str, ) -> Result<()> { - let db = app_handle.db(); + // Files first, so the write transaction never waits on the filesystem + let pending = apply_fs_sync_ops(workspace_id, sync_dir, sync_ops)?; let blobs = app_handle.blob_manager(); - let sync_state_ops = apply_sync_ops(&db, &blobs, workspace_id, sync_dir, sync_ops)?; - apply_sync_state_ops(&db, workspace_id, sync_dir, sync_state_ops)?; - Ok(()) + app_handle.db_manager().with_tx(|tx| { + let sync_state_ops = apply_db_sync_ops(tx, &blobs, workspace_id, sync_dir, pending)?; + apply_sync_state_ops(tx, workspace_id, sync_dir, sync_state_ops)?; + Ok(()) + }) } pub(crate) async fn sync_watch( diff --git a/crates-tauri/yaak-app-client/src/ws_ext.rs b/crates-tauri/yaak-app-client/src/ws_ext.rs index 46419678..45980139 100644 --- a/crates-tauri/yaak-app-client/src/ws_ext.rs +++ b/crates-tauri/yaak-app-client/src/ws_ext.rs @@ -43,18 +43,20 @@ pub async fn cmd_ws_send( { Ok(connection) => Ok(connection), Err(e) => { - app_handle.db().upsert_websocket_event( - &WebsocketEvent { - connection_id: connection.id.clone(), - request_id: connection.request_id.clone(), - workspace_id: connection.workspace_id.clone(), - is_server: false, - message_type: WebsocketEventType::Error, - message: e.to_string().into(), - ..Default::default() - }, - &UpdateSource::from_window_label(window.label()), - )?; + app_handle.with_tx(|tx| { + tx.upsert_websocket_event( + &WebsocketEvent { + connection_id: connection.id.clone(), + request_id: connection.request_id.clone(), + workspace_id: connection.workspace_id.clone(), + is_server: false, + message_type: WebsocketEventType::Error, + message: e.to_string().into(), + ..Default::default() + }, + &UpdateSource::from_window_label(window.label()), + ) + })?; Ok(connection) } @@ -96,18 +98,20 @@ async fn send_websocket_message( let mut ws_manager = ws_manager.lock().await; ws_manager.send(&connection.id, Message::Text(message.clone().into())).await?; - app_handle.db().upsert_websocket_event( - &WebsocketEvent { - connection_id: connection.id.clone(), - request_id: request.id.clone(), - workspace_id: connection.workspace_id.clone(), - is_server: false, - message_type: WebsocketEventType::Text, - message: message.into(), - ..Default::default() - }, - &UpdateSource::from_window_label(window.label()), - )?; + app_handle.with_tx(|tx| { + tx.upsert_websocket_event( + &WebsocketEvent { + connection_id: connection.id.clone(), + request_id: request.id.clone(), + workspace_id: connection.workspace_id.clone(), + is_server: false, + message_type: WebsocketEventType::Text, + message: message.into(), + ..Default::default() + }, + &UpdateSource::from_window_label(window.label()), + ) + })?; Ok(connection.clone()) } @@ -118,14 +122,13 @@ pub async fn cmd_ws_close( window: WebviewWindow, ws_manager: State<'_, Mutex>, ) -> Result { - let connection = { - let db = app_handle.db(); - let connection = db.get_websocket_connection(connection_id)?; - db.upsert_websocket_connection( + let connection = app_handle.with_tx(|tx| { + let connection = tx.get_websocket_connection(connection_id)?; + tx.upsert_websocket_connection( &WebsocketConnection { state: WebsocketConnectionState::Closing, ..connection }, &UpdateSource::from_window_label(window.label()), - )? - }; + ) + })?; let mut ws_manager = ws_manager.lock().await; if let Err(e) = ws_manager.close(&connection.id).await { @@ -169,14 +172,16 @@ pub async fn cmd_ws_connect( ) .await?; - let connection = app_handle.db().upsert_websocket_connection( - &WebsocketConnection { - workspace_id: request.workspace_id.clone(), - request_id: request_id.to_string(), - ..Default::default() - }, - &UpdateSource::from_window_label(window.label()), - )?; + let connection = app_handle.with_tx(|tx| { + tx.upsert_websocket_connection( + &WebsocketConnection { + workspace_id: request.workspace_id.clone(), + request_id: request_id.to_string(), + ..Default::default() + }, + &UpdateSource::from_window_label(window.label()), + ) + })?; let (mut url, url_parameters) = apply_path_placeholders(&request.url, &request.url_parameters); if !url.starts_with("ws://") && !url.starts_with("wss://") { @@ -187,14 +192,16 @@ pub async fn cmd_ws_connect( let mut url = match Url::parse(&url) { Ok(url) => url, Err(e) => { - return Ok(app_handle.db().upsert_websocket_connection( - &WebsocketConnection { - error: Some(format!("Failed to parse URL {}", e.to_string())), - state: WebsocketConnectionState::Closed, - ..connection - }, - &UpdateSource::from_window_label(window.label()), - )?); + return Ok(app_handle.with_tx(|tx| { + tx.upsert_websocket_connection( + &WebsocketConnection { + error: Some(format!("Failed to parse URL {}", e.to_string())), + state: WebsocketConnectionState::Closed, + ..connection + }, + &UpdateSource::from_window_label(window.label()), + ) + })?); } }; @@ -321,28 +328,32 @@ pub async fn cmd_ws_connect( { Ok(r) => r, Err(e) => { - return Ok(app_handle.db().upsert_websocket_connection( - &WebsocketConnection { - error: Some(e.to_string()), - state: WebsocketConnectionState::Closed, - ..connection - }, - &UpdateSource::from_window_label(window.label()), - )?); + return Ok(app_handle.with_tx(|tx| { + tx.upsert_websocket_connection( + &WebsocketConnection { + error: Some(e.to_string()), + state: WebsocketConnectionState::Closed, + ..connection + }, + &UpdateSource::from_window_label(window.label()), + ) + })?); } }; - app_handle.db().upsert_websocket_event( - &WebsocketEvent { - connection_id: connection.id.clone(), - request_id: request.id.clone(), - workspace_id: connection.workspace_id.clone(), - is_server: false, - message_type: WebsocketEventType::Open, - ..Default::default() - }, - &UpdateSource::from_window_label(window.label()), - )?; + app_handle.with_tx(|tx| { + tx.upsert_websocket_event( + &WebsocketEvent { + connection_id: connection.id.clone(), + request_id: request.id.clone(), + workspace_id: connection.workspace_id.clone(), + is_server: false, + message_type: WebsocketEventType::Open, + ..Default::default() + }, + &UpdateSource::from_window_label(window.label()), + ) + })?; let response_headers = response .headers() @@ -366,20 +377,22 @@ pub async fn cmd_ws_connect( if !set_cookie_headers.is_empty() { store.store_cookies_from_response(&convert_ws_url_to_http(&url), &set_cookie_headers); cookie_jar.cookies = store.get_all_cookies(); - app_handle.db().upsert_cookie_jar(cookie_jar, &UpdateSource::Background)?; + app_handle.with_tx(|tx| tx.upsert_cookie_jar(cookie_jar, &UpdateSource::Background))?; } } - let connection = app_handle.db().upsert_websocket_connection( - &WebsocketConnection { - state: WebsocketConnectionState::Connected, - headers: response_headers, - status: response.status().as_u16() as i32, - url: request.url.clone(), - ..connection - }, - &UpdateSource::from_window_label(window.label()), - )?; + let connection = app_handle.with_tx(|tx| { + tx.upsert_websocket_connection( + &WebsocketConnection { + state: WebsocketConnectionState::Connected, + headers: response_headers, + status: response.status().as_u16() as i32, + url: request.url.clone(), + ..connection + }, + &UpdateSource::from_window_label(window.label()), + ) + })?; { let connection_id = connection.id.clone(); @@ -395,57 +408,60 @@ pub async fn cmd_ws_connect( } app_handle - .db() - .upsert_websocket_event( - &WebsocketEvent { - connection_id: connection_id.clone(), - request_id: request_id.clone(), - workspace_id: workspace_id.clone(), - is_server: true, - message_type: match message { - Message::Text(_) => WebsocketEventType::Text, - Message::Binary(_) => WebsocketEventType::Binary, - Message::Ping(_) => WebsocketEventType::Ping, - Message::Pong(_) => WebsocketEventType::Pong, - Message::Close(_) => WebsocketEventType::Close, - // Raw frame will never happen during a read - Message::Frame(_) => WebsocketEventType::Frame, + .with_tx(|tx| { + tx.upsert_websocket_event( + &WebsocketEvent { + connection_id: connection_id.clone(), + request_id: request_id.clone(), + workspace_id: workspace_id.clone(), + is_server: true, + message_type: match message { + Message::Text(_) => WebsocketEventType::Text, + Message::Binary(_) => WebsocketEventType::Binary, + Message::Ping(_) => WebsocketEventType::Ping, + Message::Pong(_) => WebsocketEventType::Pong, + Message::Close(_) => WebsocketEventType::Close, + // Raw frame will never happen during a read + Message::Frame(_) => WebsocketEventType::Frame, + }, + message: message.into_data().into(), + ..Default::default() }, - message: message.into_data().into(), - ..Default::default() - }, - &UpdateSource::from_window_label(&window_label), - ) + &UpdateSource::from_window_label(&window_label), + ) + }) .unwrap(); } info!("Websocket connection closed"); if !has_written_close { app_handle - .db() - .upsert_websocket_event( - &WebsocketEvent { - connection_id: connection_id.clone(), - request_id: request_id.clone(), - workspace_id: workspace_id.clone(), - is_server: true, - message_type: WebsocketEventType::Close, - ..Default::default() - }, - &UpdateSource::from_window_label(&window_label), - ) + .with_tx(|tx| { + tx.upsert_websocket_event( + &WebsocketEvent { + connection_id: connection_id.clone(), + request_id: request_id.clone(), + workspace_id: workspace_id.clone(), + is_server: true, + message_type: WebsocketEventType::Close, + ..Default::default() + }, + &UpdateSource::from_window_label(&window_label), + ) + }) .unwrap(); } app_handle - .db() - .upsert_websocket_connection( - &WebsocketConnection { - workspace_id: request.workspace_id.clone(), - request_id: request_id.to_string(), - state: WebsocketConnectionState::Closed, - ..connection - }, - &UpdateSource::from_window_label(&window_label), - ) + .with_tx(|tx| { + tx.upsert_websocket_connection( + &WebsocketConnection { + workspace_id: request.workspace_id.clone(), + request_id: request_id.to_string(), + state: WebsocketConnectionState::Closed, + ..connection + }, + &UpdateSource::from_window_label(&window_label), + ) + }) .unwrap(); }); } diff --git a/crates-tauri/yaak-license/src/license.rs b/crates-tauri/yaak-license/src/license.rs index dcc3fd12..de4c3f6a 100644 --- a/crates-tauri/yaak-license/src/license.rs +++ b/crates-tauri/yaak-license/src/license.rs @@ -9,7 +9,7 @@ use tauri::{AppHandle, Emitter, Manager, Runtime, WebviewWindow, is_dev}; use ts_rs::TS; use yaak_api::{ApiClientKind, yaak_api_client}; use yaak_common::platform::get_os_str; -use yaak_models::client_db::ClientDb; +use yaak_models::client_db::{ClientDb, WriteDb}; use yaak_models::query_manager::QueryManager; use yaak_models::util::UpdateSource; @@ -17,6 +17,10 @@ use yaak_models::util::UpdateSource; /// This is needed temporarily until all crates are refactored to not use Tauri. trait QueryManagerExt<'a, R> { fn db(&'a self) -> ClientDb<'a>; + fn with_tx( + &'a self, + func: impl FnOnce(&WriteDb) -> yaak_models::error::Result, + ) -> yaak_models::error::Result; } impl<'a, R: Runtime, M: Manager> QueryManagerExt<'a, R> for M { @@ -24,6 +28,14 @@ impl<'a, R: Runtime, M: Manager> QueryManagerExt<'a, R> for M { let qm = self.state::(); qm.inner().connect() } + + fn with_tx( + &'a self, + func: impl FnOnce(&WriteDb) -> yaak_models::error::Result, + ) -> yaak_models::error::Result { + let qm = self.state::(); + qm.inner().with_tx(func) + } } const KV_NAMESPACE: &str = "license"; @@ -137,12 +149,17 @@ pub async fn activate_license( } let body: ActivateLicenseResponsePayload = response.json().await?; - window.app_handle().db().set_key_value_str( - KV_ACTIVATION_ID_KEY, - KV_NAMESPACE, - body.activation_id.as_str(), - &UpdateSource::from_window_label(window.label()), - ); + if let Err(e) = window.app_handle().with_tx(|tx| { + tx.set_key_value_str( + KV_ACTIVATION_ID_KEY, + KV_NAMESPACE, + body.activation_id.as_str(), + &UpdateSource::from_window_label(window.label()), + ); + Ok(()) + }) { + warn!("Failed to store license activation: {e}"); + } if let Err(e) = window.emit("license-activated", true) { warn!("Failed to emit check-license event: {}", e); @@ -172,11 +189,13 @@ pub async fn deactivate_license(window: &WebviewWindow) -> Result return Err(ServerError); } - app_handle.db().delete_key_value( - KV_ACTIVATION_ID_KEY, - KV_NAMESPACE, - &UpdateSource::from_window_label(window.label()), - )?; + app_handle.with_tx(|tx| { + tx.delete_key_value( + KV_ACTIVATION_ID_KEY, + KV_NAMESPACE, + &UpdateSource::from_window_label(window.label()), + ) + })?; if let Err(e) = app_handle.emit("license-deactivated", true) { warn!("Failed to emit deactivate-license event: {}", e); diff --git a/crates/common/yaak-database/src/pool.rs b/crates/common/yaak-database/src/pool.rs index 863087da..b5054416 100644 --- a/crates/common/yaak-database/src/pool.rs +++ b/crates/common/yaak-database/src/pool.rs @@ -21,7 +21,8 @@ //! the second ask. Sharing the handle instead makes nested *reads* work the way //! they do on the desktop; nested *write transactions* fail on both, only //! differently (here SQLite refuses the inner `BEGIN`; natively the inner -//! connection blocks on `busy_timeout` and then fails). +//! call waits for the one writer connection, which the outer call holds, and +//! times out). #[cfg(not(target_arch = "wasm32"))] mod imp { diff --git a/crates/yaak-commands/src/host.rs b/crates/yaak-commands/src/host.rs index 5d8a54b1..b5ff79f4 100644 --- a/crates/yaak-commands/src/host.rs +++ b/crates/yaak-commands/src/host.rs @@ -72,7 +72,7 @@ pub trait Host: Clone { self.query_manager().connect() } - fn blobs(&self) -> BlobContext { + fn blobs(&self) -> BlobContext<'_> { self.blob_manager().connect() } } diff --git a/crates/yaak-commands/src/models.rs b/crates/yaak-commands/src/models.rs index 86ec67d7..99e266cb 100644 --- a/crates/yaak-commands/src/models.rs +++ b/crates/yaak-commands/src/models.rs @@ -11,10 +11,11 @@ use yaak_models::queries::workspaces::default_headers; use yaak_rpc_schema::*; pub async fn models_upsert(host: H, req: ModelsUpsertReq) -> Result { - let db = host.db(); let blobs = host.blob_manager(); let source = host.update_source(); - Ok(yaak_models::models_ops::upsert_model(&db, blobs, req.model, &source)?) + Ok(host + .query_manager() + .with_tx(|tx| yaak_models::models_ops::upsert_model(tx, blobs, req.model, &source))?) } /// Deletes cascade — a workspace can hold thousands of requests — and run in a @@ -75,12 +76,9 @@ pub async fn models_upsert_graphql_introspection( req: ModelsUpsertGraphqlIntrospectionReq, ) -> Result { let source = host.update_source(); - Ok(host.db().upsert_graphql_introspection( - &req.workspace_id, - &req.request_id, - req.content, - &source, - )?) + Ok(host.query_manager().with_tx(|tx| { + tx.upsert_graphql_introspection(&req.workspace_id, &req.request_id, req.content, &source) + })?) } /// Everything the frontend's model store needs to boot, as one JSON string. @@ -112,9 +110,16 @@ pub async fn models_workspace_models( // Add the workspace children if let Some(wid) = req.workspace_id.as_deref() { + // Opening a workspace is where the rows it is assumed to have get created + host.query_manager().with_tx(|tx| { + tx.ensure_base_environment(wid)?; + tx.ensure_default_cookie_jar(wid)?; + tx.ensure_workspace_meta(wid)?; + Ok::<(), yaak_models::error::Error>(()) + })?; let db = host.db(); l.append(&mut db.list_cookie_jars(wid)?.into_iter().map(Into::into).collect()); - l.append(&mut db.list_environments_ensure_base(wid)?.into_iter().map(Into::into).collect()); + l.append(&mut db.list_environments(wid)?.into_iter().map(Into::into).collect()); l.append(&mut db.list_folders(wid)?.into_iter().map(Into::into).collect()); l.append(&mut db.list_grpc_connections(wid)?.into_iter().map(Into::into).collect()); l.append(&mut db.list_grpc_requests(wid)?.into_iter().map(Into::into).collect()); @@ -132,23 +137,26 @@ pub async fn cmd_get_workspace_meta( host: H, req: CmdGetWorkspaceMetaReq, ) -> Result { - let db = host.db(); - let workspace = db.get_workspace(&req.workspace_id)?; - Ok(db.get_or_create_workspace_meta(&workspace.id)?) + let workspace = host.db().get_workspace(&req.workspace_id)?; + Ok(host.query_manager().with_tx(|tx| tx.ensure_workspace_meta(&workspace.id))?) } pub async fn cmd_delete_all_grpc_connections( host: H, req: CmdDeleteAllGrpcConnectionsReq, ) -> Result<()> { - Ok(host.db().delete_all_grpc_connections_for_request(&req.request_id, &host.update_source())?) + Ok(host.query_manager().with_tx(|tx| { + tx.delete_all_grpc_connections_for_request(&req.request_id, &host.update_source()) + })?) } pub async fn cmd_delete_all_http_responses( host: H, req: CmdDeleteAllHttpResponsesReq, ) -> Result<()> { - host.db().delete_all_http_responses_for_request(&req.request_id, &host.update_source())?; + host.query_manager().with_tx(|tx| { + tx.delete_all_http_responses_for_request(&req.request_id, &host.update_source()) + })?; Ok(()) } @@ -156,9 +164,9 @@ pub async fn cmd_ws_delete_connections( host: H, req: CmdWsDeleteConnectionsReq, ) -> Result<()> { - Ok(host - .db() - .delete_all_websocket_connections_for_request(&req.request_id, &host.update_source())?) + Ok(host.query_manager().with_tx(|tx| { + tx.delete_all_websocket_connections_for_request(&req.request_id, &host.update_source()) + })?) } pub async fn cmd_delete_send_history(host: H, req: CmdDeleteSendHistoryReq) -> Result<()> { diff --git a/crates/yaak-commands/tests/test_host.rs b/crates/yaak-commands/tests/test_host.rs index 15440c8c..de31d9f2 100644 --- a/crates/yaak-commands/tests/test_host.rs +++ b/crates/yaak-commands/tests/test_host.rs @@ -118,6 +118,7 @@ impl Host for TestHost { #[tokio::test(flavor = "multi_thread")] async fn writes_carry_the_client_id() { let host = TestHost::new(); + host.drain_writes(); // the rows startup creates let workspace = Workspace { name: "From a test".to_string(), ..Default::default() }; let id = models_upsert(host.clone(), ModelsUpsertReq { model: AnyModel::Workspace(workspace) }) @@ -406,21 +407,23 @@ async fn a_single_threaded_host_can_implement_the_trait() { // shared code; only the callback came from the host. Rendering a real // variable is what proves the chain was resolved rather than skipped. let environment = host - .db() - .upsert_environment( - &Environment { - workspace_id: id.clone(), - name: "Test env".to_string(), - variables: vec![EnvironmentVariable { - enabled: true, - name: "greeting".to_string(), - value: "hello".to_string(), - id: None, - }], - ..Default::default() - }, - &host.update_source(), - ) + .query_manager() + .with_tx(|tx| { + tx.upsert_environment( + &Environment { + workspace_id: id.clone(), + name: "Test env".to_string(), + variables: vec![EnvironmentVariable { + enabled: true, + name: "greeting".to_string(), + value: "hello".to_string(), + id: None, + }], + ..Default::default() + }, + &host.update_source(), + ) + }) .expect("seed environment"); let rendered = cmd_render_template( @@ -459,30 +462,33 @@ async fn auth_values_are_rendered_before_the_host_sees_them() { }; let workspace = host - .db() - .upsert_workspace( - &Workspace { name: "Auth".to_string(), ..Default::default() }, - &host.update_source(), - ) + .query_manager() + .with_tx(|tx| { + tx.upsert_workspace( + &Workspace { name: "Auth".to_string(), ..Default::default() }, + &host.update_source(), + ) + }) .expect("workspace"); - host.db() - .upsert_environment( - &Environment { - workspace_id: workspace.id.clone(), - name: "Env".to_string(), - variables: vec![EnvironmentVariable { - enabled: true, - name: "token".to_string(), - value: "s3cret".to_string(), - id: None, - }], - ..Default::default() - }, - &host.update_source(), - ) + host.query_manager() + .with_tx(|tx| { + tx.upsert_environment( + &Environment { + workspace_id: workspace.id.clone(), + name: "Env".to_string(), + variables: vec![EnvironmentVariable { + enabled: true, + name: "token".to_string(), + value: "s3cret".to_string(), + id: None, + }], + ..Default::default() + }, + &host.update_source(), + ) + }) .expect("environment"); - let environment = - host.db().list_environments_ensure_base(&workspace.id).expect("list").remove(0); + let environment = host.db().list_environments(&workspace.id).expect("list").remove(0); let mut values = HashMap::new(); values.insert("password".to_string(), JsonPrimitive::String("${[ token ]}".to_string())); @@ -521,30 +527,33 @@ async fn template_function_values_are_rendered_before_the_host_sees_them() { }; let workspace = host - .db() - .upsert_workspace( - &Workspace { name: "Functions".to_string(), ..Default::default() }, - &host.update_source(), - ) + .query_manager() + .with_tx(|tx| { + tx.upsert_workspace( + &Workspace { name: "Functions".to_string(), ..Default::default() }, + &host.update_source(), + ) + }) .expect("workspace"); - host.db() - .upsert_environment( - &Environment { - workspace_id: workspace.id.clone(), - name: "Env".to_string(), - variables: vec![EnvironmentVariable { - enabled: true, - name: "1PASSWORD_TOKEN".to_string(), - value: "ops_abc123".to_string(), - id: None, - }], - ..Default::default() - }, - &host.update_source(), - ) + host.query_manager() + .with_tx(|tx| { + tx.upsert_environment( + &Environment { + workspace_id: workspace.id.clone(), + name: "Env".to_string(), + variables: vec![EnvironmentVariable { + enabled: true, + name: "1PASSWORD_TOKEN".to_string(), + value: "ops_abc123".to_string(), + id: None, + }], + ..Default::default() + }, + &host.update_source(), + ) + }) .expect("environment"); - let environment = - host.db().list_environments_ensure_base(&workspace.id).expect("list").remove(0); + let environment = host.db().list_environments(&workspace.id).expect("list").remove(0); let mut values = HashMap::new(); values.insert("token".to_string(), JsonPrimitive::String("${[1PASSWORD_TOKEN]}".to_string())); diff --git a/crates/yaak-crypto/src/manager.rs b/crates/yaak-crypto/src/manager.rs index 3eb63e95..b17111cb 100644 --- a/crates/yaak-crypto/src/manager.rs +++ b/crates/yaak-crypto/src/manager.rs @@ -83,7 +83,7 @@ impl EncryptionManager { let workspace_meta = self.query_manager.with_tx::(|tx| { let workspace = tx.get_workspace(workspace_id)?; - let workspace_meta = tx.get_or_create_workspace_meta(workspace_id)?; + let workspace_meta = tx.ensure_workspace_meta(workspace_id)?; tx.upsert_workspace( &Workspace { encryption_key_challenge, ..workspace }, &UpdateSource::Background, @@ -103,7 +103,7 @@ impl EncryptionManager { pub fn ensure_workspace_key(&self, workspace_id: &str) -> Result { let workspace_meta = - self.query_manager.connect().get_or_create_workspace_meta(workspace_id)?; + self.query_manager.with_tx(|tx| tx.ensure_workspace_meta(workspace_id))?; // Already exists if let Some(_) = workspace_meta.encryption_key { @@ -120,7 +120,7 @@ impl EncryptionManager { self.query_manager.with_tx::<(), Error>(|tx| { let workspace = tx.get_workspace(workspace_id)?; - let workspace_meta = tx.get_or_create_workspace_meta(workspace_id)?; + let workspace_meta = tx.ensure_workspace_meta(workspace_id)?; // Clear encryption challenge on workspace tx.upsert_workspace( @@ -152,10 +152,12 @@ impl EncryptionManager { } }; - let db = self.query_manager.connect(); - let workspace_meta = db.get_or_create_workspace_meta(workspace_id)?; - - let key = match workspace_meta.encryption_key { + let key = match self + .query_manager + .connect() + .get_workspace_meta(workspace_id) + .and_then(|m| m.encryption_key) + { None => return Err(MissingWorkspaceKey), Some(k) => k, }; diff --git a/crates/yaak-lifecycle/src/lib.rs b/crates/yaak-lifecycle/src/lib.rs index 33ad68f7..bbcc42b2 100644 --- a/crates/yaak-lifecycle/src/lib.rs +++ b/crates/yaak-lifecycle/src/lib.rs @@ -7,7 +7,7 @@ use log::info; use std::path::PathBuf; use yaak_models::blob_manager::BlobManager; -use yaak_models::client_db::ClientDb; +use yaak_models::client_db::WriteDb; use yaak_models::error::Result; const MODEL_CHANGES_RETENTION_HOURS: i64 = 1; @@ -44,7 +44,8 @@ impl Host { } /// Run once after the database is open, before the host answers anything. -pub fn on_launch(host: &Host, db: &ClientDb, blobs: &BlobManager) -> Result<()> { +/// Takes the write handle: a launch closes what the last session left open. +pub fn on_launch(host: &Host, db: &WriteDb, blobs: &BlobManager) -> Result<()> { db.prune_model_changes_older_than_hours(MODEL_CHANGES_RETENTION_HOURS)?; if host.role == Role::Owner { @@ -77,53 +78,48 @@ mod tests { #[test] fn only_the_owner_closes_what_the_last_session_left_open() { let (query_manager, blob_manager, _rx) = init_in_memory().expect("Failed to init DB"); - let db = query_manager.connect(); let source = &UpdateSource::Background; - let workspace = db - .upsert_workspace( - &Workspace { name: "Hooks".to_string(), ..Default::default() }, - source, - ) - .unwrap(); - let request = db - .upsert_http_request( - &HttpRequest { workspace_id: workspace.id.clone(), ..Default::default() }, - source, - ) - .unwrap(); - let pending = db - .upsert_http_response( - &HttpResponse { - request_id: request.id.clone(), - workspace_id: workspace.id.clone(), - state: HttpResponseState::Connected, - ..Default::default() - }, - source, - &blob_manager, - ) + let pending = query_manager + .with_tx(|db| { + let workspace = db.upsert_workspace( + &Workspace { name: "Hooks".to_string(), ..Default::default() }, + source, + )?; + let request = db.upsert_http_request( + &HttpRequest { workspace_id: workspace.id.clone(), ..Default::default() }, + source, + )?; + db.upsert_http_response( + &HttpResponse { + request_id: request.id.clone(), + workspace_id: workspace.id.clone(), + state: HttpResponseState::Connected, + ..Default::default() + }, + source, + &blob_manager, + ) + }) .unwrap(); - on_launch(&Host::guest(), &db, &blob_manager).unwrap(); - let response = db.get_http_response(&pending.id).unwrap(); + query_manager.with_tx(|db| on_launch(&Host::guest(), db, &blob_manager)).unwrap(); + let response = query_manager.connect().get_http_response(&pending.id).unwrap(); assert!(matches!(response.state, HttpResponseState::Connected)); - on_launch(&Host::owner(), &db, &blob_manager).unwrap(); - let response = db.get_http_response(&pending.id).unwrap(); + query_manager.with_tx(|db| on_launch(&Host::owner(), db, &blob_manager)).unwrap(); + let response = query_manager.connect().get_http_response(&pending.id).unwrap(); assert!(matches!(response.state, HttpResponseState::Closed)); } #[test] fn owner_without_a_filesystem_still_sweeps_blobs() { let (query_manager, blob_manager, _rx) = init_in_memory().expect("Failed to init DB"); - let db = query_manager.connect(); - { - let blob_ctx = blob_manager.connect(); - blob_ctx.insert_chunk(&BodyChunk::new("rs_gone", 0, b"dead".to_vec())).unwrap(); - } + blob_manager + .with_tx(|b| b.insert_chunk(&BodyChunk::new("rs_gone", 0, b"dead".to_vec()))) + .unwrap(); - on_launch(&Host::owner(), &db, &blob_manager).unwrap(); + query_manager.with_tx(|db| on_launch(&Host::owner(), db, &blob_manager)).unwrap(); assert!(!blob_manager.connect().body_exists("rs_gone").unwrap()); } diff --git a/crates/yaak-models/src/blob_manager.rs b/crates/yaak-models/src/blob_manager.rs index 78243227..41ba8bc5 100644 --- a/crates/yaak-models/src/blob_manager.rs +++ b/crates/yaak-models/src/blob_manager.rs @@ -1,9 +1,11 @@ +use crate::error::Error::GenericError; use crate::error::Result; use crate::util::generate_prefixed_id; use include_dir::{Dir, include_dir}; use log::{debug, info}; -use rusqlite::{OptionalExtension, params}; -use yaak_database::{SqliteConn, SqlitePool}; +use rusqlite::{OptionalExtension, Transaction, TransactionBehavior, params}; +use std::ops::Deref; +use yaak_database::{ConnectionOrTx, SqlitePool}; static BLOB_MIGRATIONS_DIR: Dir = include_dir!("$CARGO_MANIFEST_DIR/blob_migrations"); @@ -22,41 +24,64 @@ impl BodyChunk { } } -/// Manages the blob database connection pool. -// Pool is internally synchronized — don't wrap it in a Mutex. A Mutex held across the +/// Manages the blob database: a reader pool and a single writer, for the +/// same reason as [`crate::query_manager::QueryManager`]. +// Pools are internally synchronized — don't wrap them in a Mutex. A Mutex held across the // blocking `get()` serializes every blob access behind the slowest waiter, freezing the // whole app whenever the pool is exhausted. #[derive(Debug, Clone)] pub struct BlobManager { - pool: SqlitePool, + readers: SqlitePool, + writer: SqlitePool, } impl BlobManager { - pub fn new(pool: SqlitePool) -> Self { - Self { pool } + /// `writer` must be a pool with a single connection. + pub fn new(readers: SqlitePool, writer: SqlitePool) -> Self { + Self { readers, writer } } - pub fn connect(&self) -> BlobContext { - let conn = self.pool.get().expect("Failed to get blob DB connection from pool"); - BlobContext { conn } + /// A read handle from the reader pool. + pub fn connect(&self) -> BlobContext<'_> { + let conn = self.readers.get().expect("Failed to get blob DB connection from pool"); + BlobContext { conn: ConnectionOrTx::Connection(conn) } + } + + /// Run `func` in a transaction on the writer connection. + pub fn with_tx( + &self, + func: impl FnOnce(&BlobWriter) -> std::result::Result, + ) -> std::result::Result + where + E: From, + { + let conn = self.writer.get().map_err(crate::error::Error::SqlPoolError)?; + let tx = Transaction::new_unchecked(&conn, TransactionBehavior::Immediate) + .map_err(crate::error::Error::SqlError)?; + let writer = BlobWriter { ctx: BlobContext { conn: ConnectionOrTx::Transaction(&tx) } }; + match func(&writer) { + Ok(val) => { + tx.commit().map_err(|e| { + GenericError(format!("Failed to commit blob transaction {e:?}")) + })?; + Ok(val) + } + Err(e) => { + tx.rollback().map_err(|e| { + GenericError(format!("Failed to rollback blob transaction {e:?}")) + })?; + Err(e) + } + } } } -/// Context for blob database operations. -pub struct BlobContext { - conn: SqliteConn, +/// Read handle for the blob database. +pub struct BlobContext<'a> { + conn: ConnectionOrTx<'a>, } -impl BlobContext { - /// Insert a single chunk. - pub fn insert_chunk(&self, chunk: &BodyChunk) -> Result<()> { - self.conn.execute( - "INSERT INTO body_chunks (id, body_id, chunk_index, data) VALUES (?1, ?2, ?3, ?4)", - params![chunk.id, chunk.body_id, chunk.chunk_index, chunk.data], - )?; - Ok(()) - } - +impl<'a> BlobContext<'a> { /// Get all chunks for a body, ordered by chunk_index. pub fn get_chunks(&self, body_id: &str) -> Result> { let mut stmt = self.conn.prepare( @@ -87,25 +112,11 @@ impl BlobContext { Ok(ids) } - /// Delete all chunks for a body. - pub fn delete_chunks(&self, body_id: &str) -> Result<()> { - self.conn.execute("DELETE FROM body_chunks WHERE body_id = ?1", params![body_id])?; - Ok(()) - } - - /// Delete all chunks matching a body_id prefix (e.g., "rs_abc123.%" to delete all bodies for a response). - pub fn delete_chunks_like(&self, body_id_prefix: &str) -> Result<()> { - self.conn - .execute("DELETE FROM body_chunks WHERE body_id LIKE ?1", params![body_id_prefix])?; - Ok(()) - } -} - -/// Get total size of a body without loading data. -impl BlobContext { + /// Get total size of a body without loading data. pub fn get_body_size(&self, body_id: &str) -> Result { let size: i64 = self .conn + .resolve() .query_row( "SELECT COALESCE(SUM(LENGTH(data)), 0) FROM body_chunks WHERE body_id = ?1", params![body_id], @@ -119,6 +130,7 @@ impl BlobContext { pub fn body_exists(&self, body_id: &str) -> Result { let count: i64 = self .conn + .resolve() .query_row( "SELECT COUNT(*) FROM body_chunks WHERE body_id = ?1", params![body_id], @@ -129,6 +141,44 @@ impl BlobContext { } } +/// Write handle for the blob database. Derefs to [`BlobContext`] for reads. +pub struct BlobWriter<'a> { + ctx: BlobContext<'a>, +} + +impl<'a> Deref for BlobWriter<'a> { + type Target = BlobContext<'a>; + + fn deref(&self) -> &BlobContext<'a> { + &self.ctx + } +} + +impl<'a> BlobWriter<'a> { + /// Insert a single chunk. + pub fn insert_chunk(&self, chunk: &BodyChunk) -> Result<()> { + self.conn.execute( + "INSERT INTO body_chunks (id, body_id, chunk_index, data) VALUES (?1, ?2, ?3, ?4)", + params![chunk.id, chunk.body_id, chunk.chunk_index, chunk.data], + )?; + Ok(()) + } + + /// Delete all chunks for a body. + pub fn delete_chunks(&self, body_id: &str) -> Result<()> { + self.conn.execute("DELETE FROM body_chunks WHERE body_id = ?1", params![body_id])?; + Ok(()) + } + + /// Delete all chunks matching a body_id prefix (e.g., "rs_abc123.%" to delete all + /// bodies for a response). + pub fn delete_chunks_like(&self, body_id_prefix: &str) -> Result<()> { + self.conn + .execute("DELETE FROM body_chunks WHERE body_id LIKE ?1", params![body_id_prefix])?; + Ok(()) + } +} + /// Run migrations for the blob database. pub fn migrate_blob_db(pool: &SqlitePool) -> Result<()> { info!("Running blob database migrations"); @@ -196,28 +246,39 @@ pub fn migrate_blob_db(pool: &SqlitePool) -> Result<()> { #[cfg(test)] mod tests { use super::*; + use crate::error::Error; - fn create_test_pool() -> SqlitePool { + fn create_test_manager() -> BlobManager { let manager = r2d2_sqlite::SqliteConnectionManager::memory(); let pool = r2d2::Pool::builder().max_size(1).build(manager).unwrap(); migrate_blob_db(&pool).unwrap(); - pool + BlobManager::new(pool.clone(), pool) + } + + fn insert(manager: &BlobManager, chunks: &[BodyChunk]) { + manager + .with_tx(|b| { + for c in chunks { + b.insert_chunk(c)?; + } + Ok::<_, Error>(()) + }) + .unwrap(); } #[test] fn test_insert_and_get_chunks() { - let pool = create_test_pool(); - let manager = BlobManager::new(pool); - let ctx = manager.connect(); - + let manager = create_test_manager(); let body_id = "rs_test123.request"; - let chunk1 = BodyChunk::new(body_id, 0, b"Hello, ".to_vec()); - let chunk2 = BodyChunk::new(body_id, 1, b"World!".to_vec()); + insert( + &manager, + &[ + BodyChunk::new(body_id, 0, b"Hello, ".to_vec()), + BodyChunk::new(body_id, 1, b"World!".to_vec()), + ], + ); - ctx.insert_chunk(&chunk1).unwrap(); - ctx.insert_chunk(&chunk2).unwrap(); - - let chunks = ctx.get_chunks(body_id).unwrap(); + let chunks = manager.connect().get_chunks(body_id).unwrap(); assert_eq!(chunks.len(), 2); assert_eq!(chunks[0].chunk_index, 0); assert_eq!(chunks[0].data, b"Hello, "); @@ -227,18 +288,19 @@ mod tests { #[test] fn test_get_chunks_ordered_by_index() { - let pool = create_test_pool(); - let manager = BlobManager::new(pool); - let ctx = manager.connect(); - + let manager = create_test_manager(); let body_id = "rs_test123.request"; - // Insert out of order - ctx.insert_chunk(&BodyChunk::new(body_id, 2, b"C".to_vec())).unwrap(); - ctx.insert_chunk(&BodyChunk::new(body_id, 0, b"A".to_vec())).unwrap(); - ctx.insert_chunk(&BodyChunk::new(body_id, 1, b"B".to_vec())).unwrap(); + insert( + &manager, + &[ + BodyChunk::new(body_id, 2, b"C".to_vec()), + BodyChunk::new(body_id, 0, b"A".to_vec()), + BodyChunk::new(body_id, 1, b"B".to_vec()), + ], + ); - let chunks = ctx.get_chunks(body_id).unwrap(); + let chunks = manager.connect().get_chunks(body_id).unwrap(); assert_eq!(chunks.len(), 3); assert_eq!(chunks[0].data, b"A"); assert_eq!(chunks[1].data, b"B"); @@ -247,89 +309,85 @@ mod tests { #[test] fn test_delete_chunks() { - let pool = create_test_pool(); - let manager = BlobManager::new(pool); - let ctx = manager.connect(); - + let manager = create_test_manager(); let body_id = "rs_test123.request"; - ctx.insert_chunk(&BodyChunk::new(body_id, 0, b"data".to_vec())).unwrap(); + insert(&manager, &[BodyChunk::new(body_id, 0, b"data".to_vec())]); + assert!(manager.connect().body_exists(body_id).unwrap()); - assert!(ctx.body_exists(body_id).unwrap()); - - ctx.delete_chunks(body_id).unwrap(); + manager.with_tx(|b| b.delete_chunks(body_id)).unwrap(); + let ctx = manager.connect(); assert!(!ctx.body_exists(body_id).unwrap()); assert_eq!(ctx.get_chunks(body_id).unwrap().len(), 0); } #[test] fn test_delete_chunks_like() { - let pool = create_test_pool(); - let manager = BlobManager::new(pool); - let ctx = manager.connect(); - + let manager = create_test_manager(); // Insert chunks for same response but different body types - ctx.insert_chunk(&BodyChunk::new("rs_abc.request", 0, b"req".to_vec())).unwrap(); - ctx.insert_chunk(&BodyChunk::new("rs_abc.response", 0, b"resp".to_vec())).unwrap(); - ctx.insert_chunk(&BodyChunk::new("rs_other.request", 0, b"other".to_vec())).unwrap(); + insert( + &manager, + &[ + BodyChunk::new("rs_abc.request", 0, b"req".to_vec()), + BodyChunk::new("rs_abc.response", 0, b"resp".to_vec()), + BodyChunk::new("rs_other.request", 0, b"other".to_vec()), + ], + ); // Delete all bodies for rs_abc - ctx.delete_chunks_like("rs_abc.%").unwrap(); + manager.with_tx(|b| b.delete_chunks_like("rs_abc.%")).unwrap(); - // rs_abc bodies should be gone + let ctx = manager.connect(); assert!(!ctx.body_exists("rs_abc.request").unwrap()); assert!(!ctx.body_exists("rs_abc.response").unwrap()); - - // rs_other should still exist assert!(ctx.body_exists("rs_other.request").unwrap()); } #[test] fn test_get_body_size() { - let pool = create_test_pool(); - let manager = BlobManager::new(pool); - let ctx = manager.connect(); - + let manager = create_test_manager(); let body_id = "rs_test123.request"; - ctx.insert_chunk(&BodyChunk::new(body_id, 0, b"Hello".to_vec())).unwrap(); - ctx.insert_chunk(&BodyChunk::new(body_id, 1, b"World".to_vec())).unwrap(); + insert( + &manager, + &[ + BodyChunk::new(body_id, 0, b"Hello".to_vec()), + BodyChunk::new(body_id, 1, b"World".to_vec()), + ], + ); - let size = ctx.get_body_size(body_id).unwrap(); + let size = manager.connect().get_body_size(body_id).unwrap(); assert_eq!(size, 10); // "Hello" + "World" = 10 bytes } #[test] fn test_get_body_size_empty() { - let pool = create_test_pool(); - let manager = BlobManager::new(pool); - let ctx = manager.connect(); - - let size = ctx.get_body_size("nonexistent").unwrap(); + let manager = create_test_manager(); + let size = manager.connect().get_body_size("nonexistent").unwrap(); assert_eq!(size, 0); } #[test] fn test_body_exists() { - let pool = create_test_pool(); - let manager = BlobManager::new(pool); - let ctx = manager.connect(); + let manager = create_test_manager(); + assert!(!manager.connect().body_exists("rs_test.request").unwrap()); - assert!(!ctx.body_exists("rs_test.request").unwrap()); + insert(&manager, &[BodyChunk::new("rs_test.request", 0, b"data".to_vec())]); - ctx.insert_chunk(&BodyChunk::new("rs_test.request", 0, b"data".to_vec())).unwrap(); - - assert!(ctx.body_exists("rs_test.request").unwrap()); + assert!(manager.connect().body_exists("rs_test.request").unwrap()); } #[test] fn test_multiple_bodies_isolated() { - let pool = create_test_pool(); - let manager = BlobManager::new(pool); + let manager = create_test_manager(); + insert( + &manager, + &[ + BodyChunk::new("body1", 0, b"data1".to_vec()), + BodyChunk::new("body2", 0, b"data2".to_vec()), + ], + ); + let ctx = manager.connect(); - - ctx.insert_chunk(&BodyChunk::new("body1", 0, b"data1".to_vec())).unwrap(); - ctx.insert_chunk(&BodyChunk::new("body2", 0, b"data2".to_vec())).unwrap(); - let chunks1 = ctx.get_chunks("body1").unwrap(); let chunks2 = ctx.get_chunks("body2").unwrap(); @@ -341,16 +399,13 @@ mod tests { #[test] fn test_large_chunk() { - let pool = create_test_pool(); - let manager = BlobManager::new(pool); - let ctx = manager.connect(); - + let manager = create_test_manager(); // 1MB chunk let large_data: Vec = (0..1024 * 1024).map(|i| (i % 256) as u8).collect(); let body_id = "rs_large.request"; + insert(&manager, &[BodyChunk::new(body_id, 0, large_data.clone())]); - ctx.insert_chunk(&BodyChunk::new(body_id, 0, large_data.clone())).unwrap(); - + let ctx = manager.connect(); let chunks = ctx.get_chunks(body_id).unwrap(); assert_eq!(chunks.len(), 1); assert_eq!(chunks[0].data, large_data); diff --git a/crates/yaak-models/src/client_db.rs b/crates/yaak-models/src/client_db.rs index e4023e42..7208334f 100644 --- a/crates/yaak-models/src/client_db.rs +++ b/crates/yaak-models/src/client_db.rs @@ -3,18 +3,25 @@ use crate::models::{AnyModel, UpsertModelInfo}; use crate::util::{ModelChangeEvent, ModelPayload, UpdateSource}; use rusqlite::params; use sea_query::{IntoColumnRef, IntoIden, SimpleExpr}; +use std::cell::RefCell; use std::fmt::Debug; +use std::ops::Deref; use std::sync::mpsc; use yaak_database::DbContext; +/// A read handle. Comes from the reader pool and can only query. +/// +/// Anything that changes a row lives on [`WriteDb`], which is only ever handed +/// out inside a transaction on the single writer connection. That split is +/// what keeps the pool from filling with writers waiting on each other: there +/// is one writer, so there is never a second one to wait for. pub struct ClientDb<'a> { pub(crate) ctx: DbContext<'a>, - pub(crate) events_tx: mpsc::Sender, } impl<'a> ClientDb<'a> { - pub fn new(ctx: DbContext<'a>, events_tx: mpsc::Sender) -> Self { - Self { ctx, events_tx } + pub fn new(ctx: DbContext<'a>) -> Self { + Self { ctx } } /// Access the underlying connection for custom queries. @@ -22,8 +29,6 @@ impl<'a> ClientDb<'a> { self.ctx.conn() } - // --- Read delegates (thin wrappers over DbContext) --- - pub(crate) fn find_one( &self, col: impl IntoColumnRef + IntoIden + Clone, @@ -64,6 +69,39 @@ impl<'a> ClientDb<'a> { { Ok(self.ctx.find_many(col, value, limit)?) } +} + +/// A write handle: a [`ClientDb`] on the writer connection, inside a +/// transaction, that can also change rows. Derefs to [`ClientDb`] so every +/// query is available while writing, and reads inside the transaction see +/// its own uncommitted writes. +/// +/// Model events are held back until the transaction commits; a rollback +/// discards them along with the rows. +pub struct WriteDb<'a> { + db: ClientDb<'a>, + events_tx: mpsc::Sender, + pending_events: RefCell>, +} + +impl<'a> Deref for WriteDb<'a> { + type Target = ClientDb<'a>; + + fn deref(&self) -> &ClientDb<'a> { + &self.db + } +} + +impl<'a> WriteDb<'a> { + pub fn new(ctx: DbContext<'a>, events_tx: mpsc::Sender) -> Self { + Self { db: ClientDb::new(ctx), events_tx, pending_events: RefCell::new(Vec::new()) } + } + + /// The events for everything written so far, to send once the + /// transaction has committed. + pub(crate) fn into_events(self) -> Vec { + self.pending_events.into_inner() + } /// Bulk-delete all rows matching a column value WITHOUT recording model /// changes or emitting events. Only use for cascades whose deletion is @@ -80,8 +118,6 @@ impl<'a> ClientDb<'a> { Ok(self.ctx.delete_many::(col, value)?) } - // --- Write operations (with event recording) --- - pub(crate) fn upsert(&self, model: &M, source: &UpdateSource) -> Result where M: Into + UpsertModelInfo + Clone, @@ -95,7 +131,7 @@ impl<'a> ClientDb<'a> { }; self.record_model_change(&payload)?; - let _ = self.events_tx.send(payload); + self.pending_events.borrow_mut().push(payload); Ok(m) } @@ -113,7 +149,7 @@ impl<'a> ClientDb<'a> { }; self.record_model_change(&payload)?; - let _ = self.events_tx.send(payload); + self.pending_events.borrow_mut().push(payload); Ok(m.clone()) } diff --git a/crates/yaak-models/src/lib.rs b/crates/yaak-models/src/lib.rs index 755f1289..4e0af568 100644 --- a/crates/yaak-models/src/lib.rs +++ b/crates/yaak-models/src/lib.rs @@ -64,6 +64,17 @@ mod open { .map_err(|e| Error::Database(e.to_string())) } + /// `(readers, writer)` over one file: a pool of `max_size` readers and a + /// pool of exactly one writer. + pub fn file_pools( + path: impl Into, + max_size: u32, + min_idle: u32, + ) -> Result<(SqlitePool, SqlitePool)> { + let path: PathBuf = path.into(); + Ok((file_pool(&path, max_size, min_idle)?, file_pool(&path, 1, 1)?)) + } + pub fn memory_pool() -> Result { let manager = SqliteConnectionManager::memory().with_init(|c| init_connection(c)); // In-memory DB doesn't support multiple connections @@ -90,6 +101,16 @@ mod open { Ok(SqlitePool::single(conn)) } + /// One connection is all a browser VFS allows, so it reads and writes. + pub fn file_pools( + path: impl Into, + max_size: u32, + min_idle: u32, + ) -> Result<(SqlitePool, SqlitePool)> { + let pool = file_pool(path, max_size, min_idle)?; + Ok((pool.clone(), pool)) + } + pub fn memory_pool() -> Result { let conn = Connection::open_in_memory()?; init_connection(&conn)?; @@ -108,27 +129,33 @@ pub fn init_standalone( let db_path = db_path.as_ref(); let blob_path = blob_path.as_ref(); - // Main database pool. Sized for concurrent in-flight queries, not concurrent app - // features — connections are held per-statement, so even heavy fan-out (e.g. many - // gRPC streams) only needs a handful at once. Keep max_size modest: WAL connections - // hold ~3 file descriptors each, and macOS GUI apps get a 256 fd soft limit. + // Each database gets a reader pool and a one-connection writer pool; see + // `QueryManager` for why. Reader pools are sized for concurrent in-flight + // queries, not concurrent app features — connections are held per-statement, + // so even heavy fan-out (e.g. many gRPC streams) only needs a handful at once. + // Keep them modest: WAL connections hold ~3 file descriptors each, and macOS + // GUI apps get a 256 fd soft limit. info!("Initializing app database {db_path:?}"); - let pool = open::file_pool(db_path, 20, 2)?; - migrate_db(&pool)?; + let (readers, writer) = open::file_pools(db_path, 20, 2)?; + migrate_db(&writer)?; info!("Initializing blobs database {blob_path:?}"); - let blob_pool = open::file_pool(blob_path, 10, 1)?; - migrate_blob_db(&blob_pool)?; + let (blob_readers, blob_writer) = open::file_pools(blob_path, 10, 1)?; + migrate_blob_db(&blob_writer)?; let (tx, rx) = mpsc::channel(); - let query_manager = QueryManager::new(pool, tx); - let blob_manager = BlobManager::new(blob_pool); + let query_manager = QueryManager::new(readers, writer, tx); + let blob_manager = BlobManager::new(blob_readers, blob_writer); + bootstrap(&query_manager)?; Ok((query_manager, blob_manager, rx)) } /// Initialize the database managers with in-memory SQLite databases. /// Useful for testing and CI environments. +/// +/// An in-memory database is private to its connection, so the one connection +/// is both the reader pool and the writer. pub fn init_in_memory() -> Result<(QueryManager, BlobManager, mpsc::Receiver)> { let pool = open::memory_pool()?; migrate_db(&pool)?; @@ -137,8 +164,18 @@ pub fn init_in_memory() -> Result<(QueryManager, BlobManager, mpsc::Receiver Result<()> { + query_manager.with_tx(|tx| { + tx.ensure_settings()?; + tx.ensure_default_workspace()?; + Ok(()) + }) +} diff --git a/crates/yaak-models/src/models_ops.rs b/crates/yaak-models/src/models_ops.rs index 3964b678..2dfe808c 100644 --- a/crates/yaak-models/src/models_ops.rs +++ b/crates/yaak-models/src/models_ops.rs @@ -8,14 +8,14 @@ //! caller is a desktop window or an HTTP request. use crate::blob_manager::BlobManager; -use crate::client_db::ClientDb; +use crate::client_db::WriteDb; use crate::error::Error::GenericError; use crate::error::Result; use crate::models::AnyModel; use crate::util::UpdateSource; pub fn upsert_model( - db: &ClientDb, + db: &WriteDb, blobs: &BlobManager, model: AnyModel, source: &UpdateSource, @@ -41,7 +41,7 @@ pub fn upsert_model( /// Deletes cascade, so callers run this inside a transaction. pub fn delete_model( - tx: &ClientDb, + tx: &WriteDb, blobs: &BlobManager, model: AnyModel, source: &UpdateSource, @@ -69,7 +69,7 @@ pub fn delete_model( /// The model is re-read from the database rather than taken from the caller, so /// a duplicate never comes from a stale frontend snapshot. pub fn duplicate_model( - tx: &ClientDb, + tx: &WriteDb, model_type: &str, model_id: &str, source: &UpdateSource, diff --git a/crates/yaak-models/src/queries/batch.rs b/crates/yaak-models/src/queries/batch.rs index f4163724..9c96b241 100644 --- a/crates/yaak-models/src/queries/batch.rs +++ b/crates/yaak-models/src/queries/batch.rs @@ -1,10 +1,10 @@ -use crate::client_db::ClientDb; +use crate::client_db::WriteDb; use crate::error::Result; use crate::models::{Environment, Folder, GrpcRequest, HttpRequest, WebsocketRequest, Workspace}; use crate::util::{BatchUpsertResult, UpdateSource}; use log::info; -impl<'a> ClientDb<'a> { +impl<'a> WriteDb<'a> { pub fn batch_upsert( &self, workspaces: Vec, diff --git a/crates/yaak-models/src/queries/cookie_jars.rs b/crates/yaak-models/src/queries/cookie_jars.rs index 3c4f6695..2b9aa166 100644 --- a/crates/yaak-models/src/queries/cookie_jars.rs +++ b/crates/yaak-models/src/queries/cookie_jars.rs @@ -1,4 +1,4 @@ -use crate::client_db::ClientDb; +use crate::client_db::{ClientDb, WriteDb}; use crate::error::Result; use crate::models::{CookieJar, CookieJarIden}; use crate::util::UpdateSource; @@ -9,18 +9,22 @@ impl<'a> ClientDb<'a> { } pub fn list_cookie_jars(&self, workspace_id: &str) -> Result> { - let mut cookie_jars = self.find_many(CookieJarIden::WorkspaceId, workspace_id, None)?; + self.find_many(CookieJarIden::WorkspaceId, workspace_id, None) + } +} - if cookie_jars.is_empty() { +impl<'a> WriteDb<'a> { + /// A workspace with no cookie jar gets a default one. + pub fn ensure_default_cookie_jar(&self, workspace_id: &str) -> Result<()> { + if self.list_cookie_jars(workspace_id)?.is_empty() { let jar = CookieJar { name: "Default".to_string(), workspace_id: workspace_id.to_string(), ..Default::default() }; - cookie_jars.push(self.upsert_cookie_jar(&jar, &UpdateSource::Background)?); + self.upsert_cookie_jar(&jar, &UpdateSource::Background)?; } - - Ok(cookie_jars) + Ok(()) } pub fn delete_cookie_jar( diff --git a/crates/yaak-models/src/queries/environments.rs b/crates/yaak-models/src/queries/environments.rs index 9432d0f4..5163fce2 100644 --- a/crates/yaak-models/src/queries/environments.rs +++ b/crates/yaak-models/src/queries/environments.rs @@ -1,5 +1,5 @@ use super::conflict_free_name; -use crate::client_db::ClientDb; +use crate::client_db::{ClientDb, WriteDb}; use crate::error::Error::{MissingBaseEnvironment, MultipleBaseEnvironments}; use crate::error::Result; use crate::models::{Environment, EnvironmentIden, EnvironmentVariable}; @@ -20,7 +20,7 @@ impl<'a> ClientDb<'a> { } pub fn get_base_environment(&self, workspace_id: &str) -> Result { - let environments = self.list_environments_ensure_base(workspace_id)?; + let environments = self.list_environments(workspace_id)?; let base_environments = environments .into_iter() .filter(|e| e.parent_model == "workspace") @@ -30,38 +30,96 @@ impl<'a> ClientDb<'a> { return Err(MultipleBaseEnvironments(workspace_id.to_string())); } - Ok(base_environments.first().cloned().ok_or( - // Should never happen because one should be created above if it does not exist - MissingBaseEnvironment(workspace_id.to_string()), - )?) + Ok(base_environments + .first() + .cloned() + .ok_or(MissingBaseEnvironment(workspace_id.to_string()))?) } - /// Lists environments and will create a base environment if one doesn't exist - pub fn list_environments_ensure_base(&self, workspace_id: &str) -> Result> { - let mut environments = self.list_environments_dangerous(workspace_id)?; + pub fn list_environments(&self, workspace_id: &str) -> Result> { + Ok(self.find_many::(EnvironmentIden::WorkspaceId, workspace_id, None)?) + } - let base_environment = environments.iter().find(|e| e.parent_model == "workspace"); + /// Find other environments with the same parent folder + fn list_duplicate_folder_environments(&self, environment: &Environment) -> Vec { + if environment.parent_model != "folder" { + return Vec::new(); + } - if let None = base_environment { - let e = self.upsert_environment( - &Environment { - workspace_id: workspace_id.to_string(), - name: "Global Variables".to_string(), - parent_model: "workspace".to_string(), - ..Default::default() - }, - &UpdateSource::Background, + self.list_environments(&environment.workspace_id) + .unwrap_or_default() + .into_iter() + .filter(|e| { + e.id != environment.id + && e.parent_model == "folder" + && e.parent_id == environment.parent_id + }) + .collect() + } + + pub fn resolve_environments( + &self, + workspace_id: &str, + folder_id: Option<&str>, + active_environment_id: Option<&str>, + ) -> Result> { + let mut environments = Vec::new(); + + if let Some(folder_id) = folder_id { + let folder = self.get_folder(folder_id)?; + + // Add current folder's environment + if let Some(e) = self.get_environment_by_folder_id(folder_id)? { + environments.push(e); + }; + + // Recurse up + let ancestors = self.resolve_environments( + workspace_id, + folder.folder_id.as_deref(), + active_environment_id, )?; - info!("Created base environment {} for {workspace_id}", e.id); - environments.push(e); + environments.extend(ancestors); + } else { + // Add active and base environments + if let Some(id) = active_environment_id { + if let Ok(e) = self.get_environment(&id) { + // Add active sub environment + environments.push(e); + }; + }; + + // Add the base environment. A workspace that has never been + // opened has none yet; it simply contributes no variables. + match self.get_base_environment(workspace_id) { + Ok(e) => environments.push(e), + Err(MissingBaseEnvironment(_)) => {} + Err(e) => return Err(e), + } } Ok(environments) } +} - /// List environments for a workspace. Prefer list_environments_ensure_base() - fn list_environments_dangerous(&self, workspace_id: &str) -> Result> { - Ok(self.find_many::(EnvironmentIden::WorkspaceId, workspace_id, None)?) +impl<'a> WriteDb<'a> { + /// The workspace's base environment, created if it does not exist. + pub fn ensure_base_environment(&self, workspace_id: &str) -> Result { + match self.get_base_environment(workspace_id) { + Err(MissingBaseEnvironment(_)) => {} + other => return other, + } + let e = self.upsert_environment( + &Environment { + workspace_id: workspace_id.to_string(), + name: "Global Variables".to_string(), + parent_model: "workspace".to_string(), + ..Default::default() + }, + &UpdateSource::Background, + )?; + info!("Created base environment {} for {workspace_id}", e.id); + Ok(e) } pub fn delete_environment( @@ -72,7 +130,7 @@ impl<'a> ClientDb<'a> { let deleted_environment = self.delete(environment, source)?; // Recreate the base environment if we happened to delete it - self.list_environments_ensure_base(&environment.workspace_id)?; + self.ensure_base_environment(&environment.workspace_id)?; Ok(deleted_environment) } @@ -90,7 +148,7 @@ impl<'a> ClientDb<'a> { let mut environment = environment.clone(); environment.id = "".to_string(); let sibling_names = self - .list_environments_dangerous(&environment.workspace_id)? + .list_environments(&environment.workspace_id)? .into_iter() .map(|e| e.name) .collect::>(); @@ -98,23 +156,6 @@ impl<'a> ClientDb<'a> { self.upsert_environment(&environment, source) } - /// Find other environments with the same parent folder - fn list_duplicate_folder_environments(&self, environment: &Environment) -> Vec { - if environment.parent_model != "folder" { - return Vec::new(); - } - - self.list_environments_dangerous(&environment.workspace_id) - .unwrap_or_default() - .into_iter() - .filter(|e| { - e.id != environment.id - && e.parent_model == "folder" - && e.parent_id == environment.parent_id - }) - .collect() - } - pub fn upsert_environment( &self, environment: &Environment, @@ -154,43 +195,4 @@ impl<'a> ClientDb<'a> { source, ) } - - pub fn resolve_environments( - &self, - workspace_id: &str, - folder_id: Option<&str>, - active_environment_id: Option<&str>, - ) -> Result> { - let mut environments = Vec::new(); - - if let Some(folder_id) = folder_id { - let folder = self.get_folder(folder_id)?; - - // Add current folder's environment - if let Some(e) = self.get_environment_by_folder_id(folder_id)? { - environments.push(e); - }; - - // Recurse up - let ancestors = self.resolve_environments( - workspace_id, - folder.folder_id.as_deref(), - active_environment_id, - )?; - environments.extend(ancestors); - } else { - // Add active and base environments - if let Some(id) = active_environment_id { - if let Ok(e) = self.get_environment(&id) { - // Add active sub environment - environments.push(e); - }; - }; - - // Add the base environment - environments.push(self.get_base_environment(workspace_id)?); - } - - Ok(environments) - } } diff --git a/crates/yaak-models/src/queries/folders.rs b/crates/yaak-models/src/queries/folders.rs index ce6d0b63..df325a48 100644 --- a/crates/yaak-models/src/queries/folders.rs +++ b/crates/yaak-models/src/queries/folders.rs @@ -1,5 +1,5 @@ use super::{conflict_free_name, merge_headers}; -use crate::client_db::ClientDb; +use crate::client_db::{ClientDb, WriteDb}; use crate::connection_or_tx::ConnectionOrTx; use crate::error::Result; use crate::models::{ @@ -20,99 +20,6 @@ impl<'a> ClientDb<'a> { self.find_many(FolderIden::WorkspaceId, workspace_id, None) } - pub fn delete_folder(&self, folder: &Folder, source: &UpdateSource) -> Result { - match self.conn() { - ConnectionOrTx::Connection(_) => {} - ConnectionOrTx::Transaction(_) => {} - } - - let fid = &folder.id; - for m in self.find_many::(HttpRequestIden::FolderId, fid, None)? { - self.delete_http_request(&m, source)?; - } - - for m in self.find_many::(GrpcRequestIden::FolderId, fid, None)? { - self.delete_grpc_request(&m, source)?; - } - - for m in self.find_many::(WebsocketRequestIden::FolderId, fid, None)? { - self.delete_websocket_request(&m, source)?; - } - - for e in self.find_many(EnvironmentIden::ParentId, fid, None)? { - self.delete_environment(&e, source)?; - } - - // Recurse down into child folders - for folder in self.find_many::(FolderIden::FolderId, fid, None)? { - self.delete_folder(&folder, source)?; - } - - self.delete(folder, source) - } - - pub fn delete_folder_by_id(&self, id: &str, source: &UpdateSource) -> Result { - let folder = self.get_folder(id)?; - self.delete_folder(&folder, source) - } - - pub fn upsert_folder(&self, folder: &Folder, source: &UpdateSource) -> Result { - self.upsert(folder, source) - } - - pub fn duplicate_folder(&self, src_folder: &Folder, source: &UpdateSource) -> Result { - let fid = &src_folder.id; - - let mut folder = Folder { - id: "".into(), - sort_priority: src_folder.sort_priority + 0.001, - ..src_folder.clone() - }; - let sibling_names = self - .list_folders(&folder.workspace_id)? - .into_iter() - .filter(|f| f.folder_id == folder.folder_id) - .map(|f| f.name) - .collect::>(); - folder.name = conflict_free_name(&folder.name, &sibling_names); - let new_folder = self.upsert_folder(&folder, source)?; - - for m in self.find_many::(HttpRequestIden::FolderId, fid, None)? { - self.upsert_http_request( - &HttpRequest { id: "".into(), folder_id: Some(new_folder.id.clone()), ..m }, - source, - )?; - } - - for m in self.find_many::(WebsocketRequestIden::FolderId, fid, None)? { - self.upsert_websocket_request( - &WebsocketRequest { id: "".into(), folder_id: Some(new_folder.id.clone()), ..m }, - source, - )?; - } - - for m in self.find_many::(GrpcRequestIden::FolderId, fid, None)? { - self.upsert_grpc_request( - &GrpcRequest { id: "".into(), folder_id: Some(new_folder.id.clone()), ..m }, - source, - )?; - } - - for m in self.find_many::(EnvironmentIden::ParentId, fid, None)? { - self.upsert_environment( - &Environment { id: "".into(), parent_id: Some(new_folder.id.clone()), ..m }, - source, - )?; - } - - for m in self.find_many::(FolderIden::FolderId, fid, None)? { - // Recurse down - self.duplicate_folder(&Folder { folder_id: Some(new_folder.id.clone()), ..m }, source)?; - } - - Ok(new_folder) - } - pub fn resolve_auth_for_folder( &self, folder: &Folder, @@ -219,3 +126,98 @@ impl<'a> ClientDb<'a> { }) } } + +impl<'a> WriteDb<'a> { + pub fn delete_folder(&self, folder: &Folder, source: &UpdateSource) -> Result { + match self.conn() { + ConnectionOrTx::Connection(_) => {} + ConnectionOrTx::Transaction(_) => {} + } + + let fid = &folder.id; + for m in self.find_many::(HttpRequestIden::FolderId, fid, None)? { + self.delete_http_request(&m, source)?; + } + + for m in self.find_many::(GrpcRequestIden::FolderId, fid, None)? { + self.delete_grpc_request(&m, source)?; + } + + for m in self.find_many::(WebsocketRequestIden::FolderId, fid, None)? { + self.delete_websocket_request(&m, source)?; + } + + for e in self.find_many(EnvironmentIden::ParentId, fid, None)? { + self.delete_environment(&e, source)?; + } + + // Recurse down into child folders + for folder in self.find_many::(FolderIden::FolderId, fid, None)? { + self.delete_folder(&folder, source)?; + } + + self.delete(folder, source) + } + + pub fn delete_folder_by_id(&self, id: &str, source: &UpdateSource) -> Result { + let folder = self.get_folder(id)?; + self.delete_folder(&folder, source) + } + + pub fn upsert_folder(&self, folder: &Folder, source: &UpdateSource) -> Result { + self.upsert(folder, source) + } + + pub fn duplicate_folder(&self, src_folder: &Folder, source: &UpdateSource) -> Result { + let fid = &src_folder.id; + + let mut folder = Folder { + id: "".into(), + sort_priority: src_folder.sort_priority + 0.001, + ..src_folder.clone() + }; + let sibling_names = self + .list_folders(&folder.workspace_id)? + .into_iter() + .filter(|f| f.folder_id == folder.folder_id) + .map(|f| f.name) + .collect::>(); + folder.name = conflict_free_name(&folder.name, &sibling_names); + let new_folder = self.upsert_folder(&folder, source)?; + + for m in self.find_many::(HttpRequestIden::FolderId, fid, None)? { + self.upsert_http_request( + &HttpRequest { id: "".into(), folder_id: Some(new_folder.id.clone()), ..m }, + source, + )?; + } + + for m in self.find_many::(WebsocketRequestIden::FolderId, fid, None)? { + self.upsert_websocket_request( + &WebsocketRequest { id: "".into(), folder_id: Some(new_folder.id.clone()), ..m }, + source, + )?; + } + + for m in self.find_many::(GrpcRequestIden::FolderId, fid, None)? { + self.upsert_grpc_request( + &GrpcRequest { id: "".into(), folder_id: Some(new_folder.id.clone()), ..m }, + source, + )?; + } + + for m in self.find_many::(EnvironmentIden::ParentId, fid, None)? { + self.upsert_environment( + &Environment { id: "".into(), parent_id: Some(new_folder.id.clone()), ..m }, + source, + )?; + } + + for m in self.find_many::(FolderIden::FolderId, fid, None)? { + // Recurse down + self.duplicate_folder(&Folder { folder_id: Some(new_folder.id.clone()), ..m }, source)?; + } + + Ok(new_folder) + } +} diff --git a/crates/yaak-models/src/queries/graphql_introspections.rs b/crates/yaak-models/src/queries/graphql_introspections.rs index 8ece36de..ffaa5a1b 100644 --- a/crates/yaak-models/src/queries/graphql_introspections.rs +++ b/crates/yaak-models/src/queries/graphql_introspections.rs @@ -1,4 +1,4 @@ -use crate::client_db::ClientDb; +use crate::client_db::{ClientDb, WriteDb}; use crate::error::Result; use crate::models::{GraphQlIntrospection, GraphQlIntrospectionIden}; use crate::util::UpdateSource; @@ -11,7 +11,9 @@ impl<'a> ClientDb<'a> { pub fn get_graphql_introspection(&self, request_id: &str) -> Option { self.find_optional(GraphQlIntrospectionIden::RequestId, request_id) } +} +impl<'a> WriteDb<'a> { pub fn upsert_graphql_introspection( &self, workspace_id: &str, diff --git a/crates/yaak-models/src/queries/grpc_connections.rs b/crates/yaak-models/src/queries/grpc_connections.rs index 7de34ee2..f8e15397 100644 --- a/crates/yaak-models/src/queries/grpc_connections.rs +++ b/crates/yaak-models/src/queries/grpc_connections.rs @@ -1,4 +1,4 @@ -use crate::client_db::ClientDb; +use crate::client_db::{ClientDb, WriteDb}; use crate::error::Result; use crate::models::{GrpcConnection, GrpcConnectionIden, GrpcConnectionState}; use crate::queries::MAX_HISTORY_ITEMS; @@ -13,6 +13,20 @@ impl<'a> ClientDb<'a> { self.find_one(GrpcConnectionIden::Id, id) } + pub fn list_grpc_connections_for_request( + &self, + request_id: &str, + limit: Option, + ) -> Result> { + self.find_many(GrpcConnectionIden::RequestId, request_id, limit) + } + + pub fn list_grpc_connections(&self, workspace_id: &str) -> Result> { + self.find_many(GrpcConnectionIden::WorkspaceId, workspace_id, None) + } +} + +impl<'a> WriteDb<'a> { pub fn delete_all_grpc_connections_for_request( &self, request_id: &str, @@ -53,18 +67,6 @@ impl<'a> ClientDb<'a> { self.delete_grpc_connection(&grpc_connection, source) } - pub fn list_grpc_connections_for_request( - &self, - request_id: &str, - limit: Option, - ) -> Result> { - self.find_many(GrpcConnectionIden::RequestId, request_id, limit) - } - - pub fn list_grpc_connections(&self, workspace_id: &str) -> Result> { - self.find_many(GrpcConnectionIden::WorkspaceId, workspace_id, None) - } - pub fn cancel_pending_grpc_connections(&self) -> Result<()> { let closed = serde_json::to_value(&GrpcConnectionState::Closed)?; let (sql, params) = Query::update() diff --git a/crates/yaak-models/src/queries/grpc_events.rs b/crates/yaak-models/src/queries/grpc_events.rs index 3e1d1479..353ac77b 100644 --- a/crates/yaak-models/src/queries/grpc_events.rs +++ b/crates/yaak-models/src/queries/grpc_events.rs @@ -1,4 +1,4 @@ -use crate::client_db::ClientDb; +use crate::client_db::{ClientDb, WriteDb}; use crate::error::Result; use crate::models::{GrpcEvent, GrpcEventIden}; use crate::util::UpdateSource; @@ -11,7 +11,9 @@ impl<'a> ClientDb<'a> { pub fn list_grpc_events(&self, connection_id: &str) -> Result> { self.find_many(GrpcEventIden::ConnectionId, connection_id, None) } +} +impl<'a> WriteDb<'a> { pub fn upsert_grpc_event( &self, grpc_event: &GrpcEvent, diff --git a/crates/yaak-models/src/queries/grpc_requests.rs b/crates/yaak-models/src/queries/grpc_requests.rs index 593e128f..b7f43b1e 100644 --- a/crates/yaak-models/src/queries/grpc_requests.rs +++ b/crates/yaak-models/src/queries/grpc_requests.rs @@ -1,5 +1,5 @@ use super::{conflict_free_name, merge_headers}; -use crate::client_db::ClientDb; +use crate::client_db::{ClientDb, WriteDb}; use crate::error::Result; use crate::models::{ AnyModel, Folder, FolderIden, GrpcRequest, GrpcRequestIden, HttpRequestHeader, @@ -32,50 +32,6 @@ impl<'a> ClientDb<'a> { Ok(children) } - pub fn delete_grpc_request( - &self, - m: &GrpcRequest, - source: &UpdateSource, - ) -> Result { - self.delete_all_grpc_connections_for_request(m.id.as_str(), source)?; - self.delete(m, source) - } - - pub fn delete_grpc_request_by_id( - &self, - id: &str, - source: &UpdateSource, - ) -> Result { - let request = self.get_grpc_request(id)?; - self.delete_grpc_request(&request, source) - } - - pub fn duplicate_grpc_request( - &self, - grpc_request: &GrpcRequest, - source: &UpdateSource, - ) -> Result { - let mut request = grpc_request.clone(); - request.id = "".to_string(); - request.sort_priority = request.sort_priority + 0.001; - let sibling_names = self - .list_grpc_requests(&request.workspace_id)? - .into_iter() - .filter(|m| m.folder_id == request.folder_id) - .map(|m| m.name) - .collect::>(); - request.name = conflict_free_name(&request.name, &sibling_names); - self.upsert(&request, source) - } - - pub fn upsert_grpc_request( - &self, - grpc_request: &GrpcRequest, - source: &UpdateSource, - ) -> Result { - self.upsert(grpc_request, source) - } - pub fn resolve_auth_for_grpc_request( &self, grpc_request: &GrpcRequest, @@ -146,3 +102,49 @@ impl<'a> ClientDb<'a> { }) } } + +impl<'a> WriteDb<'a> { + pub fn delete_grpc_request( + &self, + m: &GrpcRequest, + source: &UpdateSource, + ) -> Result { + self.delete_all_grpc_connections_for_request(m.id.as_str(), source)?; + self.delete(m, source) + } + + pub fn delete_grpc_request_by_id( + &self, + id: &str, + source: &UpdateSource, + ) -> Result { + let request = self.get_grpc_request(id)?; + self.delete_grpc_request(&request, source) + } + + pub fn duplicate_grpc_request( + &self, + grpc_request: &GrpcRequest, + source: &UpdateSource, + ) -> Result { + let mut request = grpc_request.clone(); + request.id = "".to_string(); + request.sort_priority = request.sort_priority + 0.001; + let sibling_names = self + .list_grpc_requests(&request.workspace_id)? + .into_iter() + .filter(|m| m.folder_id == request.folder_id) + .map(|m| m.name) + .collect::>(); + request.name = conflict_free_name(&request.name, &sibling_names); + self.upsert(&request, source) + } + + pub fn upsert_grpc_request( + &self, + grpc_request: &GrpcRequest, + source: &UpdateSource, + ) -> Result { + self.upsert(grpc_request, source) + } +} diff --git a/crates/yaak-models/src/queries/http_requests.rs b/crates/yaak-models/src/queries/http_requests.rs index 18fe9171..dc042df8 100644 --- a/crates/yaak-models/src/queries/http_requests.rs +++ b/crates/yaak-models/src/queries/http_requests.rs @@ -1,5 +1,5 @@ use super::{conflict_free_name, merge_headers}; -use crate::client_db::ClientDb; +use crate::client_db::{ClientDb, WriteDb}; use crate::error::Result; use crate::models::{ AnyModel, Folder, FolderIden, HttpRequest, HttpRequestHeader, HttpRequestIden, @@ -18,50 +18,6 @@ impl<'a> ClientDb<'a> { self.find_many(HttpRequestIden::WorkspaceId, workspace_id, None) } - pub fn delete_http_request( - &self, - m: &HttpRequest, - source: &UpdateSource, - ) -> Result { - self.delete_all_http_responses_for_request(m.id.as_str(), source)?; - self.delete(m, source) - } - - pub fn delete_http_request_by_id( - &self, - id: &str, - source: &UpdateSource, - ) -> Result { - let http_request = self.get_http_request(id)?; - self.delete_http_request(&http_request, source) - } - - pub fn duplicate_http_request( - &self, - http_request: &HttpRequest, - source: &UpdateSource, - ) -> Result { - let mut http_request = http_request.clone(); - http_request.id = "".to_string(); - http_request.sort_priority = http_request.sort_priority + 0.001; - let sibling_names = self - .list_http_requests(&http_request.workspace_id)? - .into_iter() - .filter(|m| m.folder_id == http_request.folder_id) - .map(|m| m.name) - .collect::>(); - http_request.name = conflict_free_name(&http_request.name, &sibling_names); - self.upsert(&http_request, source) - } - - pub fn upsert_http_request( - &self, - http_request: &HttpRequest, - source: &UpdateSource, - ) -> Result { - self.upsert(http_request, source) - } - pub fn resolve_auth_for_http_request( &self, http_request: &HttpRequest, @@ -179,6 +135,52 @@ impl<'a> ClientDb<'a> { } } +impl<'a> WriteDb<'a> { + pub fn delete_http_request( + &self, + m: &HttpRequest, + source: &UpdateSource, + ) -> Result { + self.delete_all_http_responses_for_request(m.id.as_str(), source)?; + self.delete(m, source) + } + + pub fn delete_http_request_by_id( + &self, + id: &str, + source: &UpdateSource, + ) -> Result { + let http_request = self.get_http_request(id)?; + self.delete_http_request(&http_request, source) + } + + pub fn duplicate_http_request( + &self, + http_request: &HttpRequest, + source: &UpdateSource, + ) -> Result { + let mut http_request = http_request.clone(); + http_request.id = "".to_string(); + http_request.sort_priority = http_request.sort_priority + 0.001; + let sibling_names = self + .list_http_requests(&http_request.workspace_id)? + .into_iter() + .filter(|m| m.folder_id == http_request.folder_id) + .map(|m| m.name) + .collect::>(); + http_request.name = conflict_free_name(&http_request.name, &sibling_names); + self.upsert(&http_request, source) + } + + pub fn upsert_http_request( + &self, + http_request: &HttpRequest, + source: &UpdateSource, + ) -> Result { + self.upsert(http_request, source) + } +} + #[cfg(test)] mod tests { use crate::init_in_memory; @@ -225,72 +227,83 @@ mod tests { #[test] fn http_version_resolves_through_the_inheritance_chain() { let (query_manager, _blob_manager, _rx) = init_in_memory().expect("Failed to init DB"); - let db = query_manager.connect(); + let source = &UpdateSource::Background; - let workspace = db - .upsert_workspace( - &Workspace { - name: "Test".to_string(), - setting_http_version: HttpVersion::Http2, - ..Default::default() - }, - &UpdateSource::Background, - ) - .expect("Failed to upsert workspace"); - - let folder = db - .upsert_folder( - &Folder { workspace_id: workspace.id.clone(), ..Default::default() }, - &UpdateSource::Background, - ) - .expect("Failed to upsert folder"); - - let request = db - .upsert_http_request( - &HttpRequest { - workspace_id: workspace.id.clone(), - folder_id: Some(folder.id.clone()), - ..Default::default() - }, - &UpdateSource::Background, - ) - .expect("Failed to upsert request"); + let (folder, request) = query_manager + .with_tx(|db| { + let workspace = db.upsert_workspace( + &Workspace { + name: "Test".to_string(), + setting_http_version: HttpVersion::Http2, + ..Default::default() + }, + source, + )?; + let folder = db.upsert_folder( + &Folder { workspace_id: workspace.id.clone(), ..Default::default() }, + source, + )?; + let request = db.upsert_http_request( + &HttpRequest { + workspace_id: workspace.id.clone(), + folder_id: Some(folder.id.clone()), + ..Default::default() + }, + source, + )?; + Ok::<_, crate::error::Error>((folder, request)) + }) + .expect("Failed to seed"); // No overrides, so the workspace base value applies - let resolved = db.resolve_settings_for_http_request(&request).expect("Failed to resolve"); + let resolved = query_manager + .connect() + .resolve_settings_for_http_request(&request) + .expect("Failed to resolve"); assert_eq!(resolved.http_version.value, HttpVersion::Http2); assert_eq!(resolved.http_version.source_model, "workspace"); // A folder override beats the workspace base - db.upsert_folder( - &Folder { - setting_http_version: InheritedHttpVersionSetting { - enabled: true, - value: HttpVersion::Http1, - }, - ..folder - }, - &UpdateSource::Background, - ) - .expect("Failed to update folder"); - let resolved = db.resolve_settings_for_http_request(&request).expect("Failed to resolve"); + query_manager + .with_tx(|db| { + db.upsert_folder( + &Folder { + setting_http_version: InheritedHttpVersionSetting { + enabled: true, + value: HttpVersion::Http1, + }, + ..folder + }, + source, + ) + }) + .expect("Failed to update folder"); + let resolved = query_manager + .connect() + .resolve_settings_for_http_request(&request) + .expect("Failed to resolve"); assert_eq!(resolved.http_version.value, HttpVersion::Http1); assert_eq!(resolved.http_version.source_model, "folder"); // A request override beats them both - let request = db - .upsert_http_request( - &HttpRequest { - setting_http_version: InheritedHttpVersionSetting { - enabled: true, - value: HttpVersion::Auto, + let request = query_manager + .with_tx(|db| { + db.upsert_http_request( + &HttpRequest { + setting_http_version: InheritedHttpVersionSetting { + enabled: true, + value: HttpVersion::Auto, + }, + ..request }, - ..request - }, - &UpdateSource::Background, - ) + source, + ) + }) .expect("Failed to update request"); - let resolved = db.resolve_settings_for_http_request(&request).expect("Failed to resolve"); + let resolved = query_manager + .connect() + .resolve_settings_for_http_request(&request) + .expect("Failed to resolve"); assert_eq!(resolved.http_version.value, HttpVersion::Auto); assert_eq!(resolved.http_version.source_model, "http_request"); } diff --git a/crates/yaak-models/src/queries/http_response_events.rs b/crates/yaak-models/src/queries/http_response_events.rs index 2c873ecc..9e5bb8a1 100644 --- a/crates/yaak-models/src/queries/http_response_events.rs +++ b/crates/yaak-models/src/queries/http_response_events.rs @@ -1,4 +1,4 @@ -use crate::client_db::ClientDb; +use crate::client_db::{ClientDb, WriteDb}; use crate::error::Result; use crate::models::{HttpResponseEvent, HttpResponseEventIden}; use crate::util::UpdateSource; @@ -7,7 +7,9 @@ impl<'a> ClientDb<'a> { pub fn list_http_response_events(&self, response_id: &str) -> Result> { self.find_many(HttpResponseEventIden::ResponseId, response_id, None) } +} +impl<'a> WriteDb<'a> { pub fn upsert_http_response_event( &self, http_response_event: &HttpResponseEvent, diff --git a/crates/yaak-models/src/queries/http_responses.rs b/crates/yaak-models/src/queries/http_responses.rs index f77c0c4f..5151d285 100644 --- a/crates/yaak-models/src/queries/http_responses.rs +++ b/crates/yaak-models/src/queries/http_responses.rs @@ -1,5 +1,5 @@ use crate::blob_manager::BlobManager; -use crate::client_db::ClientDb; +use crate::client_db::{ClientDb, WriteDb}; use crate::error::Result; use crate::models::{HttpResponse, HttpResponseIden, HttpResponseState}; use crate::queries::MAX_HISTORY_ITEMS; @@ -31,20 +31,6 @@ impl<'a> ClientDb<'a> { self.find_many(HttpResponseIden::WorkspaceId, workspace_id, limit) } - /// Returns the number of responses deleted. - pub fn delete_all_http_responses_for_request( - &self, - request_id: &str, - source: &UpdateSource, - ) -> Result { - let responses = self.list_http_responses_for_request(request_id, None)?; - let count = responses.len(); - for m in responses { - self.delete(&m, source)?; - } - Ok(count) - } - /// Delete blob-stored response bodies whose owning HTTP response row no /// longer exists. Blob ids are keyed by the response that owns them — /// "{response_id}" for a response body, "{response_id}.request" for the @@ -55,19 +41,24 @@ impl<'a> ClientDb<'a> { /// /// Returns the number of orphaned bodies deleted. pub fn delete_orphaned_response_body_blobs(&self, blobs: &BlobManager) -> Result { - let mut deleted = 0; + let orphaned = blobs + .connect() + .list_body_ids()? + .into_iter() + .filter(|body_id| { + let response_id = body_id.split('.').next().unwrap_or_default(); + self.find_optional::(HttpResponseIden::Id, response_id).is_none() + }) + .collect::>(); - let blob_ctx = blobs.connect(); - for body_id in blob_ctx.list_body_ids()? { - let response_id = body_id.split('.').next().unwrap_or_default(); - if self.find_optional::(HttpResponseIden::Id, response_id).is_some() { - continue; + blobs.with_tx(|b| { + for body_id in &orphaned { + b.delete_chunks(body_id)?; } - blob_ctx.delete_chunks(&body_id)?; - deleted += 1; - } + Ok::<_, crate::error::Error>(()) + })?; - Ok(deleted) + Ok(orphaned.len()) } /// Delete response body data (blob chunks and body files) whose owning HTTP @@ -107,6 +98,22 @@ impl<'a> ClientDb<'a> { Ok(deleted) } +} + +impl<'a> WriteDb<'a> { + /// Returns the number of responses deleted. + pub fn delete_all_http_responses_for_request( + &self, + request_id: &str, + source: &UpdateSource, + ) -> Result { + let responses = self.list_http_responses_for_request(request_id, None)?; + let count = responses.len(); + for m in responses { + self.delete(&m, source)?; + } + Ok(count) + } /// Returns the number of responses deleted. pub fn delete_all_http_responses_for_workspace( @@ -137,9 +144,8 @@ impl<'a> ClientDb<'a> { } // Delete request body blobs (pattern: {response_id}.request) - let blob_ctx = blob_manager.connect(); let body_id = format!("{}.request", http_response.id); - if let Err(e) = blob_ctx.delete_chunks(&body_id) { + if let Err(e) = blob_manager.with_tx(|b| b.delete_chunks(&body_id)) { error!("Failed to delete request body blobs: {}", e); } @@ -186,36 +192,50 @@ impl<'a> ClientDb<'a> { #[cfg(test)] mod tests { use crate::blob_manager::{BlobManager, BodyChunk}; - use crate::client_db::ClientDb; + use crate::error::Error; use crate::init_in_memory; use crate::models::{HttpRequest, HttpResponse, Workspace}; + use crate::query_manager::QueryManager; use crate::util::UpdateSource; /// A workspace, a request, and one response that still exists. - fn seed_live_response(db: &ClientDb, blob_manager: &BlobManager) -> HttpResponse { + fn seed_live_response( + query_manager: &QueryManager, + blob_manager: &BlobManager, + ) -> HttpResponse { let source = &UpdateSource::Background; - let workspace = db - .upsert_workspace( - &Workspace { name: "GC Test".to_string(), ..Default::default() }, - source, - ) - .expect("Failed to upsert workspace"); - let request = db - .upsert_http_request( - &HttpRequest { workspace_id: workspace.id.clone(), ..Default::default() }, - source, - ) - .expect("Failed to upsert request"); - db.upsert_http_response( - &HttpResponse { - request_id: request.id.clone(), - workspace_id: workspace.id.clone(), - ..Default::default() - }, - source, - blob_manager, - ) - .expect("Failed to upsert response") + query_manager + .with_tx(|db| { + let workspace = db.upsert_workspace( + &Workspace { name: "GC Test".to_string(), ..Default::default() }, + source, + )?; + let request = db.upsert_http_request( + &HttpRequest { workspace_id: workspace.id.clone(), ..Default::default() }, + source, + )?; + db.upsert_http_response( + &HttpResponse { + request_id: request.id.clone(), + workspace_id: workspace.id.clone(), + ..Default::default() + }, + source, + blob_manager, + ) + }) + .expect("Failed to seed response") + } + + fn insert_bodies(blob_manager: &BlobManager, body_ids: &[&str]) { + blob_manager + .with_tx(|b| { + for id in body_ids { + b.insert_chunk(&BodyChunk::new(*id, 0, b"data".to_vec()))?; + } + Ok::<_, Error>(()) + }) + .expect("Failed to insert chunks"); } /// What a browser host runs: no filesystem, so bodies exist only as blob @@ -223,23 +243,21 @@ mod tests { #[test] fn deletes_orphaned_response_body_blobs() { let (query_manager, blob_manager, _rx) = init_in_memory().expect("Failed to init DB"); - let db = query_manager.connect(); - let live = seed_live_response(&db, &blob_manager); + let live = seed_live_response(&query_manager, &blob_manager); let live_request_body_id = format!("{}.request", live.id); - { - // Scope the connection: the in-memory pool only has one, and the GC - // needs to take it - let blob_ctx = blob_manager.connect(); - blob_ctx.insert_chunk(&BodyChunk::new(&live.id, 0, b"live".to_vec())).unwrap(); - blob_ctx - .insert_chunk(&BodyChunk::new(&live_request_body_id, 0, b"live".to_vec())) - .unwrap(); - blob_ctx.insert_chunk(&BodyChunk::new("rs_gone", 0, b"dead".to_vec())).unwrap(); - blob_ctx.insert_chunk(&BodyChunk::new("rs_gone.request", 0, b"dead".to_vec())).unwrap(); - } + insert_bodies( + &blob_manager, + &[ + &live.id, + &live_request_body_id, + "rs_gone", + "rs_gone.request", + ], + ); - let deleted = db + let deleted = query_manager + .connect() .delete_orphaned_response_body_blobs(&blob_manager) .expect("Failed to GC response body blobs"); assert_eq!(deleted, 2); @@ -254,24 +272,18 @@ mod tests { #[test] fn deletes_orphaned_response_bodies() { let (query_manager, blob_manager, _rx) = init_in_memory().expect("Failed to init DB"); - let db = query_manager.connect(); - let live = seed_live_response(&db, &blob_manager); + let live = seed_live_response(&query_manager, &blob_manager); let live_body_id = format!("{}.request", live.id); - { - // Scope the connection: the in-memory pool only has one, and the GC - // needs to take it - let blob_ctx = blob_manager.connect(); - blob_ctx.insert_chunk(&BodyChunk::new(&live_body_id, 0, b"live".to_vec())).unwrap(); - blob_ctx.insert_chunk(&BodyChunk::new("rs_gone.request", 0, b"dead".to_vec())).unwrap(); - } + insert_bodies(&blob_manager, &[&live_body_id, "rs_gone.request"]); let dir = std::env::temp_dir().join(format!("yaak-blob-gc-test-{}", live.id)); std::fs::create_dir_all(&dir).unwrap(); std::fs::write(dir.join(&live.id), b"live").unwrap(); std::fs::write(dir.join("rs_gone"), b"dead").unwrap(); - let deleted = db + let deleted = query_manager + .connect() .delete_orphaned_response_bodies(&blob_manager, &dir) .expect("Failed to GC response bodies"); assert_eq!(deleted, 2); diff --git a/crates/yaak-models/src/queries/import_source_resources.rs b/crates/yaak-models/src/queries/import_source_resources.rs index 6c0abb63..91fa190f 100644 --- a/crates/yaak-models/src/queries/import_source_resources.rs +++ b/crates/yaak-models/src/queries/import_source_resources.rs @@ -1,4 +1,4 @@ -use crate::client_db::ClientDb; +use crate::client_db::{ClientDb, WriteDb}; use crate::error::Result; use crate::models::{ImportSourceResource, ImportSourceResourceIden}; use sea_query::ExprTrait; @@ -20,7 +20,9 @@ impl<'a> ClientDb<'a> { let items = stmt.query_map(&*params.as_params(), |row| row.try_into())?; Ok(items.filter_map(|v| v.ok()).collect()) } +} +impl<'a> WriteDb<'a> { pub fn upsert_import_source_resource( &self, resource: &ImportSourceResource, diff --git a/crates/yaak-models/src/queries/import_sources.rs b/crates/yaak-models/src/queries/import_sources.rs index c077cac4..342a7cc6 100644 --- a/crates/yaak-models/src/queries/import_sources.rs +++ b/crates/yaak-models/src/queries/import_sources.rs @@ -1,4 +1,4 @@ -use crate::client_db::ClientDb; +use crate::client_db::{ClientDb, WriteDb}; use crate::error::Result; use crate::models::{ImportSource, ImportSourceIden}; use crate::util::UpdateSource; @@ -25,7 +25,9 @@ impl<'a> ClientDb<'a> { let sources = self.list_import_sources(workspace_id)?; Ok(sources.into_iter().find(|s| s.importer == importer && s.origin == origin)) } +} +impl<'a> WriteDb<'a> { pub fn upsert_import_source( &self, import_source: &ImportSource, diff --git a/crates/yaak-models/src/queries/key_values.rs b/crates/yaak-models/src/queries/key_values.rs index 1a528eeb..ff5ad186 100644 --- a/crates/yaak-models/src/queries/key_values.rs +++ b/crates/yaak-models/src/queries/key_values.rs @@ -1,4 +1,4 @@ -use crate::client_db::ClientDb; +use crate::client_db::{ClientDb, WriteDb}; use crate::error::Result; use crate::models::{KeyValue, KeyValueIden, UpsertModelInfo}; use crate::util::UpdateSource; @@ -89,7 +89,9 @@ impl<'a> ClientDb<'a> { .build_rusqlite(SqliteQueryBuilder); self.conn().resolve().query_row(sql.as_str(), &*params.as_params(), KeyValue::from_row).ok() } +} +impl<'a> WriteDb<'a> { pub fn set_key_value_dte( &self, namespace: &str, diff --git a/crates/yaak-models/src/queries/model_changes.rs b/crates/yaak-models/src/queries/model_changes.rs index 987118ad..57e62fd3 100644 --- a/crates/yaak-models/src/queries/model_changes.rs +++ b/crates/yaak-models/src/queries/model_changes.rs @@ -1,4 +1,4 @@ -use crate::client_db::ClientDb; +use crate::client_db::{ClientDb, WriteDb}; use crate::error::Result; use crate::util::ModelPayload; use rusqlite::params; @@ -69,7 +69,9 @@ impl<'a> ClientDb<'a> { Ok(items.collect::, rusqlite::Error>>()?) } +} +impl<'a> WriteDb<'a> { pub fn prune_model_changes_older_than_days(&self, days: i64) -> Result { let offset = format!("-{days} days"); Ok(self.conn().resolve().execute( @@ -101,23 +103,36 @@ mod tests { use crate::util::{ModelChangeEvent, UpdateSource}; use serde_json::json; + /// Startup bootstraps rows of its own; these tests count only their own. + fn clear_changes(query_manager: &crate::query_manager::QueryManager) { + query_manager + .with_tx(|db| { + db.conn().resolve().execute("DELETE FROM model_changes", [])?; + Ok::<_, crate::error::Error>(()) + }) + .expect("Failed to clear model changes"); + } + #[test] fn records_model_changes_for_upsert_and_delete() { let (query_manager, blob_manager, _rx) = init_in_memory().expect("Failed to init DB"); - let db = query_manager.connect(); + clear_changes(&query_manager); - let workspace = db - .upsert_workspace( - &Workspace { - name: "Changes Test".to_string(), - setting_follow_redirects: true, - setting_validate_certificates: true, - ..Default::default() - }, - &UpdateSource::Sync, - ) + let workspace = query_manager + .with_tx(|db| { + db.upsert_workspace( + &Workspace { + name: "Changes Test".to_string(), + setting_follow_redirects: true, + setting_validate_certificates: true, + ..Default::default() + }, + &UpdateSource::Sync, + ) + }) .expect("Failed to upsert workspace"); + let db = query_manager.connect(); let created_changes = db.list_model_changes_after(0, 10).expect("Failed to list changes"); assert_eq!(created_changes.len(), 1); assert_eq!(created_changes[0].payload.model.id(), workspace.id); @@ -128,9 +143,14 @@ mod tests { )); assert!(matches!(created_changes[0].payload.update_source, UpdateSource::Sync)); - db.delete_workspace_by_id(&workspace.id, &UpdateSource::Sync, &blob_manager) + drop(db); + query_manager + .with_tx(|db| { + db.delete_workspace_by_id(&workspace.id, &UpdateSource::Sync, &blob_manager) + }) .expect("Failed to delete workspace"); + let db = query_manager.connect(); let all_changes = db.list_model_changes_after(0, 10).expect("Failed to list changes"); assert_eq!(all_changes.len(), 2); assert!(matches!(all_changes[1].payload.change, ModelChangeEvent::Delete)); @@ -146,19 +166,23 @@ mod tests { #[test] fn prunes_old_model_changes() { let (query_manager, _blob_manager, _rx) = init_in_memory().expect("Failed to init DB"); + clear_changes(&query_manager); + + query_manager + .with_tx(|db| { + db.upsert_workspace( + &Workspace { + name: "Prune Test".to_string(), + setting_follow_redirects: true, + setting_validate_certificates: true, + ..Default::default() + }, + &UpdateSource::Sync, + ) + }) + .expect("Failed to upsert workspace"); + let db = query_manager.connect(); - - db.upsert_workspace( - &Workspace { - name: "Prune Test".to_string(), - setting_follow_redirects: true, - setting_validate_certificates: true, - ..Default::default() - }, - &UpdateSource::Sync, - ) - .expect("Failed to upsert workspace"); - let changes = db.list_model_changes_after(0, 10).expect("Failed to list changes"); assert_eq!(changes.len(), 1); @@ -170,31 +194,41 @@ mod tests { ) .expect("Failed to age model change row"); - let pruned = - db.prune_model_changes_older_than_days(30).expect("Failed to prune model changes"); + drop(db); + let pruned = query_manager + .with_tx(|db| db.prune_model_changes_older_than_days(30)) + .expect("Failed to prune model changes"); assert_eq!(pruned, 1); - assert!(db.list_model_changes_after(0, 10).expect("Failed to list changes").is_empty()); + assert!( + query_manager + .connect() + .list_model_changes_after(0, 10) + .expect("Failed to list changes") + .is_empty() + ); } #[test] fn list_model_changes_since_uses_timestamp_with_id_tiebreaker() { let (query_manager, blob_manager, _rx) = init_in_memory().expect("Failed to init DB"); + clear_changes(&query_manager); + + query_manager + .with_tx(|db| { + let workspace = db.upsert_workspace( + &Workspace { + name: "Cursor Test".to_string(), + setting_follow_redirects: true, + setting_validate_certificates: true, + ..Default::default() + }, + &UpdateSource::Sync, + )?; + db.delete_workspace_by_id(&workspace.id, &UpdateSource::Sync, &blob_manager) + }) + .expect("Failed to seed changes"); + let db = query_manager.connect(); - - let workspace = db - .upsert_workspace( - &Workspace { - name: "Cursor Test".to_string(), - setting_follow_redirects: true, - setting_validate_certificates: true, - ..Default::default() - }, - &UpdateSource::Sync, - ) - .expect("Failed to upsert workspace"); - db.delete_workspace_by_id(&workspace.id, &UpdateSource::Sync, &blob_manager) - .expect("Failed to delete workspace"); - let all = db.list_model_changes_after(0, 10).expect("Failed to list changes"); assert_eq!(all.len(), 2); @@ -213,19 +247,23 @@ mod tests { #[test] fn prunes_old_model_changes_by_hours() { let (query_manager, _blob_manager, _rx) = init_in_memory().expect("Failed to init DB"); + clear_changes(&query_manager); + + query_manager + .with_tx(|db| { + db.upsert_workspace( + &Workspace { + name: "Prune Hour Test".to_string(), + setting_follow_redirects: true, + setting_validate_certificates: true, + ..Default::default() + }, + &UpdateSource::Sync, + ) + }) + .expect("Failed to upsert workspace"); + let db = query_manager.connect(); - - db.upsert_workspace( - &Workspace { - name: "Prune Hour Test".to_string(), - setting_follow_redirects: true, - setting_validate_certificates: true, - ..Default::default() - }, - &UpdateSource::Sync, - ) - .expect("Failed to upsert workspace"); - let changes = db.list_model_changes_after(0, 10).expect("Failed to list changes"); assert_eq!(changes.len(), 1); @@ -237,14 +275,17 @@ mod tests { ) .expect("Failed to age model change row"); - let pruned = - db.prune_model_changes_older_than_hours(1).expect("Failed to prune model changes"); + drop(db); + let pruned = query_manager + .with_tx(|db| db.prune_model_changes_older_than_hours(1)) + .expect("Failed to prune model changes"); assert_eq!(pruned, 1); } #[test] fn list_model_changes_deserializes_http_response_event_payload() { let (query_manager, _blob_manager, _rx) = init_in_memory().expect("Failed to init DB"); + clear_changes(&query_manager); let db = query_manager.connect(); let payload = json!({ diff --git a/crates/yaak-models/src/queries/plugin_key_values.rs b/crates/yaak-models/src/queries/plugin_key_values.rs index 979fb11c..dabc8bbb 100644 --- a/crates/yaak-models/src/queries/plugin_key_values.rs +++ b/crates/yaak-models/src/queries/plugin_key_values.rs @@ -1,4 +1,4 @@ -use crate::client_db::ClientDb; +use crate::client_db::{ClientDb, WriteDb}; use crate::error::Result; use crate::models::{PluginKeyValue, PluginKeyValueIden}; use sea_query::ExprTrait; @@ -22,7 +22,9 @@ impl<'a> ClientDb<'a> { .query_row(sql.as_str(), &*params.as_params(), |row| row.try_into()) .ok() } +} +impl<'a> WriteDb<'a> { pub fn set_plugin_key_value( &self, plugin_name: &str, diff --git a/crates/yaak-models/src/queries/plugins.rs b/crates/yaak-models/src/queries/plugins.rs index 48dd1998..007c9210 100644 --- a/crates/yaak-models/src/queries/plugins.rs +++ b/crates/yaak-models/src/queries/plugins.rs @@ -1,4 +1,4 @@ -use crate::client_db::ClientDb; +use crate::client_db::{ClientDb, WriteDb}; use crate::error::Result; use crate::models::{Plugin, PluginIden}; use crate::util::UpdateSource; @@ -15,7 +15,9 @@ impl<'a> ClientDb<'a> { pub fn list_plugins(&self) -> Result> { self.find_all() } +} +impl<'a> WriteDb<'a> { pub fn delete_plugin(&self, plugin: &Plugin, source: &UpdateSource) -> Result { self.delete(plugin, source) } diff --git a/crates/yaak-models/src/queries/settings.rs b/crates/yaak-models/src/queries/settings.rs index b2d70832..12fd175b 100644 --- a/crates/yaak-models/src/queries/settings.rs +++ b/crates/yaak-models/src/queries/settings.rs @@ -1,52 +1,61 @@ use std::collections::HashMap; -use crate::client_db::ClientDb; +use crate::client_db::{ClientDb, WriteDb}; use crate::error::Result; use crate::models::{EditorKeymap, Settings, SettingsIden}; use crate::util::UpdateSource; impl<'a> ClientDb<'a> { + /// The settings row, or the defaults if it has not been written yet. + /// [`WriteDb::ensure_settings`] persists it at startup. pub fn get_settings(&self) -> Settings { - let id = "default".to_string(); + self.find_optional::(SettingsIden::Id, "default").unwrap_or_else(default_settings) + } +} - if let Some(s) = self.find_optional::(SettingsIden::Id, &id) { - return s; - }; - - let settings = Settings { - model: "settings".to_string(), - id, - created_at: Default::default(), - updated_at: Default::default(), - - appearance: "system".to_string(), - client_certificates: Vec::new(), - editor_font_size: 12, - editor_font: None, - editor_keymap: EditorKeymap::Default, - editor_soft_wrap: true, - interface_font_size: 14, - interface_scale: 1.0, - interface_font: None, - hide_window_controls: false, - use_native_titlebar: false, - open_workspace_new_window: None, - proxy: None, - theme_dark: "yaak-dark".to_string(), - theme_light: "yaak-light".to_string(), - update_channel: "stable".to_string(), - autoupdate: true, - colored_methods: false, - hide_license_badge: false, - prompt_feedback: true, - auto_download_updates: true, - check_notifications: true, - hotkeys: HashMap::new(), - }; - self.upsert(&settings, &UpdateSource::Background).expect("Failed to upsert settings") +impl<'a> WriteDb<'a> { + /// Create the settings row if it does not exist. + pub fn ensure_settings(&self) -> Result { + if let Some(s) = self.find_optional::(SettingsIden::Id, "default") { + return Ok(s); + } + self.upsert(&default_settings(), &UpdateSource::Background) } pub fn upsert_settings(&self, settings: &Settings, source: &UpdateSource) -> Result { self.upsert(settings, source) } } + +fn default_settings() -> Settings { + Settings { + model: "settings".to_string(), + id: "default".to_string(), + created_at: Default::default(), + updated_at: Default::default(), + + appearance: "system".to_string(), + client_certificates: Vec::new(), + editor_font_size: 12, + editor_font: None, + editor_keymap: EditorKeymap::Default, + editor_soft_wrap: true, + interface_font_size: 14, + interface_scale: 1.0, + interface_font: None, + hide_window_controls: false, + use_native_titlebar: false, + open_workspace_new_window: None, + proxy: None, + theme_dark: "yaak-dark".to_string(), + theme_light: "yaak-light".to_string(), + update_channel: "stable".to_string(), + autoupdate: true, + colored_methods: false, + hide_license_badge: false, + prompt_feedback: true, + auto_download_updates: true, + check_notifications: true, + hotkeys: HashMap::new(), + } +} diff --git a/crates/yaak-models/src/queries/sync_states.rs b/crates/yaak-models/src/queries/sync_states.rs index b3053912..3a999013 100644 --- a/crates/yaak-models/src/queries/sync_states.rs +++ b/crates/yaak-models/src/queries/sync_states.rs @@ -1,4 +1,4 @@ -use crate::client_db::ClientDb; +use crate::client_db::{ClientDb, WriteDb}; use crate::error::Result; use crate::models::{SyncState, SyncStateIden, UpsertModelInfo}; use crate::util::UpdateSource; @@ -12,10 +12,6 @@ impl<'a> ClientDb<'a> { self.find_one(SyncStateIden::Id, id) } - pub fn upsert_sync_state(&self, sync_state: &SyncState) -> Result { - self.upsert(sync_state, &UpdateSource::Sync) - } - pub fn list_sync_states_for_workspace( &self, workspace_id: &str, @@ -34,6 +30,12 @@ impl<'a> ClientDb<'a> { let items = stmt.query_map(&*params.as_params(), SyncState::from_row)?; Ok(items.map(|v| v.unwrap()).collect()) } +} + +impl<'a> WriteDb<'a> { + pub fn upsert_sync_state(&self, sync_state: &SyncState) -> Result { + self.upsert(sync_state, &UpdateSource::Sync) + } pub fn delete_sync_state(&self, sync_state: &SyncState) -> Result { self.delete(sync_state, &UpdateSource::Sync) diff --git a/crates/yaak-models/src/queries/websocket_connections.rs b/crates/yaak-models/src/queries/websocket_connections.rs index d1dbd37a..87bda32f 100644 --- a/crates/yaak-models/src/queries/websocket_connections.rs +++ b/crates/yaak-models/src/queries/websocket_connections.rs @@ -1,4 +1,4 @@ -use crate::client_db::ClientDb; +use crate::client_db::{ClientDb, WriteDb}; use crate::error::Result; use crate::models::{WebsocketConnection, WebsocketConnectionIden, WebsocketConnectionState}; use crate::queries::MAX_HISTORY_ITEMS; @@ -13,6 +13,22 @@ impl<'a> ClientDb<'a> { self.find_one(WebsocketConnectionIden::Id, id) } + pub fn list_websocket_connections( + &self, + workspace_id: &str, + ) -> Result> { + self.find_many(WebsocketConnectionIden::WorkspaceId, workspace_id, None) + } + + pub fn list_websocket_connections_for_request( + &self, + request_id: &str, + ) -> Result> { + self.find_many(WebsocketConnectionIden::RequestId, request_id, None) + } +} + +impl<'a> WriteDb<'a> { pub fn delete_all_websocket_connections_for_request( &self, request_id: &str, @@ -37,20 +53,6 @@ impl<'a> ClientDb<'a> { Ok(()) } - pub fn list_websocket_connections( - &self, - workspace_id: &str, - ) -> Result> { - self.find_many(WebsocketConnectionIden::WorkspaceId, workspace_id, None) - } - - pub fn list_websocket_connections_for_request( - &self, - request_id: &str, - ) -> Result> { - self.find_many(WebsocketConnectionIden::RequestId, request_id, None) - } - pub fn delete_websocket_connection( &self, websocket_connection: &WebsocketConnection, diff --git a/crates/yaak-models/src/queries/websocket_events.rs b/crates/yaak-models/src/queries/websocket_events.rs index 458f516c..d1779996 100644 --- a/crates/yaak-models/src/queries/websocket_events.rs +++ b/crates/yaak-models/src/queries/websocket_events.rs @@ -1,4 +1,4 @@ -use crate::client_db::ClientDb; +use crate::client_db::{ClientDb, WriteDb}; use crate::error::Result; use crate::models::{WebsocketEvent, WebsocketEventIden}; use crate::util::UpdateSource; @@ -11,7 +11,9 @@ impl<'a> ClientDb<'a> { pub fn list_websocket_events(&self, connection_id: &str) -> Result> { self.find_many(WebsocketEventIden::ConnectionId, connection_id, None) } +} +impl<'a> WriteDb<'a> { pub fn upsert_websocket_event( &self, websocket_event: &WebsocketEvent, diff --git a/crates/yaak-models/src/queries/websocket_requests.rs b/crates/yaak-models/src/queries/websocket_requests.rs index 1dcf752f..7b3a66a1 100644 --- a/crates/yaak-models/src/queries/websocket_requests.rs +++ b/crates/yaak-models/src/queries/websocket_requests.rs @@ -1,5 +1,5 @@ use super::{conflict_free_name, merge_headers}; -use crate::client_db::ClientDb; +use crate::client_db::{ClientDb, WriteDb}; use crate::error::Result; use crate::models::{ AnyModel, Folder, FolderIden, HttpRequestHeader, ResolvedHttpRequestSettings, ResolvedSetting, @@ -34,50 +34,6 @@ impl<'a> ClientDb<'a> { Ok(children) } - pub fn delete_websocket_request( - &self, - websocket_request: &WebsocketRequest, - source: &UpdateSource, - ) -> Result { - self.delete_all_websocket_connections_for_request(websocket_request.id.as_str(), source)?; - self.delete(websocket_request, source) - } - - pub fn delete_websocket_request_by_id( - &self, - id: &str, - source: &UpdateSource, - ) -> Result { - let request = self.get_websocket_request(id)?; - self.delete_websocket_request(&request, source) - } - - pub fn duplicate_websocket_request( - &self, - websocket_request: &WebsocketRequest, - source: &UpdateSource, - ) -> Result { - let mut websocket_request = websocket_request.clone(); - websocket_request.id = "".to_string(); - websocket_request.sort_priority = websocket_request.sort_priority + 0.001; - let sibling_names = self - .list_websocket_requests(&websocket_request.workspace_id)? - .into_iter() - .filter(|m| m.folder_id == websocket_request.folder_id) - .map(|m| m.name) - .collect::>(); - websocket_request.name = conflict_free_name(&websocket_request.name, &sibling_names); - self.upsert(&websocket_request, source) - } - - pub fn upsert_websocket_request( - &self, - websocket_request: &WebsocketRequest, - source: &UpdateSource, - ) -> Result { - self.upsert(websocket_request, source) - } - pub fn resolve_auth_for_websocket_request( &self, websocket_request: &WebsocketRequest, @@ -168,3 +124,49 @@ impl<'a> ClientDb<'a> { }) } } + +impl<'a> WriteDb<'a> { + pub fn delete_websocket_request( + &self, + websocket_request: &WebsocketRequest, + source: &UpdateSource, + ) -> Result { + self.delete_all_websocket_connections_for_request(websocket_request.id.as_str(), source)?; + self.delete(websocket_request, source) + } + + pub fn delete_websocket_request_by_id( + &self, + id: &str, + source: &UpdateSource, + ) -> Result { + let request = self.get_websocket_request(id)?; + self.delete_websocket_request(&request, source) + } + + pub fn duplicate_websocket_request( + &self, + websocket_request: &WebsocketRequest, + source: &UpdateSource, + ) -> Result { + let mut websocket_request = websocket_request.clone(); + websocket_request.id = "".to_string(); + websocket_request.sort_priority = websocket_request.sort_priority + 0.001; + let sibling_names = self + .list_websocket_requests(&websocket_request.workspace_id)? + .into_iter() + .filter(|m| m.folder_id == websocket_request.folder_id) + .map(|m| m.name) + .collect::>(); + websocket_request.name = conflict_free_name(&websocket_request.name, &sibling_names); + self.upsert(&websocket_request, source) + } + + pub fn upsert_websocket_request( + &self, + websocket_request: &WebsocketRequest, + source: &UpdateSource, + ) -> Result { + self.upsert(websocket_request, source) + } +} diff --git a/crates/yaak-models/src/queries/workspace_metas.rs b/crates/yaak-models/src/queries/workspace_metas.rs index a977c524..01018b8a 100644 --- a/crates/yaak-models/src/queries/workspace_metas.rs +++ b/crates/yaak-models/src/queries/workspace_metas.rs @@ -1,4 +1,4 @@ -use crate::client_db::ClientDb; +use crate::client_db::{ClientDb, WriteDb}; use crate::error::Result; use crate::models::{WorkspaceMeta, WorkspaceMetaIden}; use crate::util::UpdateSource; @@ -10,18 +10,12 @@ impl<'a> ClientDb<'a> { } pub fn list_workspace_metas(&self, workspace_id: &str) -> Result> { - let mut workspace_metas = - self.find_many(WorkspaceMetaIden::WorkspaceId, workspace_id, None)?; - - if workspace_metas.is_empty() { - let wm = WorkspaceMeta { workspace_id: workspace_id.to_string(), ..Default::default() }; - workspace_metas.push(self.upsert_workspace_meta(&wm, &UpdateSource::Background)?) - } - - Ok(workspace_metas) + self.find_many(WorkspaceMetaIden::WorkspaceId, workspace_id, None) } +} - pub fn get_or_create_workspace_meta(&self, workspace_id: &str) -> Result { +impl<'a> WriteDb<'a> { + pub fn ensure_workspace_meta(&self, workspace_id: &str) -> Result { let workspace_meta = self.get_workspace_meta(workspace_id); if let Some(workspace_meta) = workspace_meta { return Ok(workspace_meta); diff --git a/crates/yaak-models/src/queries/workspaces.rs b/crates/yaak-models/src/queries/workspaces.rs index 797f7eaa..f74fb099 100644 --- a/crates/yaak-models/src/queries/workspaces.rs +++ b/crates/yaak-models/src/queries/workspaces.rs @@ -1,6 +1,6 @@ use super::merge_headers; use crate::blob_manager::BlobManager; -use crate::client_db::ClientDb; +use crate::client_db::{ClientDb, WriteDb}; use crate::error::Result; use crate::models::{ AnyModel, CookieJar, CookieJarIden, Environment, EnvironmentIden, Folder, FolderIden, @@ -23,119 +23,7 @@ impl<'a> ClientDb<'a> { } pub fn list_workspaces(&self) -> Result> { - let mut workspaces = self.find_all()?; - - if workspaces.is_empty() { - workspaces.push(self.upsert_workspace( - &Workspace { name: "Yaak".to_string(), ..Default::default() }, - &UpdateSource::Background, - )?) - } - - Ok(workspaces) - } - - /// Delete a workspace and everything in it. - /// - /// Children are bulk-deleted with one statement per table and are NOT - /// individually recorded in model_changes or emitted as events — the single - /// workspace delete event implies the subtree (see [`ModelChangeEvent::Delete`]). - /// This keeps huge workspaces (thousands of requests) fast and avoids - /// flooding event consumers. - pub fn delete_workspace( - &self, - workspace: &Workspace, - source: &UpdateSource, - blobs: &BlobManager, - ) -> Result { - let wid = workspace.id.as_str(); - - // Collect response cleanup targets before their rows disappear. The actual - // cleanup runs at the end: response bodies live on disk and in the blob DB, - // which don't participate in this transaction, so removing them must wait - // until every statement that could fail (and roll back the rows) is done. - let responses = self.find_many::(HttpResponseIden::WorkspaceId, wid, None)?; - - // Sync and the CLI call this on a plain connection where each statement - // would otherwise commit on its own, leaving a partially-deleted workspace - // if one fails. A savepoint makes the cascade atomic there, and nests - // harmlessly inside the interactive path's transaction. - let conn = self.conn().resolve(); - conn.execute_batch("SAVEPOINT delete_workspace")?; - - let result: Result = (|| { - self.delete_many_untracked::( - HttpResponseEventIden::WorkspaceId, - wid, - )?; - self.delete_many_untracked::(HttpResponseIden::WorkspaceId, wid)?; - self.delete_many_untracked::(HttpRequestIden::WorkspaceId, wid)?; - self.delete_many_untracked::(GrpcEventIden::WorkspaceId, wid)?; - self.delete_many_untracked::(GrpcConnectionIden::WorkspaceId, wid)?; - self.delete_many_untracked::(GrpcRequestIden::WorkspaceId, wid)?; - self.delete_many_untracked::(WebsocketEventIden::WorkspaceId, wid)?; - self.delete_many_untracked::( - WebsocketConnectionIden::WorkspaceId, - wid, - )?; - self.delete_many_untracked::(WebsocketRequestIden::WorkspaceId, wid)?; - self.delete_many_untracked::( - GraphQlIntrospectionIden::WorkspaceId, - wid, - )?; - self.delete_many_untracked::(FolderIden::WorkspaceId, wid)?; - self.delete_many_untracked::(EnvironmentIden::WorkspaceId, wid)?; - self.delete_many_untracked::(CookieJarIden::WorkspaceId, wid)?; - for import_source in self.list_import_sources(wid)? { - self.delete_import_source_resources(&import_source.id)?; - } - self.delete_many_untracked::(ImportSourceIden::WorkspaceId, wid)?; - self.delete_many_untracked::(SyncStateIden::WorkspaceId, wid)?; - self.delete_many_untracked::(WorkspaceMetaIden::WorkspaceId, wid)?; - self.delete(workspace, source) - })(); - - let deleted = match result { - Ok(deleted) => { - conn.execute_batch("RELEASE delete_workspace")?; - deleted - } - Err(e) => { - let _ = - conn.execute_batch("ROLLBACK TO delete_workspace; RELEASE delete_workspace"); - return Err(e); - } - }; - - // Best-effort cleanup of response bodies (disk files and blob chunks). - // Failures only orphan unreferenced data, and are logged. - let blob_ctx = blobs.connect(); - for m in responses { - if let Some(p) = m.body_path { - if let Err(e) = std::fs::remove_file(&p) { - warn!("Failed to delete response body file {p:?}: {e}"); - } - } - if let Err(e) = blob_ctx.delete_chunks_like(&format!("{}.%", m.id)) { - warn!("Failed to delete blobs for response {}: {e}", m.id); - } - } - - Ok(deleted) - } - - pub fn delete_workspace_by_id( - &self, - id: &str, - source: &UpdateSource, - blobs: &BlobManager, - ) -> Result { - let workspace = self.get_workspace(id)?; - self.delete_workspace(&workspace, source, blobs) - } - - pub fn upsert_workspace(&self, w: &Workspace, source: &UpdateSource) -> Result { - self.upsert(w, source) + self.find_all() } pub fn resolve_auth_for_workspace( @@ -190,6 +78,100 @@ impl<'a> ClientDb<'a> { } } +impl<'a> WriteDb<'a> { + /// There is always at least one workspace. Called at startup and after a + /// workspace is deleted. + pub fn ensure_default_workspace(&self) -> Result<()> { + if self.find_all::()?.is_empty() { + self.upsert_workspace( + &Workspace { name: "Yaak".to_string(), ..Default::default() }, + &UpdateSource::Background, + )?; + } + Ok(()) + } + + /// Delete a workspace and everything in it. + /// + /// Children are bulk-deleted with one statement per table and are NOT + /// individually recorded in model_changes or emitted as events — the single + /// workspace delete event implies the subtree (see [`ModelChangeEvent::Delete`]). + /// This keeps huge workspaces (thousands of requests) fast and avoids + /// flooding event consumers. + pub fn delete_workspace( + &self, + workspace: &Workspace, + source: &UpdateSource, + blobs: &BlobManager, + ) -> Result { + let wid = workspace.id.as_str(); + + // Collect response cleanup targets before their rows disappear. The actual + // cleanup runs at the end: response bodies live on disk and in the blob DB, + // which don't participate in this transaction, so removing them must wait + // until every statement that could fail (and roll back the rows) is done. + let responses = self.find_many::(HttpResponseIden::WorkspaceId, wid, None)?; + + self.delete_many_untracked::(HttpResponseEventIden::WorkspaceId, wid)?; + self.delete_many_untracked::(HttpResponseIden::WorkspaceId, wid)?; + self.delete_many_untracked::(HttpRequestIden::WorkspaceId, wid)?; + self.delete_many_untracked::(GrpcEventIden::WorkspaceId, wid)?; + self.delete_many_untracked::(GrpcConnectionIden::WorkspaceId, wid)?; + self.delete_many_untracked::(GrpcRequestIden::WorkspaceId, wid)?; + self.delete_many_untracked::(WebsocketEventIden::WorkspaceId, wid)?; + self.delete_many_untracked::( + WebsocketConnectionIden::WorkspaceId, + wid, + )?; + self.delete_many_untracked::(WebsocketRequestIden::WorkspaceId, wid)?; + self.delete_many_untracked::( + GraphQlIntrospectionIden::WorkspaceId, + wid, + )?; + self.delete_many_untracked::(FolderIden::WorkspaceId, wid)?; + self.delete_many_untracked::(EnvironmentIden::WorkspaceId, wid)?; + self.delete_many_untracked::(CookieJarIden::WorkspaceId, wid)?; + for import_source in self.list_import_sources(wid)? { + self.delete_import_source_resources(&import_source.id)?; + } + self.delete_many_untracked::(ImportSourceIden::WorkspaceId, wid)?; + self.delete_many_untracked::(SyncStateIden::WorkspaceId, wid)?; + self.delete_many_untracked::(WorkspaceMetaIden::WorkspaceId, wid)?; + let deleted = self.delete(workspace, source)?; + self.ensure_default_workspace()?; + + // Best-effort cleanup of response bodies (disk files and blob chunks). + // Failures only orphan unreferenced data, and are logged. + for m in responses { + if let Some(p) = m.body_path { + if let Err(e) = std::fs::remove_file(&p) { + warn!("Failed to delete response body file {p:?}: {e}"); + } + } + let pattern = format!("{}.%", m.id); + if let Err(e) = blobs.with_tx(|b| b.delete_chunks_like(&pattern)) { + warn!("Failed to delete blobs for response {}: {e}", m.id); + } + } + + Ok(deleted) + } + + pub fn delete_workspace_by_id( + &self, + id: &str, + source: &UpdateSource, + blobs: &BlobManager, + ) -> Result { + let workspace = self.get_workspace(id)?; + self.delete_workspace(&workspace, source, blobs) + } + + pub fn upsert_workspace(&self, w: &Workspace, source: &UpdateSource) -> Result { + self.upsert(w, source) + } +} + /// Global default headers that are always sent with requests unless overridden. /// These are prepended to the inheritance chain so workspace/folder/request headers /// can override or disable them. diff --git a/crates/yaak-models/src/query_manager.rs b/crates/yaak-models/src/query_manager.rs index 8b2c86b8..b08f149b 100644 --- a/crates/yaak-models/src/query_manager.rs +++ b/crates/yaak-models/src/query_manager.rs @@ -1,65 +1,82 @@ -use crate::client_db::ClientDb; +use crate::client_db::{ClientDb, WriteDb}; use crate::error::Error::GenericError; use crate::util::ModelPayload; use rusqlite::{Transaction, TransactionBehavior}; use std::sync::mpsc; use yaak_database::{ConnectionOrTx, DbContext, SqlitePool}; -// Pool is internally synchronized — don't wrap it in a Mutex. A Mutex held across the -// blocking `get()` serializes every DB access behind the slowest waiter, freezing the -// whole app whenever the pool is exhausted. +/// Reads come from a pool; writes go through one connection. +/// +/// SQLite in WAL mode lets many readers run alongside a single writer, and +/// never more than one writer. A second in-process writer can only wait, and +/// while it waits in the busy handler it sleeps, retries, and keeps its pool +/// slot. Enough of those and the pool is full of writers that are all asleep, +/// and every read in the app queues behind them. Giving writes exactly one +/// connection turns that into a plain queue: the next write starts the moment +/// the previous one commits, and it never takes a slot a read could use. +/// +/// The pools are internally synchronized — don't wrap them in a Mutex. A Mutex +/// held across the blocking `get()` serializes every DB access behind the +/// slowest waiter. #[derive(Debug, Clone)] pub struct QueryManager { - pool: SqlitePool, + readers: SqlitePool, + writer: SqlitePool, events_tx: mpsc::Sender, } impl QueryManager { - pub fn new(pool: SqlitePool, events_tx: mpsc::Sender) -> Self { - QueryManager { pool, events_tx } + /// `writer` must be a pool with a single connection; see [`crate::init_standalone`]. + pub fn new( + readers: SqlitePool, + writer: SqlitePool, + events_tx: mpsc::Sender, + ) -> Self { + QueryManager { readers, writer, events_tx } } + /// A read handle from the reader pool. pub fn connect(&self) -> ClientDb<'_> { - let conn = self.pool.get().expect("Failed to get a new DB connection from the pool"); - let ctx = DbContext::new(ConnectionOrTx::Connection(conn)); - ClientDb::new(ctx, self.events_tx.clone()) - } - - pub fn with_conn(&self, func: F) -> T - where - F: FnOnce(&ClientDb) -> T, - { - let conn = self.pool.get().expect("Failed to get new DB connection from the pool"); - - let ctx = DbContext::new(ConnectionOrTx::Connection(conn)); - let db = ClientDb::new(ctx, self.events_tx.clone()); - - func(&db) + let conn = self.readers.get().expect("Failed to get a new DB connection from the pool"); + ClientDb::new(DbContext::new(ConnectionOrTx::Connection(conn))) } + /// Run `func` in a transaction on the writer connection. + /// + /// Waits for any write in progress to commit first, and fails with a pool + /// error if that takes longer than the pool's timeout. Do not call this + /// from inside another `with_tx` closure: the inner call would wait for + /// the outer transaction, which is waiting on it. + /// + /// Model events for the writes are sent once the transaction commits. pub fn with_tx( &self, - func: impl FnOnce(&ClientDb) -> std::result::Result, + func: impl FnOnce(&WriteDb) -> std::result::Result, ) -> std::result::Result where E: From, { - let conn = self.pool.get().expect("Failed to get new DB connection from the pool"); + let conn = self.writer.get().map_err(crate::error::Error::SqlPoolError)?; // `new_unchecked` takes `&Connection`; see yaak_database::pool for why // the pool never hands out `&mut`. let tx = Transaction::new_unchecked(&conn, TransactionBehavior::Immediate) - .expect("Failed to start DB transaction"); + .map_err(crate::error::Error::SqlError)?; - let ctx = DbContext::new(ConnectionOrTx::Transaction(&tx)); - let db = ClientDb::new(ctx, self.events_tx.clone()); + let db = + WriteDb::new(DbContext::new(ConnectionOrTx::Transaction(&tx)), self.events_tx.clone()); match func(&db) { Ok(val) => { + let events = db.into_events(); tx.commit() .map_err(|e| GenericError(format!("Failed to commit transaction {e:?}")))?; + for payload in events { + let _ = self.events_tx.send(payload); + } Ok(val) } Err(e) => { + drop(db); tx.rollback() .map_err(|e| GenericError(format!("Failed to rollback transaction {e:?}")))?; Err(e) diff --git a/crates/yaak-models/src/util.rs b/crates/yaak-models/src/util.rs index 0729ef62..e90dcc5b 100644 --- a/crates/yaak-models/src/util.rs +++ b/crates/yaak-models/src/util.rs @@ -277,7 +277,7 @@ pub fn get_workspace_export_resources( data.resources.workspaces.push(db.find_one(WorkspaceIden::Id, workspace_id)?); data.resources.environments.append( &mut db - .list_environments_ensure_base(workspace_id)? + .list_environments(workspace_id)? .into_iter() .filter(|e| include_private_environments || e.public) .collect(), diff --git a/crates/yaak-plugins/src/install.rs b/crates/yaak-plugins/src/install.rs index c97117ad..fe677650 100644 --- a/crates/yaak-plugins/src/install.rs +++ b/crates/yaak-plugins/src/install.rs @@ -24,11 +24,7 @@ pub async fn delete_and_uninstall( Some(label) => UpdateSource::from_window_label(label), None => UpdateSource::Background, }; - // Scope the db connection so it doesn't live across await - let plugin = { - let db = query_manager.connect(); - db.delete_plugin_by_id(plugin_id, &update_source)? - }; + let plugin = query_manager.with_tx(|db| db.delete_plugin_by_id(plugin_id, &update_source))?; if let Err(err) = plugin_manager.uninstall(plugin_context, plugin.directory.as_str()).await { if !matches!(err, PluginNotFoundErr(_)) { return Err(err); @@ -72,9 +68,7 @@ pub async fn download_and_install( zip_extract::extract(Cursor::new(&bytes), &plugin_dir, true)?; info!("Extracted plugin {} to {}", plugin_version.id, plugin_dir_str); - // Scope the db connection so it doesn't live across await - let plugin = { - let db = query_manager.connect(); + let plugin = query_manager.with_tx(|db| { db.upsert_plugin( &Plugin { id: plugin_version.id.clone(), @@ -86,8 +80,8 @@ pub async fn download_and_install( ..Default::default() }, &UpdateSource::Background, - )? - }; + ) + })?; plugin_manager.add_plugin(plugin_context, &plugin).await?; diff --git a/crates/yaak-plugins/src/manager.rs b/crates/yaak-plugins/src/manager.rs index 2e05dd8b..310a1a5b 100644 --- a/crates/yaak-plugins/src/manager.rs +++ b/crates/yaak-plugins/src/manager.rs @@ -187,9 +187,7 @@ impl PluginManager { } let bundled_dirs = plugin_manager.list_bundled_plugin_dirs().await?; - // Scope the db connection so the future stays Send across the await below - let plugins = { - let db = query_manager.connect(); + let plugins = query_manager.with_tx(|db| { for dir in &bundled_dirs { if db.get_plugin_by_directory(dir).is_none() { db.upsert_plugin( @@ -204,8 +202,8 @@ impl PluginManager { )?; } } - db.list_plugins()? - }; + db.list_plugins() + })?; let init_errors = plugin_manager.initialize_all_plugins(plugins, plugin_context).await; if !init_errors.is_empty() { diff --git a/crates/yaak-sync/src/sync.rs b/crates/yaak-sync/src/sync.rs index 5aa745fb..cc9faa1b 100644 --- a/crates/yaak-sync/src/sync.rs +++ b/crates/yaak-sync/src/sync.rs @@ -11,8 +11,11 @@ use std::io::Write; use std::path::{Path, PathBuf}; use ts_rs::TS; use yaak_models::blob_manager::BlobManager; -use yaak_models::client_db::ClientDb; -use yaak_models::models::{SyncState, WorkspaceMeta}; +use yaak_models::client_db::{ClientDb, WriteDb}; +use yaak_models::models::{ + Environment, Folder, GrpcRequest, HttpRequest, SyncState, WebsocketRequest, Workspace, + WorkspaceMeta, +}; use yaak_models::util::{UpdateSource, get_workspace_export_resources}; #[derive(Debug, Clone, Serialize, Deserialize, TS)] @@ -336,17 +339,40 @@ fn workspace_models(db: &ClientDb, version: &str, workspace_id: &str) -> Result< Ok(sync_models) } -/// Apply sync operations to the filesystem and database. -/// Returns a list of SyncStateOps that should be applied afterward. -pub fn apply_sync_ops( - db: &ClientDb, - blobs: &BlobManager, +/// The database half of a sync apply, ready to run once the files are on disk. +pub struct PendingDbSyncOps { + sync_state_ops: Vec, + deletes: Vec, + workspaces: Vec, + environments: Vec, + folders: Vec, + http_requests: Vec, + grpc_requests: Vec, + websocket_requests: Vec, +} + +/// Apply the filesystem half of the sync operations: create, rewrite and +/// delete files. Returns the database half, for [`apply_db_sync_ops`]. +/// +/// Split this way so the file work, which can be slow, happens before the +/// write transaction is opened rather than inside it. +pub fn apply_fs_sync_ops( workspace_id: &str, sync_dir: &Path, sync_ops: Vec, -) -> Result> { +) -> Result { + let mut pending = PendingDbSyncOps { + sync_state_ops: Vec::new(), + deletes: Vec::new(), + workspaces: Vec::new(), + environments: Vec::new(), + folders: Vec::new(), + http_requests: Vec::new(), + grpc_requests: Vec::new(), + websocket_requests: Vec::new(), + }; if sync_ops.is_empty() { - return Ok(Vec::new()); + return Ok(pending); } info!( @@ -354,21 +380,13 @@ pub fn apply_sync_ops( sync_ops.iter().map(|op| op.to_string()).collect::>().join(", ") ); - let mut sync_state_ops = Vec::new(); - let mut workspaces_to_upsert = Vec::new(); - let mut environments_to_upsert = Vec::new(); - let mut folders_to_upsert = Vec::new(); - let mut http_requests_to_upsert = Vec::new(); - let mut grpc_requests_to_upsert = Vec::new(); - let mut websocket_requests_to_upsert = Vec::new(); - for op in sync_ops { // Only apply things if workspace ID matches if op.workspace_id() != workspace_id { continue; } - sync_state_ops.push(match op { + let state_op = match op { SyncOp::FsCreate { model } => { let rel_path = derive_model_filename(&model); let abs_path = sync_dir.join(rel_path.clone()); @@ -402,17 +420,7 @@ pub fn apply_sync_ops( }, SyncOp::DbCreate { fs } => { let model_id = fs.model.id(); - - // Push updates to arrays so we can do them all in a single - // batch upsert to make foreign keys happy - match fs.model { - SyncModel::Environment(m) => environments_to_upsert.push(m), - SyncModel::Folder(m) => folders_to_upsert.push(m), - SyncModel::GrpcRequest(m) => grpc_requests_to_upsert.push(m), - SyncModel::HttpRequest(m) => http_requests_to_upsert.push(m), - SyncModel::WebsocketRequest(m) => websocket_requests_to_upsert.push(m), - SyncModel::Workspace(m) => workspaces_to_upsert.push(m), - }; + pending.push_upsert(fs.model); SyncStateOp::Create { model_id, checksum: fs.checksum.to_owned(), @@ -420,16 +428,7 @@ pub fn apply_sync_ops( } } SyncOp::DbUpdate { state, fs } => { - // Push updates to arrays so we can do them all in a single - // batch upsert to make foreign keys happy - match fs.model { - SyncModel::Environment(m) => environments_to_upsert.push(m), - SyncModel::Folder(m) => folders_to_upsert.push(m), - SyncModel::GrpcRequest(m) => grpc_requests_to_upsert.push(m), - SyncModel::HttpRequest(m) => http_requests_to_upsert.push(m), - SyncModel::WebsocketRequest(m) => websocket_requests_to_upsert.push(m), - SyncModel::Workspace(m) => workspaces_to_upsert.push(m), - } + pending.push_upsert(fs.model); SyncStateOp::Update { state: state.to_owned(), checksum: fs.checksum.to_owned(), @@ -437,20 +436,52 @@ pub fn apply_sync_ops( } } SyncOp::DbDelete { model, state } => { - delete_model(db, blobs, &model)?; + pending.deletes.push(model); SyncStateOp::Delete { state: state.to_owned() } } SyncOp::IgnorePrivate { .. } => SyncStateOp::NoOp, - }); + }; + pending.sync_state_ops.push(state_op); + } + + Ok(pending) +} + +impl PendingDbSyncOps { + /// Upserts are collected per model type and written in one batch so + /// foreign keys are satisfied. + fn push_upsert(&mut self, model: SyncModel) { + match model { + SyncModel::Environment(m) => self.environments.push(m), + SyncModel::Folder(m) => self.folders.push(m), + SyncModel::GrpcRequest(m) => self.grpc_requests.push(m), + SyncModel::HttpRequest(m) => self.http_requests.push(m), + SyncModel::WebsocketRequest(m) => self.websocket_requests.push(m), + SyncModel::Workspace(m) => self.workspaces.push(m), + } + } +} + +/// Apply the database half of the sync operations. +/// Returns a list of SyncStateOps that should be applied afterward. +pub fn apply_db_sync_ops( + db: &WriteDb, + blobs: &BlobManager, + workspace_id: &str, + sync_dir: &Path, + pending: PendingDbSyncOps, +) -> Result> { + for model in &pending.deletes { + delete_model(db, blobs, model)?; } let upserted_models = db.batch_upsert( - workspaces_to_upsert, - environments_to_upsert, - folders_to_upsert, - http_requests_to_upsert, - grpc_requests_to_upsert, - websocket_requests_to_upsert, + pending.workspaces, + pending.environments, + pending.folders, + pending.http_requests, + pending.grpc_requests, + pending.websocket_requests, &UpdateSource::Sync, )?; @@ -482,7 +513,7 @@ pub fn apply_sync_ops( }?; } - Ok(sync_state_ops) + Ok(pending.sync_state_ops) } #[derive(Debug)] @@ -504,7 +535,7 @@ pub enum SyncStateOp { } pub fn apply_sync_state_ops( - db: &ClientDb, + db: &WriteDb, workspace_id: &str, sync_dir: &Path, ops: Vec, @@ -549,7 +580,7 @@ fn derive_model_filename(m: &SyncModel) -> PathBuf { Path::new(&rel).to_path_buf() } -fn delete_model(db: &ClientDb, blobs: &BlobManager, model: &SyncModel) -> Result<()> { +fn delete_model(db: &WriteDb, blobs: &BlobManager, model: &SyncModel) -> Result<()> { match model { SyncModel::Workspace(m) => { db.delete_workspace(&m, &UpdateSource::Sync, blobs)?; diff --git a/crates/yaak-wasm/src/lib.rs b/crates/yaak-wasm/src/lib.rs index 45704162..c79d2911 100644 --- a/crates/yaak-wasm/src/lib.rs +++ b/crates/yaak-wasm/src/lib.rs @@ -104,7 +104,7 @@ pub async fn boot() -> Result<()> { let (queries, blobs, events) = yaak_models::init_standalone(DB_NAME, BLOB_DB_NAME).map_err(js_error)?; - if let Err(e) = yaak_lifecycle::on_launch(&lifecycle_host(), &queries.connect(), &blobs) { + if let Err(e) = queries.with_tx(|tx| yaak_lifecycle::on_launch(&lifecycle_host(), tx, &blobs)) { web_sys::console::warn_2(&"on_launch hook failed".into(), &js_error(e)); } @@ -266,10 +266,17 @@ fn dispatch( if let Some(wid) = req.workspace_id.as_deref() { let e = js_error; + // Opening a workspace is where the rows it is assumed to have get created + host.queries + .with_tx(|tx| { + tx.ensure_base_environment(wid)?; + tx.ensure_default_cookie_jar(wid)?; + tx.ensure_workspace_meta(wid)?; + Ok::<(), yaak_models::error::Error>(()) + }) + .map_err(e)?; list.extend(db.list_cookie_jars(wid).map_err(e)?.into_iter().map(Into::into)); - list.extend( - db.list_environments_ensure_base(wid).map_err(e)?.into_iter().map(Into::into), - ); + list.extend(db.list_environments(wid).map_err(e)?.into_iter().map(Into::into)); list.extend(db.list_folders(wid).map_err(e)?.into_iter().map(Into::into)); list.extend(db.list_grpc_connections(wid).map_err(e)?.into_iter().map(Into::into)); list.extend(db.list_grpc_requests(wid).map_err(e)?.into_iter().map(Into::into)); @@ -291,9 +298,10 @@ fn dispatch( "models_upsert" => { let req: ModelReq = from_js(payload)?; - let db = host.queries.connect(); - let id = - models_ops::upsert_model(&db, &host.blobs, req.model, source).map_err(js_error)?; + let id = host + .queries + .with_tx(|tx| models_ops::upsert_model(tx, &host.blobs, req.model, source)) + .map_err(js_error)?; to_json(id) } @@ -330,13 +338,14 @@ fn dispatch( let req: UpsertIntrospectionReq = from_js(payload)?; let saved = host .queries - .connect() - .upsert_graphql_introspection( - &req.workspace_id, - &req.request_id, - req.content, - source, - ) + .with_tx(|tx| { + tx.upsert_graphql_introspection( + &req.workspace_id, + &req.request_id, + req.content, + source, + ) + }) .map_err(js_error)?; to_json(saved) } @@ -366,10 +375,14 @@ fn dispatch( if req.before == req.after { return to_json(()); } - let db = host.queries.connect(); - let jar = db.get_cookie_jar(&req.cookie_jar_id).map_err(js_error)?; - let cookies = apply_cookie_changes(jar.cookies.clone(), &req.before, &req.after); - db.upsert_cookie_jar(&CookieJar { cookies, ..jar }, source).map_err(js_error)?; + host.queries + .with_tx(|tx| { + let jar = tx.get_cookie_jar(&req.cookie_jar_id)?; + let cookies = + apply_cookie_changes(jar.cookies.clone(), &req.before, &req.after); + tx.upsert_cookie_jar(&CookieJar { cookies, ..jar }, source) + }) + .map_err(js_error)?; to_json(()) } @@ -378,26 +391,34 @@ fn dispatch( // writes fan out to every tab as `model_writes` like any other. "web_insert_http_response_events" => { let req: InsertResponseEventsReq = from_js(payload)?; - let db = host.queries.connect(); - for event in req.events { - let model = HttpResponseEvent::new(&req.response_id, &req.workspace_id, event); - db.upsert_http_response_event(&model, source).map_err(js_error)?; - } + host.queries + .with_tx(|tx| { + for event in req.events { + let model = + HttpResponseEvent::new(&req.response_id, &req.workspace_id, event); + tx.upsert_http_response_event(&model, source)?; + } + Ok::<(), yaak_models::error::Error>(()) + }) + .map_err(js_error)?; to_json(()) } "cmd_get_workspace_meta" => { let req: WorkspaceIdReq = from_js(payload)?; - let db = host.queries.connect(); - let workspace = db.get_workspace(&req.workspace_id).map_err(js_error)?; - to_json(db.get_or_create_workspace_meta(&workspace.id).map_err(js_error)?) + let workspace = + host.queries.connect().get_workspace(&req.workspace_id).map_err(js_error)?; + to_json( + host.queries + .with_tx(|tx| tx.ensure_workspace_meta(&workspace.id)) + .map_err(js_error)?, + ) } "cmd_delete_all_http_responses" => { let req: RequestIdReq = from_js(payload)?; host.queries - .connect() - .delete_all_http_responses_for_request(&req.request_id, source) + .with_tx(|tx| tx.delete_all_http_responses_for_request(&req.request_id, source)) .map_err(js_error)?; to_json(()) } @@ -569,16 +590,19 @@ pub fn blob_get(id: &str) -> Result>> { pub fn blob_put(id: &str, bytes: &[u8]) -> Result<()> { const CHUNK: usize = 512 * 1024; with_host(|host| { - let ctx = host.blobs.connect(); - ctx.delete_chunks(id).map_err(js_error)?; - for (i, part) in bytes.chunks(CHUNK).enumerate() { - ctx.insert_chunk(&BodyChunk::new(id, i as i32, part.to_vec())).map_err(js_error)?; - } - Ok(()) + host.blobs + .with_tx(|b| { + b.delete_chunks(id)?; + for (i, part) in bytes.chunks(CHUNK).enumerate() { + b.insert_chunk(&BodyChunk::new(id, i as i32, part.to_vec()))?; + } + Ok::<(), yaak_models::error::Error>(()) + }) + .map_err(js_error) }) } #[wasm_bindgen] pub fn blob_delete(id: &str) -> Result<()> { - with_host(|host| host.blobs.connect().delete_chunks(id).map_err(js_error)) + with_host(|host| host.blobs.with_tx(|b| b.delete_chunks(id)).map_err(js_error)) } diff --git a/crates/yaak/src/import.rs b/crates/yaak/src/import.rs index d5e5039b..6d66f5fc 100644 --- a/crates/yaak/src/import.rs +++ b/crates/yaak/src/import.rs @@ -4,7 +4,7 @@ use log::info; use serde_json::Value; use sha2::{Digest, Sha256}; use std::collections::{BTreeMap, BTreeSet}; -use yaak_models::client_db::ClientDb; +use yaak_models::client_db::{ClientDb, WriteDb}; use yaak_models::models::{ AnyModel, DEFAULT_REQUEST_MESSAGE_SIZE, Environment, Folder, GrpcRequest, HttpRequest, ImportSource, ImportSourceResource, UpsertModelInfo, WebsocketRequest, Workspace, @@ -351,7 +351,7 @@ pub fn commit_import_plan( }) } -fn commit_plan_in_tx(db: &ClientDb, plan: ImportPlan) -> Result { +fn commit_plan_in_tx(db: &WriteDb, plan: ImportPlan) -> Result { let items: BTreeMap = plan.items.iter().map(|item| (item.model_id.clone(), item.clone())).collect(); @@ -460,7 +460,7 @@ fn commit_plan_in_tx(db: &ClientDb, plan: ImportPlan) -> Result Result<()> { +fn delete_existing_model(db: &WriteDb, resource: ImportResourceType, id: &str) -> Result<()> { use ImportResourceType::*; let source = &UpdateSource::Import; match resource { @@ -503,7 +503,7 @@ fn delete_existing_model(db: &ClientDb, resource: ImportResourceType, id: &str) /// offered again next time. A resource the user turned down is remembered as a row without a /// model, so it is neither re-offered nor resurrected. fn record_import_source( - db: &ClientDb, + db: &WriteDb, plan: &ImportPlan, items: &BTreeMap, upserted: &BatchUpsertResult, @@ -1619,32 +1619,30 @@ mod tests { name: "Selected Folder".to_string(), ..Default::default() }; - { - let db = query_manager.connect(); - destination = db - .upsert_workspace(&destination, &UpdateSource::Import) - .expect("create destination"); - db.upsert_folder(&selected_folder, &UpdateSource::Import) - .expect("create selected folder"); - db.upsert_environment( - &Environment { - id: "ev_destination_base".to_string(), - model: "environment".to_string(), - workspace_id: destination.id.clone(), - name: "Destination Variables".to_string(), - parent_model: "workspace".to_string(), - variables: vec![EnvironmentVariable { - enabled: true, - name: "destination".to_string(), - value: "keep".to_string(), - id: None, - }], - ..Default::default() - }, - &UpdateSource::Import, - ) - .expect("create base environment"); - } + query_manager + .with_tx(|db| { + destination = db.upsert_workspace(&destination, &UpdateSource::Import)?; + db.upsert_folder(&selected_folder, &UpdateSource::Import)?; + db.upsert_environment( + &Environment { + id: "ev_destination_base".to_string(), + model: "environment".to_string(), + workspace_id: destination.id.clone(), + name: "Destination Variables".to_string(), + parent_model: "workspace".to_string(), + variables: vec![EnvironmentVariable { + enabled: true, + name: "destination".to_string(), + value: "keep".to_string(), + id: None, + }], + ..Default::default() + }, + &UpdateSource::Import, + ) + }) + .expect("seed destination"); + let workspace_count = query_manager.connect().list_workspaces().expect("list").len(); let plan = plan_import_resources( &query_manager, @@ -1662,13 +1660,10 @@ mod tests { // Planning performed only reads. { let db = query_manager.connect(); - assert_eq!(db.list_workspaces().expect("list workspaces").len(), 1); + assert_eq!(db.list_workspaces().expect("list workspaces").len(), workspace_count); assert_eq!(db.list_folders(&destination.id).expect("list folders").len(), 1); assert!(db.list_http_requests(&destination.id).expect("list requests").is_empty()); - assert_eq!( - db.list_environments_ensure_base(&destination.id).expect("list environments").len(), - 1 - ); + assert_eq!(db.list_environments(&destination.id).expect("list environments").len(), 1); assert_eq!(db.get_workspace(&destination.id).expect("get destination"), destination); } @@ -1861,8 +1856,7 @@ mod tests { yaak_models::init_in_memory().expect("initialize database"); let destination = destination_workspace(); query_manager - .connect() - .upsert_workspace(&destination, &UpdateSource::Import) + .with_tx(|tx| tx.upsert_workspace(&destination, &UpdateSource::Import)) .expect("create destination"); let resources = ImportResources { workspaces: vec![ @@ -2247,10 +2241,7 @@ mod tests { let db = query_manager.connect(); assert_eq!(db.list_http_requests(&workspace_id).expect("list requests").len(), 2); assert_eq!(db.list_folders(&workspace_id).expect("list folders").len(), 1); - assert_eq!( - db.list_environments_ensure_base(&workspace_id).expect("list environments").len(), - 1 - ); + assert_eq!(db.list_environments(&workspace_id).expect("list environments").len(), 1); let rows = db.list_import_source_resources(&source.id).expect("list resource rows"); assert_eq!(rows.len(), 4, "re-commit replaces rows instead of accumulating"); } @@ -2269,23 +2260,22 @@ mod tests { .id .clone(); - { - let db = query_manager.connect(); - let nested = db - .list_http_requests(&workspace_id) - .expect("list requests") - .into_iter() - .find(|r| r.name == "Nested Request") - .expect("nested request"); - db.upsert_http_request( - &HttpRequest { - url: "https://example.com/nested-local".to_string(), - ..nested.clone() - }, - &UpdateSource::Background, - ) + query_manager + .with_tx(|db| { + let nested = db + .list_http_requests(&workspace_id)? + .into_iter() + .find(|r| r.name == "Nested Request") + .expect("nested request"); + db.upsert_http_request( + &HttpRequest { + url: "https://example.com/nested-local".to_string(), + ..nested.clone() + }, + &UpdateSource::Background, + ) + }) .expect("edit nested request locally"); - } let mut resources = imported_resources(); resources.http_requests[0].url = "https://example.com/root-v2".to_string(); @@ -2370,15 +2360,15 @@ mod tests { .id .clone(); - { - let db = query_manager.connect(); - let root = db.get_http_request(&root_id).expect("get root"); - db.upsert_http_request( - &HttpRequest { url: "https://example.com/root-local".to_string(), ..root }, - &UpdateSource::Background, - ) + query_manager + .with_tx(|db| { + let root = db.get_http_request(&root_id)?; + db.upsert_http_request( + &HttpRequest { url: "https://example.com/root-local".to_string(), ..root }, + &UpdateSource::Background, + ) + }) .expect("edit root locally"); - } let mut resources = imported_resources(); resources.http_requests[0].url = "https://example.com/root-v2".to_string(); @@ -2507,8 +2497,7 @@ mod tests { .clone(); query_manager - .connect() - .delete_http_request_by_id(&root_id, &UpdateSource::Background) + .with_tx(|tx| tx.delete_http_request_by_id(&root_id, &UpdateSource::Background)) .expect("delete root locally"); let plan = replan(&query_manager, &workspace_id, imported_resources()); @@ -2703,10 +2692,9 @@ mod tests { let workspace_id = committed.workspaces[0].id.clone(); // A second source claiming the same keys leaves nothing to merge into safely. - { - let db = query_manager.connect(); - let other = db - .upsert_import_source( + query_manager + .with_tx(|db| { + let other = db.upsert_import_source( &ImportSource { workspace_id: workspace_id.clone(), importer: "OpenAPI".to_string(), @@ -2715,18 +2703,18 @@ mod tests { ..Default::default() }, &UpdateSource::Import, - ) - .expect("create second source"); - for key in ["env:base", "folder:src", "op:root", "op:nested"] { - db.upsert_import_source_resource(&ImportSourceResource { - import_source_id: other.id.clone(), - source_key: key.to_string(), - model_type: "http_request".to_string(), - ..Default::default() - }) - .expect("claim the same keys"); - } - } + )?; + for key in ["env:base", "folder:src", "op:root", "op:nested"] { + db.upsert_import_source_resource(&ImportSourceResource { + import_source_id: other.id.clone(), + source_key: key.to_string(), + model_type: "http_request".to_string(), + ..Default::default() + })?; + } + Ok::<_, yaak_models::error::Error>(()) + }) + .expect("claim the same keys"); let third = ImportOrigin { origin: "/tmp/third.yaml".to_string(), label: "third.yaml".to_string() }; @@ -2821,21 +2809,20 @@ mod tests { let committed = first_import(&query_manager); let workspace_id = committed.workspaces[0].id.clone(); - { - let db = query_manager.connect(); - let sources = db.list_import_sources(&workspace_id).expect("list import sources"); - let row = db - .list_import_source_resources(&sources[0].id) - .expect("list rows") - .into_iter() - .find(|r| r.source_key == "op:root") - .expect("row for the root request"); - db.upsert_import_source_resource(&ImportSourceResource { - content_hash: Some("v99:from-the-future".to_string()), - ..row + query_manager + .with_tx(|db| { + let sources = db.list_import_sources(&workspace_id)?; + let row = db + .list_import_source_resources(&sources[0].id)? + .into_iter() + .find(|r| r.source_key == "op:root") + .expect("row for the root request"); + db.upsert_import_source_resource(&ImportSourceResource { + content_hash: Some("v99:from-the-future".to_string()), + ..row + }) }) .expect("write an unreadable hash"); - } let plan = replan(&query_manager, &workspace_id, imported_resources()); assert_eq!( @@ -2874,17 +2861,20 @@ mod tests { .clone(); // Opening the request in the editor stamps a row ID onto every header it renders. - { - let db = query_manager.connect(); - let root = db.get_http_request(&root_id).expect("get root"); - let headers = root - .headers - .iter() - .map(|h| HttpRequestHeader { id: Some("hd_generated".to_string()), ..h.clone() }) - .collect(); - db.upsert_http_request(&HttpRequest { headers, ..root }, &UpdateSource::Background) - .expect("stamp row ids"); - } + query_manager + .with_tx(|db| { + let root = db.get_http_request(&root_id)?; + let headers = root + .headers + .iter() + .map(|h| HttpRequestHeader { + id: Some("hd_generated".to_string()), + ..h.clone() + }) + .collect(); + db.upsert_http_request(&HttpRequest { headers, ..root }, &UpdateSource::Background) + }) + .expect("stamp row ids"); let plan = replan(&query_manager, &workspace_id, resources); assert_eq!( @@ -2997,15 +2987,15 @@ mod tests { .id .clone(); - { - let db = query_manager.connect(); - let root = db.get_http_request(&root_id).expect("get root"); - db.upsert_http_request( - &HttpRequest { url: "https://example.com/root-local".to_string(), ..root }, - &UpdateSource::Background, - ) + query_manager + .with_tx(|db| { + let root = db.get_http_request(&root_id)?; + db.upsert_http_request( + &HttpRequest { url: "https://example.com/root-local".to_string(), ..root }, + &UpdateSource::Background, + ) + }) .expect("edit root locally"); - } let mut plan = replan(&query_manager, &workspace_id, imported_resources()); let root = item_by_name(&plan, "Root Request"); diff --git a/crates/yaak/src/plugin_events.rs b/crates/yaak/src/plugin_events.rs index 767c6550..136e697a 100644 --- a/crates/yaak/src/plugin_events.rs +++ b/crates/yaak/src/plugin_events.rs @@ -1,6 +1,7 @@ use crate::response_body::ResponseBodyStore; use base64::Engine; use base64::prelude::BASE64_STANDARD; +use log::warn; use yaak_models::models::AnyModel; use yaak_models::query_manager::QueryManager; use yaak_models::util::UpdateSource; @@ -226,11 +227,18 @@ fn build_shared_reply( InternalEventPayload::GetKeyValueResponse(GetKeyValueResponse { value }) } SharedRequest::SetKeyValue(req) => { - query_manager.connect().set_plugin_key_value(context.plugin_name, &req.key, &req.value); + if let Err(e) = query_manager.with_tx(|tx| { + tx.set_plugin_key_value(context.plugin_name, &req.key, &req.value); + Ok::<(), yaak_models::error::Error>(()) + }) { + warn!("Failed to set plugin key value: {e}"); + } InternalEventPayload::SetKeyValueResponse(yaak_plugins::events::SetKeyValueResponse {}) } SharedRequest::DeleteKeyValue(req) => { - match query_manager.connect().delete_plugin_key_value(context.plugin_name, &req.key) { + match query_manager + .with_tx(|tx| tx.delete_plugin_key_value(context.plugin_name, &req.key)) + { Ok(deleted) => { InternalEventPayload::DeleteKeyValueResponse(DeleteKeyValueResponse { deleted }) } @@ -331,7 +339,9 @@ fn build_shared_reply( let model = match &req.model { HttpRequest(m) => { - match query_manager.connect().upsert_http_request(m, &UpdateSource::Plugin) { + match query_manager + .with_tx(|tx| tx.upsert_http_request(m, &UpdateSource::Plugin)) + { Ok(model) => HttpRequest(model), Err(err) => { return InternalEventPayload::ErrorResponse(ErrorResponse { @@ -341,7 +351,9 @@ fn build_shared_reply( } } GrpcRequest(m) => { - match query_manager.connect().upsert_grpc_request(m, &UpdateSource::Plugin) { + match query_manager + .with_tx(|tx| tx.upsert_grpc_request(m, &UpdateSource::Plugin)) + { Ok(model) => GrpcRequest(model), Err(err) => { return InternalEventPayload::ErrorResponse(ErrorResponse { @@ -351,7 +363,8 @@ fn build_shared_reply( } } WebsocketRequest(m) => { - match query_manager.connect().upsert_websocket_request(m, &UpdateSource::Plugin) + match query_manager + .with_tx(|tx| tx.upsert_websocket_request(m, &UpdateSource::Plugin)) { Ok(model) => WebsocketRequest(model), Err(err) => { @@ -362,7 +375,7 @@ fn build_shared_reply( } } Folder(m) => { - match query_manager.connect().upsert_folder(m, &UpdateSource::Plugin) { + match query_manager.with_tx(|tx| tx.upsert_folder(m, &UpdateSource::Plugin)) { Ok(model) => Folder(model), Err(err) => { return InternalEventPayload::ErrorResponse(ErrorResponse { @@ -372,7 +385,9 @@ fn build_shared_reply( } } Environment(m) => { - match query_manager.connect().upsert_environment(m, &UpdateSource::Plugin) { + match query_manager + .with_tx(|tx| tx.upsert_environment(m, &UpdateSource::Plugin)) + { Ok(model) => Environment(model), Err(err) => { return InternalEventPayload::ErrorResponse(ErrorResponse { @@ -382,7 +397,8 @@ fn build_shared_reply( } } Workspace(m) => { - match query_manager.connect().upsert_workspace(m, &UpdateSource::Plugin) { + match query_manager.with_tx(|tx| tx.upsert_workspace(m, &UpdateSource::Plugin)) + { Ok(model) => Workspace(model), Err(err) => { return InternalEventPayload::ErrorResponse(ErrorResponse { @@ -404,8 +420,7 @@ fn build_shared_reply( let model = match req.model.as_str() { "http_request" => { match query_manager - .connect() - .delete_http_request_by_id(&req.id, &UpdateSource::Plugin) + .with_tx(|tx| tx.delete_http_request_by_id(&req.id, &UpdateSource::Plugin)) { Ok(model) => AnyModel::HttpRequest(model), Err(err) => { @@ -417,8 +432,7 @@ fn build_shared_reply( } "grpc_request" => { match query_manager - .connect() - .delete_grpc_request_by_id(&req.id, &UpdateSource::Plugin) + .with_tx(|tx| tx.delete_grpc_request_by_id(&req.id, &UpdateSource::Plugin)) { Ok(model) => AnyModel::GrpcRequest(model), Err(err) => { @@ -429,10 +443,9 @@ fn build_shared_reply( } } "websocket_request" => { - match query_manager - .connect() - .delete_websocket_request_by_id(&req.id, &UpdateSource::Plugin) - { + match query_manager.with_tx(|tx| { + tx.delete_websocket_request_by_id(&req.id, &UpdateSource::Plugin) + }) { Ok(model) => AnyModel::WebsocketRequest(model), Err(err) => { return InternalEventPayload::ErrorResponse(ErrorResponse { @@ -442,8 +455,7 @@ fn build_shared_reply( } } "folder" => match query_manager - .connect() - .delete_folder_by_id(&req.id, &UpdateSource::Plugin) + .with_tx(|tx| tx.delete_folder_by_id(&req.id, &UpdateSource::Plugin)) { Ok(model) => AnyModel::Folder(model), Err(err) => { @@ -454,8 +466,7 @@ fn build_shared_reply( }, "environment" => { match query_manager - .connect() - .delete_environment_by_id(&req.id, &UpdateSource::Plugin) + .with_tx(|tx| tx.delete_environment_by_id(&req.id, &UpdateSource::Plugin)) { Ok(model) => AnyModel::Environment(model), Err(err) => { @@ -508,45 +519,38 @@ mod tests { yaak_models::init_standalone(&db_path, &blob_path).expect("Failed to initialize DB"); query_manager - .connect() - .upsert_workspace( - &Workspace { - id: "wk_test".to_string(), - name: "Workspace".to_string(), - ..Default::default() - }, - &UpdateSource::Sync, - ) - .expect("Failed to seed workspace"); - - query_manager - .connect() - .upsert_folder( - &Folder { - id: "fl_test".to_string(), - workspace_id: "wk_test".to_string(), - name: "Folder".to_string(), - ..Default::default() - }, - &UpdateSource::Sync, - ) - .expect("Failed to seed folder"); - - query_manager - .connect() - .upsert_http_request( - &HttpRequest { - id: "rq_test".to_string(), - workspace_id: "wk_test".to_string(), - folder_id: Some("fl_test".to_string()), - name: "Request".to_string(), - method: "GET".to_string(), - url: "https://example.com".to_string(), - ..Default::default() - }, - &UpdateSource::Sync, - ) - .expect("Failed to seed request"); + .with_tx(|tx| { + tx.upsert_workspace( + &Workspace { + id: "wk_test".to_string(), + name: "Workspace".to_string(), + ..Default::default() + }, + &UpdateSource::Sync, + )?; + tx.upsert_folder( + &Folder { + id: "fl_test".to_string(), + workspace_id: "wk_test".to_string(), + name: "Folder".to_string(), + ..Default::default() + }, + &UpdateSource::Sync, + )?; + tx.upsert_http_request( + &HttpRequest { + id: "rq_test".to_string(), + workspace_id: "wk_test".to_string(), + folder_id: Some("fl_test".to_string()), + name: "Request".to_string(), + method: "GET".to_string(), + url: "https://example.com".to_string(), + ..Default::default() + }, + &UpdateSource::Sync, + ) + }) + .expect("Failed to seed"); (query_manager, temp_dir) } diff --git a/crates/yaak/src/response_body.rs b/crates/yaak/src/response_body.rs index 38c9daab..38762038 100644 --- a/crates/yaak/src/response_body.rs +++ b/crates/yaak/src/response_body.rs @@ -126,23 +126,25 @@ mod tests { .unwrap(); query_manager - .connect() - .upsert_workspace( - &Workspace { id: "wk_test".to_string(), ..Default::default() }, - &UpdateSource::Sync, - ) + .with_tx(|tx| { + tx.upsert_workspace( + &Workspace { id: "wk_test".to_string(), ..Default::default() }, + &UpdateSource::Sync, + ) + }) .unwrap(); query_manager - .connect() - .upsert_http_request( - &HttpRequest { - id: "rq_test".to_string(), - workspace_id: "wk_test".to_string(), - ..Default::default() - }, - &UpdateSource::Sync, - ) + .with_tx(|tx| { + tx.upsert_http_request( + &HttpRequest { + id: "rq_test".to_string(), + workspace_id: "wk_test".to_string(), + ..Default::default() + }, + &UpdateSource::Sync, + ) + }) .unwrap(); let body_path = body.map(|bytes| { @@ -153,21 +155,22 @@ mod tests { }); let response = query_manager - .connect() - .upsert_http_response( - &HttpResponse { - workspace_id: "wk_test".to_string(), - request_id: "rq_test".to_string(), - body_path, - headers: vec![HttpResponseHeader { - name: "Content-Type".to_string(), - value: "application/json; charset=utf-8".to_string(), - }], - ..Default::default() - }, - &UpdateSource::Sync, - &blob_manager, - ) + .with_tx(|tx| { + tx.upsert_http_response( + &HttpResponse { + workspace_id: "wk_test".to_string(), + request_id: "rq_test".to_string(), + body_path, + headers: vec![HttpResponseHeader { + name: "Content-Type".to_string(), + value: "application/json; charset=utf-8".to_string(), + }], + ..Default::default() + }, + &UpdateSource::Sync, + &blob_manager, + ) + }) .unwrap(); let id = response.id.clone(); @@ -209,7 +212,7 @@ mod tests { let mut response = qm.connect().get_http_response(&id).unwrap(); response.state = HttpResponseState::Closed; - qm.connect().update_http_response_if_id(&response, &UpdateSource::Sync).unwrap(); + qm.with_tx(|tx| tx.update_http_response_if_id(&response, &UpdateSource::Sync)).unwrap(); assert!(FileResponseBodyStore::new(&qm).info(&id).unwrap().complete); } diff --git a/crates/yaak/src/send.rs b/crates/yaak/src/send.rs index cc99144d..026c60d4 100644 --- a/crates/yaak/src/send.rs +++ b/crates/yaak/src/send.rs @@ -646,8 +646,9 @@ pub async fn send_http_request( if let Some(store) = store { response = store .query_manager - .connect() - .upsert_http_response(&response, &store.update_source, store.blob_manager) + .with_tx(|tx| { + tx.upsert_http_response(&response, &store.update_source, store.blob_manager) + }) .map_err(SendHttpRequestError::PersistResponse)?; } else if response.id.is_empty() { response.id = generate_prefixed_id("rs"); @@ -700,8 +701,8 @@ pub async fn send_http_request( &event_workspace_id, event.clone().into(), ); - if let Err(err) = - query_manager.connect().upsert_http_response_event(&db_event, update_source) + if let Err(err) = query_manager + .with_tx(|tx| tx.upsert_http_response_event(&db_event, update_source)) { warn!("Failed to persist HTTP response event: {}", err); } @@ -799,8 +800,13 @@ pub async fn send_http_request( if let Some(store) = store { response = store .query_manager - .connect() - .upsert_http_response(&connected_response, &store.update_source, store.blob_manager) + .with_tx(|tx| { + tx.upsert_http_response( + &connected_response, + &store.update_source, + store.blob_manager, + ) + }) .map_err(SendHttpRequestError::PersistResponse)?; } else { response = connected_response; @@ -886,12 +892,13 @@ pub async fn send_http_request( if let Some(store) = store { response = store .query_manager - .connect() - .upsert_http_response( - &progress_response, - &store.update_source, - store.blob_manager, - ) + .with_tx(|tx| { + tx.upsert_http_response( + &progress_response, + &store.update_source, + store.blob_manager, + ) + }) .map_err(SendHttpRequestError::PersistResponse)?; } else { response = progress_response; @@ -960,8 +967,9 @@ pub async fn send_http_request( if let Some(store) = store { response = store .query_manager - .connect() - .upsert_http_response(&final_response, &store.update_source, store.blob_manager) + .with_tx(|tx| { + tx.upsert_http_response(&final_response, &store.update_source, store.blob_manager) + }) .map_err(SendHttpRequestError::PersistResponse)?; } else { response = final_response; @@ -998,8 +1006,9 @@ pub async fn send_http_request( if update_response && let Some(store) = store { response = store .query_manager - .connect() - .upsert_http_response(&response, &store.update_source, store.blob_manager) + .with_tx(|tx| { + tx.upsert_http_response(&response, &store.update_source, store.blob_manager) + }) .map_err(SendHttpRequestError::PersistResponse)?; } } @@ -1027,17 +1036,14 @@ fn persist_request_body_bytes( return Ok(()); } - let blob_ctx = blob_manager.connect(); - let mut offset = 0; - let mut chunk_index: i32 = 0; - while offset < bytes.len() { - let end = std::cmp::min(offset + REQUEST_BODY_CHUNK_SIZE, bytes.len()); - let chunk = BodyChunk::new(body_id, chunk_index, bytes[offset..end].to_vec()); - blob_ctx.insert_chunk(&chunk).map_err(|e| e.to_string())?; - chunk_index += 1; - offset = end; - } - Ok(()) + blob_manager + .with_tx(|b| { + for (chunk_index, data) in bytes.chunks(REQUEST_BODY_CHUNK_SIZE).enumerate() { + b.insert_chunk(&BodyChunk::new(body_id, chunk_index as i32, data.to_vec()))?; + } + Ok::<_, yaak_models::error::Error>(()) + }) + .map_err(|e| e.to_string()) } async fn persist_request_body_stream( @@ -1057,14 +1063,14 @@ async fn persist_request_body_stream( while buf.len() >= REQUEST_BODY_CHUNK_SIZE { let data = buf.drain(..REQUEST_BODY_CHUNK_SIZE).collect(); let chunk = BodyChunk::new(&body_id, chunk_index, data); - blob_manager.connect().insert_chunk(&chunk).map_err(|e| e.to_string())?; + blob_manager.with_tx(|b| b.insert_chunk(&chunk)).map_err(|e| e.to_string())?; chunk_index += 1; } } if !buf.is_empty() { let chunk = BodyChunk::new(&body_id, chunk_index, buf); - blob_manager.connect().insert_chunk(&chunk).map_err(|e| e.to_string())?; + blob_manager.with_tx(|b| b.insert_chunk(&chunk)).map_err(|e| e.to_string())?; } Ok(total_bytes) @@ -1114,8 +1120,7 @@ pub fn persist_cookies_after_send( cookie_jar.cookies = cookies; query_manager - .connect() - .upsert_cookie_jar(cookie_jar, &UpdateSource::Background) + .with_tx(|tx| tx.upsert_cookie_jar(cookie_jar, &UpdateSource::Background)) .map_err(SendHttpRequestError::PersistCookieJar)?; Ok(()) } @@ -1209,23 +1214,24 @@ fn persist_response_error( let elapsed = duration_to_i32(started_at.elapsed()); store .query_manager - .connect() - .upsert_http_response( - &HttpResponse { - state: HttpResponseState::Closed, - elapsed, - elapsed_headers: if response.elapsed_headers == 0 { - elapsed - } else { - response.elapsed_headers + .with_tx(|tx| { + tx.upsert_http_response( + &HttpResponse { + state: HttpResponseState::Closed, + elapsed, + elapsed_headers: if response.elapsed_headers == 0 { + elapsed + } else { + response.elapsed_headers + }, + error: Some(error), + url: if response.url.is_empty() { fallback_url } else { response.url.clone() }, + ..response.clone() }, - error: Some(error), - url: if response.url.is_empty() { fallback_url } else { response.url.clone() }, - ..response.clone() - }, - &store.update_source, - store.blob_manager, - ) + &store.update_source, + store.blob_manager, + ) + }) .map_err(SendHttpRequestError::PersistResponse) } @@ -1444,23 +1450,25 @@ mod tests { .expect("Failed to initialize DB"); query_manager - .connect() - .upsert_workspace( - &Workspace { id: "wk_test".to_string(), ..Default::default() }, - &UpdateSource::Sync, - ) + .with_tx(|tx| { + tx.upsert_workspace( + &Workspace { id: "wk_test".to_string(), ..Default::default() }, + &UpdateSource::Sync, + ) + }) .expect("Failed to seed workspace"); let cookie_jar = query_manager - .connect() - .upsert_cookie_jar( - &CookieJar { - id: "cj_test".to_string(), - workspace_id: "wk_test".to_string(), - name: "Default".to_string(), - ..Default::default() - }, - &UpdateSource::Sync, - ) + .with_tx(|tx| { + tx.upsert_cookie_jar( + &CookieJar { + id: "cj_test".to_string(), + workspace_id: "wk_test".to_string(), + name: "Default".to_string(), + ..Default::default() + }, + &UpdateSource::Sync, + ) + }) .expect("Failed to seed cookie jar"); (query_manager, cookie_jar, temp_dir) @@ -1507,18 +1515,18 @@ mod tests { let (query_manager, mut cookie_jar, _temp_dir) = seed_cookie_jar(); cookie_jar.cookies = vec![cookie("original")]; cookie_jar = query_manager - .connect() - .upsert_cookie_jar(&cookie_jar, &UpdateSource::Sync) + .with_tx(|tx| tx.upsert_cookie_jar(&cookie_jar, &UpdateSource::Sync)) .expect("Failed to seed cookies"); let store = CookieStore::from_cookies(cookie_jar.cookies.clone()); // Someone else updates the jar while the send is in flight. query_manager - .connect() - .upsert_cookie_jar( - &CookieJar { cookies: vec![cookie("newer")], ..cookie_jar.clone() }, - &UpdateSource::Sync, - ) + .with_tx(|tx| { + tx.upsert_cookie_jar( + &CookieJar { cookies: vec![cookie("newer")], ..cookie_jar.clone() }, + &UpdateSource::Sync, + ) + }) .expect("Failed to update cookie jar"); persist_cookies_after_send(&query_manager, Some(&mut cookie_jar), Some(&store))