starred/yaak-mountain-loop
1
0
Fork 0
mirror of https://github.com/mountain-loop/yaak.git synced 2026-09-14 13:52:01 +02:00
Code Issues 8 Packages Projects Releases 10 Wiki Activity

Route all database writes through one connection (#642)

Browse Source 
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
This commit is contained in:
Gregory Schier
2026-09-13 08:12:41 -07:00
committed by GitHub
co-authored by Claude Fable 5.1
parent ff7eebf3cd
commit 90cb578e26
62 changed files with 2231 additions and 1852 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/yaak/src/import.rs
+105 -115
View File
@@ -4,7 +4,7 @@ use log::info;
use serde_json::Value;
use sha2::{Digest, Sha256};
use std::collections::{BTreeMap, BTreeSet};
use yaak_models::client_db::ClientDb;
use yaak_models::client_db::{ClientDb, WriteDb};
use yaak_models::models::{
AnyModel, DEFAULT_REQUEST_MESSAGE_SIZE, Environment, Folder, GrpcRequest, HttpRequest,
ImportSource, ImportSourceResource, UpsertModelInfo, WebsocketRequest, Workspace,
@@ -351,7 +351,7 @@ pub fn commit_import_plan(
})
}
fn commit_plan_in_tx(db: &ClientDb, plan: ImportPlan) -> Result<BatchUpsertResult> {
fn commit_plan_in_tx(db: &WriteDb, plan: ImportPlan) -> Result<BatchUpsertResult> {
let items: BTreeMap<String, ImportPlanItem> =
plan.items.iter().map(|item| (item.model_id.clone(), item.clone())).collect();
@@ -460,7 +460,7 @@ fn commit_plan_in_tx(db: &ClientDb, plan: ImportPlan) -> Result<BatchUpsertResul
}
/// A folder deletion may have already cascaded over the model, so absent models are skipped.
fn delete_existing_model(db: &ClientDb, resource: ImportResourceType, id: &str) -> Result<()> {
fn delete_existing_model(db: &WriteDb, resource: ImportResourceType, id: &str) -> Result<()> {
use ImportResourceType::*;
let source = &UpdateSource::Import;
match resource {
@@ -503,7 +503,7 @@ fn delete_existing_model(db: &ClientDb, resource: ImportResourceType, id: &str)
/// offered again next time. A resource the user turned down is remembered as a row without a
/// model, so it is neither re-offered nor resurrected.
fn record_import_source(
db: &ClientDb,
db: &WriteDb,
plan: &ImportPlan,
items: &BTreeMap<String, ImportPlanItem>,
upserted: &BatchUpsertResult,
@@ -1619,32 +1619,30 @@ mod tests {
name: "Selected Folder".to_string(),
..Default::default()
};
{
let db = query_manager.connect();
destination = db
.upsert_workspace(&destination, &UpdateSource::Import)
.expect("create destination");
db.upsert_folder(&selected_folder, &UpdateSource::Import)
.expect("create selected folder");
db.upsert_environment(
&Environment {
id: "ev_destination_base".to_string(),
model: "environment".to_string(),
workspace_id: destination.id.clone(),
name: "Destination Variables".to_string(),
parent_model: "workspace".to_string(),
variables: vec![EnvironmentVariable {
enabled: true,
name: "destination".to_string(),
value: "keep".to_string(),
id: None,
}],
..Default::default()
},
&UpdateSource::Import,
)
.expect("create base environment");
}
query_manager
.with_tx(|db| {
destination = db.upsert_workspace(&destination, &UpdateSource::Import)?;
db.upsert_folder(&selected_folder, &UpdateSource::Import)?;
db.upsert_environment(
&Environment {
id: "ev_destination_base".to_string(),
model: "environment".to_string(),
workspace_id: destination.id.clone(),
name: "Destination Variables".to_string(),
parent_model: "workspace".to_string(),
variables: vec![EnvironmentVariable {
enabled: true,
name: "destination".to_string(),
value: "keep".to_string(),
id: None,
}],
..Default::default()
},
&UpdateSource::Import,
)
})
.expect("seed destination");
let workspace_count = query_manager.connect().list_workspaces().expect("list").len();
let plan = plan_import_resources(
&query_manager,
@@ -1662,13 +1660,10 @@ mod tests {
// Planning performed only reads.
{
let db = query_manager.connect();
assert_eq!(db.list_workspaces().expect("list workspaces").len(), 1);
assert_eq!(db.list_workspaces().expect("list workspaces").len(), workspace_count);
assert_eq!(db.list_folders(&destination.id).expect("list folders").len(), 1);
assert!(db.list_http_requests(&destination.id).expect("list requests").is_empty());
assert_eq!(
db.list_environments_ensure_base(&destination.id).expect("list environments").len(),
1
);
assert_eq!(db.list_environments(&destination.id).expect("list environments").len(), 1);
assert_eq!(db.get_workspace(&destination.id).expect("get destination"), destination);
}
@@ -1861,8 +1856,7 @@ mod tests {
yaak_models::init_in_memory().expect("initialize database");
let destination = destination_workspace();
query_manager
.connect()
.upsert_workspace(&destination, &UpdateSource::Import)
.with_tx(|tx| tx.upsert_workspace(&destination, &UpdateSource::Import))
.expect("create destination");
let resources = ImportResources {
workspaces: vec![
@@ -2247,10 +2241,7 @@ mod tests {
let db = query_manager.connect();
assert_eq!(db.list_http_requests(&workspace_id).expect("list requests").len(), 2);
assert_eq!(db.list_folders(&workspace_id).expect("list folders").len(), 1);
assert_eq!(
db.list_environments_ensure_base(&workspace_id).expect("list environments").len(),
1
);
assert_eq!(db.list_environments(&workspace_id).expect("list environments").len(), 1);
let rows = db.list_import_source_resources(&source.id).expect("list resource rows");
assert_eq!(rows.len(), 4, "re-commit replaces rows instead of accumulating");
}
@@ -2269,23 +2260,22 @@ mod tests {
.id
.clone();
{
let db = query_manager.connect();
let nested = db
.list_http_requests(&workspace_id)
.expect("list requests")
.into_iter()
.find(|r| r.name == "Nested Request")
.expect("nested request");
db.upsert_http_request(
&HttpRequest {
url: "https://example.com/nested-local".to_string(),
..nested.clone()
},
&UpdateSource::Background,
)
query_manager
.with_tx(|db| {
let nested = db
.list_http_requests(&workspace_id)?
.into_iter()
.find(|r| r.name == "Nested Request")
.expect("nested request");
db.upsert_http_request(
&HttpRequest {
url: "https://example.com/nested-local".to_string(),
..nested.clone()
},
&UpdateSource::Background,
)
})
.expect("edit nested request locally");
}
let mut resources = imported_resources();
resources.http_requests[0].url = "https://example.com/root-v2".to_string();
@@ -2370,15 +2360,15 @@ mod tests {
.id
.clone();
{
let db = query_manager.connect();
let root = db.get_http_request(&root_id).expect("get root");
db.upsert_http_request(
&HttpRequest { url: "https://example.com/root-local".to_string(), ..root },
&UpdateSource::Background,
)
query_manager
.with_tx(|db| {
let root = db.get_http_request(&root_id)?;
db.upsert_http_request(
&HttpRequest { url: "https://example.com/root-local".to_string(), ..root },
&UpdateSource::Background,
)
})
.expect("edit root locally");
}
let mut resources = imported_resources();
resources.http_requests[0].url = "https://example.com/root-v2".to_string();
@@ -2507,8 +2497,7 @@ mod tests {
.clone();
query_manager
.connect()
.delete_http_request_by_id(&root_id, &UpdateSource::Background)
.with_tx(|tx| tx.delete_http_request_by_id(&root_id, &UpdateSource::Background))
.expect("delete root locally");
let plan = replan(&query_manager, &workspace_id, imported_resources());
@@ -2703,10 +2692,9 @@ mod tests {
let workspace_id = committed.workspaces[0].id.clone();
// A second source claiming the same keys leaves nothing to merge into safely.
{
let db = query_manager.connect();
let other = db
.upsert_import_source(
query_manager
.with_tx(|db| {
let other = db.upsert_import_source(
&ImportSource {
workspace_id: workspace_id.clone(),
importer: "OpenAPI".to_string(),
@@ -2715,18 +2703,18 @@ mod tests {
..Default::default()
},
&UpdateSource::Import,
)
.expect("create second source");
for key in ["env:base", "folder:src", "op:root", "op:nested"] {
db.upsert_import_source_resource(&ImportSourceResource {
import_source_id: other.id.clone(),
source_key: key.to_string(),
model_type: "http_request".to_string(),
..Default::default()
})
.expect("claim the same keys");
}
}
)?;
for key in ["env:base", "folder:src", "op:root", "op:nested"] {
db.upsert_import_source_resource(&ImportSourceResource {
import_source_id: other.id.clone(),
source_key: key.to_string(),
model_type: "http_request".to_string(),
..Default::default()
})?;
}
Ok::<_, yaak_models::error::Error>(())
})
.expect("claim the same keys");
let third =
ImportOrigin { origin: "/tmp/third.yaml".to_string(), label: "third.yaml".to_string() };
@@ -2821,21 +2809,20 @@ mod tests {
let committed = first_import(&query_manager);
let workspace_id = committed.workspaces[0].id.clone();
{
let db = query_manager.connect();
let sources = db.list_import_sources(&workspace_id).expect("list import sources");
let row = db
.list_import_source_resources(&sources[0].id)
.expect("list rows")
.into_iter()
.find(|r| r.source_key == "op:root")
.expect("row for the root request");
db.upsert_import_source_resource(&ImportSourceResource {
content_hash: Some("v99:from-the-future".to_string()),
..row
query_manager
.with_tx(|db| {
let sources = db.list_import_sources(&workspace_id)?;
let row = db
.list_import_source_resources(&sources[0].id)?
.into_iter()
.find(|r| r.source_key == "op:root")
.expect("row for the root request");
db.upsert_import_source_resource(&ImportSourceResource {
content_hash: Some("v99:from-the-future".to_string()),
..row
})
})
.expect("write an unreadable hash");
}
let plan = replan(&query_manager, &workspace_id, imported_resources());
assert_eq!(
@@ -2874,17 +2861,20 @@ mod tests {
.clone();
// Opening the request in the editor stamps a row ID onto every header it renders.
{
let db = query_manager.connect();
let root = db.get_http_request(&root_id).expect("get root");
let headers = root
.headers
.iter()
.map(|h| HttpRequestHeader { id: Some("hd_generated".to_string()), ..h.clone() })
.collect();
db.upsert_http_request(&HttpRequest { headers, ..root }, &UpdateSource::Background)
.expect("stamp row ids");
}
query_manager
.with_tx(|db| {
let root = db.get_http_request(&root_id)?;
let headers = root
.headers
.iter()
.map(|h| HttpRequestHeader {
id: Some("hd_generated".to_string()),
..h.clone()
})
.collect();
db.upsert_http_request(&HttpRequest { headers, ..root }, &UpdateSource::Background)
})
.expect("stamp row ids");
let plan = replan(&query_manager, &workspace_id, resources);
assert_eq!(
@@ -2997,15 +2987,15 @@ mod tests {
.id
.clone();
{
let db = query_manager.connect();
let root = db.get_http_request(&root_id).expect("get root");
db.upsert_http_request(
&HttpRequest { url: "https://example.com/root-local".to_string(), ..root },
&UpdateSource::Background,
)
query_manager
.with_tx(|db| {
let root = db.get_http_request(&root_id)?;
db.upsert_http_request(
&HttpRequest { url: "https://example.com/root-local".to_string(), ..root },
&UpdateSource::Background,
)
})
.expect("edit root locally");
}
let mut plan = replan(&query_manager, &workspace_id, imported_resources());
let root = item_by_name(&plan, "Root Request");
crates/yaak/src/plugin_events.rs
+63 -59
View File
@@ -1,6 +1,7 @@
use crate::response_body::ResponseBodyStore;
use base64::Engine;
use base64::prelude::BASE64_STANDARD;
use log::warn;
use yaak_models::models::AnyModel;
use yaak_models::query_manager::QueryManager;
use yaak_models::util::UpdateSource;
@@ -226,11 +227,18 @@ fn build_shared_reply(
InternalEventPayload::GetKeyValueResponse(GetKeyValueResponse { value })
}
SharedRequest::SetKeyValue(req) => {
query_manager.connect().set_plugin_key_value(context.plugin_name, &req.key, &req.value);
if let Err(e) = query_manager.with_tx(|tx| {
tx.set_plugin_key_value(context.plugin_name, &req.key, &req.value);
Ok::<(), yaak_models::error::Error>(())
}) {
warn!("Failed to set plugin key value: {e}");
}
InternalEventPayload::SetKeyValueResponse(yaak_plugins::events::SetKeyValueResponse {})
}
SharedRequest::DeleteKeyValue(req) => {
match query_manager.connect().delete_plugin_key_value(context.plugin_name, &req.key) {
match query_manager
.with_tx(|tx| tx.delete_plugin_key_value(context.plugin_name, &req.key))
{
Ok(deleted) => {
InternalEventPayload::DeleteKeyValueResponse(DeleteKeyValueResponse { deleted })
}
@@ -331,7 +339,9 @@ fn build_shared_reply(
let model = match &req.model {
HttpRequest(m) => {
match query_manager.connect().upsert_http_request(m, &UpdateSource::Plugin) {
match query_manager
.with_tx(|tx| tx.upsert_http_request(m, &UpdateSource::Plugin))
{
Ok(model) => HttpRequest(model),
Err(err) => {
return InternalEventPayload::ErrorResponse(ErrorResponse {
@@ -341,7 +351,9 @@ fn build_shared_reply(
}
}
GrpcRequest(m) => {
match query_manager.connect().upsert_grpc_request(m, &UpdateSource::Plugin) {
match query_manager
.with_tx(|tx| tx.upsert_grpc_request(m, &UpdateSource::Plugin))
{
Ok(model) => GrpcRequest(model),
Err(err) => {
return InternalEventPayload::ErrorResponse(ErrorResponse {
@@ -351,7 +363,8 @@ fn build_shared_reply(
}
}
WebsocketRequest(m) => {
match query_manager.connect().upsert_websocket_request(m, &UpdateSource::Plugin)
match query_manager
.with_tx(|tx| tx.upsert_websocket_request(m, &UpdateSource::Plugin))
{
Ok(model) => WebsocketRequest(model),
Err(err) => {
@@ -362,7 +375,7 @@ fn build_shared_reply(
}
}
Folder(m) => {
match query_manager.connect().upsert_folder(m, &UpdateSource::Plugin) {
match query_manager.with_tx(|tx| tx.upsert_folder(m, &UpdateSource::Plugin)) {
Ok(model) => Folder(model),
Err(err) => {
return InternalEventPayload::ErrorResponse(ErrorResponse {
@@ -372,7 +385,9 @@ fn build_shared_reply(
}
}
Environment(m) => {
match query_manager.connect().upsert_environment(m, &UpdateSource::Plugin) {
match query_manager
.with_tx(|tx| tx.upsert_environment(m, &UpdateSource::Plugin))
{
Ok(model) => Environment(model),
Err(err) => {
return InternalEventPayload::ErrorResponse(ErrorResponse {
@@ -382,7 +397,8 @@ fn build_shared_reply(
}
}
Workspace(m) => {
match query_manager.connect().upsert_workspace(m, &UpdateSource::Plugin) {
match query_manager.with_tx(|tx| tx.upsert_workspace(m, &UpdateSource::Plugin))
{
Ok(model) => Workspace(model),
Err(err) => {
return InternalEventPayload::ErrorResponse(ErrorResponse {
@@ -404,8 +420,7 @@ fn build_shared_reply(
let model = match req.model.as_str() {
"http_request" => {
match query_manager
.connect()
.delete_http_request_by_id(&req.id, &UpdateSource::Plugin)
.with_tx(|tx| tx.delete_http_request_by_id(&req.id, &UpdateSource::Plugin))
{
Ok(model) => AnyModel::HttpRequest(model),
Err(err) => {
@@ -417,8 +432,7 @@ fn build_shared_reply(
}
"grpc_request" => {
match query_manager
.connect()
.delete_grpc_request_by_id(&req.id, &UpdateSource::Plugin)
.with_tx(|tx| tx.delete_grpc_request_by_id(&req.id, &UpdateSource::Plugin))
{
Ok(model) => AnyModel::GrpcRequest(model),
Err(err) => {
@@ -429,10 +443,9 @@ fn build_shared_reply(
}
}
"websocket_request" => {
match query_manager
.connect()
.delete_websocket_request_by_id(&req.id, &UpdateSource::Plugin)
{
match query_manager.with_tx(|tx| {
tx.delete_websocket_request_by_id(&req.id, &UpdateSource::Plugin)
}) {
Ok(model) => AnyModel::WebsocketRequest(model),
Err(err) => {
return InternalEventPayload::ErrorResponse(ErrorResponse {
@@ -442,8 +455,7 @@ fn build_shared_reply(
}
}
"folder" => match query_manager
.connect()
.delete_folder_by_id(&req.id, &UpdateSource::Plugin)
.with_tx(|tx| tx.delete_folder_by_id(&req.id, &UpdateSource::Plugin))
{
Ok(model) => AnyModel::Folder(model),
Err(err) => {
@@ -454,8 +466,7 @@ fn build_shared_reply(
},
"environment" => {
match query_manager
.connect()
.delete_environment_by_id(&req.id, &UpdateSource::Plugin)
.with_tx(|tx| tx.delete_environment_by_id(&req.id, &UpdateSource::Plugin))
{
Ok(model) => AnyModel::Environment(model),
Err(err) => {
@@ -508,45 +519,38 @@ mod tests {
yaak_models::init_standalone(&db_path, &blob_path).expect("Failed to initialize DB");
query_manager
.connect()
.upsert_workspace(
&Workspace {
id: "wk_test".to_string(),
name: "Workspace".to_string(),
..Default::default()
},
&UpdateSource::Sync,
)
.expect("Failed to seed workspace");
query_manager
.connect()
.upsert_folder(
&Folder {
id: "fl_test".to_string(),
workspace_id: "wk_test".to_string(),
name: "Folder".to_string(),
..Default::default()
},
&UpdateSource::Sync,
)
.expect("Failed to seed folder");
query_manager
.connect()
.upsert_http_request(
&HttpRequest {
id: "rq_test".to_string(),
workspace_id: "wk_test".to_string(),
folder_id: Some("fl_test".to_string()),
name: "Request".to_string(),
method: "GET".to_string(),
url: "https://example.com".to_string(),
..Default::default()
},
&UpdateSource::Sync,
)
.expect("Failed to seed request");
.with_tx(|tx| {
tx.upsert_workspace(
&Workspace {
id: "wk_test".to_string(),
name: "Workspace".to_string(),
..Default::default()
},
&UpdateSource::Sync,
)?;
tx.upsert_folder(
&Folder {
id: "fl_test".to_string(),
workspace_id: "wk_test".to_string(),
name: "Folder".to_string(),
..Default::default()
},
&UpdateSource::Sync,
)?;
tx.upsert_http_request(
&HttpRequest {
id: "rq_test".to_string(),
workspace_id: "wk_test".to_string(),
folder_id: Some("fl_test".to_string()),
name: "Request".to_string(),
method: "GET".to_string(),
url: "https://example.com".to_string(),
..Default::default()
},
&UpdateSource::Sync,
)
})
.expect("Failed to seed");
(query_manager, temp_dir)
}
crates/yaak/src/response_body.rs
+33 -30
View File
@@ -126,23 +126,25 @@ mod tests {
.unwrap();
query_manager
.connect()
.upsert_workspace(
&Workspace { id: "wk_test".to_string(), ..Default::default() },
&UpdateSource::Sync,
)
.with_tx(|tx| {
tx.upsert_workspace(
&Workspace { id: "wk_test".to_string(), ..Default::default() },
&UpdateSource::Sync,
)
})
.unwrap();
query_manager
.connect()
.upsert_http_request(
&HttpRequest {
id: "rq_test".to_string(),
workspace_id: "wk_test".to_string(),
..Default::default()
},
&UpdateSource::Sync,
)
.with_tx(|tx| {
tx.upsert_http_request(
&HttpRequest {
id: "rq_test".to_string(),
workspace_id: "wk_test".to_string(),
..Default::default()
},
&UpdateSource::Sync,
)
})
.unwrap();
let body_path = body.map(|bytes| {
@@ -153,21 +155,22 @@ mod tests {
});
let response = query_manager
.connect()
.upsert_http_response(
&HttpResponse {
workspace_id: "wk_test".to_string(),
request_id: "rq_test".to_string(),
body_path,
headers: vec![HttpResponseHeader {
name: "Content-Type".to_string(),
value: "application/json; charset=utf-8".to_string(),
}],
..Default::default()
},
&UpdateSource::Sync,
&blob_manager,
)
.with_tx(|tx| {
tx.upsert_http_response(
&HttpResponse {
workspace_id: "wk_test".to_string(),
request_id: "rq_test".to_string(),
body_path,
headers: vec![HttpResponseHeader {
name: "Content-Type".to_string(),
value: "application/json; charset=utf-8".to_string(),
}],
..Default::default()
},
&UpdateSource::Sync,
&blob_manager,
)
})
.unwrap();
let id = response.id.clone();
@@ -209,7 +212,7 @@ mod tests {
let mut response = qm.connect().get_http_response(&id).unwrap();
response.state = HttpResponseState::Closed;
qm.connect().update_http_response_if_id(&response, &UpdateSource::Sync).unwrap();
qm.with_tx(|tx| tx.update_http_response_if_id(&response, &UpdateSource::Sync)).unwrap();
assert!(FileResponseBodyStore::new(&qm).info(&id).unwrap().complete);
}
crates/yaak/src/send.rs
+77 -69
View File
@@ -646,8 +646,9 @@ pub async fn send_http_request<T: TemplateCallback>(
if let Some(store) = store {
response = store
.query_manager
.connect()
.upsert_http_response(&response, &store.update_source, store.blob_manager)
.with_tx(|tx| {
tx.upsert_http_response(&response, &store.update_source, store.blob_manager)
})
.map_err(SendHttpRequestError::PersistResponse)?;
} else if response.id.is_empty() {
response.id = generate_prefixed_id("rs");
@@ -700,8 +701,8 @@ pub async fn send_http_request<T: TemplateCallback>(
&event_workspace_id,
event.clone().into(),
);
if let Err(err) =
query_manager.connect().upsert_http_response_event(&db_event, update_source)
if let Err(err) = query_manager
.with_tx(|tx| tx.upsert_http_response_event(&db_event, update_source))
{
warn!("Failed to persist HTTP response event: {}", err);
}
@@ -799,8 +800,13 @@ pub async fn send_http_request<T: TemplateCallback>(
if let Some(store) = store {
response = store
.query_manager
.connect()
.upsert_http_response(&connected_response, &store.update_source, store.blob_manager)
.with_tx(|tx| {
tx.upsert_http_response(
&connected_response,
&store.update_source,
store.blob_manager,
)
})
.map_err(SendHttpRequestError::PersistResponse)?;
} else {
response = connected_response;
@@ -886,12 +892,13 @@ pub async fn send_http_request<T: TemplateCallback>(
if let Some(store) = store {
response = store
.query_manager
.connect()
.upsert_http_response(
&progress_response,
&store.update_source,
store.blob_manager,
)
.with_tx(|tx| {
tx.upsert_http_response(
&progress_response,
&store.update_source,
store.blob_manager,
)
})
.map_err(SendHttpRequestError::PersistResponse)?;
} else {
response = progress_response;
@@ -960,8 +967,9 @@ pub async fn send_http_request<T: TemplateCallback>(
if let Some(store) = store {
response = store
.query_manager
.connect()
.upsert_http_response(&final_response, &store.update_source, store.blob_manager)
.with_tx(|tx| {
tx.upsert_http_response(&final_response, &store.update_source, store.blob_manager)
})
.map_err(SendHttpRequestError::PersistResponse)?;
} else {
response = final_response;
@@ -998,8 +1006,9 @@ pub async fn send_http_request<T: TemplateCallback>(
if update_response && let Some(store) = store {
response = store
.query_manager
.connect()
.upsert_http_response(&response, &store.update_source, store.blob_manager)
.with_tx(|tx| {
tx.upsert_http_response(&response, &store.update_source, store.blob_manager)
})
.map_err(SendHttpRequestError::PersistResponse)?;
}
}
@@ -1027,17 +1036,14 @@ fn persist_request_body_bytes(
return Ok(());
}
let blob_ctx = blob_manager.connect();
let mut offset = 0;
let mut chunk_index: i32 = 0;
while offset < bytes.len() {
let end = std::cmp::min(offset + REQUEST_BODY_CHUNK_SIZE, bytes.len());
let chunk = BodyChunk::new(body_id, chunk_index, bytes[offset..end].to_vec());
blob_ctx.insert_chunk(&chunk).map_err(|e| e.to_string())?;
chunk_index += 1;
offset = end;
}
Ok(())
blob_manager
.with_tx(|b| {
for (chunk_index, data) in bytes.chunks(REQUEST_BODY_CHUNK_SIZE).enumerate() {
b.insert_chunk(&BodyChunk::new(body_id, chunk_index as i32, data.to_vec()))?;
}
Ok::<_, yaak_models::error::Error>(())
})
.map_err(|e| e.to_string())
}
async fn persist_request_body_stream(
@@ -1057,14 +1063,14 @@ async fn persist_request_body_stream(
while buf.len() >= REQUEST_BODY_CHUNK_SIZE {
let data = buf.drain(..REQUEST_BODY_CHUNK_SIZE).collect();
let chunk = BodyChunk::new(&body_id, chunk_index, data);
blob_manager.connect().insert_chunk(&chunk).map_err(|e| e.to_string())?;
blob_manager.with_tx(|b| b.insert_chunk(&chunk)).map_err(|e| e.to_string())?;
chunk_index += 1;
}
}
if !buf.is_empty() {
let chunk = BodyChunk::new(&body_id, chunk_index, buf);
blob_manager.connect().insert_chunk(&chunk).map_err(|e| e.to_string())?;
blob_manager.with_tx(|b| b.insert_chunk(&chunk)).map_err(|e| e.to_string())?;
}
Ok(total_bytes)
@@ -1114,8 +1120,7 @@ pub fn persist_cookies_after_send(
cookie_jar.cookies = cookies;
query_manager
.connect()
.upsert_cookie_jar(cookie_jar, &UpdateSource::Background)
.with_tx(|tx| tx.upsert_cookie_jar(cookie_jar, &UpdateSource::Background))
.map_err(SendHttpRequestError::PersistCookieJar)?;
Ok(())
}
@@ -1209,23 +1214,24 @@ fn persist_response_error(
let elapsed = duration_to_i32(started_at.elapsed());
store
.query_manager
.connect()
.upsert_http_response(
&HttpResponse {
state: HttpResponseState::Closed,
elapsed,
elapsed_headers: if response.elapsed_headers == 0 {
elapsed
} else {
response.elapsed_headers
.with_tx(|tx| {
tx.upsert_http_response(
&HttpResponse {
state: HttpResponseState::Closed,
elapsed,
elapsed_headers: if response.elapsed_headers == 0 {
elapsed
} else {
response.elapsed_headers
},
error: Some(error),
url: if response.url.is_empty() { fallback_url } else { response.url.clone() },
..response.clone()
},
error: Some(error),
url: if response.url.is_empty() { fallback_url } else { response.url.clone() },
..response.clone()
},
&store.update_source,
store.blob_manager,
)
&store.update_source,
store.blob_manager,
)
})
.map_err(SendHttpRequestError::PersistResponse)
}
@@ -1444,23 +1450,25 @@ mod tests {
.expect("Failed to initialize DB");
query_manager
.connect()
.upsert_workspace(
&Workspace { id: "wk_test".to_string(), ..Default::default() },
&UpdateSource::Sync,
)
.with_tx(|tx| {
tx.upsert_workspace(
&Workspace { id: "wk_test".to_string(), ..Default::default() },
&UpdateSource::Sync,
)
})
.expect("Failed to seed workspace");
let cookie_jar = query_manager
.connect()
.upsert_cookie_jar(
&CookieJar {
id: "cj_test".to_string(),
workspace_id: "wk_test".to_string(),
name: "Default".to_string(),
..Default::default()
},
&UpdateSource::Sync,
)
.with_tx(|tx| {
tx.upsert_cookie_jar(
&CookieJar {
id: "cj_test".to_string(),
workspace_id: "wk_test".to_string(),
name: "Default".to_string(),
..Default::default()
},
&UpdateSource::Sync,
)
})
.expect("Failed to seed cookie jar");
(query_manager, cookie_jar, temp_dir)
@@ -1507,18 +1515,18 @@ mod tests {
let (query_manager, mut cookie_jar, _temp_dir) = seed_cookie_jar();
cookie_jar.cookies = vec![cookie("original")];
cookie_jar = query_manager
.connect()
.upsert_cookie_jar(&cookie_jar, &UpdateSource::Sync)
.with_tx(|tx| tx.upsert_cookie_jar(&cookie_jar, &UpdateSource::Sync))
.expect("Failed to seed cookies");
let store = CookieStore::from_cookies(cookie_jar.cookies.clone());
// Someone else updates the jar while the send is in flight.
query_manager
.connect()
.upsert_cookie_jar(
&CookieJar { cookies: vec![cookie("newer")], ..cookie_jar.clone() },
&UpdateSource::Sync,
)
.with_tx(|tx| {
tx.upsert_cookie_jar(
&CookieJar { cookies: vec![cookie("newer")], ..cookie_jar.clone() },
&UpdateSource::Sync,
)
})
.expect("Failed to update cookie jar");
persist_cookies_after_send(&query_manager, Some(&mut cookie_jar), Some(&store))