From c9e884d4fe237b2401740f123178e64fc3800a14 Mon Sep 17 00:00:00 2001 From: Gregory Schier Date: Wed, 9 Sep 2026 21:46:15 -0700 Subject: [PATCH] fix(models): return a database error instead of panicking when the pool is exhausted MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two users on 2026.7.x reported the app going half-dead: local editing kept working, but Send did nothing and every encrypted secret rendered blank until they restarted. The webview console carried a Rust task panic — Delete task failed: task 1314 panicked with message "Failed to get new DB connection from the pool: Error(None)" `Error(None)` is r2d2's acquire timeout. `QueryManager::connect` and `with_tx` turned that timeout into `expect`, so a pool that was momentarily full took down whatever task asked for a connection instead of reporting a failure the caller could show, retry, or ignore. Every DB-backed path shares those two functions, which is why one stalled pool reads as "Send is dead and the secrets are gone". Acquisition is now fallible end to end: `connect()` and `BlobManager::connect()` return `Result`, and so do the `db()`/`blobs()` accessors on every host. Call sites propagate with `?` where they already returned a result; the fire-and-forget writers (streaming events, the model-change poller) log and carry on. `with_conn` had no callers and only added a third panic, so it is gone. Two places held a pooled connection across work that has nothing to do with the database and can block for a long time — the OS keyring in `EncryptionManager::get_workspace_key`, and the sync directory walk in `cmd_sync_calculate`. Both now hand the connection back first. Those are the paths behind the two visible symptoms. Not fixed here: the gRPC streaming block in yaak-app-client still `.unwrap()`s its database writes, as it did before, so acquisition failures panic there too. `ProxyQueryManager` in the separate proxy app keeps its own `expect`. Co-Authored-By: Claude Opus 5 --- Cargo.lock | 1 + .../yaak-cli/src/commands/cookie_jar.rs | 1 + .../yaak-cli/src/commands/environment.rs | 7 + crates-cli/yaak-cli/src/commands/folder.rs | 25 +- .../yaak-cli/src/commands/import_export.rs | 9 +- crates-cli/yaak-cli/src/commands/plugin.rs | 1 + crates-cli/yaak-cli/src/commands/request.rs | 21 +- crates-cli/yaak-cli/src/commands/response.rs | 13 +- crates-cli/yaak-cli/src/commands/send.rs | 12 +- crates-cli/yaak-cli/src/commands/workspace.rs | 13 +- crates-cli/yaak-cli/src/context.rs | 10 +- crates-cli/yaak-cli/src/main.rs | 8 +- crates-cli/yaak-cli/src/plugin_events.rs | 121 +++-- crates-cli/yaak-cli/src/utils/workspace.rs | 7 +- crates-cli/yaak-cli/tests/common/mod.rs | 5 + .../yaak-cli/tests/environment_commands.rs | 2 +- crates-cli/yaak-cli/tests/folder_commands.rs | 2 +- .../yaak-cli/tests/import_export_commands.rs | 12 +- crates-cli/yaak-cli/tests/request_commands.rs | 11 +- crates-cli/yaak-cli/tests/send_commands.rs | 1 + .../yaak-cli/tests/workspace_commands.rs | 2 +- crates-tauri/yaak-app-client/src/history.rs | 9 +- crates-tauri/yaak-app-client/src/lib.rs | 466 ++++++++++-------- .../yaak-app-client/src/models_ext.rs | 12 +- .../yaak-app-client/src/notifications.rs | 10 +- .../yaak-app-client/src/plugin_events.rs | 18 +- .../yaak-app-client/src/plugins_ext.rs | 8 +- crates-tauri/yaak-app-client/src/rpc_ext.rs | 22 +- crates-tauri/yaak-app-client/src/sync_ext.rs | 7 +- crates-tauri/yaak-app-client/src/updates.rs | 2 +- crates-tauri/yaak-app-client/src/ws_ext.rs | 70 +-- crates-tauri/yaak-license/src/license.rs | 15 +- crates/yaak-commands/src/actions.rs | 12 +- crates/yaak-commands/src/host.rs | 4 +- crates/yaak-commands/src/models.rs | 28 +- crates/yaak-commands/src/plugins.rs | 2 +- crates/yaak-commands/src/render.rs | 2 +- crates/yaak-commands/src/responses.rs | 8 +- crates/yaak-commands/src/templates.rs | 2 +- crates/yaak-commands/tests/test_host.rs | 15 +- crates/yaak-crypto/src/manager.rs | 10 +- crates/yaak-lifecycle/src/lib.rs | 8 +- crates/yaak-models/Cargo.toml | 3 + crates/yaak-models/src/blob_manager.rs | 24 +- .../yaak-models/src/queries/http_requests.rs | 4 +- .../yaak-models/src/queries/http_responses.rs | 16 +- .../yaak-models/src/queries/model_changes.rs | 10 +- crates/yaak-models/src/queries/workspaces.rs | 4 +- crates/yaak-models/src/query_manager.rs | 115 ++++- crates/yaak-plugins/src/install.rs | 4 +- crates/yaak-plugins/src/manager.rs | 2 +- crates/yaak-wasm/src/lib.rs | 42 +- crates/yaak/src/export.rs | 2 +- crates/yaak/src/import.rs | 76 +-- crates/yaak/src/plugin_events.rs | 144 +++--- crates/yaak/src/response_body.rs | 11 +- crates/yaak/src/send.rs | 55 ++- 57 files changed, 899 insertions(+), 627 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index c81eef004..1849c4f69 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11595,6 +11595,7 @@ dependencies = [ "serde", "serde_json", "sha2", + "tempfile", "thiserror 2.0.17", "ts-rs", "urlencoding", diff --git a/crates-cli/yaak-cli/src/commands/cookie_jar.rs b/crates-cli/yaak-cli/src/commands/cookie_jar.rs index 3494cde28..31a890f37 100644 --- a/crates-cli/yaak-cli/src/commands/cookie_jar.rs +++ b/crates-cli/yaak-cli/src/commands/cookie_jar.rs @@ -22,6 +22,7 @@ fn list(ctx: &CliContext, workspace_id: Option<&str>) -> CommandResult { let workspace_id = resolve_workspace_id(ctx, workspace_id, "cookie-jar list")?; let cookie_jars = ctx .db() + .map_err(|e| e.to_string())? .list_cookie_jars(&workspace_id) .map_err(|e| format!("Failed to list cookie jars: {e}"))?; diff --git a/crates-cli/yaak-cli/src/commands/environment.rs b/crates-cli/yaak-cli/src/commands/environment.rs index 2c060a920..7222a8471 100644 --- a/crates-cli/yaak-cli/src/commands/environment.rs +++ b/crates-cli/yaak-cli/src/commands/environment.rs @@ -50,6 +50,7 @@ fn list(ctx: &CliContext, workspace_id: Option<&str>) -> CommandResult { let workspace_id = resolve_workspace_id(ctx, workspace_id, "environment list")?; let environments = ctx .db() + .map_err(|e| e.to_string())? .list_environments_ensure_base(&workspace_id) .map_err(|e| format!("Failed to list environments: {e}"))?; @@ -66,6 +67,7 @@ fn list(ctx: &CliContext, workspace_id: Option<&str>) -> CommandResult { fn show(ctx: &CliContext, environment_id: &str) -> CommandResult { let environment = ctx .db() + .map_err(|e| e.to_string())? .get_environment(environment_id) .map_err(|e| format!("Failed to get environment: {e}"))?; let output = serde_json::to_string_pretty(&environment) @@ -112,6 +114,7 @@ fn create( let created = ctx .db() + .map_err(|e| e.to_string())? .upsert_environment(&environment, &UpdateSource::Sync) .map_err(|e| format!("Failed to create environment: {e}"))?; @@ -134,6 +137,7 @@ fn create( let created = ctx .db() + .map_err(|e| e.to_string())? .upsert_environment(&environment, &UpdateSource::Sync) .map_err(|e| format!("Failed to create environment: {e}"))?; @@ -147,12 +151,14 @@ fn update(ctx: &CliContext, json: Option, json_input: Option) -> let existing = ctx .db() + .map_err(|e| e.to_string())? .get_environment(&id) .map_err(|e| format!("Failed to get environment for update: {e}"))?; let updated = apply_merge_patch(&existing, &patch, &id, "environment update")?; let saved = ctx .db() + .map_err(|e| e.to_string())? .upsert_environment(&updated, &UpdateSource::Sync) .map_err(|e| format!("Failed to update environment: {e}"))?; @@ -168,6 +174,7 @@ fn delete(ctx: &CliContext, environment_id: &str, yes: bool) -> CommandResult { let deleted = ctx .db() + .map_err(|e| e.to_string())? .delete_environment_by_id(environment_id, &UpdateSource::Sync) .map_err(|e| format!("Failed to delete environment: {e}"))?; diff --git a/crates-cli/yaak-cli/src/commands/folder.rs b/crates-cli/yaak-cli/src/commands/folder.rs index 71eab72cc..5a37c3e57 100644 --- a/crates-cli/yaak-cli/src/commands/folder.rs +++ b/crates-cli/yaak-cli/src/commands/folder.rs @@ -48,8 +48,11 @@ fn schema(pretty: bool) -> CommandResult { fn list(ctx: &CliContext, workspace_id: Option<&str>) -> CommandResult { let workspace_id = resolve_workspace_id(ctx, workspace_id, "folder list")?; - let folders = - ctx.db().list_folders(&workspace_id).map_err(|e| format!("Failed to list folders: {e}"))?; + let folders = ctx + .db() + .map_err(|e| e.to_string())? + .list_folders(&workspace_id) + .map_err(|e| format!("Failed to list folders: {e}"))?; if folders.is_empty() { println!("No folders found in workspace {}", workspace_id); } else { @@ -61,8 +64,11 @@ fn list(ctx: &CliContext, workspace_id: Option<&str>) -> CommandResult { } fn show(ctx: &CliContext, folder_id: &str) -> CommandResult { - let folder = - ctx.db().get_folder(folder_id).map_err(|e| format!("Failed to get folder: {e}"))?; + let folder = ctx + .db() + .map_err(|e| e.to_string())? + .get_folder(folder_id) + .map_err(|e| format!("Failed to get folder: {e}"))?; let output = serde_json::to_string_pretty(&folder) .map_err(|e| format!("Failed to serialize folder: {e}"))?; println!("{output}"); @@ -103,6 +109,7 @@ fn create( let created = ctx .db() + .map_err(|e| e.to_string())? .upsert_folder(&folder, &UpdateSource::Sync) .map_err(|e| format!("Failed to create folder: {e}"))?; @@ -119,6 +126,7 @@ fn create( let created = ctx .db() + .map_err(|e| e.to_string())? .upsert_folder(&folder, &UpdateSource::Sync) .map_err(|e| format!("Failed to create folder: {e}"))?; @@ -130,12 +138,16 @@ fn update(ctx: &CliContext, json: Option, json_input: Option) -> let patch = parse_required_json(json, json_input, "folder update")?; let id = require_id(&patch, "folder update")?; - let existing = - ctx.db().get_folder(&id).map_err(|e| format!("Failed to get folder for update: {e}"))?; + let existing = ctx + .db() + .map_err(|e| e.to_string())? + .get_folder(&id) + .map_err(|e| format!("Failed to get folder for update: {e}"))?; let updated = apply_merge_patch(&existing, &patch, &id, "folder update")?; let saved = ctx .db() + .map_err(|e| e.to_string())? .upsert_folder(&updated, &UpdateSource::Sync) .map_err(|e| format!("Failed to update folder: {e}"))?; @@ -151,6 +163,7 @@ fn delete(ctx: &CliContext, folder_id: &str, yes: bool) -> CommandResult { let deleted = ctx .db() + .map_err(|e| e.to_string())? .delete_folder_by_id(folder_id, &UpdateSource::Sync) .map_err(|e| format!("Failed to delete folder: {e}"))?; diff --git a/crates-cli/yaak-cli/src/commands/import_export.rs b/crates-cli/yaak-cli/src/commands/import_export.rs index e2b9e935e..53ed4693e 100644 --- a/crates-cli/yaak-cli/src/commands/import_export.rs +++ b/crates-cli/yaak-cli/src/commands/import_export.rs @@ -47,6 +47,7 @@ async fn import( ) -> CommandResult<(BatchUpsertResult, Vec)> { if let Some(workspace_id) = args.workspace_id.as_deref() { ctx.db() + .map_err(|e| e.to_string())? .get_workspace(workspace_id) .map_err(|e| format!("Failed to get workspace '{workspace_id}': {e}"))?; } @@ -146,8 +147,11 @@ fn resolve_export_workspace_ids( all: bool, ) -> CommandResult> { if all { - let workspaces = - ctx.db().list_workspaces().map_err(|e| format!("Failed to list workspaces: {e}"))?; + let workspaces = ctx + .db() + .map_err(|e| e.to_string())? + .list_workspaces() + .map_err(|e| format!("Failed to list workspaces: {e}"))?; if workspaces.is_empty() { return Err("No workspaces found to export".to_string()); } @@ -160,6 +164,7 @@ fn resolve_export_workspace_ids( for workspace_id in &workspace_ids { ctx.db() + .map_err(|e| e.to_string())? .get_workspace(workspace_id) .map_err(|e| format!("Failed to get workspace '{workspace_id}': {e}"))?; } diff --git a/crates-cli/yaak-cli/src/commands/plugin.rs b/crates-cli/yaak-cli/src/commands/plugin.rs index 951042c5a..54e64f786 100644 --- a/crates-cli/yaak-cli/src/commands/plugin.rs +++ b/crates-cli/yaak-cli/src/commands/plugin.rs @@ -349,6 +349,7 @@ async fn install_from_directory(context: &CliContext, source: &str) -> CommandRe let plugin = context .db() + .map_err(|e| e.to_string())? .upsert_plugin( &Plugin { directory: plugin_dir_str, diff --git a/crates-cli/yaak-cli/src/commands/request.rs b/crates-cli/yaak-cli/src/commands/request.rs index ec37789ab..9d0d1f242 100644 --- a/crates-cli/yaak-cli/src/commands/request.rs +++ b/crates-cli/yaak-cli/src/commands/request.rs @@ -71,6 +71,7 @@ fn list(ctx: &CliContext, workspace_id: Option<&str>) -> CommandResult { let workspace_id = resolve_workspace_id(ctx, workspace_id, "request list")?; let requests = ctx .db() + .map_err(|e| e.to_string())? .list_http_requests(&workspace_id) .map_err(|e| format!("Failed to list requests: {e}"))?; if requests.is_empty() { @@ -425,6 +426,7 @@ fn create( let created = ctx .db() + .map_err(|e| e.to_string())? .upsert_http_request(&request, &UpdateSource::Sync) .map_err(|e| format!("Failed to create request: {e}"))?; @@ -444,6 +446,7 @@ fn create( let created = ctx .db() + .map_err(|e| e.to_string())? .upsert_http_request(&request, &UpdateSource::Sync) .map_err(|e| format!("Failed to create request: {e}"))?; @@ -458,12 +461,14 @@ fn update(ctx: &CliContext, json: Option, json_input: Option) -> let existing = ctx .db() + .map_err(|e| e.to_string())? .get_http_request(&id) .map_err(|e| format!("Failed to get request for update: {e}"))?; let updated = apply_merge_patch(&existing, &patch, &id, "request update")?; let saved = ctx .db() + .map_err(|e| e.to_string())? .upsert_http_request(&updated, &UpdateSource::Sync) .map_err(|e| format!("Failed to update request: {e}"))?; @@ -472,8 +477,11 @@ fn update(ctx: &CliContext, json: Option, json_input: Option) -> } fn show(ctx: &CliContext, request_id: &str) -> CommandResult { - let request = - ctx.db().get_http_request(request_id).map_err(|e| format!("Failed to get request: {e}"))?; + let request = ctx + .db() + .map_err(|e| e.to_string())? + .get_http_request(request_id) + .map_err(|e| format!("Failed to get request: {e}"))?; let output = serde_json::to_string_pretty(&request) .map_err(|e| format!("Failed to serialize request: {e}"))?; println!("{output}"); @@ -488,6 +496,7 @@ fn delete(ctx: &CliContext, request_id: &str, yes: bool) -> CommandResult { let deleted = ctx .db() + .map_err(|e| e.to_string())? .delete_http_request_by_id(request_id, &UpdateSource::Sync) .map_err(|e| format!("Failed to delete request: {e}"))?; println!("Deleted request: {}", deleted.id); @@ -502,8 +511,11 @@ pub async fn send_request_by_id( cookie_jar_id: Option<&str>, verbose: bool, ) -> Result<(), String> { - let request = - ctx.db().get_any_request(request_id).map_err(|e| format!("Failed to get request: {e}"))?; + let request = ctx + .db() + .map_err(|e| e.to_string())? + .get_any_request(request_id) + .map_err(|e| format!("Failed to get request: {e}"))?; match request { AnyRequest::HttpRequest(http_request) => { send_http_request_by_id( @@ -611,6 +623,7 @@ pub(crate) fn resolve_cookie_jar_id( let default_cookie_jar = ctx .db() + .map_err(|e| e.to_string())? .list_cookie_jars(workspace_id) .map_err(|e| format!("Failed to list cookie jars: {e}"))? .into_iter() diff --git a/crates-cli/yaak-cli/src/commands/response.rs b/crates-cli/yaak-cli/src/commands/response.rs index 3a8b3ba82..c3c475391 100644 --- a/crates-cli/yaak-cli/src/commands/response.rs +++ b/crates-cli/yaak-cli/src/commands/response.rs @@ -27,13 +27,14 @@ pub fn run(ctx: &CliContext, args: ResponseArgs) -> i32 { /// Accepts a response ID, or a request ID meaning "the latest response for that request". fn resolve_response(ctx: &CliContext, id: &str) -> CommandResult { - if let Ok(response) = ctx.db().get_http_response(id) { + if let Ok(response) = ctx.db().map_err(|e| e.to_string())?.get_http_response(id) { return Ok(response); } - if ctx.db().get_http_request(id).is_ok() { + if ctx.db().map_err(|e| e.to_string())?.get_http_request(id).is_ok() { return ctx .db() + .map_err(|e| e.to_string())? .list_http_responses_for_request(id, Some(1)) .map_err(|e| format!("Failed to list responses: {e}"))? .into_iter() @@ -47,13 +48,15 @@ fn resolve_response(ctx: &CliContext, id: &str) -> CommandResult { fn list(ctx: &CliContext, id: Option<&str>, limit: Option) -> CommandResult { // An explicit request ID scopes to that request; anything else lists the workspace. let responses = match id { - Some(id) if ctx.db().get_http_request(id).is_ok() => ctx + Some(id) if ctx.db().map_err(|e| e.to_string())?.get_http_request(id).is_ok() => ctx .db() + .map_err(|e| e.to_string())? .list_http_responses_for_request(id, limit) .map_err(|e| format!("Failed to list responses: {e}"))?, other => { let workspace_id = resolve_workspace_id(ctx, other, "response list")?; ctx.db() + .map_err(|e| e.to_string())? .list_http_responses(&workspace_id, limit) .map_err(|e| format!("Failed to list responses: {e}"))? } @@ -112,13 +115,14 @@ fn body(ctx: &CliContext, id: &str) -> CommandResult { fn delete(ctx: &CliContext, id: &str, yes: bool) -> CommandResult { // Deleting is scoped to a whole request or workspace, since a single stored // response is rarely the thing someone wants to remove. - if ctx.db().get_http_request(id).is_ok() { + if ctx.db().map_err(|e| e.to_string())?.get_http_request(id).is_ok() { if !yes && !confirm_delete("responses for request", id) { println!("Aborted"); return Ok(()); } let count = ctx .db() + .map_err(|e| e.to_string())? .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}"); @@ -132,6 +136,7 @@ fn delete(ctx: &CliContext, id: &str, yes: bool) -> CommandResult { } let count = ctx .db() + .map_err(|e| e.to_string())? .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}"); diff --git a/crates-cli/yaak-cli/src/commands/send.rs b/crates-cli/yaak-cli/src/commands/send.rs index ba96235c1..c3860e600 100644 --- a/crates-cli/yaak-cli/src/commands/send.rs +++ b/crates-cli/yaak-cli/src/commands/send.rs @@ -34,7 +34,7 @@ async fn send_target( ) -> Result<(), String> { let mode = if args.parallel { ExecutionMode::Parallel } else { ExecutionMode::Sequential }; - if let Ok(request) = ctx.db().get_any_request(&args.id) { + if let Ok(request) = ctx.db().map_err(|e| e.to_string())?.get_any_request(&args.id) { let workspace_id = match &request { AnyRequest::HttpRequest(r) => r.workspace_id.clone(), AnyRequest::GrpcRequest(r) => r.workspace_id.clone(), @@ -53,7 +53,7 @@ async fn send_target( .await; } - if let Ok(folder) = ctx.db().get_folder(&args.id) { + if let Ok(folder) = ctx.db().map_err(|e| e.to_string())?.get_folder(&args.id) { let resolved_cookie_jar_id = request::resolve_cookie_jar_id(ctx, &folder.workspace_id, cookie_jar_id)?; @@ -74,7 +74,7 @@ async fn send_target( .await; } - if let Ok(workspace) = ctx.db().get_workspace(&args.id) { + if let Ok(workspace) = ctx.db().map_err(|e| e.to_string())?.get_workspace(&args.id) { let resolved_cookie_jar_id = request::resolve_cookie_jar_id(ctx, &workspace.id, cookie_jar_id)?; @@ -103,6 +103,7 @@ fn collect_folder_request_ids(ctx: &CliContext, folder_id: &str) -> Result Result Result CommandResult { } fn list(ctx: &CliContext) -> CommandResult { - let workspaces = - ctx.db().list_workspaces().map_err(|e| format!("Failed to list workspaces: {e}"))?; + let workspaces = ctx + .db() + .map_err(|e| e.to_string())? + .list_workspaces() + .map_err(|e| format!("Failed to list workspaces: {e}"))?; if workspaces.is_empty() { println!("No workspaces found"); } else { @@ -58,6 +61,7 @@ fn list(ctx: &CliContext) -> CommandResult { fn show(ctx: &CliContext, workspace_id: &str) -> CommandResult { let workspace = ctx .db() + .map_err(|e| e.to_string())? .get_workspace(workspace_id) .map_err(|e| format!("Failed to get workspace: {e}"))?; let output = serde_json::to_string_pretty(&workspace) @@ -85,6 +89,7 @@ fn create( let created = ctx .db() + .map_err(|e| e.to_string())? .upsert_workspace(&workspace, &UpdateSource::Sync) .map_err(|e| format!("Failed to create workspace: {e}"))?; println!("Created workspace: {}", created.id); @@ -98,6 +103,7 @@ fn create( let workspace = Workspace { name, ..Default::default() }; let created = ctx .db() + .map_err(|e| e.to_string())? .upsert_workspace(&workspace, &UpdateSource::Sync) .map_err(|e| format!("Failed to create workspace: {e}"))?; println!("Created workspace: {}", created.id); @@ -110,12 +116,14 @@ fn update(ctx: &CliContext, json: Option, json_input: Option) -> let existing = ctx .db() + .map_err(|e| e.to_string())? .get_workspace(&id) .map_err(|e| format!("Failed to get workspace for update: {e}"))?; let updated = apply_merge_patch(&existing, &patch, &id, "workspace update")?; let saved = ctx .db() + .map_err(|e| e.to_string())? .upsert_workspace(&updated, &UpdateSource::Sync) .map_err(|e| format!("Failed to update workspace: {e}"))?; @@ -131,6 +139,7 @@ fn delete(ctx: &CliContext, workspace_id: &str, yes: bool) -> CommandResult { let deleted = ctx .db() + .map_err(|e| e.to_string())? .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); diff --git a/crates-cli/yaak-cli/src/context.rs b/crates-cli/yaak-cli/src/context.rs index be30439f5..98be3fbe6 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, - ); + if let Ok(db) = query_manager.connect() { + let _ = yaak_lifecycle::on_launch(&yaak_lifecycle::Host::guest(), &db, &blob_manager); + } let encryption_manager = Arc::new(EncryptionManager::new(query_manager.clone(), app_id)); @@ -120,7 +118,7 @@ impl CliContext { &self.data_dir } - pub fn db(&self) -> ClientDb<'_> { + pub fn db(&self) -> yaak_models::error::Result> { self.query_manager.connect() } diff --git a/crates-cli/yaak-cli/src/main.rs b/crates-cli/yaak-cli/src/main.rs index b6020378a..8d92a8132 100644 --- a/crates-cli/yaak-cli/src/main.rs +++ b/crates-cli/yaak-cli/src/main.rs @@ -203,7 +203,7 @@ fn resolve_send_execution_context( environment: Option<&str>, explicit_cookie_jar_id: Option<&str>, ) -> Result { - if let Ok(request) = context.db().get_any_request(id) { + if let Ok(request) = context.db().map_err(|e| e.to_string())?.get_any_request(id) { let (request_id, workspace_id) = match request { AnyRequest::HttpRequest(r) => (Some(r.id), r.workspace_id), AnyRequest::GrpcRequest(r) => (Some(r.id), r.workspace_id), @@ -218,7 +218,7 @@ fn resolve_send_execution_context( }); } - if let Ok(folder) = context.db().get_folder(id) { + if let Ok(folder) = context.db().map_err(|e| e.to_string())?.get_folder(id) { let cookie_jar_id = resolve_cookie_jar_id(context, &folder.workspace_id, explicit_cookie_jar_id)?; return Ok(CliExecutionContext { @@ -229,7 +229,7 @@ fn resolve_send_execution_context( }); } - if let Ok(workspace) = context.db().get_workspace(id) { + if let Ok(workspace) = context.db().map_err(|e| e.to_string())?.get_workspace(id) { let cookie_jar_id = resolve_cookie_jar_id(context, &workspace.id, explicit_cookie_jar_id)?; return Ok(CliExecutionContext { request_id: None, @@ -250,6 +250,7 @@ fn resolve_request_execution_context( ) -> Result { let request = context .db() + .map_err(|e| e.to_string())? .get_any_request(request_id) .map_err(|e| format!("Failed to get request: {e}"))?; @@ -279,6 +280,7 @@ fn resolve_cookie_jar_id( let default_cookie_jar = context .db() + .map_err(|e| e.to_string())? .list_cookie_jars(workspace_id) .map_err(|e| format!("Failed to list cookie jars: {e}"))? .into_iter() diff --git a/crates-cli/yaak-cli/src/plugin_events.rs b/crates-cli/yaak-cli/src/plugin_events.rs index d59bbdd91..0e2e76b70 100644 --- a/crates-cli/yaak-cli/src/plugin_events.rs +++ b/crates-cli/yaak-cli/src/plugin_events.rs @@ -13,7 +13,6 @@ use tokio::task::JoinHandle; use yaak::plugin_events::{ GroupedPluginEvent, HostRequest, SharedPluginEventContext, handle_shared_plugin_event, }; -use yaak_models::render::{render_grpc_request, render_http_request}; use yaak::response_body::FileResponseBodyStore; use yaak::send::{SendHttpRequestWithPluginsParams, send_http_request_with_plugins}; use yaak_crypto::manager::EncryptionManager; @@ -24,6 +23,7 @@ use yaak_models::models::Environment; use yaak_models::queries::any_request::AnyRequest; use yaak_models::query_manager::QueryManager; use yaak_models::render::make_vars_hashmap; +use yaak_models::render::{render_grpc_request, render_http_request}; use yaak_models::util::UpdateSource; use yaak_plugins::events::{ EmptyPayload, ErrorResponse, FormInput, GetCookieValueResponse, InternalEvent, @@ -150,7 +150,11 @@ async fn build_plugin_reply( Some(InternalEventPayload::ShowToastResponse(EmptyPayload {})) } HostRequest::ListOpenWorkspaces(_) => { - let workspaces = match host_context.query_manager.connect().list_workspaces() { + let workspaces = match host_context + .query_manager + .connect() + .and_then(|db| db.list_workspaces()) + { Ok(workspaces) => workspaces .into_iter() .map(|w| WorkspaceInfo { id: w.id.clone(), name: w.name, label: w.id }) @@ -189,7 +193,7 @@ async fn build_plugin_reply( match host_context .query_manager .connect() - .list_cookie_jars(http_request.workspace_id.as_str()) + .and_then(|db| db.list_cookie_jars(http_request.workspace_id.as_str())) { Ok(cookie_jars) => cookie_jars .into_iter() @@ -264,19 +268,20 @@ async fn build_plugin_reply( ..event.context.clone() }; - let environment_chain = - match host_context.query_manager.connect().resolve_environments( + let environment_chain = match host_context.query_manager.connect().and_then(|db| { + db.resolve_environments( &grpc_request.workspace_id, grpc_request.folder_id.as_deref(), execution_context.environment_id.as_deref(), - ) { - Ok(chain) => chain, - Err(err) => { - return Some(InternalEventPayload::ErrorResponse(ErrorResponse { - error: format!("Failed to resolve environments in CLI: {err}"), - })); - } - }; + ) + }) { + Ok(chain) => chain, + Err(err) => { + return Some(InternalEventPayload::ErrorResponse(ErrorResponse { + error: format!("Failed to resolve environments in CLI: {err}"), + })); + } + }; let template_callback = PluginTemplateCallback::new( host_context.plugin_manager.clone(), @@ -324,19 +329,20 @@ async fn build_plugin_reply( ..event.context.clone() }; - let environment_chain = - match host_context.query_manager.connect().resolve_environments( + let environment_chain = match host_context.query_manager.connect().and_then(|db| { + db.resolve_environments( &http_request.workspace_id, http_request.folder_id.as_deref(), execution_context.environment_id.as_deref(), - ) { - Ok(chain) => chain, - Err(err) => { - return Some(InternalEventPayload::ErrorResponse(ErrorResponse { - error: format!("Failed to resolve environments in CLI: {err}"), - })); - } - }; + ) + }) { + Ok(chain) => chain, + Err(err) => { + return Some(InternalEventPayload::ErrorResponse(ErrorResponse { + error: format!("Failed to resolve environments in CLI: {err}"), + })); + } + }; let template_callback = PluginTemplateCallback::new( host_context.plugin_manager.clone(), @@ -380,7 +386,11 @@ async fn build_plugin_reply( }; let folder_id = execution_context.request_id.as_ref().and_then(|rid| { - match host_context.query_manager.connect().get_any_request(rid) { + match host_context + .query_manager + .connect() + .and_then(|db| db.get_any_request(rid)) + { Ok(AnyRequest::HttpRequest(r)) => r.folder_id, Ok(AnyRequest::GrpcRequest(r)) => r.folder_id, Ok(AnyRequest::WebsocketRequest(r)) => r.folder_id, @@ -388,19 +398,20 @@ async fn build_plugin_reply( } }); - let environment_chain = - match host_context.query_manager.connect().resolve_environments( + let environment_chain = match host_context.query_manager.connect().and_then(|db| { + db.resolve_environments( &workspace_id, folder_id.as_deref(), execution_context.environment_id.as_deref(), - ) { - Ok(chain) => chain, - Err(err) => { - return Some(InternalEventPayload::ErrorResponse(ErrorResponse { - error: format!("Failed to resolve environments in CLI: {err}"), - })); - } - }; + ) + }) { + Ok(chain) => chain, + Err(err) => { + return Some(InternalEventPayload::ErrorResponse(ErrorResponse { + error: format!("Failed to resolve environments in CLI: {err}"), + })); + } + }; let template_callback = PluginTemplateCallback::new( host_context.plugin_manager.clone(), @@ -476,15 +487,18 @@ async fn build_plugin_reply( )); }; - let cookie_jar = - match host_context.query_manager.connect().get_cookie_jar(cookie_jar_id) { - Ok(cookie_jar) => cookie_jar, - Err(err) => { - return Some(InternalEventPayload::ErrorResponse(ErrorResponse { - error: format!("Failed to load cookie jar in CLI: {err}"), - })); - } - }; + let cookie_jar = match host_context + .query_manager + .connect() + .and_then(|db| db.get_cookie_jar(cookie_jar_id)) + { + Ok(cookie_jar) => cookie_jar, + Err(err) => { + return Some(InternalEventPayload::ErrorResponse(ErrorResponse { + error: format!("Failed to load cookie jar in CLI: {err}"), + })); + } + }; let names = cookie_jar.cookies.into_iter().map(|c| c.name).collect(); @@ -499,15 +513,18 @@ async fn build_plugin_reply( )); }; - let cookie_jar = - match host_context.query_manager.connect().get_cookie_jar(cookie_jar_id) { - Ok(cookie_jar) => cookie_jar, - Err(err) => { - return Some(InternalEventPayload::ErrorResponse(ErrorResponse { - error: format!("Failed to load cookie jar in CLI: {err}"), - })); - } - }; + let cookie_jar = match host_context + .query_manager + .connect() + .and_then(|db| db.get_cookie_jar(cookie_jar_id)) + { + Ok(cookie_jar) => cookie_jar, + Err(err) => { + return Some(InternalEventPayload::ErrorResponse(ErrorResponse { + error: format!("Failed to load cookie jar in CLI: {err}"), + })); + } + }; let value = get_cookie_value_from_jar(cookie_jar.cookies, &req.name, req.domain.as_deref()); diff --git a/crates-cli/yaak-cli/src/utils/workspace.rs b/crates-cli/yaak-cli/src/utils/workspace.rs index a1f54693d..881f58de6 100644 --- a/crates-cli/yaak-cli/src/utils/workspace.rs +++ b/crates-cli/yaak-cli/src/utils/workspace.rs @@ -9,8 +9,11 @@ pub fn resolve_workspace_id( return Ok(workspace_id.to_string()); } - let workspaces = - ctx.db().list_workspaces().map_err(|e| format!("Failed to list workspaces: {e}"))?; + let workspaces = ctx + .db() + .map_err(|e| e.to_string())? + .list_workspaces() + .map_err(|e| format!("Failed to list workspaces: {e}"))?; match workspaces.as_slice() { [] => Err(format!("No workspaces found. {command_name} requires a workspace ID.")), [workspace] => Ok(workspace.id.clone()), diff --git a/crates-cli/yaak-cli/tests/common/mod.rs b/crates-cli/yaak-cli/tests/common/mod.rs index 2a8749399..9b569fcb0 100644 --- a/crates-cli/yaak-cli/tests/common/mod.rs +++ b/crates-cli/yaak-cli/tests/common/mod.rs @@ -41,6 +41,7 @@ pub fn seed_workspace(data_dir: &Path, workspace_id: &str) { query_manager(data_dir) .connect() + .unwrap() .upsert_workspace(&workspace, &UpdateSource::Sync) .expect("Failed to seed workspace"); } @@ -57,6 +58,7 @@ pub fn seed_request(data_dir: &Path, workspace_id: &str, request_id: &str) { query_manager(data_dir) .connect() + .unwrap() .upsert_http_request(&request, &UpdateSource::Sync) .expect("Failed to seed request"); } @@ -71,6 +73,7 @@ pub fn seed_folder(data_dir: &Path, workspace_id: &str, folder_id: &str) { query_manager(data_dir) .connect() + .unwrap() .upsert_folder(&folder, &UpdateSource::Sync) .expect("Failed to seed folder"); } @@ -86,6 +89,7 @@ pub fn seed_grpc_request(data_dir: &Path, workspace_id: &str, request_id: &str) query_manager(data_dir) .connect() + .unwrap() .upsert_grpc_request(&request, &UpdateSource::Sync) .expect("Failed to seed gRPC request"); } @@ -101,6 +105,7 @@ pub fn seed_websocket_request(data_dir: &Path, workspace_id: &str, request_id: & query_manager(data_dir) .connect() + .unwrap() .upsert_websocket_request(&request, &UpdateSource::Sync) .expect("Failed to seed WebSocket request"); } diff --git a/crates-cli/yaak-cli/tests/environment_commands.rs b/crates-cli/yaak-cli/tests/environment_commands.rs index 30e411dba..2fd404700 100644 --- a/crates-cli/yaak-cli/tests/environment_commands.rs +++ b/crates-cli/yaak-cli/tests/environment_commands.rs @@ -42,7 +42,7 @@ fn create_list_show_delete_round_trip() { .success() .stdout(contains(format!("Deleted environment: {environment_id}"))); - assert!(query_manager(data_dir).connect().get_environment(&environment_id).is_err()); + assert!(query_manager(data_dir).connect().unwrap().get_environment(&environment_id).is_err()); } #[test] diff --git a/crates-cli/yaak-cli/tests/folder_commands.rs b/crates-cli/yaak-cli/tests/folder_commands.rs index 8b250b22b..7b1c903fe 100644 --- a/crates-cli/yaak-cli/tests/folder_commands.rs +++ b/crates-cli/yaak-cli/tests/folder_commands.rs @@ -36,7 +36,7 @@ fn create_list_show_delete_round_trip() { .success() .stdout(contains(format!("Deleted folder: {folder_id}"))); - assert!(query_manager(data_dir).connect().get_folder(&folder_id).is_err()); + assert!(query_manager(data_dir).connect().unwrap().get_folder(&folder_id).is_err()); } #[test] diff --git a/crates-cli/yaak-cli/tests/import_export_commands.rs b/crates-cli/yaak-cli/tests/import_export_commands.rs index f05e55615..23a7cb304 100644 --- a/crates-cli/yaak-cli/tests/import_export_commands.rs +++ b/crates-cli/yaak-cli/tests/import_export_commands.rs @@ -81,7 +81,7 @@ fn import_reads_yaak_workspace_file() { .stdout(contains("Imported 1 workspace, 1 HTTP request")); let query_manager = query_manager(data_dir); - let db = query_manager.connect(); + let db = query_manager.connect().unwrap(); let workspaces = db.list_workspaces().expect("list imported workspaces"); let workspace = workspaces .iter() @@ -160,7 +160,7 @@ fn import_postman_environment_uses_workspace_id() { .stdout(contains("Imported 1 environment")); let query_manager = query_manager(data_dir); - let db = query_manager.connect(); + let db = query_manager.connect().unwrap(); let environments = db.list_environments_ensure_base(&workspace_id).expect("list imported environments"); @@ -217,7 +217,7 @@ fn re_import_merges_into_linked_workspace() { let workspace_id = { let query_manager = query_manager(data_dir); - let db = query_manager.connect(); + let db = query_manager.connect().unwrap(); db.list_workspaces() .expect("list workspaces") .into_iter() @@ -248,7 +248,7 @@ fn re_import_merges_into_linked_workspace() { .stdout(contains("Skipped 1 removed from source")); let query_manager = query_manager(data_dir); - let db = query_manager.connect(); + let db = query_manager.connect().unwrap(); let requests = db.list_http_requests(&workspace_id).expect("list requests"); assert_eq!(requests.len(), 3, "merge must not duplicate: {requests:?}"); assert_eq!( @@ -279,7 +279,7 @@ fn re_import_leaves_deleted_resources_alone() { let workspace_id = { let query_manager = query_manager(data_dir); - let db = query_manager.connect(); + let db = query_manager.connect().unwrap(); let workspace_id = db .list_workspaces() .expect("list workspaces") @@ -311,7 +311,7 @@ fn re_import_leaves_deleted_resources_alone() { let query_manager = query_manager(data_dir); let requests = - query_manager.connect().list_http_requests(&workspace_id).expect("list requests"); + query_manager.connect().unwrap().list_http_requests(&workspace_id).expect("list requests"); assert_eq!(requests.len(), 1, "a deleted request must not come back: {requests:?}"); assert_eq!(requests[0].name, "Request A"); } diff --git a/crates-cli/yaak-cli/tests/request_commands.rs b/crates-cli/yaak-cli/tests/request_commands.rs index cdb78c1fb..62ab8ed24 100644 --- a/crates-cli/yaak-cli/tests/request_commands.rs +++ b/crates-cli/yaak-cli/tests/request_commands.rs @@ -43,7 +43,7 @@ fn show_and_delete_yes_round_trip() { .success() .stdout(contains(format!("Deleted request: {request_id}"))); - assert!(query_manager(data_dir).connect().get_http_request(&request_id).is_err()); + assert!(query_manager(data_dir).connect().unwrap().get_http_request(&request_id).is_err()); } #[test] @@ -61,7 +61,11 @@ fn delete_without_yes_fails_in_non_interactive_mode() { .stderr(contains("Refusing to delete in non-interactive mode without --yes")); assert!( - query_manager(data_dir).connect().get_http_request("rq_seed_delete_noninteractive").is_ok() + query_manager(data_dir) + .connect() + .unwrap() + .get_http_request("rq_seed_delete_noninteractive") + .is_ok() ); } @@ -122,6 +126,7 @@ fn create_allows_workspace_only_with_empty_defaults() { let request = query_manager(data_dir) .connect() + .unwrap() .get_http_request(&request_id) .expect("Failed to load created request"); assert_eq!(request.workspace_id, "wk_test"); @@ -207,7 +212,7 @@ fn request_send_persists_response_body_and_events() { .stdout(contains("hello from integration test")); let qm = query_manager(data_dir); - let db = qm.connect(); + let db = qm.connect().unwrap(); let responses = db.list_http_responses_for_request(&request_id, None).expect("Failed to load responses"); assert_eq!(responses.len(), 1, "expected exactly one persisted response"); diff --git a/crates-cli/yaak-cli/tests/send_commands.rs b/crates-cli/yaak-cli/tests/send_commands.rs index 6ed033823..212f06015 100644 --- a/crates-cli/yaak-cli/tests/send_commands.rs +++ b/crates-cli/yaak-cli/tests/send_commands.rs @@ -26,6 +26,7 @@ fn top_level_send_folder_sends_http_requests_and_prints_summary() { }; query_manager(data_dir) .connect() + .unwrap() .upsert_http_request(&request, &UpdateSource::Sync) .expect("Failed to seed folder request"); diff --git a/crates-cli/yaak-cli/tests/workspace_commands.rs b/crates-cli/yaak-cli/tests/workspace_commands.rs index c2bb64086..ab8731561 100644 --- a/crates-cli/yaak-cli/tests/workspace_commands.rs +++ b/crates-cli/yaak-cli/tests/workspace_commands.rs @@ -26,7 +26,7 @@ fn create_show_delete_round_trip() { .success() .stdout(contains(format!("Deleted workspace: {workspace_id}"))); - assert!(query_manager(data_dir).connect().get_workspace(&workspace_id).is_err()); + assert!(query_manager(data_dir).connect().unwrap().get_workspace(&workspace_id).is_err()); } #[test] diff --git a/crates-tauri/yaak-app-client/src/history.rs b/crates-tauri/yaak-app-client/src/history.rs index a56149193..f6bf290b1 100644 --- a/crates-tauri/yaak-app-client/src/history.rs +++ b/crates-tauri/yaak-app-client/src/history.rs @@ -27,12 +27,9 @@ pub fn get_or_upsert_launch_info(app_handle: &AppHandle) -> &Laun LAUNCH_INFO.get_or_init(|| { let now = Utc::now().naive_utc(); let mut info = LaunchEventInfo { - version_since: app_handle.db().get_key_value_dte(NAMESPACE, VERSION_SINCE_KEY, now), current_version: app_handle.package_info().version.to_string(), - user_since: app_handle.db().get_settings().created_at, - num_launches: app_handle.db().get_key_value_int(NAMESPACE, NUM_LAUNCHES_KEY, 0) + 1, - // The rest will be set below + // The rest is read and written inside the transaction below ..Default::default() }; @@ -42,6 +39,10 @@ pub fn get_or_upsert_launch_info(app_handle: &AppHandle) -> &Laun app_handle .with_tx(|tx| { + info.version_since = tx.get_key_value_dte(NAMESPACE, VERSION_SINCE_KEY, now); + info.user_since = tx.get_settings().created_at; + info.num_launches = tx.get_key_value_int(NAMESPACE, NUM_LAUNCHES_KEY, 0) + 1; + // Load the previously tracked version let curr_db = tx.get_key_value_str(NAMESPACE, LAST_VERSION_KEY, ""); let prev_db = tx.get_key_value_str(NAMESPACE, PREV_VERSION_KEY, ""); diff --git a/crates-tauri/yaak-app-client/src/lib.rs b/crates-tauri/yaak-app-client/src/lib.rs index dec37c7f9..0fe683017 100644 --- a/crates-tauri/yaak-app-client/src/lib.rs +++ b/crates-tauri/yaak-app-client/src/lib.rs @@ -237,17 +237,17 @@ async fn cmd_grpc_reflect( app_handle: AppHandle, grpc_handle: State<'_, Mutex>, ) -> YaakResult> { - let unrendered_request = app_handle.db().get_grpc_request(request_id)?; + let unrendered_request = app_handle.db()?.get_grpc_request(request_id)?; let (resolved_request, auth_context_id) = - resolve_grpc_request(&window.db(), &unrendered_request)?; + resolve_grpc_request(&window.db()?, &unrendered_request)?; - let environment_chain = app_handle.db().resolve_environments( + let environment_chain = app_handle.db()?.resolve_environments( &unrendered_request.workspace_id, unrendered_request.folder_id.as_deref(), environment_id, )?; let resolved_settings = - app_handle.db().resolve_settings_for_grpc_request(&unrendered_request)?; + app_handle.db()?.resolve_settings_for_grpc_request(&unrendered_request)?; let plugin_manager = Arc::new(crate::plugins_ext::plugin_manager(&app_handle).await?); let encryption_manager = Arc::new((*app_handle.state::()).clone()); @@ -266,7 +266,7 @@ async fn cmd_grpc_reflect( let uri = safe_uri(&req.url); let metadata = build_metadata(&window, &req, &auth_context_id).await?; - let settings = window.db().get_settings(); + let settings = window.db()?.get_settings(); let client_certificate = find_client_certificate(req.url.as_str(), &settings.client_certificates); let proto_files: Vec = @@ -298,16 +298,16 @@ async fn cmd_grpc_go( window: WebviewWindow, grpc_handle: State<'_, Mutex>, ) -> YaakResult { - let unrendered_request = app_handle.db().get_grpc_request(request_id)?; + let unrendered_request = app_handle.db()?.get_grpc_request(request_id)?; let (resolved_request, auth_context_id) = - resolve_grpc_request(&window.db(), &unrendered_request)?; - let environment_chain = app_handle.db().resolve_environments( + resolve_grpc_request(&window.db()?, &unrendered_request)?; + let environment_chain = app_handle.db()?.resolve_environments( &unrendered_request.workspace_id, unrendered_request.folder_id.as_deref(), environment_id, )?; let resolved_settings = - app_handle.db().resolve_settings_for_grpc_request(&unrendered_request)?; + app_handle.db()?.resolve_settings_for_grpc_request(&unrendered_request)?; let plugin_manager = Arc::new(crate::plugins_ext::plugin_manager(&app_handle).await?); let encryption_manager = Arc::new((*app_handle.state::()).clone()); @@ -327,10 +327,10 @@ async fn cmd_grpc_go( let metadata = build_metadata(&window, &request, &auth_context_id).await?; // Find matching client certificate for this URL - let settings = app_handle.db().get_settings(); + 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( + let conn = app_handle.db()?.upsert_grpc_connection( &GrpcConnection { workspace_id: request.workspace_id.clone(), request_id: request.id.clone(), @@ -386,7 +386,7 @@ async fn cmd_grpc_go( let connection = match connection { Ok(c) => c, Err(err) => { - app_handle.db().upsert_grpc_connection( + app_handle.db()?.upsert_grpc_connection( &GrpcConnection { elapsed: start.elapsed().as_millis() as i32, error: Some(err.to_string()), @@ -495,7 +495,7 @@ async fn cmd_grpc_go( .await?; let msg = strip_json_comments(&msg); - app_handle.db().upsert_grpc_event( + app_handle.db()?.upsert_grpc_event( &GrpcEvent { content: format!("Connecting to {}", req.url), event_type: GrpcEventType::ConnectionStart, @@ -513,24 +513,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.db().and_then(|db| { + db.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.db().and_then(|db| { + 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), + ) + }); } } }; @@ -584,14 +588,16 @@ 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()), - ) + .and_then(|db| { + db.upsert_grpc_event( + &GrpcEvent { + event_type: GrpcEventType::ClientMessage, + content: msg, + ..base_event.clone() + }, + &UpdateSource::from_window_label(window.label()), + ) + }) .unwrap(); } @@ -599,20 +605,22 @@ async fn cmd_grpc_go( 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()), - ) + .and_then(|db| { + 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()), + ) + }) .unwrap(); let response_message = msg.into_inner(); let content = match connection @@ -623,82 +631,92 @@ async fn cmd_grpc_go( 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()), - ) + .and_then(|db| { + 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()), + ) + }) .unwrap(); return; } }; app_handle .db() - .upsert_grpc_event( - &GrpcEvent { - content, - event_type: GrpcEventType::ServerMessage, - ..base_event.clone() - }, - &UpdateSource::from_window_label(window.label()), - ) + .and_then(|db| { + db.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()), - ) + .and_then(|db| { + 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()), + ) + }) .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), - status: Some(Code::Unknown as i32), - content: "Request failed".to_string(), - event_type: GrpcEventType::ConnectionEnd, - ..base_event.clone() - }, - }), - &UpdateSource::from_window_label(window.label()), - ) + .and_then(|db| { + 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), + 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()), - ) + .and_then(|db| { + 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()), + ) + }) .unwrap(); } None => { @@ -710,20 +728,22 @@ async fn cmd_grpc_go( 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()), - ) + .and_then(|db| { + 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()), + ) + }) .unwrap(); stream.into_inner() } @@ -731,42 +751,46 @@ async fn cmd_grpc_go( 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()), - ) + .and_then(|db| { + 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()), + ) + }) .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()), - ) + .and_then(|db| { + 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()), + ) + }) .unwrap(); return; } @@ -784,30 +808,34 @@ async fn cmd_grpc_go( 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()), - ) + .and_then(|db| { + 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()), + ) + }) .unwrap(); break; } }; app_handle .db() - .upsert_grpc_event( - &GrpcEvent { - content: message, - event_type: GrpcEventType::ServerMessage, - ..base_event.clone() - }, - &UpdateSource::from_window_label(window.label()), - ) + .and_then(|db| { + db.upsert_grpc_event( + &GrpcEvent { + content: message, + event_type: GrpcEventType::ServerMessage, + ..base_event.clone() + }, + &UpdateSource::from_window_label(window.label()), + ) + }) .unwrap(); } Ok(None) => { @@ -815,33 +843,37 @@ async fn cmd_grpc_go( 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()), - ) + .and_then(|db| { + 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()), + ) + }) .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()), - ) + .and_then(|db| { + 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()), + ) + }) .unwrap(); break; } @@ -856,7 +888,7 @@ async fn cmd_grpc_go( let w = app_handle.clone(); tokio::select! { _ = grpc_listen => { - let events = w.db().list_grpc_events(&conn_id).unwrap(); + let events = w.db().and_then(|db| db.list_grpc_events(&conn_id)).unwrap(); let closed_event = events .iter() .find(|e| GrpcEventType::ConnectionEnd == e.event_type); @@ -874,7 +906,7 @@ async fn cmd_grpc_go( }).unwrap(); }, _ = cancelled_rx.changed() => { - w.db().upsert_grpc_event( + w.db().and_then(|db| db.upsert_grpc_event( &GrpcEvent { content: "Cancelled".to_string(), event_type: GrpcEventType::ConnectionEnd, @@ -882,7 +914,7 @@ async fn cmd_grpc_go( ..base_msg.clone() }, &UpdateSource::from_window_label(window.label()), - ).unwrap(); + )).unwrap(); w.with_tx(|c| { c.upsert_grpc_connection( &GrpcConnection{ @@ -923,11 +955,11 @@ async fn cmd_send_ephemeral_request( let response = HttpResponse::default(); request.id = "".to_string(); let environment = match environment_id { - Some(id) => Some(app_handle.db().get_environment(id)?), + Some(id) => Some(app_handle.db()?.get_environment(id)?), None => None, }; let cookie_jar = match cookie_jar_id { - Some(id) => Some(app_handle.db().get_cookie_jar(id)?), + Some(id) => Some(app_handle.db()?.get_cookie_jar(id)?), None => None, }; @@ -964,7 +996,7 @@ async fn cmd_http_response_body( response_id: &str, filter: Option<&str>, ) -> YaakResult { - let location = locate_response_body(&window.db(), response_id)?; + let location = locate_response_body(&window.db()?, response_id)?; let Some(body_path) = location.path else { return Ok(FilterResponse { content: String::new(), error: None }); }; @@ -987,7 +1019,7 @@ async fn cmd_get_sse_events( app_handle: AppHandle, response_id: &str, ) -> YaakResult> { - let Some(body_path) = locate_response_body(&app_handle.db(), response_id)?.path else { + let Some(body_path) = locate_response_body(&app_handle.db()?, response_id)?.path else { return Ok(Vec::new()); }; @@ -1077,10 +1109,10 @@ async fn cmd_send_http_request( cookie_jar_id: Option<&str>, request_id: String, ) -> YaakResult { - let request = app_handle.db().get_http_request(&request_id)?; + let request = app_handle.db()?.get_http_request(&request_id)?; let blobs = app_handle.blob_manager(); - let response = app_handle.db().upsert_http_response( + let response = app_handle.db()?.upsert_http_response( &HttpResponse { request_id: request.id.clone(), workspace_id: request.workspace_id.clone(), @@ -1098,7 +1130,7 @@ async fn cmd_send_http_request( }); let environment = match environment_id { - Some(id) => match app_handle.db().get_environment(id) { + Some(id) => match app_handle.db()?.get_environment(id) { Ok(env) => Some(env), Err(e) => { warn!("Failed to find environment by id {id} {}", e); @@ -1109,7 +1141,7 @@ async fn cmd_send_http_request( }; let cookie_jar = match cookie_jar_id { - Some(id) => Some(app_handle.db().get_cookie_jar(id)?), + Some(id) => Some(app_handle.db()?.get_cookie_jar(id)?), None => None, }; @@ -1125,8 +1157,8 @@ 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( + 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()), @@ -1149,7 +1181,7 @@ async fn cmd_new_child_window( title: &str, inner_size: (f64, f64), ) -> YaakResult<()> { - let use_native_titlebar = parent_window.app_handle().db().get_settings().use_native_titlebar; + let use_native_titlebar = parent_window.app_handle().db()?.get_settings().use_native_titlebar; let initialization_script = initial_appearance_script(&parent_window.app_handle()); let win = yaak_window::window::create_child_window( &parent_window, @@ -1165,7 +1197,7 @@ async fn cmd_new_child_window( } async fn cmd_new_main_window(app_handle: AppHandle, url: &str) -> YaakResult<()> { - let use_native_titlebar = app_handle.db().get_settings().use_native_titlebar; + let use_native_titlebar = app_handle.db()?.get_settings().use_native_titlebar; let initialization_script = initial_appearance_script(&app_handle); let win = yaak_window::window::create_main_window( &app_handle, @@ -1182,7 +1214,7 @@ async fn cmd_check_for_updates( yaak_updater: State<'_, Mutex>, ) -> YaakResult { let update_mode = get_update_mode(&window).await?; - let settings = window.db().get_settings(); + let settings = window.db()?.get_settings(); Ok(yaak_updater .lock() .await @@ -1270,7 +1302,7 @@ pub fn run() { 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()) + yaak_lifecycle::on_launch(&lifecycle_host, &app.db()?, &app.blob_manager()) { error!("on_launch hook failed: {e:?}"); } @@ -1346,7 +1378,7 @@ pub fn run() { app.manage(Mutex::new(ws_manager)); // Specific settings - let settings = app.db().get_settings(); + let settings = app.db()?.get_settings(); app.app_handle().set_native_titlebar(settings.use_native_titlebar); monitor_plugin_events(&app.app_handle().clone()); @@ -1359,7 +1391,8 @@ pub fn run() { .run(|app_handle, event| { match event { RunEvent::Ready => { - let use_native_titlebar = app_handle.db().get_settings().use_native_titlebar; + let settings = app_handle.db().map(|db| db.get_settings()).unwrap_or_default(); + let use_native_titlebar = settings.use_native_titlebar; let initialization_script = initial_appearance_script(app_handle); if let Ok(win) = yaak_window::window::create_main_window( app_handle, @@ -1398,7 +1431,10 @@ pub fn run() { let w = app_handle.get_webview_window(&label).unwrap(); let h = app_handle.clone(); tauri::async_runtime::spawn(async move { - let settings = w.db().get_settings(); + let Ok(settings) = w.db().map(|db| db.get_settings()) else { + warn!("Failed to read settings for update check"); + return; + }; if settings.autoupdate { time::sleep(Duration::from_secs(3)).await; // Wait a bit so it's not so jarring let val: State<'_, Mutex> = h.state(); @@ -1434,7 +1470,7 @@ pub fn run() { } async fn get_update_mode(window: &WebviewWindow) -> YaakResult { - let settings = window.db().get_settings(); + let settings = window.db()?.get_settings(); Ok(UpdateMode::new(settings.update_channel.as_str())) } @@ -1561,13 +1597,13 @@ fn get_window_from_plugin_context( } fn workspace_from_window(window: &WebviewWindow) -> Option { - window.workspace_id().and_then(|id| window.db().get_workspace(&id).ok()) + window.workspace_id().and_then(|id| window.db().and_then(|db| db.get_workspace(&id)).ok()) } fn environment_from_window(window: &WebviewWindow) -> Option { - window.environment_id().and_then(|id| window.db().get_environment(&id).ok()) + window.environment_id().and_then(|id| window.db().and_then(|db| db.get_environment(&id)).ok()) } fn cookie_jar_from_window(window: &WebviewWindow) -> Option { - window.cookie_jar_id().and_then(|id| window.db().get_cookie_jar(&id).ok()) + window.cookie_jar_id().and_then(|id| window.db().and_then(|db| db.get_cookie_jar(&id)).ok()) } diff --git a/crates-tauri/yaak-app-client/src/models_ext.rs b/crates-tauri/yaak-app-client/src/models_ext.rs index 835547dd0..3e3b6188d 100644 --- a/crates-tauri/yaak-app-client/src/models_ext.rs +++ b/crates-tauri/yaak-app-client/src/models_ext.rs @@ -37,11 +37,9 @@ fn drain_model_changes_batch( app_handle: &tauri::AppHandle, cursor: &mut ModelChangeCursor, ) -> bool { - let changes = match query_manager.connect().list_model_changes_since( - &cursor.created_at, - cursor.id, - MODEL_CHANGES_POLL_BATCH_SIZE, - ) { + let changes = match query_manager.connect().and_then(|db| { + db.list_model_changes_since(&cursor.created_at, cursor.id, MODEL_CHANGES_POLL_BATCH_SIZE) + }) { Ok(changes) => changes, Err(err) => { error!("Failed to poll model_changes rows: {err:?}"); @@ -92,7 +90,7 @@ async fn run_model_change_poller( /// Extension trait for accessing the QueryManager from Tauri Manager types. pub trait QueryManagerExt<'a, R> { fn db_manager(&'a self) -> State<'a, QueryManager>; - fn db(&'a self) -> ClientDb<'a>; + fn db(&'a self) -> Result>; fn with_tx(&'a self, func: F) -> Result where F: FnOnce(&ClientDb) -> Result; @@ -103,7 +101,7 @@ impl<'a, R: Runtime, M: Manager> QueryManagerExt<'a, R> for M { self.state::() } - fn db(&'a self) -> ClientDb<'a> { + fn db(&'a self) -> Result> { let qm = self.state::(); qm.inner().connect() } diff --git a/crates-tauri/yaak-app-client/src/notifications.rs b/crates-tauri/yaak-app-client/src/notifications.rs index c653c55e4..dbea12ef7 100644 --- a/crates-tauri/yaak-app-client/src/notifications.rs +++ b/crates-tauri/yaak-app-client/src/notifications.rs @@ -55,7 +55,7 @@ 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( + window.db()?.set_key_value_raw( KV_NAMESPACE, KV_KEY, seen_json.as_str(), @@ -74,7 +74,7 @@ impl YaakNotifier { self.last_check = Some(Instant::now()); - if !app_handle.db().get_settings().check_notifications { + if !app_handle.db()?.get_settings().check_notifications { info!("Notifications are disabled. Skipping check."); return Ok(()); } @@ -139,7 +139,7 @@ impl YaakNotifier { } async fn get_kv(app_handle: &AppHandle) -> Result> { - match app_handle.db().get_key_value_raw("notifications", "seen") { + match app_handle.db()?.get_key_value_raw("notifications", "seen") { None => Ok(Vec::new()), Some(v) => Ok(serde_json::from_str(&v.value)?), } @@ -155,7 +155,7 @@ fn get_updater_status(app_handle: &AppHandle) -> &'static str { #[cfg(all(feature = "updater", target_os = "linux"))] { - let settings = app_handle.db().get_settings(); + let settings = app_handle.db()?.get_settings(); if !settings.autoupdate { // Updates are explicitly disabled "disabled" @@ -170,7 +170,7 @@ fn get_updater_status(app_handle: &AppHandle) -> &'static str { #[cfg(all(feature = "updater", not(target_os = "linux")))] { - let settings = app_handle.db().get_settings(); + let settings = app_handle.db()?.get_settings(); if settings.autoupdate { "enabled" } else { "disabled" } } } diff --git a/crates-tauri/yaak-app-client/src/plugin_events.rs b/crates-tauri/yaak-app-client/src/plugin_events.rs index 9071313c6..1cae34f64 100644 --- a/crates-tauri/yaak-app-client/src/plugin_events.rs +++ b/crates-tauri/yaak-app-client/src/plugin_events.rs @@ -104,14 +104,14 @@ async fn handle_host_plugin_request( Box::pin(handle_plugin_event(app_handle, &toast_event, plugin_handle)).await } HostRequest::ReloadResponse(req) => { - let plugins = app_handle.db().list_plugins()?; + let plugins = app_handle.db()?.list_plugins()?; for plugin in plugins { if plugin.directory != plugin_handle.dir { continue; } let new_plugin = Plugin { updated_at: Utc::now().naive_utc(), ..plugin }; - app_handle.db().upsert_plugin(&new_plugin, &UpdateSource::Plugin)?; + app_handle.db()?.upsert_plugin(&new_plugin, &UpdateSource::Plugin)?; } if !req.silent { @@ -199,7 +199,7 @@ async fn handle_host_plugin_request( let workspace = workspace_from_window(&window).expect("Failed to get workspace_id from window URL"); let environment_id = environment_from_window(&window).map(|e| e.id); - let environment_chain = window.db().resolve_environments( + let environment_chain = window.db()?.resolve_environments( &workspace.id, req.grpc_request.folder_id.as_deref(), environment_id.as_deref(), @@ -225,7 +225,7 @@ async fn handle_host_plugin_request( let workspace = workspace_from_window(&window).expect("Failed to get workspace_id from window URL"); let environment_id = environment_from_window(&window).map(|e| e.id); - let environment_chain = window.db().resolve_environments( + let environment_chain = window.db()?.resolve_environments( &workspace.id, req.http_request.folder_id.as_deref(), environment_id.as_deref(), @@ -252,7 +252,7 @@ async fn handle_host_plugin_request( workspace_from_window(&window).expect("Failed to get workspace_id from window URL"); let environment_id = environment_from_window(&window).map(|e| e.id); let folder_id = if let Some(id) = window.request_id() { - match window.db().get_any_request(&id) { + match window.db()?.get_any_request(&id) { Ok(AnyRequest::HttpRequest(r)) => r.folder_id, Ok(AnyRequest::GrpcRequest(r)) => r.folder_id, Ok(AnyRequest::WebsocketRequest(r)) => r.folder_id, @@ -261,7 +261,7 @@ async fn handle_host_plugin_request( } else { None }; - let environment_chain = window.db().resolve_environments( + let environment_chain = window.db()?.resolve_environments( &workspace.id, folder_id.as_deref(), environment_id.as_deref(), @@ -294,7 +294,7 @@ async fn handle_host_plugin_request( HttpResponse::default() } else { let blobs = window.blob_manager(); - window.db().upsert_http_response( + window.db()?.upsert_http_response( &HttpResponse { request_id: http_request.id.clone(), workspace_id: http_request.workspace_id.clone(), @@ -328,7 +328,7 @@ async fn handle_host_plugin_request( HostRequest::OpenWindow(req) => { let (navigation_tx, mut navigation_rx) = tokio::sync::mpsc::channel(128); let (close_tx, mut close_rx) = tokio::sync::mpsc::channel(128); - let use_native_titlebar = app_handle.db().get_settings().use_native_titlebar; + let use_native_titlebar = app_handle.db()?.get_settings().use_native_titlebar; let win_config = CreateWindowConfig { url: &req.url, label: &req.label, @@ -438,7 +438,7 @@ async fn handle_host_plugin_request( let environment_id = environment_from_window(&w).map(|m| m.id); let workspace_id = workspace_from_window(&w).map(|m| m.id); let request_id = - match app_handle.db().get_any_request(&w.request_id().unwrap_or_default()) { + match app_handle.db()?.get_any_request(&w.request_id().unwrap_or_default()) { Ok(AnyRequest::HttpRequest(r)) => Some(r.id), Ok(AnyRequest::WebsocketRequest(r)) => Some(r.id), Ok(AnyRequest::GrpcRequest(r)) => Some(r.id), diff --git a/crates-tauri/yaak-app-client/src/plugins_ext.rs b/crates-tauri/yaak-app-client/src/plugins_ext.rs index bff3c6ed8..49b1af298 100644 --- a/crates-tauri/yaak-app-client/src/plugins_ext.rs +++ b/crates-tauri/yaak-app-client/src/plugins_ext.rs @@ -107,7 +107,7 @@ impl PluginUpdater { let app_version = window.app_handle().package_info().version.to_string(); let http_client = yaak_api_client(ApiClientKind::App, &app_version)?; - let plugins = window.app_handle().db().list_plugins()?; + let plugins = window.app_handle().db()?.list_plugins()?; let updates = check_plugin_updates(&http_client, plugins.clone()).await?; if updates.plugins.is_empty() { @@ -203,7 +203,7 @@ 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( + let plugin = window.db()?.upsert_plugin( &Plugin { directory: directory.into(), url: None, @@ -234,7 +234,7 @@ pub async fn cmd_plugins_updates( ) -> Result { let app_version = app_handle.package_info().version.to_string(); let http_client = yaak_api_client(ApiClientKind::App, &app_version)?; - let plugins = app_handle.db().list_plugins()?; + let plugins = app_handle.db()?.list_plugins()?; Ok(check_plugin_updates(&http_client, plugins).await?) } @@ -243,7 +243,7 @@ pub async fn cmd_plugins_update_all( ) -> Result> { let app_version = window.app_handle().package_info().version.to_string(); let http_client = yaak_api_client(ApiClientKind::App, &app_version)?; - let plugins = window.db().list_plugins()?; + let plugins = window.db()?.list_plugins()?; // Get list of available updates (already filtered to only registry plugins) let updates = check_plugin_updates(&http_client, plugins).await?; diff --git a/crates-tauri/yaak-app-client/src/rpc_ext.rs b/crates-tauri/yaak-app-client/src/rpc_ext.rs index 4d1f5f26d..95b3da391 100644 --- a/crates-tauri/yaak-app-client/src/rpc_ext.rs +++ b/crates-tauri/yaak-app-client/src/rpc_ext.rs @@ -42,26 +42,26 @@ use yaak_models::models::{ }; use yaak_models::query_manager::QueryManager; use yaak_models::util::{BatchUpsertResult, ImportPlan}; +use yaak_plugins::api::{PluginNameVersion, PluginSearchResponse, PluginUpdatesResponse}; use yaak_plugins::events::{ CallFolderActionRequest, CallGrpcRequestActionRequest, CallHttpRequestActionRequest, - CallWebsocketRequestActionRequest, CallWorkspaceActionRequest, FilterResponse, ImportResponse, - JsonPrimitive, RenderPurpose, GetFolderActionsResponse, GetGrpcRequestActionsResponse, - GetHttpAuthenticationConfigResponse, GetHttpAuthenticationSummaryResponse, - GetHttpRequestActionsResponse, GetTemplateFunctionConfigResponse, - GetTemplateFunctionSummaryResponse, GetThemesResponse, GetWebsocketRequestActionsResponse, - GetWorkspaceActionsResponse, + CallWebsocketRequestActionRequest, CallWorkspaceActionRequest, FilterResponse, + GetFolderActionsResponse, GetGrpcRequestActionsResponse, GetHttpAuthenticationConfigResponse, + GetHttpAuthenticationSummaryResponse, GetHttpRequestActionsResponse, + GetTemplateFunctionConfigResponse, GetTemplateFunctionSummaryResponse, GetThemesResponse, + GetWebsocketRequestActionsResponse, GetWorkspaceActionsResponse, ImportResponse, JsonPrimitive, + RenderPurpose, }; -use yaak_plugins::api::{PluginNameVersion, PluginSearchResponse, PluginUpdatesResponse}; use yaak_plugins::manager::PluginManager; use yaak_plugins::native_template_functions::encrypt_secure_template_function; -use yaak_plugins::template_callback::PluginTemplateCallback; use yaak_plugins::plugin_meta::PluginMetadata; +use yaak_plugins::template_callback::PluginTemplateCallback; use yaak_rpc::RpcRouter; use yaak_rpc_schema::*; use yaak_sse::sse::ServerSentEvent; use yaak_sync::sync::SyncOp; -use yaak_templates::TemplateCallback; use yaak_tauri_utils::window::WorkspaceWindowTrait; +use yaak_templates::TemplateCallback; use yaak_ws::WebsocketManager; /// Per-call context: the window a command was invoked from. @@ -467,7 +467,7 @@ async fn cmd_commit_import(ctx: ClientCtx, req: CmdCommitImportRe async fn cmd_list_import_sources(ctx: ClientCtx, req: CmdListImportSourcesReq) -> Result> { use crate::models_ext::QueryManagerExt; - Ok(ctx.window.db().list_import_sources(&req.workspace_id)?) + Ok(ctx.window.db()?.list_import_sources(&req.workspace_id)?) } async fn cmd_import_sources_for_origin(ctx: ClientCtx, req: CmdImportSourcesForOriginReq) -> Result> { @@ -480,7 +480,7 @@ async fn cmd_import_sources_for_origin(ctx: ClientCtx, req: CmdIm }, (None, None) => return Ok(Vec::new()), }; - Ok(ctx.window.db().list_import_sources_by_origin(&origin)?) + Ok(ctx.window.db()?.list_import_sources_by_origin(&origin)?) } async fn cmd_http_request_actions(ctx: ClientCtx, req: CmdHttpRequestActionsReq) -> Result> { diff --git a/crates-tauri/yaak-app-client/src/sync_ext.rs b/crates-tauri/yaak-app-client/src/sync_ext.rs index 223573ce0..c9db6f1dd 100644 --- a/crates-tauri/yaak-app-client/src/sync_ext.rs +++ b/crates-tauri/yaak-app-client/src/sync_ext.rs @@ -26,9 +26,10 @@ pub(crate) async fn cmd_sync_calculate( return Err(InvalidSyncDirectory(sync_dir.to_string_lossy().to_string()).into()); } - let db = app_handle.db(); let version = app_handle.package_info().version.to_string(); - let db_candidates = get_db_candidates(&db, &version, workspace_id, sync_dir)?; + // Scoped so the pooled connection is back before the directory walk below, which is + // as slow as the sync directory is large. + let db_candidates = get_db_candidates(&app_handle.db()?, &version, workspace_id, sync_dir)?; let fs_candidates = get_fs_candidates(sync_dir)? .into_iter() // Only keep items in the same workspace @@ -49,7 +50,7 @@ pub(crate) async fn cmd_sync_apply( sync_dir: &Path, workspace_id: &str, ) -> Result<()> { - let db = app_handle.db(); + let db = app_handle.db()?; 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)?; diff --git a/crates-tauri/yaak-app-client/src/updates.rs b/crates-tauri/yaak-app-client/src/updates.rs index ffef3a004..cd8ce8a07 100644 --- a/crates-tauri/yaak-app-client/src/updates.rs +++ b/crates-tauri/yaak-app-client/src/updates.rs @@ -75,7 +75,7 @@ impl YaakUpdater { auto_download: bool, update_trigger: UpdateTrigger, ) -> Result { - let settings = window.db().get_settings(); + let settings = window.db()?.get_settings(); let update_key = format!("{:x}", md5::compute(settings.id)); self.last_check = Some(Instant::now()); diff --git a/crates-tauri/yaak-app-client/src/ws_ext.rs b/crates-tauri/yaak-app-client/src/ws_ext.rs index 7ccba1fe2..634de4ba1 100644 --- a/crates-tauri/yaak-app-client/src/ws_ext.rs +++ b/crates-tauri/yaak-app-client/src/ws_ext.rs @@ -13,6 +13,7 @@ use tauri::{AppHandle, Manager, Runtime, State, WebviewWindow}; use tokio::sync::{Mutex, mpsc}; use tokio_tungstenite::tungstenite::Message; use url::Url; +use yaak_commands::resolve::resolve_websocket_request; use yaak_crypto::manager::EncryptionManager; use yaak_http::cookies::CookieStore; use yaak_http::path_placeholders::apply_path_placeholders; @@ -26,7 +27,6 @@ use yaak_plugins::template_callback::PluginTemplateCallback; use yaak_templates::strip_json_comments::maybe_strip_json_comments; use yaak_templates::{RenderErrorBehavior, RenderOptions}; use yaak_tls::find_client_certificate; -use yaak_commands::resolve::resolve_websocket_request; use yaak_ws::{WebsocketManager, render_websocket_request}; pub async fn cmd_ws_send( @@ -36,14 +36,14 @@ pub async fn cmd_ws_send( window: WebviewWindow, ws_manager: State<'_, Mutex>, ) -> Result { - let connection = app_handle.db().get_websocket_connection(connection_id)?; + let connection = app_handle.db()?.get_websocket_connection(connection_id)?; match send_websocket_message(&connection, environment_id, &app_handle, &window, &ws_manager) .await { Ok(connection) => Ok(connection), Err(e) => { - app_handle.db().upsert_websocket_event( + app_handle.db()?.upsert_websocket_event( &WebsocketEvent { connection_id: connection.id.clone(), request_id: connection.request_id.clone(), @@ -68,14 +68,14 @@ async fn send_websocket_message( window: &WebviewWindow, ws_manager: &Mutex, ) -> Result { - let unrendered_request = app_handle.db().get_websocket_request(&connection.request_id)?; - let environment_chain = app_handle.db().resolve_environments( + let unrendered_request = app_handle.db()?.get_websocket_request(&connection.request_id)?; + let environment_chain = app_handle.db()?.resolve_environments( &unrendered_request.workspace_id, unrendered_request.folder_id.as_deref(), environment_id, )?; let (resolved_request, _auth_context_id) = - resolve_websocket_request(&window.db(), &unrendered_request)?; + resolve_websocket_request(&window.db()?, &unrendered_request)?; let plugin_manager = Arc::new(crate::plugins_ext::plugin_manager(app_handle).await?); let encryption_manager = Arc::new((*app_handle.state::()).clone()); let request = render_websocket_request( @@ -96,7 +96,7 @@ 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( + app_handle.db()?.upsert_websocket_event( &WebsocketEvent { connection_id: connection.id.clone(), request_id: request.id.clone(), @@ -119,7 +119,7 @@ pub async fn cmd_ws_close( ws_manager: State<'_, Mutex>, ) -> Result { let connection = { - let db = app_handle.db(); + let db = app_handle.db()?; let connection = db.get_websocket_connection(connection_id)?; db.upsert_websocket_connection( &WebsocketConnection { state: WebsocketConnectionState::Closing, ..connection }, @@ -143,17 +143,17 @@ pub async fn cmd_ws_connect( window: WebviewWindow, ws_manager: State<'_, Mutex>, ) -> Result { - let unrendered_request = app_handle.db().get_websocket_request(request_id)?; - let environment_chain = app_handle.db().resolve_environments( + let unrendered_request = app_handle.db()?.get_websocket_request(request_id)?; + let environment_chain = app_handle.db()?.resolve_environments( &unrendered_request.workspace_id, unrendered_request.folder_id.as_deref(), environment_id, )?; let resolved_settings = - app_handle.db().resolve_settings_for_websocket_request(&unrendered_request)?; - let settings = app_handle.db().get_settings(); + app_handle.db()?.resolve_settings_for_websocket_request(&unrendered_request)?; + let settings = app_handle.db()?.get_settings(); let (resolved_request, auth_context_id) = - resolve_websocket_request(&window.db(), &unrendered_request)?; + resolve_websocket_request(&window.db()?, &unrendered_request)?; let plugin_manager = Arc::new(crate::plugins_ext::plugin_manager(&app_handle).await?); let encryption_manager = Arc::new((*app_handle.state::()).clone()); let request = render_websocket_request( @@ -169,7 +169,7 @@ pub async fn cmd_ws_connect( ) .await?; - let connection = app_handle.db().upsert_websocket_connection( + let connection = app_handle.db()?.upsert_websocket_connection( &WebsocketConnection { workspace_id: request.workspace_id.clone(), request_id: request_id.to_string(), @@ -187,7 +187,7 @@ 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( + return Ok(app_handle.db()?.upsert_websocket_connection( &WebsocketConnection { error: Some(format!("Failed to parse URL {}", e.to_string())), state: WebsocketConnectionState::Closed, @@ -268,7 +268,7 @@ pub async fn cmd_ws_connect( resolved_settings.send_cookies.value || resolved_settings.store_cookies.value, cookie_jar_id, ) { - (true, Some(id)) => Some(app_handle.db().get_cookie_jar(id)?), + (true, Some(id)) => Some(app_handle.db()?.get_cookie_jar(id)?), _ => None, }; let cookie_store = @@ -321,7 +321,7 @@ pub async fn cmd_ws_connect( { Ok(r) => r, Err(e) => { - return Ok(app_handle.db().upsert_websocket_connection( + return Ok(app_handle.db()?.upsert_websocket_connection( &WebsocketConnection { error: Some(e.to_string()), state: WebsocketConnectionState::Closed, @@ -332,7 +332,7 @@ pub async fn cmd_ws_connect( } }; - app_handle.db().upsert_websocket_event( + app_handle.db()?.upsert_websocket_event( &WebsocketEvent { connection_id: connection.id.clone(), request_id: request.id.clone(), @@ -366,11 +366,11 @@ 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.db()?.upsert_cookie_jar(cookie_jar, &UpdateSource::Background)?; } } - let connection = app_handle.db().upsert_websocket_connection( + let connection = app_handle.db()?.upsert_websocket_connection( &WebsocketConnection { state: WebsocketConnectionState::Connected, headers: response_headers, @@ -394,9 +394,8 @@ pub async fn cmd_ws_connect( has_written_close = true; } - app_handle - .db() - .upsert_websocket_event( + let event = || { + app_handle.db()?.upsert_websocket_event( &WebsocketEvent { connection_id: connection_id.clone(), request_id: request_id.clone(), @@ -416,13 +415,15 @@ pub async fn cmd_ws_connect( }, &UpdateSource::from_window_label(&window_label), ) - .unwrap(); + }; + if let Err(e) = event() { + warn!("Failed to store WebSocket event: {e}"); + } } info!("Websocket connection closed"); if !has_written_close { - app_handle - .db() - .upsert_websocket_event( + let close_event = app_handle.db().and_then(|db| { + db.upsert_websocket_event( &WebsocketEvent { connection_id: connection_id.clone(), request_id: request_id.clone(), @@ -433,11 +434,13 @@ pub async fn cmd_ws_connect( }, &UpdateSource::from_window_label(&window_label), ) - .unwrap(); + }); + if let Err(e) = close_event { + warn!("Failed to store WebSocket close event: {e}"); + } } - app_handle - .db() - .upsert_websocket_connection( + let closed = app_handle.db().and_then(|db| { + db.upsert_websocket_connection( &WebsocketConnection { workspace_id: request.workspace_id.clone(), request_id: request_id.to_string(), @@ -446,7 +449,10 @@ pub async fn cmd_ws_connect( }, &UpdateSource::from_window_label(&window_label), ) - .unwrap(); + }); + if let Err(e) = closed { + warn!("Failed to mark WebSocket connection closed: {e}"); + } }); } diff --git a/crates-tauri/yaak-license/src/license.rs b/crates-tauri/yaak-license/src/license.rs index dcc3fd122..349b18192 100644 --- a/crates-tauri/yaak-license/src/license.rs +++ b/crates-tauri/yaak-license/src/license.rs @@ -16,11 +16,11 @@ use yaak_models::util::UpdateSource; /// Extension trait for accessing the QueryManager from Tauri Manager types. /// This is needed temporarily until all crates are refactored to not use Tauri. trait QueryManagerExt<'a, R> { - fn db(&'a self) -> ClientDb<'a>; + fn db(&'a self) -> yaak_models::error::Result>; } impl<'a, R: Runtime, M: Manager> QueryManagerExt<'a, R> for M { - fn db(&'a self) -> ClientDb<'a> { + fn db(&'a self) -> yaak_models::error::Result> { let qm = self.state::(); qm.inner().connect() } @@ -137,7 +137,7 @@ pub async fn activate_license( } let body: ActivateLicenseResponsePayload = response.json().await?; - window.app_handle().db().set_key_value_str( + window.app_handle().db()?.set_key_value_str( KV_ACTIVATION_ID_KEY, KV_NAMESPACE, body.activation_id.as_str(), @@ -172,7 +172,7 @@ pub async fn deactivate_license(window: &WebviewWindow) -> Result return Err(ServerError); } - app_handle.db().delete_key_value( + app_handle.db()?.delete_key_value( KV_ACTIVATION_ID_KEY, KV_NAMESPACE, &UpdateSource::from_window_label(window.label()), @@ -191,7 +191,7 @@ pub async fn check_license(window: &WebviewWindow) -> Result String { } pub async fn get_activation_id(app_handle: &AppHandle) -> String { - app_handle.db().get_key_value_str(KV_ACTIVATION_ID_KEY, KV_NAMESPACE, "") + app_handle + .db() + .map(|db| db.get_key_value_str(KV_ACTIVATION_ID_KEY, KV_NAMESPACE, "")) + .unwrap_or_default() } diff --git a/crates/yaak-commands/src/actions.rs b/crates/yaak-commands/src/actions.rs index 97b30e79c..b2154263a 100644 --- a/crates/yaak-commands/src/actions.rs +++ b/crates/yaak-commands/src/actions.rs @@ -63,7 +63,7 @@ pub async fn cmd_call_http_request_action( req: CmdCallHttpRequestActionReq, ) -> Result<()> { let inner = req.req; - let http_request = resolve_http_request(&host.db(), &inner.args.http_request)?.0; + let http_request = resolve_http_request(&host.db()?, &inner.args.http_request)?.0; host.call_http_request_action(CallHttpRequestActionRequest { args: CallHttpRequestActionArgs { http_request }, ..inner @@ -76,7 +76,7 @@ pub async fn cmd_call_grpc_request_action( req: CmdCallGrpcRequestActionReq, ) -> Result<()> { let inner = req.req; - let grpc_request = resolve_grpc_request(&host.db(), &inner.args.grpc_request)?.0; + let grpc_request = resolve_grpc_request(&host.db()?, &inner.args.grpc_request)?.0; host.call_grpc_request_action(CallGrpcRequestActionRequest { args: CallGrpcRequestActionArgs { grpc_request, ..inner.args }, ..inner @@ -89,7 +89,7 @@ pub async fn cmd_call_websocket_request_action( req: CmdCallWebsocketRequestActionReq, ) -> Result<()> { let inner = req.req; - let websocket_request = host.db().get_websocket_request(&inner.args.websocket_request.id)?; + let websocket_request = host.db()?.get_websocket_request(&inner.args.websocket_request.id)?; host.call_websocket_request_action(CallWebsocketRequestActionRequest { args: CallWebsocketRequestActionArgs { websocket_request }, ..inner @@ -102,7 +102,7 @@ pub async fn cmd_call_workspace_action( req: CmdCallWorkspaceActionReq, ) -> Result<()> { let inner = req.req; - let workspace = host.db().get_workspace(&inner.args.workspace.id)?; + let workspace = host.db()?.get_workspace(&inner.args.workspace.id)?; host.call_workspace_action(CallWorkspaceActionRequest { args: CallWorkspaceActionArgs { workspace }, ..inner @@ -115,7 +115,7 @@ pub async fn cmd_call_folder_action( req: CmdCallFolderActionReq, ) -> Result<()> { let inner = req.req; - let folder = host.db().get_folder(&inner.args.folder.id)?; + let folder = host.db()?.get_folder(&inner.args.folder.id)?; host.call_folder_action(CallFolderActionRequest { args: CallFolderActionArgs { folder }, ..inner @@ -152,6 +152,6 @@ pub async fn cmd_reload_plugins( host: H, _req: CmdReloadPluginsReq, ) -> Result> { - let plugins = host.db().list_plugins()?; + let plugins = host.db()?.list_plugins()?; Ok(host.reload_plugins(plugins).await) } diff --git a/crates/yaak-commands/src/host.rs b/crates/yaak-commands/src/host.rs index 5d8a54b10..3c9eb419a 100644 --- a/crates/yaak-commands/src/host.rs +++ b/crates/yaak-commands/src/host.rs @@ -68,11 +68,11 @@ pub trait Host: Clone { PluginContext::new(Some(self.client_id().to_string()), self.session().workspace_id) } - fn db(&self) -> ClientDb<'_> { + fn db(&self) -> yaak_models::error::Result> { self.query_manager().connect() } - fn blobs(&self) -> BlobContext { + fn blobs(&self) -> yaak_models::error::Result { self.blob_manager().connect() } } diff --git a/crates/yaak-commands/src/models.rs b/crates/yaak-commands/src/models.rs index 86ec67d7e..58dfbbac1 100644 --- a/crates/yaak-commands/src/models.rs +++ b/crates/yaak-commands/src/models.rs @@ -11,7 +11,7 @@ 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 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)?) @@ -49,25 +49,25 @@ pub async fn models_websocket_events( host: H, req: ModelsWebsocketEventsReq, ) -> Result> { - Ok(host.db().list_websocket_events(&req.connection_id)?) + Ok(host.db()?.list_websocket_events(&req.connection_id)?) } pub async fn models_grpc_events( host: H, req: ModelsGrpcEventsReq, ) -> Result> { - Ok(host.db().list_grpc_events(&req.connection_id)?) + Ok(host.db()?.list_grpc_events(&req.connection_id)?) } pub async fn models_get_settings(host: H, _req: ModelsGetSettingsReq) -> Result { - Ok(host.db().get_settings()) + Ok(host.db()?.get_settings()) } pub async fn models_get_graphql_introspection( host: H, req: ModelsGetGraphqlIntrospectionReq, ) -> Result> { - Ok(host.db().get_graphql_introspection(&req.request_id)) + Ok(host.db()?.get_graphql_introspection(&req.request_id)) } pub async fn models_upsert_graphql_introspection( @@ -75,7 +75,7 @@ pub async fn models_upsert_graphql_introspection( req: ModelsUpsertGraphqlIntrospectionReq, ) -> Result { let source = host.update_source(); - Ok(host.db().upsert_graphql_introspection( + Ok(host.db()?.upsert_graphql_introspection( &req.workspace_id, &req.request_id, req.content, @@ -96,14 +96,14 @@ pub async fn models_workspace_models( // Add the global models { - let db = host.db(); + let db = host.db()?; l.push(db.get_settings().into()); l.append(&mut db.list_workspaces()?.into_iter().map(Into::into).collect()); l.append(&mut db.list_key_values()?.into_iter().map(Into::into).collect()); } let plugins = { - let db = host.db(); + let db = host.db()?; db.list_plugins()? }; @@ -112,7 +112,7 @@ pub async fn models_workspace_models( // Add the workspace children if let Some(wid) = req.workspace_id.as_deref() { - let db = host.db(); + 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_folders(wid)?.into_iter().map(Into::into).collect()); @@ -132,7 +132,7 @@ pub async fn cmd_get_workspace_meta( host: H, req: CmdGetWorkspaceMetaReq, ) -> Result { - let db = host.db(); + let db = host.db()?; let workspace = db.get_workspace(&req.workspace_id)?; Ok(db.get_or_create_workspace_meta(&workspace.id)?) } @@ -141,14 +141,16 @@ 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 + .db()? + .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.db()?.delete_all_http_responses_for_request(&req.request_id, &host.update_source())?; Ok(()) } @@ -157,7 +159,7 @@ pub async fn cmd_ws_delete_connections( req: CmdWsDeleteConnectionsReq, ) -> Result<()> { Ok(host - .db() + .db()? .delete_all_websocket_connections_for_request(&req.request_id, &host.update_source())?) } diff --git a/crates/yaak-commands/src/plugins.rs b/crates/yaak-commands/src/plugins.rs index 0971f930a..149332558 100644 --- a/crates/yaak-commands/src/plugins.rs +++ b/crates/yaak-commands/src/plugins.rs @@ -10,7 +10,7 @@ pub async fn cmd_plugin_info( host: H, req: CmdPluginInfoReq, ) -> Result { - let plugin = host.db().get_plugin(&req.id)?; + let plugin = host.db()?.get_plugin(&req.id)?; if let Some(metadata) = host.loaded_plugin_metadata(&plugin.directory).await { return Ok(metadata); } diff --git a/crates/yaak-commands/src/render.rs b/crates/yaak-commands/src/render.rs index 230fc7ac2..0be82316a 100644 --- a/crates/yaak-commands/src/render.rs +++ b/crates/yaak-commands/src/render.rs @@ -64,7 +64,7 @@ pub(crate) async fn render_form_values( }; let environment_chain = - host.db().resolve_environments(&workspace_id, folder_id.as_deref(), environment_id)?; + host.db()?.resolve_environments(&workspace_id, folder_id.as_deref(), environment_id)?; let cb = host.template_callback(purpose).await?; let rendered = diff --git a/crates/yaak-commands/src/responses.rs b/crates/yaak-commands/src/responses.rs index 6602b6152..af3953269 100644 --- a/crates/yaak-commands/src/responses.rs +++ b/crates/yaak-commands/src/responses.rs @@ -41,7 +41,7 @@ pub async fn cmd_get_http_response_events( host: H, req: CmdGetHttpResponseEventsReq, ) -> Result> { - let events: Vec = host.db().list_http_response_events(&req.response_id)?; + let events: Vec = host.db()?.list_http_response_events(&req.response_id)?; Ok(events) } @@ -55,7 +55,7 @@ pub async fn cmd_http_response_body_path( host: H, req: CmdHttpResponseBodyPathReq, ) -> Result> { - let location = locate_response_body(&host.db(), &req.response_id)?; + let location = locate_response_body(&host.db()?, &req.response_id)?; Ok(location.path.map(|p| p.to_string_lossy().to_string())) } @@ -64,7 +64,7 @@ pub async fn cmd_http_request_body( req: CmdHttpRequestBodyReq, ) -> Result>> { let body_id = format!("{}.request", req.response_id); - let chunks = host.blobs().get_chunks(&body_id)?; + let chunks = host.blobs()?.get_chunks(&body_id)?; if chunks.is_empty() { return Ok(None); @@ -76,7 +76,7 @@ pub async fn cmd_http_request_body( } pub async fn cmd_save_response(host: H, req: CmdSaveResponseReq) -> Result<()> { - let response = host.db().get_http_response(&req.response_id)?; + let response = host.db()?.get_http_response(&req.response_id)?; let body_path = response.body_path.ok_or(Error::Generic("Response does not have a body".to_string()))?; diff --git a/crates/yaak-commands/src/templates.rs b/crates/yaak-commands/src/templates.rs index 9044c8934..e3a8c2db5 100644 --- a/crates/yaak-commands/src/templates.rs +++ b/crates/yaak-commands/src/templates.rs @@ -20,7 +20,7 @@ pub async fn cmd_render_template( req: CmdRenderTemplateReq, ) -> Result { let environment_chain = - host.db().resolve_environments(&req.workspace_id, None, req.environment_id.as_deref())?; + host.db()?.resolve_environments(&req.workspace_id, None, req.environment_id.as_deref())?; let cb = host.template_callback(req.purpose.unwrap_or(RenderPurpose::Preview)).await?; let options = RenderOptions { // A preview that throws would show the user an error where they expect diff --git a/crates/yaak-commands/tests/test_host.rs b/crates/yaak-commands/tests/test_host.rs index 15440c8c7..b95af7b3e 100644 --- a/crates/yaak-commands/tests/test_host.rs +++ b/crates/yaak-commands/tests/test_host.rs @@ -141,13 +141,13 @@ async fn writes_carry_the_client_id() { // Deletes cascade inside a transaction; make sure that path works with no // host doing anything special around it. - let workspace = host.db().get_workspace(&id).expect("get workspace"); + let workspace = host.db().unwrap().get_workspace(&id).expect("get workspace"); let deleted = models_delete(host.clone(), ModelsDeleteReq { model: AnyModel::Workspace(workspace) }) .await .expect("delete"); assert_eq!(deleted, id); - assert!(host.db().get_workspace(&id).is_err(), "workspace should be gone"); + assert!(host.db().unwrap().get_workspace(&id).is_err(), "workspace should be gone"); } #[tokio::test] @@ -407,6 +407,7 @@ async fn a_single_threaded_host_can_implement_the_trait() { // variable is what proves the chain was resolved rather than skipped. let environment = host .db() + .unwrap() .upsert_environment( &Environment { workspace_id: id.clone(), @@ -439,7 +440,7 @@ async fn a_single_threaded_host_can_implement_the_trait() { // The delete path too, since it is the one that used to reach for a // blocking thread this host does not have. - let workspace = host.db().get_workspace(&id).expect("get workspace"); + let workspace = host.db().unwrap().get_workspace(&id).expect("get workspace"); let deleted = models_delete(host, ModelsDeleteReq { model: AnyModel::Workspace(workspace) }) .await .expect("delete"); @@ -460,12 +461,14 @@ async fn auth_values_are_rendered_before_the_host_sees_them() { let workspace = host .db() + .unwrap() .upsert_workspace( &Workspace { name: "Auth".to_string(), ..Default::default() }, &host.update_source(), ) .expect("workspace"); host.db() + .unwrap() .upsert_environment( &Environment { workspace_id: workspace.id.clone(), @@ -482,7 +485,7 @@ async fn auth_values_are_rendered_before_the_host_sees_them() { ) .expect("environment"); let environment = - host.db().list_environments_ensure_base(&workspace.id).expect("list").remove(0); + host.db().unwrap().list_environments_ensure_base(&workspace.id).expect("list").remove(0); let mut values = HashMap::new(); values.insert("password".to_string(), JsonPrimitive::String("${[ token ]}".to_string())); @@ -522,12 +525,14 @@ async fn template_function_values_are_rendered_before_the_host_sees_them() { let workspace = host .db() + .unwrap() .upsert_workspace( &Workspace { name: "Functions".to_string(), ..Default::default() }, &host.update_source(), ) .expect("workspace"); host.db() + .unwrap() .upsert_environment( &Environment { workspace_id: workspace.id.clone(), @@ -544,7 +549,7 @@ async fn template_function_values_are_rendered_before_the_host_sees_them() { ) .expect("environment"); let environment = - host.db().list_environments_ensure_base(&workspace.id).expect("list").remove(0); + host.db().unwrap().list_environments_ensure_base(&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 3eb63e956..a4ded202e 100644 --- a/crates/yaak-crypto/src/manager.rs +++ b/crates/yaak-crypto/src/manager.rs @@ -51,7 +51,7 @@ impl EncryptionManager { pub fn set_human_key(&self, workspace_id: &str, human_key: &str) -> Result { let wkey = WorkspaceKey::from_human(human_key)?; - let workspace = self.query_manager.connect().get_workspace(workspace_id)?; + let workspace = self.query_manager.connect()?.get_workspace(workspace_id)?; let encryption_key_challenge = match workspace.encryption_key_challenge { None => return self.set_workspace_key(workspace_id, &wkey), Some(c) => c, @@ -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.connect()?.get_or_create_workspace_meta(workspace_id)?; // Already exists if let Some(_) = workspace_meta.encryption_key { @@ -152,8 +152,10 @@ impl EncryptionManager { } }; - let db = self.query_manager.connect(); - let workspace_meta = db.get_or_create_workspace_meta(workspace_id)?; + // Scoped so the pooled connection is back before `get_master_key` below, which + // can block on the OS keyring for as long as it likes. + let workspace_meta = + self.query_manager.connect()?.get_or_create_workspace_meta(workspace_id)?; let key = match workspace_meta.encryption_key { None => return Err(MissingWorkspaceKey), diff --git a/crates/yaak-lifecycle/src/lib.rs b/crates/yaak-lifecycle/src/lib.rs index 33ad68f72..303195b39 100644 --- a/crates/yaak-lifecycle/src/lib.rs +++ b/crates/yaak-lifecycle/src/lib.rs @@ -77,7 +77,7 @@ 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 db = query_manager.connect().unwrap(); let source = &UpdateSource::Background; let workspace = db @@ -117,14 +117,14 @@ mod tests { #[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 db = query_manager.connect().unwrap(); { - let blob_ctx = blob_manager.connect(); + let blob_ctx = blob_manager.connect().unwrap(); blob_ctx.insert_chunk(&BodyChunk::new("rs_gone", 0, b"dead".to_vec())).unwrap(); } on_launch(&Host::owner(), &db, &blob_manager).unwrap(); - assert!(!blob_manager.connect().body_exists("rs_gone").unwrap()); + assert!(!blob_manager.connect().unwrap().body_exists("rs_gone").unwrap()); } } diff --git a/crates/yaak-models/Cargo.toml b/crates/yaak-models/Cargo.toml index e9cff8cbf..126c6588c 100644 --- a/crates/yaak-models/Cargo.toml +++ b/crates/yaak-models/Cargo.toml @@ -27,3 +27,6 @@ yaak-templates = { path = "../yaak-templates", default-features = false } [target.'cfg(not(target_arch = "wasm32"))'.dependencies] r2d2 = "0.8.10" r2d2_sqlite = { version = "0.32" } + +[dev-dependencies] +tempfile = "3" diff --git a/crates/yaak-models/src/blob_manager.rs b/crates/yaak-models/src/blob_manager.rs index 782432278..29f607d53 100644 --- a/crates/yaak-models/src/blob_manager.rs +++ b/crates/yaak-models/src/blob_manager.rs @@ -36,9 +36,9 @@ impl BlobManager { Self { pool } } - pub fn connect(&self) -> BlobContext { - let conn = self.pool.get().expect("Failed to get blob DB connection from pool"); - BlobContext { conn } + pub fn connect(&self) -> Result { + let conn = self.pool.get()?; + Ok(BlobContext { conn }) } } @@ -208,7 +208,7 @@ mod tests { fn test_insert_and_get_chunks() { let pool = create_test_pool(); let manager = BlobManager::new(pool); - let ctx = manager.connect(); + let ctx = manager.connect().unwrap(); let body_id = "rs_test123.request"; let chunk1 = BodyChunk::new(body_id, 0, b"Hello, ".to_vec()); @@ -229,7 +229,7 @@ mod tests { fn test_get_chunks_ordered_by_index() { let pool = create_test_pool(); let manager = BlobManager::new(pool); - let ctx = manager.connect(); + let ctx = manager.connect().unwrap(); let body_id = "rs_test123.request"; @@ -249,7 +249,7 @@ mod tests { fn test_delete_chunks() { let pool = create_test_pool(); let manager = BlobManager::new(pool); - let ctx = manager.connect(); + let ctx = manager.connect().unwrap(); let body_id = "rs_test123.request"; ctx.insert_chunk(&BodyChunk::new(body_id, 0, b"data".to_vec())).unwrap(); @@ -266,7 +266,7 @@ mod tests { fn test_delete_chunks_like() { let pool = create_test_pool(); let manager = BlobManager::new(pool); - let ctx = manager.connect(); + let ctx = manager.connect().unwrap(); // Insert chunks for same response but different body types ctx.insert_chunk(&BodyChunk::new("rs_abc.request", 0, b"req".to_vec())).unwrap(); @@ -288,7 +288,7 @@ mod tests { fn test_get_body_size() { let pool = create_test_pool(); let manager = BlobManager::new(pool); - let ctx = manager.connect(); + let ctx = manager.connect().unwrap(); let body_id = "rs_test123.request"; ctx.insert_chunk(&BodyChunk::new(body_id, 0, b"Hello".to_vec())).unwrap(); @@ -302,7 +302,7 @@ mod tests { fn test_get_body_size_empty() { let pool = create_test_pool(); let manager = BlobManager::new(pool); - let ctx = manager.connect(); + let ctx = manager.connect().unwrap(); let size = ctx.get_body_size("nonexistent").unwrap(); assert_eq!(size, 0); @@ -312,7 +312,7 @@ mod tests { fn test_body_exists() { let pool = create_test_pool(); let manager = BlobManager::new(pool); - let ctx = manager.connect(); + let ctx = manager.connect().unwrap(); assert!(!ctx.body_exists("rs_test.request").unwrap()); @@ -325,7 +325,7 @@ mod tests { fn test_multiple_bodies_isolated() { let pool = create_test_pool(); let manager = BlobManager::new(pool); - let ctx = manager.connect(); + let ctx = manager.connect().unwrap(); ctx.insert_chunk(&BodyChunk::new("body1", 0, b"data1".to_vec())).unwrap(); ctx.insert_chunk(&BodyChunk::new("body2", 0, b"data2".to_vec())).unwrap(); @@ -343,7 +343,7 @@ mod tests { fn test_large_chunk() { let pool = create_test_pool(); let manager = BlobManager::new(pool); - let ctx = manager.connect(); + let ctx = manager.connect().unwrap(); // 1MB chunk let large_data: Vec = (0..1024 * 1024).map(|i| (i % 256) as u8).collect(); diff --git a/crates/yaak-models/src/queries/http_requests.rs b/crates/yaak-models/src/queries/http_requests.rs index 18fe91714..8b8c9ee65 100644 --- a/crates/yaak-models/src/queries/http_requests.rs +++ b/crates/yaak-models/src/queries/http_requests.rs @@ -190,7 +190,7 @@ mod tests { #[test] fn request_resolution_preserves_duplicate_request_headers() { let (query_manager, _blob_manager, _rx) = init_in_memory().expect("Failed to init DB"); - let db = query_manager.connect(); + let db = query_manager.connect().unwrap(); let workspace = db.list_workspaces().expect("Failed to list workspaces").remove(0); let request = HttpRequest { workspace_id: workspace.id, @@ -225,7 +225,7 @@ 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 db = query_manager.connect().unwrap(); let workspace = db .upsert_workspace( diff --git a/crates/yaak-models/src/queries/http_responses.rs b/crates/yaak-models/src/queries/http_responses.rs index 2edc299d6..479113f7d 100644 --- a/crates/yaak-models/src/queries/http_responses.rs +++ b/crates/yaak-models/src/queries/http_responses.rs @@ -57,7 +57,7 @@ impl<'a> ClientDb<'a> { pub fn delete_orphaned_response_body_blobs(&self, blobs: &BlobManager) -> Result { let mut deleted = 0; - let blob_ctx = blobs.connect(); + 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() { @@ -138,7 +138,7 @@ impl<'a> ClientDb<'a> { } // Delete request body blobs (pattern: {response_id}.request) - let blob_ctx = blob_manager.connect(); + let blob_ctx = blob_manager.connect()?; let body_id = format!("{}.request", http_response.id); if let Err(e) = blob_ctx.delete_chunks(&body_id) { error!("Failed to delete request body blobs: {}", e); @@ -224,14 +224,14 @@ 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 db = query_manager.connect().unwrap(); let live = seed_live_response(&db, &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(); + let blob_ctx = blob_manager.connect().unwrap(); 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())) @@ -245,7 +245,7 @@ mod tests { .expect("Failed to GC response body blobs"); assert_eq!(deleted, 2); - let blob_ctx = blob_manager.connect(); + let blob_ctx = blob_manager.connect().unwrap(); assert!(blob_ctx.body_exists(&live.id).unwrap()); assert!(blob_ctx.body_exists(&live_request_body_id).unwrap()); assert!(!blob_ctx.body_exists("rs_gone").unwrap()); @@ -255,14 +255,14 @@ 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 db = query_manager.connect().unwrap(); let live = seed_live_response(&db, &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(); + let blob_ctx = blob_manager.connect().unwrap(); 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(); } @@ -278,7 +278,7 @@ mod tests { assert_eq!(deleted, 2); // Live data survives, orphans are gone - let blob_ctx = blob_manager.connect(); + let blob_ctx = blob_manager.connect().unwrap(); assert!(blob_ctx.body_exists(&live_body_id).unwrap()); assert!(!blob_ctx.body_exists("rs_gone.request").unwrap()); assert!(dir.join(&live.id).exists()); diff --git a/crates/yaak-models/src/queries/model_changes.rs b/crates/yaak-models/src/queries/model_changes.rs index 987118ad4..0a508b677 100644 --- a/crates/yaak-models/src/queries/model_changes.rs +++ b/crates/yaak-models/src/queries/model_changes.rs @@ -104,7 +104,7 @@ mod tests { #[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(); + let db = query_manager.connect().unwrap(); let workspace = db .upsert_workspace( @@ -146,7 +146,7 @@ mod tests { #[test] fn prunes_old_model_changes() { let (query_manager, _blob_manager, _rx) = init_in_memory().expect("Failed to init DB"); - let db = query_manager.connect(); + let db = query_manager.connect().unwrap(); db.upsert_workspace( &Workspace { @@ -179,7 +179,7 @@ mod tests { #[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"); - let db = query_manager.connect(); + let db = query_manager.connect().unwrap(); let workspace = db .upsert_workspace( @@ -213,7 +213,7 @@ mod tests { #[test] fn prunes_old_model_changes_by_hours() { let (query_manager, _blob_manager, _rx) = init_in_memory().expect("Failed to init DB"); - let db = query_manager.connect(); + let db = query_manager.connect().unwrap(); db.upsert_workspace( &Workspace { @@ -245,7 +245,7 @@ mod tests { #[test] fn list_model_changes_deserializes_http_response_event_payload() { let (query_manager, _blob_manager, _rx) = init_in_memory().expect("Failed to init DB"); - let db = query_manager.connect(); + let db = query_manager.connect().unwrap(); let payload = json!({ "model": { diff --git a/crates/yaak-models/src/queries/workspaces.rs b/crates/yaak-models/src/queries/workspaces.rs index 3b9934a1f..7de48abd4 100644 --- a/crates/yaak-models/src/queries/workspaces.rs +++ b/crates/yaak-models/src/queries/workspaces.rs @@ -109,7 +109,7 @@ impl<'a> ClientDb<'a> { // Best-effort cleanup of response bodies (disk files and blob chunks). // Failures only orphan unreferenced data, and are logged. - let blob_ctx = blobs.connect(); + 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) { @@ -215,7 +215,7 @@ mod tests { #[test] fn bootstraps_first_workspace_with_real_defaults() { let (query_manager, _blob_manager, _rx) = init_in_memory().expect("Failed to init DB"); - let db = query_manager.connect(); + let db = query_manager.connect().unwrap(); let workspaces = db.list_workspaces().expect("Failed to list workspaces"); let workspace = workspaces.first().expect("No workspace was bootstrapped"); diff --git a/crates/yaak-models/src/query_manager.rs b/crates/yaak-models/src/query_manager.rs index 8b2c86b86..fb6e7b242 100644 --- a/crates/yaak-models/src/query_manager.rs +++ b/crates/yaak-models/src/query_manager.rs @@ -19,22 +19,10 @@ impl QueryManager { QueryManager { pool, events_tx } } - pub fn connect(&self) -> ClientDb<'_> { - let conn = self.pool.get().expect("Failed to get a new DB connection from the pool"); + pub fn connect(&self) -> crate::error::Result> { + let conn = self.pool.get()?; 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) + Ok(ClientDb::new(ctx, self.events_tx.clone())) } pub fn with_tx( @@ -44,11 +32,11 @@ impl QueryManager { where E: From, { - let conn = self.pool.get().expect("Failed to get new DB connection from the pool"); + let conn = self.pool.get().map_err(crate::error::Error::from)?; // `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(|e| GenericError(format!("Failed to start DB transaction: {e:?}")))?; let ctx = DbContext::new(ConnectionOrTx::Transaction(&tx)); let db = ClientDb::new(ctx, self.events_tx.clone()); @@ -67,3 +55,96 @@ impl QueryManager { } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::error::Error; + use crate::migrate::migrate_db; + use crate::models::Workspace; + use crate::util::UpdateSource; + use r2d2_sqlite::SqliteConnectionManager; + use std::time::Duration; + + /// A one-connection pool that gives up quickly, so exhaustion is reachable in a test + /// without waiting out the app's acquire timeout. + fn one_connection_pool( + dir: &std::path::Path, + ) -> (QueryManager, mpsc::Receiver) { + let manager = SqliteConnectionManager::file(dir.join("db.sqlite")); + let pool = r2d2::Pool::builder() + .max_size(1) + .connection_timeout(Duration::from_millis(100)) + .build(manager) + .expect("build pool"); + migrate_db(&pool).expect("migrate"); + let (tx, rx) = mpsc::channel(); + (QueryManager::new(pool, tx), rx) + } + + fn workspace(name: &str) -> Workspace { + Workspace { name: name.to_string(), ..Default::default() } + } + + #[test] + fn exhausted_pool_reports_an_error_instead_of_panicking() { + let dir = tempfile::tempdir().expect("tempdir"); + let (qm, _events) = one_connection_pool(dir.path()); + + let held = qm.connect().expect("first connection"); + + match qm.connect() { + Ok(_) => panic!("pool should be exhausted"), + Err(e) => assert!(matches!(e, Error::SqlPoolError(_)), "got {e:?}"), + } + + let err = qm + .with_tx::<(), Error>(|db| { + db.upsert_workspace(&workspace("never runs"), &UpdateSource::Background)?; + Ok(()) + }) + .expect_err("pool is exhausted"); + assert!(matches!(err, Error::SqlPoolError(_)), "got {err:?}"); + + drop(held); + qm.connect().expect("pool recovers once the connection is back"); + } + + #[test] + fn with_tx_returns_its_connection_on_every_exit() { + let dir = tempfile::tempdir().expect("tempdir"); + let (qm, _events) = one_connection_pool(dir.path()); + + qm.with_tx::<(), Error>(|db| { + db.upsert_workspace(&workspace("committed"), &UpdateSource::Background)?; + Ok(()) + }) + .expect("commit"); + + let rolled_back = qm.with_tx::<(), Error>(|db| { + db.upsert_workspace(&workspace("rolled back"), &UpdateSource::Background)?; + Err(Error::Unknown) + }); + assert!(rolled_back.is_err()); + + let panicked = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + qm.with_tx::<(), Error>(|db| { + db.upsert_workspace(&workspace("panicked"), &UpdateSource::Background)?; + panic!("closure blew up"); + }) + })); + assert!(panicked.is_err()); + + // The pool holds exactly one connection, so this only succeeds if all three + // transactions above gave theirs back. + let names = qm + .connect() + .expect("connection is back in the pool") + .list_workspaces() + .expect("list") + .into_iter() + .map(|w| w.name) + .collect::>(); + assert_eq!(names, vec!["committed".to_string()]); + } +} diff --git a/crates/yaak-plugins/src/install.rs b/crates/yaak-plugins/src/install.rs index c97117ada..4aa3cfdf1 100644 --- a/crates/yaak-plugins/src/install.rs +++ b/crates/yaak-plugins/src/install.rs @@ -26,7 +26,7 @@ pub async fn delete_and_uninstall( }; // Scope the db connection so it doesn't live across await let plugin = { - let db = query_manager.connect(); + let db = query_manager.connect()?; db.delete_plugin_by_id(plugin_id, &update_source)? }; if let Err(err) = plugin_manager.uninstall(plugin_context, plugin.directory.as_str()).await { @@ -74,7 +74,7 @@ pub async fn download_and_install( // Scope the db connection so it doesn't live across await let plugin = { - let db = query_manager.connect(); + let db = query_manager.connect()?; db.upsert_plugin( &Plugin { id: plugin_version.id.clone(), diff --git a/crates/yaak-plugins/src/manager.rs b/crates/yaak-plugins/src/manager.rs index 2e05dd8ba..374a445c8 100644 --- a/crates/yaak-plugins/src/manager.rs +++ b/crates/yaak-plugins/src/manager.rs @@ -189,7 +189,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 db = query_manager.connect()?; for dir in &bundled_dirs { if db.get_plugin_by_directory(dir).is_none() { db.upsert_plugin( diff --git a/crates/yaak-wasm/src/lib.rs b/crates/yaak-wasm/src/lib.rs index 45704162c..86140e4ea 100644 --- a/crates/yaak-wasm/src/lib.rs +++ b/crates/yaak-wasm/src/lib.rs @@ -104,7 +104,9 @@ 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) = + yaak_lifecycle::on_launch(&lifecycle_host(), &queries.connect().map_err(js_error)?, &blobs) + { web_sys::console::warn_2(&"on_launch hook failed".into(), &js_error(e)); } @@ -254,7 +256,7 @@ fn dispatch( // because that is what the desktop returns and what the store parses. "models_workspace_models" => { let req: WorkspaceModelsReq = from_js(payload)?; - let db = host.queries.connect(); + let db = host.queries.connect().map_err(js_error)?; let mut list: Vec = Vec::new(); list.push(db.get_settings().into()); @@ -291,7 +293,7 @@ fn dispatch( "models_upsert" => { let req: ModelReq = from_js(payload)?; - let db = host.queries.connect(); + let db = host.queries.connect().map_err(js_error)?; let id = models_ops::upsert_model(&db, &host.blobs, req.model, source).map_err(js_error)?; to_json(id) @@ -319,11 +321,16 @@ fn dispatch( to_json(id) } - "models_get_settings" => to_json(host.queries.connect().get_settings()), + "models_get_settings" => to_json(host.queries.connect().map_err(js_error)?.get_settings()), "models_get_graphql_introspection" => { let req: RequestIdReq = from_js(payload)?; - to_json(host.queries.connect().get_graphql_introspection(&req.request_id)) + to_json( + host.queries + .connect() + .map_err(js_error)? + .get_graphql_introspection(&req.request_id), + ) } "models_upsert_graphql_introspection" => { @@ -331,6 +338,7 @@ fn dispatch( let saved = host .queries .connect() + .map_err(js_error)? .upsert_graphql_introspection( &req.workspace_id, &req.request_id, @@ -346,7 +354,13 @@ fn dispatch( "web_get_http_request" => { let req: RequestIdReq = from_js(payload)?; - to_json(host.queries.connect().get_http_request(&req.request_id).map_err(js_error)?) + to_json( + host.queries + .connect() + .map_err(js_error)? + .get_http_request(&req.request_id) + .map_err(js_error)?, + ) } "cmd_get_http_response_events" => { @@ -354,6 +368,7 @@ fn dispatch( to_json( host.queries .connect() + .map_err(js_error)? .list_http_response_events(&req.response_id) .map_err(js_error)?, ) @@ -366,7 +381,7 @@ fn dispatch( if req.before == req.after { return to_json(()); } - let db = host.queries.connect(); + let db = host.queries.connect().map_err(js_error)?; 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)?; @@ -378,7 +393,7 @@ 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(); + let db = host.queries.connect().map_err(js_error)?; 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)?; @@ -388,7 +403,7 @@ fn dispatch( "cmd_get_workspace_meta" => { let req: WorkspaceIdReq = from_js(payload)?; - let db = host.queries.connect(); + let db = host.queries.connect().map_err(js_error)?; 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)?) } @@ -397,6 +412,7 @@ fn dispatch( let req: RequestIdReq = from_js(payload)?; host.queries .connect() + .map_err(js_error)? .delete_all_http_responses_for_request(&req.request_id, source) .map_err(js_error)?; to_json(()) @@ -488,7 +504,7 @@ pub async fn prepare_http_send(payload: JsValue) -> Result { // Everything from the database first, then release the host borrow before rendering. let (request, environment_chain, settings, cookie_jar) = with_host(|host| { - let db = host.queries.connect(); + let db = host.queries.connect().map_err(js_error)?; let request = db.get_http_request(&req.request_id).map_err(js_error)?; let environment_chain = db .resolve_environments( @@ -554,7 +570,7 @@ pub async fn prepare_http_send(payload: JsValue) -> Result { #[wasm_bindgen] pub fn blob_get(id: &str) -> Result>> { with_host(|host| { - let chunks = host.blobs.connect().get_chunks(id).map_err(js_error)?; + let chunks = host.blobs.connect().map_err(js_error)?.get_chunks(id).map_err(js_error)?; if chunks.is_empty() { return Ok(None); } @@ -569,7 +585,7 @@ 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(); + let ctx = host.blobs.connect().map_err(js_error)?; 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)?; @@ -580,5 +596,5 @@ pub fn blob_put(id: &str, bytes: &[u8]) -> Result<()> { #[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.connect().map_err(js_error)?.delete_chunks(id).map_err(js_error)) } diff --git a/crates/yaak/src/export.rs b/crates/yaak/src/export.rs index 473f9d56d..be3e5d53a 100644 --- a/crates/yaak/src/export.rs +++ b/crates/yaak/src/export.rs @@ -13,7 +13,7 @@ pub struct ExportDataParams<'a> { } pub fn export_data(params: ExportDataParams<'_>) -> Result<()> { - let db = params.query_manager.connect(); + let db = params.query_manager.connect()?; let export_data = get_workspace_export_resources( &db, params.yaak_version, diff --git a/crates/yaak/src/import.rs b/crates/yaak/src/import.rs index 29841661f..36f04ffa5 100644 --- a/crates/yaak/src/import.rs +++ b/crates/yaak/src/import.rs @@ -97,7 +97,7 @@ pub fn plan_import_resources( // Two workspaces of the same name are indistinguishable in the sidebar, and importing // a document a second time is the usual way to end up with a pair. let mut taken = query_manager - .connect() + .connect()? .list_workspaces()? .into_iter() .map(|w| w.name) @@ -115,7 +115,7 @@ pub fn plan_import_resources( workspace_ids.insert(source.id.clone(), workspace_id.clone()); } if !resources.workspaces.is_empty() { - let destination_workspace = query_manager.connect().get_workspace(workspace_id)?; + let destination_workspace = query_manager.connect()?.get_workspace(workspace_id)?; let skipped_fields = resources .workspaces .iter() @@ -296,7 +296,7 @@ pub fn plan_import_resources( origin, }; merge_with_linked_source(query_manager, &mut plan, &original)?; - warn_if_imported_elsewhere(&query_manager.connect(), &mut plan)?; + warn_if_imported_elsewhere(&query_manager.connect()?, &mut plan)?; Ok(plan) } @@ -708,7 +708,7 @@ fn merge_with_linked_source( plan: &mut ImportPlan, original: &ImportResources, ) -> Result<()> { - let db = query_manager.connect(); + let db = query_manager.connect()?; let incoming_keys: BTreeSet = plan.source_keys.values().cloned().collect(); let linked = match (&plan.origin, &plan.destination) { @@ -1151,7 +1151,7 @@ fn validate_destination( query_manager: &QueryManager, destination: &ImportDestination, ) -> Result<()> { - let db = query_manager.connect(); + let db = query_manager.connect()?; validate_destination_db(&db, destination) } @@ -1591,7 +1591,7 @@ mod tests { ..Default::default() }; { - let db = query_manager.connect(); + let db = query_manager.connect().unwrap(); destination = db .upsert_workspace(&destination, &UpdateSource::Import) .expect("create destination"); @@ -1632,7 +1632,7 @@ mod tests { // Planning performed only reads. { - let db = query_manager.connect(); + let db = query_manager.connect().unwrap(); assert_eq!(db.list_workspaces().expect("list workspaces").len(), 1); assert_eq!(db.list_folders(&destination.id).expect("list folders").len(), 1); assert!(db.list_http_requests(&destination.id).expect("list requests").is_empty()); @@ -1682,7 +1682,7 @@ mod tests { assert_eq!(committed.http_requests.len(), 2); assert_eq!( query_manager - .connect() + .connect().unwrap() .get_workspace(&destination.id) .expect("get destination after commit"), destination @@ -1729,7 +1729,7 @@ mod tests { assert_eq!(third.resources.workspaces[0].name, "Imported (3)"); let names = query_manager - .connect() + .connect().unwrap() .list_workspaces() .expect("list workspaces") .into_iter() @@ -1832,7 +1832,7 @@ mod tests { yaak_models::init_in_memory().expect("initialize database"); let destination = destination_workspace(); query_manager - .connect() + .connect().unwrap() .upsert_workspace(&destination, &UpdateSource::Import) .expect("create destination"); let resources = ImportResources { @@ -1924,7 +1924,7 @@ mod tests { drop(connection); assert!(commit_import_plan(&query_manager, plan).is_err()); - let db = query_manager.connect(); + let db = query_manager.connect().unwrap(); assert!(db.get_workspace(&workspace_id).is_err(), "workspace insert must roll back"); assert!(db.get_environment(&environment_id).is_err(), "environment must not exist"); assert!( @@ -2155,7 +2155,7 @@ mod tests { } fn source_rows(query_manager: &QueryManager, workspace_id: &str) -> Vec { - let db = query_manager.connect(); + let db = query_manager.connect().unwrap(); let sources = db.list_import_sources(workspace_id).expect("list import sources"); assert_eq!(sources.len(), 1, "expected one linked source: {sources:?}"); db.list_import_source_resources(&sources[0].id).expect("list resource rows") @@ -2176,7 +2176,7 @@ mod tests { let workspace_id = committed.workspaces[0].id.clone(); let source = { - let db = query_manager.connect(); + let db = query_manager.connect().unwrap(); let source = db .find_import_source(&workspace_id, "OpenAPI", "/tmp/api.yaml") .expect("query import source") @@ -2215,7 +2215,7 @@ mod tests { assert_eq!(base.parent_model, "workspace"); commit_import_plan(&query_manager, plan).expect("commit re-import"); - let db = query_manager.connect(); + let db = query_manager.connect().unwrap(); 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!( @@ -2241,7 +2241,7 @@ mod tests { .clone(); { - let db = query_manager.connect(); + let db = query_manager.connect().unwrap(); let nested = db .list_http_requests(&workspace_id) .expect("list requests") @@ -2286,7 +2286,7 @@ mod tests { commit_import_plan(&query_manager, plan).expect("commit merge"); - let db = query_manager.connect(); + let db = query_manager.connect().unwrap(); let requests = db.list_http_requests(&workspace_id).expect("list requests"); assert_eq!(requests.len(), 3); assert_eq!( @@ -2322,7 +2322,7 @@ mod tests { commit_import_plan(&query_manager, plan).expect("commit rename"); let requests = - query_manager.connect().list_http_requests(&workspace_id).expect("list requests"); + query_manager.connect().unwrap().list_http_requests(&workspace_id).expect("list requests"); assert_eq!(requests.len(), 2, "rename must not duplicate"); assert!(requests.iter().any(|r| r.id == root_id && r.name == "Root Request Renamed")); } @@ -2342,7 +2342,7 @@ mod tests { .clone(); { - let db = query_manager.connect(); + let db = query_manager.connect().unwrap(); 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 }, @@ -2361,7 +2361,7 @@ mod tests { commit_import_plan(&query_manager, plan).expect("commit keep-mine"); assert_eq!( - query_manager.connect().get_http_request(&root_id).expect("get root").url, + query_manager.connect().unwrap().get_http_request(&root_id).expect("get root").url, "https://example.com/root-local", "keep-mine must not overwrite the local edit" ); @@ -2382,7 +2382,7 @@ mod tests { } commit_import_plan(&query_manager, plan).expect("commit take-source"); assert_eq!( - query_manager.connect().get_http_request(&root_id).expect("get root").url, + query_manager.connect().unwrap().get_http_request(&root_id).expect("get root").url, "https://example.com/root-v3" ); let plan = replan(&query_manager, &workspace_id, resources); @@ -2407,7 +2407,7 @@ mod tests { } commit_import_plan(&query_manager, plan).expect("commit with deselected update"); - let db = query_manager.connect(); + let db = query_manager.connect().unwrap(); let root = db .list_http_requests(&workspace_id) .expect("list requests") @@ -2437,7 +2437,7 @@ mod tests { assert!(!removal.selected, "deletions default to deselected"); commit_import_plan(&query_manager, plan).expect("commit with default selection"); assert_eq!( - query_manager.connect().list_http_requests(&workspace_id).expect("list").len(), + query_manager.connect().unwrap().list_http_requests(&workspace_id).expect("list").len(), 2, "deselected deletion must not delete" ); @@ -2451,7 +2451,7 @@ mod tests { } commit_import_plan(&query_manager, plan).expect("commit with deletion"); assert_eq!( - query_manager.connect().list_http_requests(&workspace_id).expect("list").len(), + query_manager.connect().unwrap().list_http_requests(&workspace_id).expect("list").len(), 1 ); @@ -2478,7 +2478,7 @@ mod tests { .clone(); query_manager - .connect() + .connect().unwrap() .delete_http_request_by_id(&root_id, &UpdateSource::Background) .expect("delete root locally"); @@ -2490,7 +2490,7 @@ mod tests { commit_import_plan(&query_manager, plan).expect("commit re-import"); assert_eq!( - query_manager.connect().list_http_requests(&workspace_id).expect("list").len(), + query_manager.connect().unwrap().list_http_requests(&workspace_id).expect("list").len(), 1, "a locally deleted resource must not come back" ); @@ -2514,7 +2514,7 @@ mod tests { select(&mut plan, "Extra Request", false); commit_import_plan(&query_manager, plan).expect("commit with deselected create"); assert_eq!( - query_manager.connect().list_http_requests(&workspace_id).expect("list").len(), + query_manager.connect().unwrap().list_http_requests(&workspace_id).expect("list").len(), 2, "a deselected create must not be created" ); @@ -2545,7 +2545,7 @@ mod tests { commit_import_plan(&query_manager, plan).expect("commit with restored item"); let extra = query_manager - .connect() + .connect().unwrap() .list_http_requests(&workspace_id) .expect("list") .into_iter() @@ -2615,7 +2615,7 @@ mod tests { let plan = replan(&query_manager, &workspace_id, with_extra_folder(true)); commit_import_plan(&query_manager, plan).expect("commit the default selection"); let root_folder = query_manager - .connect() + .connect().unwrap() .get_http_request(&root_id) .expect("root request survives") .folder_id; @@ -2627,7 +2627,7 @@ mod tests { select(&mut plan, "Root Request", true); commit_import_plan(&query_manager, plan).expect("commit the folder and the move"); - let db = query_manager.connect(); + let db = query_manager.connect().unwrap(); let folder = db .list_folders(&workspace_id) .expect("list folders") @@ -2655,7 +2655,7 @@ mod tests { assert_eq!(item_by_name(&plan, "Root Request").action, ImportPlanAction::Update); commit_import_plan(&query_manager, plan).expect("commit merge from the new path"); - let db = query_manager.connect(); + let db = query_manager.connect().unwrap(); assert_eq!( db.list_http_requests(&workspace_id).expect("list").len(), 2, @@ -2676,7 +2676,7 @@ mod tests { // A second source claiming the same keys leaves nothing to merge into safely. { - let db = query_manager.connect(); + let db = query_manager.connect().unwrap(); let other = db .upsert_import_source( &ImportSource { @@ -2753,7 +2753,7 @@ mod tests { assert_ne!(unrelated.model_id, root_id, "it must not adopt another document's model"); commit_import_plan(&query_manager, plan).expect("commit the unrelated import"); - let db = query_manager.connect(); + let db = query_manager.connect().unwrap(); assert_eq!( db.get_http_request(&root_id).expect("root request survives").url, "https://example.com/root" @@ -2794,7 +2794,7 @@ mod tests { let workspace_id = committed.workspaces[0].id.clone(); { - let db = query_manager.connect(); + let db = query_manager.connect().unwrap(); let sources = db.list_import_sources(&workspace_id).expect("list import sources"); let row = db .list_import_source_resources(&sources[0].id) @@ -2847,7 +2847,7 @@ mod tests { // Opening the request in the editor stamps a row ID onto every header it renders. { - let db = query_manager.connect(); + let db = query_manager.connect().unwrap(); let root = db.get_http_request(&root_id).expect("get root"); let headers = root .headers @@ -2914,7 +2914,7 @@ mod tests { let committed = commit_import_plan(&query_manager, plan).expect("commit import"); let workspace_id = committed.workspaces[0].id.clone(); - let db = query_manager.connect(); + let db = query_manager.connect().unwrap(); assert!(db.list_folders(&workspace_id).expect("list folders").is_empty()); let requests = db.list_http_requests(&workspace_id).expect("list requests"); assert_eq!(requests.len(), 1, "requests inside the skipped folder are skipped too"); @@ -2949,7 +2949,7 @@ fn desktop_style_json_roundtrip_records_source() { let committed = commit_import_plan(&query_manager, plan).expect("commit"); let workspace_id = committed.workspaces[0].id.clone(); let source = query_manager - .connect() + .connect().unwrap() .find_import_source(&workspace_id, "OpenAPI", "/tmp/api.yaml") .expect("query") .expect("source recorded after JSON round-trip"); @@ -2971,7 +2971,7 @@ fn selected_keep_local_reverts_the_local_edit() { .clone(); { - let db = query_manager.connect(); + let db = query_manager.connect().unwrap(); 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 }, @@ -2992,7 +2992,7 @@ fn selected_keep_local_reverts_the_local_edit() { commit_import_plan(&query_manager, plan).expect("commit revert"); assert_eq!( - query_manager.connect().get_http_request(&root_id).expect("get root").url, + query_manager.connect().unwrap().get_http_request(&root_id).expect("get root").url, "https://example.com/root", "selected keep-local must revert to the source version" ); diff --git a/crates/yaak/src/plugin_events.rs b/crates/yaak/src/plugin_events.rs index 1bb8bba9c..dccbc84ee 100644 --- a/crates/yaak/src/plugin_events.rs +++ b/crates/yaak/src/plugin_events.rs @@ -217,20 +217,33 @@ fn build_shared_reply( request: SharedRequest<'_>, context: SharedPluginEventContext<'_>, ) -> InternalEventPayload { + // Each arm takes its own connection and hands it straight back, so an arm that reads a + // response body never holds one while it does. The macro is here so failing to get one + // answers the plugin with an error instead of unwinding the event loop. + macro_rules! db { + () => { + match query_manager.connect() { + Ok(db) => db, + Err(err) => { + return InternalEventPayload::ErrorResponse(ErrorResponse { + error: format!("Failed to get a database connection: {err}"), + }); + } + } + }; + } + match request { SharedRequest::GetKeyValue(req) => { - let value = query_manager - .connect() - .get_plugin_key_value(context.plugin_name, &req.key) - .map(|v| v.value); + let value = db!().get_plugin_key_value(context.plugin_name, &req.key).map(|v| v.value); InternalEventPayload::GetKeyValueResponse(GetKeyValueResponse { value }) } SharedRequest::SetKeyValue(req) => { - query_manager.connect().set_plugin_key_value(context.plugin_name, &req.key, &req.value); + db!().set_plugin_key_value(context.plugin_name, &req.key, &req.value); InternalEventPayload::SetKeyValueResponse(yaak_plugins::events::SetKeyValueResponse {}) } SharedRequest::DeleteKeyValue(req) => { - match query_manager.connect().delete_plugin_key_value(context.plugin_name, &req.key) { + match db!().delete_plugin_key_value(context.plugin_name, &req.key) { Ok(deleted) => { InternalEventPayload::DeleteKeyValueResponse(DeleteKeyValueResponse { deleted }) } @@ -240,7 +253,7 @@ fn build_shared_reply( } } SharedRequest::GetHttpRequestById(req) => { - let http_request = query_manager.connect().get_http_request(&req.id).ok(); + let http_request = db!().get_http_request(&req.id).ok(); InternalEventPayload::GetHttpRequestByIdResponse(GetHttpRequestByIdResponse { http_request, }) @@ -251,7 +264,7 @@ fn build_shared_reply( error: "workspace_id is required for list_folders_request".to_string(), }); }; - let folders = match query_manager.connect().list_folders(workspace_id) { + let folders = match db!().list_folders(workspace_id) { Ok(folders) => folders, Err(err) => { return InternalEventPayload::ErrorResponse(ErrorResponse { @@ -263,7 +276,7 @@ fn build_shared_reply( } SharedRequest::ListHttpRequests(req) => { let http_requests = if let Some(folder_id) = req.folder_id.as_deref() { - match query_manager.connect().list_http_requests_for_folder_recursive(folder_id) { + match db!().list_http_requests_for_folder_recursive(folder_id) { Ok(http_requests) => http_requests, Err(err) => { return InternalEventPayload::ErrorResponse(ErrorResponse { @@ -279,7 +292,7 @@ fn build_shared_reply( .to_string(), }); }; - match query_manager.connect().list_http_requests(workspace_id) { + match db!().list_http_requests(workspace_id) { Ok(http_requests) => http_requests, Err(err) => { return InternalEventPayload::ErrorResponse(ErrorResponse { @@ -293,8 +306,7 @@ fn build_shared_reply( }) } SharedRequest::FindHttpResponses(req) => { - let http_responses = query_manager - .connect() + let http_responses = db!() .list_http_responses_for_request(&req.request_id, req.limit.map(|l| l as u64)) .unwrap_or_default(); InternalEventPayload::FindHttpResponsesResponse(FindHttpResponsesResponse { @@ -330,29 +342,24 @@ fn build_shared_reply( use AnyModel::*; let model = match &req.model { - HttpRequest(m) => { - match query_manager.connect().upsert_http_request(m, &UpdateSource::Plugin) { - Ok(model) => HttpRequest(model), - Err(err) => { - return InternalEventPayload::ErrorResponse(ErrorResponse { - error: format!("Failed to upsert HTTP request: {err}"), - }); - } + HttpRequest(m) => match db!().upsert_http_request(m, &UpdateSource::Plugin) { + Ok(model) => HttpRequest(model), + Err(err) => { + return InternalEventPayload::ErrorResponse(ErrorResponse { + error: format!("Failed to upsert HTTP request: {err}"), + }); } - } - GrpcRequest(m) => { - match query_manager.connect().upsert_grpc_request(m, &UpdateSource::Plugin) { - Ok(model) => GrpcRequest(model), - Err(err) => { - return InternalEventPayload::ErrorResponse(ErrorResponse { - error: format!("Failed to upsert gRPC request: {err}"), - }); - } + }, + GrpcRequest(m) => match db!().upsert_grpc_request(m, &UpdateSource::Plugin) { + Ok(model) => GrpcRequest(model), + Err(err) => { + return InternalEventPayload::ErrorResponse(ErrorResponse { + error: format!("Failed to upsert gRPC request: {err}"), + }); } } WebsocketRequest(m) => { - match query_manager.connect().upsert_websocket_request(m, &UpdateSource::Plugin) - { + match db!().upsert_websocket_request(m, &UpdateSource::Plugin) { Ok(model) => WebsocketRequest(model), Err(err) => { return InternalEventPayload::ErrorResponse(ErrorResponse { @@ -361,34 +368,28 @@ fn build_shared_reply( } } } - Folder(m) => { - match query_manager.connect().upsert_folder(m, &UpdateSource::Plugin) { - Ok(model) => Folder(model), - Err(err) => { - return InternalEventPayload::ErrorResponse(ErrorResponse { - error: format!("Failed to upsert folder: {err}"), - }); - } + Folder(m) => match db!().upsert_folder(m, &UpdateSource::Plugin) { + Ok(model) => Folder(model), + Err(err) => { + return InternalEventPayload::ErrorResponse(ErrorResponse { + error: format!("Failed to upsert folder: {err}"), + }); } - } - Environment(m) => { - match query_manager.connect().upsert_environment(m, &UpdateSource::Plugin) { - Ok(model) => Environment(model), - Err(err) => { - return InternalEventPayload::ErrorResponse(ErrorResponse { - error: format!("Failed to upsert environment: {err}"), - }); - } + }, + Environment(m) => match db!().upsert_environment(m, &UpdateSource::Plugin) { + Ok(model) => Environment(model), + Err(err) => { + return InternalEventPayload::ErrorResponse(ErrorResponse { + error: format!("Failed to upsert environment: {err}"), + }); } - } - Workspace(m) => { - match query_manager.connect().upsert_workspace(m, &UpdateSource::Plugin) { - Ok(model) => Workspace(model), - Err(err) => { - return InternalEventPayload::ErrorResponse(ErrorResponse { - error: format!("Failed to upsert workspace: {err}"), - }); - } + }, + Workspace(m) => match db!().upsert_workspace(m, &UpdateSource::Plugin) { + Ok(model) => Workspace(model), + Err(err) => { + return InternalEventPayload::ErrorResponse(ErrorResponse { + error: format!("Failed to upsert workspace: {err}"), + }); } } _ => { @@ -403,10 +404,7 @@ fn build_shared_reply( SharedRequest::DeleteModel(req) => { let model = match req.model.as_str() { "http_request" => { - match query_manager - .connect() - .delete_http_request_by_id(&req.id, &UpdateSource::Plugin) - { + match db!().delete_http_request_by_id(&req.id, &UpdateSource::Plugin) { Ok(model) => AnyModel::HttpRequest(model), Err(err) => { return InternalEventPayload::ErrorResponse(ErrorResponse { @@ -416,10 +414,7 @@ fn build_shared_reply( } } "grpc_request" => { - match query_manager - .connect() - .delete_grpc_request_by_id(&req.id, &UpdateSource::Plugin) - { + match db!().delete_grpc_request_by_id(&req.id, &UpdateSource::Plugin) { Ok(model) => AnyModel::GrpcRequest(model), Err(err) => { return InternalEventPayload::ErrorResponse(ErrorResponse { @@ -429,10 +424,7 @@ fn build_shared_reply( } } "websocket_request" => { - match query_manager - .connect() - .delete_websocket_request_by_id(&req.id, &UpdateSource::Plugin) - { + match db!().delete_websocket_request_by_id(&req.id, &UpdateSource::Plugin) { Ok(model) => AnyModel::WebsocketRequest(model), Err(err) => { return InternalEventPayload::ErrorResponse(ErrorResponse { @@ -441,10 +433,7 @@ fn build_shared_reply( } } } - "folder" => match query_manager - .connect() - .delete_folder_by_id(&req.id, &UpdateSource::Plugin) - { + "folder" => match db!().delete_folder_by_id(&req.id, &UpdateSource::Plugin) { Ok(model) => AnyModel::Folder(model), Err(err) => { return InternalEventPayload::ErrorResponse(ErrorResponse { @@ -453,10 +442,7 @@ fn build_shared_reply( } }, "environment" => { - match query_manager - .connect() - .delete_environment_by_id(&req.id, &UpdateSource::Plugin) - { + match db!().delete_environment_by_id(&req.id, &UpdateSource::Plugin) { Ok(model) => AnyModel::Environment(model), Err(err) => { return InternalEventPayload::ErrorResponse(ErrorResponse { @@ -509,6 +495,7 @@ mod tests { query_manager .connect() + .unwrap() .upsert_workspace( &Workspace { id: "wk_test".to_string(), @@ -521,6 +508,7 @@ mod tests { query_manager .connect() + .unwrap() .upsert_folder( &Folder { id: "fl_test".to_string(), @@ -534,6 +522,7 @@ mod tests { query_manager .connect() + .unwrap() .upsert_http_request( &HttpRequest { id: "rq_test".to_string(), @@ -739,6 +728,7 @@ mod tests { let existing = query_manager .connect() + .unwrap() .get_http_request("rq_test") .expect("Failed to load seeded request"); let upsert_payload = InternalEventPayload::UpsertModelRequest(UpsertModelRequest { diff --git a/crates/yaak/src/response_body.rs b/crates/yaak/src/response_body.rs index 38c9daab0..56eb116e7 100644 --- a/crates/yaak/src/response_body.rs +++ b/crates/yaak/src/response_body.rs @@ -63,13 +63,13 @@ impl<'a> FileResponseBodyStore<'a> { /// request behind it never reaches the store at all, and its bytes come /// back from the send instead — see `SendHttpRequestResponse::body`. fn body_path(&self, response_id: &str) -> Result> { - Ok(self.query_manager.connect().get_http_response(response_id)?.body_path) + Ok(self.query_manager.connect()?.get_http_response(response_id)?.body_path) } } impl ResponseBodyStore for FileResponseBodyStore<'_> { fn info(&self, response_id: &str) -> Result { - let response = self.query_manager.connect().get_http_response(response_id)?; + let response = self.query_manager.connect()?.get_http_response(response_id)?; let content_type = response .headers @@ -127,6 +127,7 @@ mod tests { query_manager .connect() + .unwrap() .upsert_workspace( &Workspace { id: "wk_test".to_string(), ..Default::default() }, &UpdateSource::Sync, @@ -135,6 +136,7 @@ mod tests { query_manager .connect() + .unwrap() .upsert_http_request( &HttpRequest { id: "rq_test".to_string(), @@ -154,6 +156,7 @@ mod tests { let response = query_manager .connect() + .unwrap() .upsert_http_response( &HttpResponse { workspace_id: "wk_test".to_string(), @@ -207,9 +210,9 @@ mod tests { // Seeded responses default to Initialized: still arriving. assert!(!FileResponseBodyStore::new(&qm).info(&id).unwrap().complete); - let mut response = qm.connect().get_http_response(&id).unwrap(); + let mut response = qm.connect().unwrap().get_http_response(&id).unwrap(); response.state = HttpResponseState::Closed; - qm.connect().update_http_response_if_id(&response, &UpdateSource::Sync).unwrap(); + qm.connect().unwrap().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 cc99144d9..124e3541e 100644 --- a/crates/yaak/src/send.rs +++ b/crates/yaak/src/send.rs @@ -45,6 +45,9 @@ const MAX_AUTH_BODY_BYTES: usize = 10 * 1024 * 1024; #[derive(Debug, Error)] pub enum SendHttpRequestError { + #[error("Failed to acquire a database connection: {0}")] + AcquireDatabase(#[source] yaak_models::error::Error), + #[error("Failed to load request: {0}")] LoadRequest(#[source] yaak_models::error::Error), @@ -409,7 +412,7 @@ pub fn resolve_send_inputs( environment_id: Option<&str>, cookies: Option>, ) -> Result { - let db = query_manager.connect(); + let db = query_manager.connect().map_err(SendHttpRequestError::AcquireDatabase)?; let environment_chain = db .resolve_environments(&request.workspace_id, request.folder_id.as_deref(), environment_id) @@ -442,7 +445,7 @@ pub fn resolve_inherited_request( query_manager: &QueryManager, request: &HttpRequest, ) -> Result { - let db = query_manager.connect(); + let db = query_manager.connect().map_err(SendHttpRequestError::AcquireDatabase)?; let (authentication_type, authentication, auth_context_id) = db .resolve_auth_for_http_request(request) .map_err(SendHttpRequestError::ResolveRequestInheritance)?; @@ -462,6 +465,7 @@ pub async fn send_http_request_by_id_with_plugins( let request = params .query_manager .connect() + .map_err(SendHttpRequestError::AcquireDatabase)? .get_http_request(params.request_id) .map_err(SendHttpRequestError::LoadRequest)?; @@ -544,6 +548,7 @@ pub async fn send_http_request_by_id( let request = params .query_manager .connect() + .map_err(SendHttpRequestError::AcquireDatabase)? .get_http_request(params.request_id) .map_err(SendHttpRequestError::LoadRequest)?; let mut cookie_jar = load_cookie_jar(params.query_manager, params.cookie_jar_id.as_deref())?; @@ -647,6 +652,7 @@ pub async fn send_http_request( response = store .query_manager .connect() + .map_err(SendHttpRequestError::AcquireDatabase)? .upsert_http_response(&response, &store.update_source, store.blob_manager) .map_err(SendHttpRequestError::PersistResponse)?; } else if response.id.is_empty() { @@ -700,8 +706,10 @@ 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 + .connect() + .unwrap() + .upsert_http_response_event(&db_event, update_source) { warn!("Failed to persist HTTP response event: {}", err); } @@ -800,6 +808,7 @@ pub async fn send_http_request( response = store .query_manager .connect() + .map_err(SendHttpRequestError::AcquireDatabase)? .upsert_http_response(&connected_response, &store.update_source, store.blob_manager) .map_err(SendHttpRequestError::PersistResponse)?; } else { @@ -887,6 +896,7 @@ pub async fn send_http_request( response = store .query_manager .connect() + .map_err(SendHttpRequestError::AcquireDatabase)? .upsert_http_response( &progress_response, &store.update_source, @@ -961,6 +971,7 @@ pub async fn send_http_request( response = store .query_manager .connect() + .map_err(SendHttpRequestError::AcquireDatabase)? .upsert_http_response(&final_response, &store.update_source, store.blob_manager) .map_err(SendHttpRequestError::PersistResponse)?; } else { @@ -999,6 +1010,7 @@ pub async fn send_http_request( response = store .query_manager .connect() + .map_err(SendHttpRequestError::AcquireDatabase)? .upsert_http_response(&response, &store.update_source, store.blob_manager) .map_err(SendHttpRequestError::PersistResponse)?; } @@ -1027,7 +1039,7 @@ fn persist_request_body_bytes( return Ok(()); } - let blob_ctx = blob_manager.connect(); + let blob_ctx = blob_manager.connect().map_err(|e| e.to_string())?; let mut offset = 0; let mut chunk_index: i32 = 0; while offset < bytes.len() { @@ -1057,14 +1069,22 @@ 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 + .connect() + .map_err(|e| e.to_string())? + .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 + .connect() + .map_err(|e| e.to_string())? + .insert_chunk(&chunk) + .map_err(|e| e.to_string())?; } Ok(total_bytes) @@ -1087,6 +1107,7 @@ pub fn load_cookie_jar( query_manager .connect() + .map_err(SendHttpRequestError::AcquireDatabase)? .get_cookie_jar(cookie_jar_id) .map(Some) .map_err(SendHttpRequestError::LoadCookieJar) @@ -1115,6 +1136,7 @@ pub fn persist_cookies_after_send( cookie_jar.cookies = cookies; query_manager .connect() + .map_err(SendHttpRequestError::AcquireDatabase)? .upsert_cookie_jar(cookie_jar, &UpdateSource::Background) .map_err(SendHttpRequestError::PersistCookieJar)?; Ok(()) @@ -1210,6 +1232,7 @@ fn persist_response_error( store .query_manager .connect() + .map_err(SendHttpRequestError::AcquireDatabase)? .upsert_http_response( &HttpResponse { state: HttpResponseState::Closed, @@ -1445,6 +1468,7 @@ mod tests { query_manager .connect() + .unwrap() .upsert_workspace( &Workspace { id: "wk_test".to_string(), ..Default::default() }, &UpdateSource::Sync, @@ -1452,6 +1476,7 @@ mod tests { .expect("Failed to seed workspace"); let cookie_jar = query_manager .connect() + .unwrap() .upsert_cookie_jar( &CookieJar { id: "cj_test".to_string(), @@ -1493,8 +1518,11 @@ mod tests { persist_cookies_after_send(&query_manager, Some(&mut cookie_jar), Some(&store)) .expect("Failed to persist cookies"); - let stored = - query_manager.connect().get_cookie_jar("cj_test").expect("Failed to load cookie jar"); + let stored = query_manager + .connect() + .unwrap() + .get_cookie_jar("cj_test") + .expect("Failed to load cookie jar"); assert_eq!(stored.cookies.len(), 1); assert_eq!(stored.cookies[0].name, "session"); assert_eq!(stored.cookies[0].value, "abc123"); @@ -1508,6 +1536,7 @@ mod tests { cookie_jar.cookies = vec![cookie("original")]; cookie_jar = query_manager .connect() + .unwrap() .upsert_cookie_jar(&cookie_jar, &UpdateSource::Sync) .expect("Failed to seed cookies"); let store = CookieStore::from_cookies(cookie_jar.cookies.clone()); @@ -1515,6 +1544,7 @@ mod tests { // Someone else updates the jar while the send is in flight. query_manager .connect() + .unwrap() .upsert_cookie_jar( &CookieJar { cookies: vec![cookie("newer")], ..cookie_jar.clone() }, &UpdateSource::Sync, @@ -1524,8 +1554,11 @@ mod tests { persist_cookies_after_send(&query_manager, Some(&mut cookie_jar), Some(&store)) .expect("Failed to persist cookies"); - let stored = - query_manager.connect().get_cookie_jar("cj_test").expect("Failed to load cookie jar"); + let stored = query_manager + .connect() + .unwrap() + .get_cookie_jar("cj_test") + .expect("Failed to load cookie jar"); assert_eq!(stored.cookies, vec![cookie("newer")]); } }