From 299a9843f3308ad798346fa2994aaf2baf5b01ab Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Wed, 1 Jul 2026 00:42:42 +0800 Subject: [PATCH] fix(elasticsearch): pass routing for document writes --- .../components/document/DocumentBrowser.vue | 68 ++++- apps/desktop/src/lib/api.ts | 5 + apps/desktop/src/lib/http.ts | 40 ++- apps/desktop/src/lib/tauri.ts | 40 ++- .../dbx-core/src/db/elasticsearch_driver.rs | 144 +++++++++- crates/dbx-core/src/document_ops.rs | 266 ++++++++++++++++++ crates/dbx-core/src/lib.rs | 1 + crates/dbx-core/src/mongo_ops.rs | 239 ++-------------- crates/dbx-web/src/main.rs | 7 +- crates/dbx-web/src/routes/document_store.rs | 195 +++++++++++++ crates/dbx-web/src/routes/mod.rs | 1 + crates/dbx-web/src/routes/mongo.rs | 37 +-- src-tauri/src/commands/document_cmd.rs | 131 +++++++++ src-tauri/src/commands/mod.rs | 1 + src-tauri/src/commands/mongo_cmd.rs | 81 ++---- src-tauri/src/lib.rs | 7 +- 16 files changed, 912 insertions(+), 351 deletions(-) create mode 100644 crates/dbx-core/src/document_ops.rs create mode 100644 crates/dbx-web/src/routes/document_store.rs create mode 100644 src-tauri/src/commands/document_cmd.rs diff --git a/apps/desktop/src/components/document/DocumentBrowser.vue b/apps/desktop/src/components/document/DocumentBrowser.vue index e1c74dc0de..bd689bb753 100644 --- a/apps/desktop/src/components/document/DocumentBrowser.vue +++ b/apps/desktop/src/components/document/DocumentBrowser.vue @@ -244,6 +244,22 @@ function documentIdFromGridValue(value: MongoInputValue | undefined): string | n return id.trim() ? id : null; } +function documentRoutingValue(value: unknown): string | undefined { + if (value === null || value === undefined) return undefined; + const routing = typeof value === "string" ? value : String(value); + const trimmed = routing.trim(); + return trimmed ? trimmed : undefined; +} + +function documentRoutingFromDocument(doc: JsonRecord | undefined): string | undefined { + return documentRoutingValue(doc?._routing); +} + +function documentRoutingFromGridRow(row: MongoInputValue[] | undefined, columns: string[]): string | undefined { + const routingColIdx = columns.indexOf("_routing"); + return routingColIdx >= 0 ? documentRoutingValue(row?.[routingColIdx]) : undefined; +} + async function gridSave(changes: DocumentGridChanges) { const cols = changes.columns; const idColIdx = cols.indexOf("_id"); @@ -258,30 +274,32 @@ async function gridSave(changes: DocumentGridChanges) { if (isEs) { const doc = documents.value[rowIdx]; if (!doc) continue; + const routing = documentRoutingFromDocument(doc); const updated = { ...doc }; for (const [colIdx, newVal] of dirtyCols) { const col = cols[colIdx]; - if (col === "_id") continue; + if (col === "_id" || col === "_routing") continue; if (newVal === null) { delete updated[col]; } else { updated[col] = parseMongoDocumentInputValue(newVal); } } - await api.mongoUpdateDocument(props.connectionId, props.database, props.collection, String(id), JSON.stringify(updated)); + await api.documentUpdateDocument(props.connectionId, props.database, props.collection, String(id), JSON.stringify(updated), routing); continue; } const updateDoc = buildMongoUpdateDocument(dirtyCols, cols); if (Object.keys(updateDoc).length === 0) continue; - await api.mongoUpdateDocument(props.connectionId, props.database, props.collection, String(id), JSON.stringify(updateDoc)); + await api.documentUpdateDocument(props.connectionId, props.database, props.collection, String(id), JSON.stringify(updateDoc)); } for (const rowIdx of changes.deletedRows) { const row = changes.rows[rowIdx]; const id = row?.[idColIdx]; if (id == null) continue; - await api.mongoDeleteDocument(props.connectionId, props.database, props.collection, String(id)); + const routing = isEs ? documentRoutingFromDocument(documents.value[rowIdx]) : undefined; + await api.documentDeleteDocument(props.connectionId, props.database, props.collection, String(id), routing); } for (const newRow of changes.newRows) { @@ -289,13 +307,13 @@ async function gridSave(changes: DocumentGridChanges) { if (isEs) { const id = documentIdFromGridValue(newRow[idColIdx]); if (id) { - await api.mongoUpdateDocument(props.connectionId, props.database, props.collection, id, JSON.stringify(doc)); + await api.documentUpdateDocument(props.connectionId, props.database, props.collection, id, JSON.stringify(doc), documentRoutingFromGridRow(newRow, cols)); } else { - await api.mongoInsertDocument(props.connectionId, props.database, props.collection, JSON.stringify(doc)); + await api.documentInsertDocument(props.connectionId, props.database, props.collection, JSON.stringify(doc)); } continue; } - await api.mongoInsertDocument(props.connectionId, props.database, props.collection, JSON.stringify(doc)); + await api.documentInsertDocument(props.connectionId, props.database, props.collection, JSON.stringify(doc)); } await load(); @@ -318,6 +336,20 @@ function elasticsearchPathIdPreview(id: string): string { return encodeURIComponent(id); } +function elasticsearchRoutingPreview(routing: string | undefined): string { + return routing ? `?routing=${encodeURIComponent(routing)}` : ""; +} + +function buildElasticsearchPartialUpdateDocument(changes: Map, columns: string[]): Record { + const filtered = new Map(); + for (const [colIdx, newVal] of changes) { + const col = columns[colIdx]; + if (col === "_id" || col === "_routing") continue; + filtered.set(colIdx, newVal); + } + return buildMongoUpdateDocument(filtered, columns); +} + async function previewDocumentChanges(changes: DocumentGridChanges): Promise { const { dirtyRows, deletedRows, newRows, columns, rows } = changes; const idColIdx = columns.indexOf("_id"); @@ -330,8 +362,9 @@ async function previewDocumentChanges(changes: DocumentGridChanges): Promise(() => ({ save: gridSave, preview: previewDocumentChanges, supportsInsert: true, + readonlyColumns: documentStoreProvider.value.kind === "elasticsearch" ? ["_routing"] : undefined, })); function stopDocumentLoadingTimer() { @@ -514,7 +549,10 @@ function startNew() { function startEdit() { const doc = selectedDoc.value; if (!doc) return; - editFields.value = Object.entries(doc).map(([name, value]) => createEditNode(name, value, name === "_id", name === "_id")); + editFields.value = Object.entries(doc).map(([name, value]) => { + const readonlyMetadata = name === "_id" || (documentStoreProvider.value.kind === "elasticsearch" && name === "_routing"); + return createEditNode(name, value, readonlyMetadata, readonlyMetadata); + }); isEditing.value = true; isNew.value = false; } @@ -603,7 +641,7 @@ function buildObjectFromNodes(nodes: EditNode[], path: string): JsonRecord { for (const field of nodes) { const name = field.keyName.trim(); - if (!name || (!path && name === "_id")) continue; + if (!name || (!path && (name === "_id" || (documentStoreProvider.value.kind === "elasticsearch" && name === "_routing")))) continue; if (seen.has(name)) throw new Error(t("mongo.duplicateField", { field: name })); seen.add(name); doc[name] = buildValueFromNode(field, path ? `${path}.${name}` : name); @@ -629,7 +667,7 @@ async function saveDoc() { try { const doc = buildDocumentFromFields(); if (isNew.value) { - await api.mongoInsertDocument(props.connectionId, props.database, props.collection, JSON.stringify(doc)); + await api.documentInsertDocument(props.connectionId, props.database, props.collection, JSON.stringify(doc)); } else if (selectedIdx.value !== null) { const current = documents.value[selectedIdx.value]; const id = current?._id; @@ -637,7 +675,7 @@ async function saveDoc() { error.value = "No _id field"; return; } - await api.mongoUpdateDocument(props.connectionId, props.database, props.collection, String(id), JSON.stringify(doc)); + await api.documentUpdateDocument(props.connectionId, props.database, props.collection, String(id), JSON.stringify(doc), documentRoutingFromDocument(current)); } isEditing.value = false; isNew.value = false; @@ -657,7 +695,7 @@ async function applyDeleteDoc(idx: number) { if (!id) return; error.value = ""; try { - await api.mongoDeleteDocument(props.connectionId, props.database, props.collection, String(id)); + await api.documentDeleteDocument(props.connectionId, props.database, props.collection, String(id), documentRoutingFromDocument(doc)); if (selectedIdx.value === idx) { selectedIdx.value = null; editJson.value = ""; diff --git a/apps/desktop/src/lib/api.ts b/apps/desktop/src/lib/api.ts index c8fd6fc04e..61d845146a 100644 --- a/apps/desktop/src/lib/api.ts +++ b/apps/desktop/src/lib/api.ts @@ -366,7 +366,9 @@ export const mqRawRequest = forward("mqRawRequest"); export const mqSendMessage = forward("mqSendMessage"); // MongoDB +export const documentListDatabases = forward("documentListDatabases"); export const mongoListDatabases = forward("mongoListDatabases"); +export const documentListCollections = forward("documentListCollections"); export const mongoListCollections = forward("mongoListCollections"); export const mongoCreateDatabase = forward("mongoCreateDatabase"); export const mongoDropDatabase = forward("mongoDropDatabase"); @@ -376,10 +378,13 @@ export const mongoFindDocuments = forward("mongoFindDocuments"); export const mongoServerVersion = forward("mongoServerVersion"); export const mongoAggregateDocuments = forward("mongoAggregateDocuments"); export const mongoCreateIndex = forward("mongoCreateIndex"); +export const documentInsertDocument = forward("documentInsertDocument"); export const mongoInsertDocument = forward("mongoInsertDocument"); export const mongoInsertDocuments = forward("mongoInsertDocuments"); +export const documentUpdateDocument = forward("documentUpdateDocument"); export const mongoUpdateDocument = forward("mongoUpdateDocument"); export const mongoUpdateDocuments = forward("mongoUpdateDocuments"); +export const documentDeleteDocument = forward("documentDeleteDocument"); export const mongoDeleteDocument = forward("mongoDeleteDocument"); export const mongoDeleteDocuments = forward("mongoDeleteDocuments"); diff --git a/apps/desktop/src/lib/http.ts b/apps/desktop/src/lib/http.ts index 620a52e0d9..6da8fe07ba 100644 --- a/apps/desktop/src/lib/http.ts +++ b/apps/desktop/src/lib/http.ts @@ -1711,12 +1711,20 @@ export async function nacosRawRequest(connectionId: string, req: NacosRawRequest // MongoDB // --------------------------------------------------------------------------- +export async function documentListDatabases(connectionId: string): Promise { + return post("/api/document-store/list-databases", { connectionId }); +} + export async function mongoListDatabases(connectionId: string): Promise { - return post("/api/mongo/list-databases", { connectionId }); + return documentListDatabases(connectionId); +} + +export async function documentListCollections(connectionId: string, database: string): Promise { + return post("/api/document-store/list-collections", { connectionId, database }); } export async function mongoListCollections(connectionId: string, database: string): Promise { - return post("/api/mongo/list-collections", { connectionId, database }); + return documentListCollections(connectionId, database); } export async function mongoCreateDatabase(connectionId: string, database: string): Promise { @@ -1732,16 +1740,16 @@ export async function mongoDropCollection(connectionId: string, database: string } export async function elasticsearchListIndices(connectionId: string): Promise { - const collections = await mongoListCollections(connectionId, "default"); + const collections = await documentListCollections(connectionId, "default"); return collections.map((c) => c.name); } export async function vectorListCollections(connectionId: string, database?: string): Promise { - return mongoListCollections(connectionId, database || "default"); + return documentListCollections(connectionId, database || "default"); } export async function mongoFindDocuments(connectionId: string, database: string, collection: string, skip: number, limit: number, filter?: string, projection?: string, sort?: string, executionId?: string): Promise { - return post("/api/mongo/find-documents", { connectionId, database, collection, skip, limit, filter, projection, sort, executionId }); + return documentFindDocuments(connectionId, database, collection, skip, limit, filter, projection, sort, executionId); } export async function documentFindDocuments(connectionId: string, database: string, collection: string, skip: number, limit: number, filter?: string, projection?: string, sort?: string, executionId?: string): Promise { @@ -1761,23 +1769,35 @@ export async function mongoCreateIndex(connectionId: string, database: string, c } export async function mongoInsertDocument(connectionId: string, database: string, collection: string, docJson: string): Promise { - return post("/api/mongo/insert-document", { connectionId, database, collection, docJson }); + return documentInsertDocument(connectionId, database, collection, docJson); +} + +export async function documentInsertDocument(connectionId: string, database: string, collection: string, docJson: string): Promise { + return post("/api/document-store/insert-document", { connectionId, database, collection, docJson }); } export async function mongoInsertDocuments(connectionId: string, database: string, collection: string, docsJson: string): Promise<{ affected_rows: number }> { return post("/api/mongo/insert-documents", { connectionId, database, collection, docsJson }); } -export async function mongoUpdateDocument(connectionId: string, database: string, collection: string, id: string, docJson: string): Promise { - return post("/api/mongo/update-document", { connectionId, database, collection, id, docJson }); +export async function mongoUpdateDocument(connectionId: string, database: string, collection: string, id: string, docJson: string, routing?: string): Promise { + return documentUpdateDocument(connectionId, database, collection, id, docJson, routing); +} + +export async function documentUpdateDocument(connectionId: string, database: string, collection: string, id: string, docJson: string, routing?: string): Promise { + return post("/api/document-store/update-document", { connectionId, database, collection, id, docJson, routing }); } export async function mongoUpdateDocuments(connectionId: string, database: string, collection: string, filterJson: string, updateJson: string, many: boolean): Promise<{ affected_rows: number }> { return post("/api/mongo/update-documents", { connectionId, database, collection, filterJson, updateJson, many }); } -export async function mongoDeleteDocument(connectionId: string, database: string, collection: string, id: string): Promise { - return post("/api/mongo/delete-document", { connectionId, database, collection, id }); +export async function mongoDeleteDocument(connectionId: string, database: string, collection: string, id: string, routing?: string): Promise { + return documentDeleteDocument(connectionId, database, collection, id, routing); +} + +export async function documentDeleteDocument(connectionId: string, database: string, collection: string, id: string, routing?: string): Promise { + return post("/api/document-store/delete-document", { connectionId, database, collection, id, routing }); } export async function mongoDeleteDocuments(connectionId: string, database: string, collection: string, filterJson: string, many: boolean): Promise<{ affected_rows: number }> { diff --git a/apps/desktop/src/lib/tauri.ts b/apps/desktop/src/lib/tauri.ts index 9dc96f1b2d..299010d97f 100644 --- a/apps/desktop/src/lib/tauri.ts +++ b/apps/desktop/src/lib/tauri.ts @@ -1441,12 +1441,20 @@ export interface MongoDocumentResult { total: number; } +export async function documentListDatabases(connectionId: string): Promise { + return invoke("document_list_databases", { connectionId }); +} + export async function mongoListDatabases(connectionId: string): Promise { - return invoke("mongo_list_databases", { connectionId }); + return documentListDatabases(connectionId); +} + +export async function documentListCollections(connectionId: string, database: string): Promise { + return invoke("document_list_collections", { connectionId, database }); } export async function mongoListCollections(connectionId: string, database: string): Promise { - return invoke("mongo_list_collections", { connectionId, database }); + return documentListCollections(connectionId, database); } export async function mongoCreateDatabase(connectionId: string, database: string): Promise { @@ -1462,16 +1470,16 @@ export async function mongoDropCollection(connectionId: string, database: string } export async function elasticsearchListIndices(connectionId: string): Promise { - const collections = await mongoListCollections(connectionId, "default"); + const collections = await documentListCollections(connectionId, "default"); return collections.map((c) => c.name); } export async function vectorListCollections(connectionId: string, database?: string): Promise { - return mongoListCollections(connectionId, database || "default"); + return documentListCollections(connectionId, database || "default"); } export async function mongoFindDocuments(connectionId: string, database: string, collection: string, skip: number, limit: number, filter?: string, projection?: string, sort?: string, executionId?: string): Promise { - return invoke("mongo_find_documents", { connectionId, database, collection, skip, limit, filter, projection, sort, executionId }); + return documentFindDocuments(connectionId, database, collection, skip, limit, filter, projection, sort, executionId); } export async function documentFindDocuments(connectionId: string, database: string, collection: string, skip: number, limit: number, filter?: string, projection?: string, sort?: string, executionId?: string): Promise { @@ -1491,7 +1499,11 @@ export async function mongoCreateIndex(connectionId: string, database: string, c } export async function mongoInsertDocument(connectionId: string, database: string, collection: string, docJson: string): Promise { - return invoke("mongo_insert_document", { connectionId, database, collection, docJson }); + return documentInsertDocument(connectionId, database, collection, docJson); +} + +export async function documentInsertDocument(connectionId: string, database: string, collection: string, docJson: string): Promise { + return invoke("document_insert_document", { connectionId, database, collection, docJson }); } export async function mongoInsertDocuments(connectionId: string, database: string, collection: string, docsJson: string): Promise<{ affected_rows: number }> { @@ -1499,8 +1511,12 @@ export async function mongoInsertDocuments(connectionId: string, database: strin return { affected_rows: affectedRows }; } -export async function mongoUpdateDocument(connectionId: string, database: string, collection: string, id: string, docJson: string): Promise { - return invoke("mongo_update_document", { connectionId, database, collection, id, docJson }); +export async function mongoUpdateDocument(connectionId: string, database: string, collection: string, id: string, docJson: string, routing?: string): Promise { + return documentUpdateDocument(connectionId, database, collection, id, docJson, routing); +} + +export async function documentUpdateDocument(connectionId: string, database: string, collection: string, id: string, docJson: string, routing?: string): Promise { + return invoke("document_update_document", { connectionId, database, collection, id, docJson, routing }); } export async function mongoUpdateDocuments(connectionId: string, database: string, collection: string, filterJson: string, updateJson: string, many: boolean): Promise<{ affected_rows: number }> { @@ -1515,8 +1531,12 @@ export async function mongoUpdateDocuments(connectionId: string, database: strin return { affected_rows: affectedRows }; } -export async function mongoDeleteDocument(connectionId: string, database: string, collection: string, id: string): Promise { - return invoke("mongo_delete_document", { connectionId, database, collection, id }); +export async function mongoDeleteDocument(connectionId: string, database: string, collection: string, id: string, routing?: string): Promise { + return documentDeleteDocument(connectionId, database, collection, id, routing); +} + +export async function documentDeleteDocument(connectionId: string, database: string, collection: string, id: string, routing?: string): Promise { + return invoke("document_delete_document", { connectionId, database, collection, id, routing }); } export async function mongoDeleteDocuments(connectionId: string, database: string, collection: string, filterJson: string, many: boolean): Promise<{ affected_rows: number }> { diff --git a/crates/dbx-core/src/db/elasticsearch_driver.rs b/crates/dbx-core/src/db/elasticsearch_driver.rs index 97cb8bba09..d9392e9981 100644 --- a/crates/dbx-core/src/db/elasticsearch_driver.rs +++ b/crates/dbx-core/src/db/elasticsearch_driver.rs @@ -27,6 +27,9 @@ const ELASTICSEARCH_PATH_SEGMENT_ENCODE_SET: &AsciiSet = &CONTROLS .add(b'|') .add(b'}'); +const ELASTICSEARCH_QUERY_VALUE_ENCODE_SET: &AsciiSet = + &CONTROLS.add(b' ').add(b'"').add(b'#').add(b'%').add(b'&').add(b'+').add(b'/').add(b'=').add(b'?'); + pub struct EsClient { http: HttpClient, base_url: String, @@ -189,6 +192,19 @@ fn elasticsearch_path_segment(value: &str) -> String { utf8_percent_encode(value, ELASTICSEARCH_PATH_SEGMENT_ENCODE_SET).to_string() } +fn elasticsearch_query_value(value: &str) -> String { + utf8_percent_encode(value, ELASTICSEARCH_QUERY_VALUE_ENCODE_SET).to_string() +} + +fn elasticsearch_document_path(index: &str, id: &str, routing: Option<&str>) -> String { + let base = format!("/{}/_doc/{}", elasticsearch_path_segment(index), elasticsearch_path_segment(id)); + if let Some(routing) = routing.map(str::trim).filter(|value| !value.is_empty()) { + format!("{base}?routing={}&refresh=true", elasticsearch_query_value(routing)) + } else { + format!("{base}?refresh=true") + } +} + fn redact_elasticsearch_url(url: &str) -> String { let Ok(mut parsed) = reqwest::Url::parse(url) else { return url.to_string(); @@ -364,6 +380,8 @@ impl<'de> Deserialize<'de> for HitsTotal { struct SearchHit { #[serde(rename = "_id")] id: String, + #[serde(rename = "_routing")] + routing: Option, #[serde(rename = "_source")] source: serde_json::Value, } @@ -398,6 +416,9 @@ pub async fn find_documents( _ => serde_json::Map::new(), }; doc.insert("_id".to_string(), serde_json::Value::String(hit.id)); + if let Some(routing) = hit.routing { + doc.insert("_routing".to_string(), serde_json::Value::String(routing)); + } serde_json::Value::Object(doc) }) .collect(); @@ -634,10 +655,16 @@ pub async fn insert_document(client: &EsClient, index: &str, doc_json: &str) -> Ok(result["_id"].as_str().unwrap_or("").to_string()) } -pub async fn update_document(client: &EsClient, index: &str, id: &str, doc_json: &str) -> Result { - let doc = elasticsearch_document_body_from_json(doc_json)?; +pub async fn update_document( + client: &EsClient, + index: &str, + id: &str, + doc_json: &str, + routing: Option<&str>, +) -> Result { + let (doc, routing) = elasticsearch_document_body_and_routing_from_json(doc_json, routing)?; - let path = format!("/{}/_doc/{}?refresh=true", elasticsearch_path_segment(index), elasticsearch_path_segment(id)); + let path = elasticsearch_document_path(index, id, routing.as_deref()); let resp = client.put(&path).json(&doc).send().await.map_err(|e| format!("Elasticsearch request failed: {e}"))?; if !resp.status().is_success() { @@ -649,15 +676,37 @@ pub async fn update_document(client: &EsClient, index: &str, id: &str, doc_json: } fn elasticsearch_document_body_from_json(doc_json: &str) -> Result { - let mut doc: serde_json::Value = serde_json::from_str(doc_json).map_err(|e| format!("Invalid JSON: {e}"))?; - if let serde_json::Value::Object(map) = &mut doc { - map.remove("_id"); - } - Ok(doc) + elasticsearch_document_body_and_routing_from_json(doc_json, None).map(|(doc, _)| doc) } -pub async fn delete_document(client: &EsClient, index: &str, id: &str) -> Result { - let path = format!("/{}/_doc/{}?refresh=true", elasticsearch_path_segment(index), elasticsearch_path_segment(id)); +fn elasticsearch_document_body_and_routing_from_json( + doc_json: &str, + routing: Option<&str>, +) -> Result<(serde_json::Value, Option), String> { + let mut doc: serde_json::Value = serde_json::from_str(doc_json).map_err(|e| format!("Invalid JSON: {e}"))?; + let mut routing = routing.map(str::trim).filter(|value| !value.is_empty()).map(str::to_string); + if let serde_json::Value::Object(map) = &mut doc { + map.remove("_id"); + if routing.is_none() { + routing = map.get("_routing").and_then(elasticsearch_routing_from_value); + } + map.remove("_routing"); + } + Ok((doc, routing)) +} + +fn elasticsearch_routing_from_value(value: &serde_json::Value) -> Option { + match value { + serde_json::Value::String(value) => { + let trimmed = value.trim(); + (!trimmed.is_empty()).then(|| trimmed.to_string()) + } + _ => None, + } +} + +pub async fn delete_document(client: &EsClient, index: &str, id: &str, routing: Option<&str>) -> Result { + let path = elasticsearch_document_path(index, id, routing); let resp = client.delete(&path).send().await.map_err(|e| format!("Elasticsearch request failed: {e}"))?; if !resp.status().is_success() { @@ -807,12 +856,15 @@ fn parse_elasticsearch_response( .iter() .map(|hit| { let mut row = serde_json::Map::new(); - row.insert("_id".to_string(), hit.get("_id").cloned().unwrap_or(serde_json::Value::Null)); if let Some(source) = hit.get("_source").and_then(|s| s.as_object()) { for (k, v) in source { row.insert(k.clone(), v.clone()); } } + row.insert("_id".to_string(), hit.get("_id").cloned().unwrap_or(serde_json::Value::Null)); + if let Some(routing) = hit.get("_routing") { + row.insert("_routing".to_string(), routing.clone()); + } for k in row.keys() { if !all_keys.contains(k) { all_keys.push(k.clone()); @@ -1475,6 +1527,15 @@ mod tests { assert_eq!(super::elasticsearch_path_segment("a%b/c"), "a%25b%2Fc"); } + #[test] + fn builds_elasticsearch_document_path_with_routing() { + assert_eq!( + super::elasticsearch_document_path("orders/2026", "a%b/c", Some("tenant/a&b")), + "/orders%2F2026/_doc/a%25b%2Fc?routing=tenant%2Fa%26b&refresh=true" + ); + assert_eq!(super::elasticsearch_document_path("orders", "1", None), "/orders/_doc/1?refresh=true"); + } + #[test] fn elasticsearch_sql_detection_does_not_treat_rest_methods_as_sql() { assert!(super::is_elasticsearch_sql_query("SELECT * FROM index_task_v1")); @@ -1594,13 +1655,72 @@ mod tests { assert_eq!(response.hits.total.value(), 5); } + #[test] + fn parses_elasticsearch_hit_routing_metadata() { + let response: SearchResponse = serde_json::from_value(json!({ + "hits": { + "total": { "value": 1, "relation": "eq" }, + "hits": [ + { "_id": "doc-1", "_routing": "tenant-1", "_source": { "name": "Alice" } } + ] + } + })) + .unwrap(); + + assert_eq!(response.hits.hits[0].routing.as_deref(), Some("tenant-1")); + } + + #[test] + fn parses_search_response_rows_with_routing_metadata() { + let result = super::parse_elasticsearch_response( + 200, + json!({ + "hits": { + "hits": [ + { "_id": "doc-1", "_routing": "tenant-1", "_source": { "name": "Alice" } } + ] + } + }), + std::time::Instant::now(), + ) + .unwrap(); + + let routing_idx = result.columns.iter().position(|column| column == "_routing").unwrap(); + assert_eq!(result.rows[0][routing_idx], json!("tenant-1")); + } + #[test] fn document_body_removes_elasticsearch_id_metadata() { - let doc = super::elasticsearch_document_body_from_json(r#"{"_id":"abc","name":"Alice"}"#).unwrap(); + let doc = super::elasticsearch_document_body_from_json(r#"{"_id":"abc","_routing":"tenant-1","name":"Alice"}"#) + .unwrap(); assert_eq!(doc, json!({ "name": "Alice" })); } + #[test] + fn document_body_extracts_elasticsearch_routing_metadata() { + let (doc, routing) = super::elasticsearch_document_body_and_routing_from_json( + r#"{"_id":"abc","_routing":"tenant-1","name":"Alice"}"#, + None, + ) + .unwrap(); + + assert_eq!(doc, json!({ "name": "Alice" })); + assert_eq!(routing.as_deref(), Some("tenant-1")); + } + + #[test] + fn explicit_elasticsearch_routing_overrides_document_metadata() { + let (doc, routing) = super::elasticsearch_document_body_and_routing_from_json( + r#"{"_id":"abc","_routing":"tenant-1","name":"Alice"}"#, + Some("tenant-2"), + ) + .unwrap(); + + assert_eq!(doc, json!({ "name": "Alice" })); + assert_eq!(routing.as_deref(), Some("tenant-2")); + } + #[test] fn document_body_preserves_user_field_order() { let doc = super::elasticsearch_document_body_from_json(r#"{"z":1,"_id":"abc","a":2}"#).unwrap(); diff --git a/crates/dbx-core/src/document_ops.rs b/crates/dbx-core/src/document_ops.rs new file mode 100644 index 0000000000..f2b2155a16 --- /dev/null +++ b/crates/dbx-core/src/document_ops.rs @@ -0,0 +1,266 @@ +use crate::connection::{AppState, PoolKind}; +use crate::db::agent_driver::mongo_document_id_params; +use crate::db::mongo_driver::MongoDocumentResult; +use crate::db::{elasticsearch_driver, mongo_driver, vector_driver}; + +pub use crate::db::vector_driver::CollectionInfo; + +fn sort_names(mut names: Vec) -> Vec { + names.sort_by(|left, right| { + let left_lower = left.to_lowercase(); + let right_lower = right.to_lowercase(); + left_lower.cmp(&right_lower).then_with(|| left.cmp(right)) + }); + names +} + +async fn ensure_document_pool(state: &AppState, connection_id: &str) -> Result<(), String> { + state.get_or_create_pool(connection_id, None).await.map(|_| ()) +} + +pub async fn list_databases_core(state: &AppState, connection_id: &str) -> Result, String> { + ensure_document_pool(state, connection_id).await?; + let fallback_database = configured_mongo_database(state, connection_id).await; + let connections = state.connections.read().await; + match connections.get(connection_id).ok_or("Not found")? { + PoolKind::MongoDb(client) => match mongo_driver::list_databases(client).await { + Ok(databases) => Ok(sort_names(databases)), + Err(error) if mongo_list_databases_unauthorized(&error) => { + fallback_mongo_database(&error, fallback_database) + } + Err(error) => Err(error), + }, + PoolKind::Elasticsearch(_) | PoolKind::VectorDb(_) => Ok(vec!["default".to_string()]), + PoolKind::Agent(client) => { + let mut client = client.lock().await; + match client.mongo_list_databases::>().await { + Ok(result) => { + Ok(sort_names(result.iter().filter_map(|v| v.get("name")?.as_str().map(String::from)).collect())) + } + Err(error) if mongo_list_databases_unauthorized(&error) => { + fallback_mongo_database(&error, fallback_database) + } + Err(error) => Err(error), + } + } + _ => Err("Not a MongoDB/Elasticsearch/vector connection".to_string()), + } +} + +async fn configured_mongo_database(state: &AppState, connection_id: &str) -> Option { + let configs = state.configs.read().await; + configs.get(connection_id).and_then(|config| config.effective_database().map(str::to_string)) +} + +fn fallback_mongo_database(error: &str, fallback_database: Option) -> Result, String> { + fallback_database.map(|database| vec![database]).ok_or_else(|| error.to_string()) +} + +fn mongo_list_databases_unauthorized(error: &str) -> bool { + let lower = error.to_lowercase(); + lower.contains("not authorized") && lower.contains("listdatabases") +} + +pub async fn list_collections_core( + state: &AppState, + connection_id: &str, + database: &str, +) -> Result, String> { + ensure_document_pool(state, connection_id).await?; + let connections = state.connections.read().await; + match connections.get(connection_id).ok_or("Not found")? { + PoolKind::MongoDb(client) => { + let names = sort_names(mongo_driver::list_collections(client, database).await?); + Ok(names.into_iter().map(|n| CollectionInfo { name: n.clone(), id: n, dimension: None }).collect()) + } + PoolKind::Elasticsearch(client) => { + let names = sort_names(elasticsearch_driver::list_indices(client).await?); + Ok(names.into_iter().map(|n| CollectionInfo { name: n.clone(), id: n, dimension: None }).collect()) + } + PoolKind::VectorDb(client) => vector_driver::list_collections_with_db(&client, database).await, + PoolKind::Agent(client) => { + let mut client = client.lock().await; + let names = sort_names(client.mongo_list_collections(database).await?); + Ok(names.into_iter().map(|n| CollectionInfo { name: n.clone(), id: n, dimension: None }).collect()) + } + _ => Err("Not a MongoDB/Elasticsearch/vector connection".to_string()), + } +} + +#[allow(clippy::too_many_arguments)] +pub async fn find_documents_core( + state: &AppState, + connection_id: &str, + database: &str, + collection: &str, + skip: u64, + limit: i64, + filter: Option<&str>, + projection: Option<&str>, + sort: Option<&str>, +) -> Result { + ensure_document_pool(state, connection_id).await?; + let connections = state.connections.read().await; + match connections.get(connection_id).ok_or("Not found")? { + PoolKind::MongoDb(client) => { + mongo_driver::find_documents(client, database, collection, skip, limit, filter, projection, sort).await + } + PoolKind::Elasticsearch(client) => { + let client = client.clone(); + drop(connections); + elasticsearch_driver::find_documents(&client, collection, skip, limit, filter, sort).await + } + PoolKind::VectorDb(client) => { + let client = client.clone(); + drop(connections); + let _ = (database, filter, sort); + vector_driver::find_documents(&client, collection, skip, limit).await + } + PoolKind::Agent(client) => { + let mut client = client.lock().await; + let mut params = serde_json::json!({ + "database": database, + "collection": collection, + "skip": skip, + "limit": limit, + "filter": filter, + "sort": sort, + }); + if let Some(projection) = projection { + params["projection"] = serde_json::json!(projection); + } + client.mongo_find_documents(params).await + } + _ => Err("Not a MongoDB/Elasticsearch/vector connection".to_string()), + } +} + +pub async fn insert_document_core( + state: &AppState, + connection_id: &str, + database: &str, + collection: &str, + doc_json: &str, +) -> Result { + ensure_document_pool(state, connection_id).await?; + let connections = state.connections.read().await; + match connections.get(connection_id).ok_or("Not found")? { + PoolKind::MongoDb(client) => mongo_driver::insert_document(client, database, collection, doc_json).await, + PoolKind::Elasticsearch(client) => { + let client = client.clone(); + drop(connections); + elasticsearch_driver::insert_document(&client, collection, doc_json).await + } + PoolKind::Agent(client) => { + let mut client = client.lock().await; + let result: serde_json::Value = client + .mongo_insert_document(serde_json::json!({ + "database": database, + "collection": collection, + "doc_json": doc_json, + })) + .await?; + Ok(result.get("inserted_id").and_then(|v| v.as_str()).unwrap_or("").to_string()) + } + _ => Err("Not a MongoDB/Elasticsearch connection".to_string()), + } +} + +pub async fn update_document_core( + state: &AppState, + connection_id: &str, + database: &str, + collection: &str, + id: &str, + doc_json: &str, + routing: Option<&str>, +) -> Result { + ensure_document_pool(state, connection_id).await?; + let connections = state.connections.read().await; + match connections.get(connection_id).ok_or("Not found")? { + PoolKind::MongoDb(client) => mongo_driver::update_document(client, database, collection, id, doc_json).await, + PoolKind::Elasticsearch(client) => { + let client = client.clone(); + drop(connections); + // Elasticsearch requires the same custom routing value for writes + // as was used to index the document. + elasticsearch_driver::update_document(&client, collection, id, doc_json, routing).await + } + PoolKind::Agent(client) => { + let mut client = client.lock().await; + let result: serde_json::Value = client + .mongo_update_document(serde_json::json!({ + "database": database, + "collection": collection, + "id": id, + "doc_json": doc_json, + })) + .await?; + Ok(result.get("modified_count").and_then(|v| v.as_u64()).unwrap_or(0)) + } + _ => Err("Not a MongoDB/Elasticsearch connection".to_string()), + } +} + +pub async fn delete_document_core( + state: &AppState, + connection_id: &str, + database: &str, + collection: &str, + id: &str, + routing: Option<&str>, +) -> Result { + ensure_document_pool(state, connection_id).await?; + let connections = state.connections.read().await; + match connections.get(connection_id).ok_or("Not found")? { + PoolKind::MongoDb(client) => mongo_driver::delete_document(client, database, collection, id).await, + PoolKind::Elasticsearch(client) => { + let client = client.clone(); + drop(connections); + // Elasticsearch requires the same custom routing value for writes + // as was used to index the document. + elasticsearch_driver::delete_document(&client, collection, id, routing).await + } + PoolKind::Agent(client) => { + let mut client = client.lock().await; + let result: serde_json::Value = + client.mongo_delete_document(mongo_document_id_params(database, collection, id)).await?; + Ok(result.get("deleted_count").and_then(|v| v.as_u64()).unwrap_or(0)) + } + _ => Err("Not a MongoDB/Elasticsearch connection".to_string()), + } +} + +#[cfg(test)] +mod tests { + use super::{fallback_mongo_database, mongo_list_databases_unauthorized, sort_names}; + + #[test] + fn sorts_names_case_insensitively() { + let sorted = sort_names(vec![ + "movies".to_string(), + "Comments".to_string(), + "users".to_string(), + "embedded_movies".to_string(), + ]); + + assert_eq!(sorted, vec!["Comments", "embedded_movies", "movies", "users"]); + } + + #[test] + fn detects_mongo_list_databases_unauthorized_errors() { + assert!(mongo_list_databases_unauthorized( + "Command failed with error 13 (Unauthorized): not authorized on admin to execute command { listDatabases: 1 }", + )); + assert!(!mongo_list_databases_unauthorized("not authorized to execute command { find: \"orders\" }")); + } + + #[test] + fn falls_back_to_configured_mongo_database() { + assert_eq!( + fallback_mongo_database("not authorized", Some("app".to_string())).unwrap(), + vec!["app".to_string()], + ); + assert_eq!(fallback_mongo_database("not authorized", None).unwrap_err(), "not authorized"); + } +} diff --git a/crates/dbx-core/src/lib.rs b/crates/dbx-core/src/lib.rs index b303746053..34c29c54d0 100644 --- a/crates/dbx-core/src/lib.rs +++ b/crates/dbx-core/src/lib.rs @@ -21,6 +21,7 @@ pub mod database_export; pub mod database_search_sql; pub mod db; pub mod db_admin_sql; +pub mod document_ops; pub mod driver_runtime; pub mod external; pub mod history; diff --git a/crates/dbx-core/src/mongo_ops.rs b/crates/dbx-core/src/mongo_ops.rs index 8e7b777a2c..615dd4322f 100644 --- a/crates/dbx-core/src/mongo_ops.rs +++ b/crates/dbx-core/src/mongo_ops.rs @@ -1,91 +1,21 @@ use crate::connection::{AppState, PoolKind}; -use crate::db::agent_driver::mongo_document_id_params; -use crate::db::elasticsearch_driver; use crate::db::mongo_driver::{self, MongoDocumentResult}; -use crate::db::vector_driver; - -fn sort_names(mut names: Vec) -> Vec { - names.sort_by(|left, right| { - let left_lower = left.to_lowercase(); - let right_lower = right.to_lowercase(); - left_lower.cmp(&right_lower).then_with(|| left.cmp(right)) - }); - names -} +use crate::document_ops::CollectionInfo; async fn ensure_document_pool(state: &AppState, connection_id: &str) -> Result<(), String> { state.get_or_create_pool(connection_id, None).await.map(|_| ()) } pub async fn mongo_list_databases_core(state: &AppState, connection_id: &str) -> Result, String> { - ensure_document_pool(state, connection_id).await?; - let fallback_database = configured_mongo_database(state, connection_id).await; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { - PoolKind::MongoDb(client) => match mongo_driver::list_databases(client).await { - Ok(databases) => Ok(sort_names(databases)), - Err(error) if mongo_list_databases_unauthorized(&error) => { - fallback_mongo_database(&error, fallback_database) - } - Err(error) => Err(error), - }, - PoolKind::Elasticsearch(_) | PoolKind::VectorDb(_) => Ok(vec!["default".to_string()]), - PoolKind::Agent(client) => { - let mut client = client.lock().await; - match client.mongo_list_databases::>().await { - Ok(result) => { - Ok(sort_names(result.iter().filter_map(|v| v.get("name")?.as_str().map(String::from)).collect())) - } - Err(error) if mongo_list_databases_unauthorized(&error) => { - fallback_mongo_database(&error, fallback_database) - } - Err(error) => Err(error), - } - } - _ => Err("Not a MongoDB/Elasticsearch/vector connection".to_string()), - } + crate::document_ops::list_databases_core(state, connection_id).await } -async fn configured_mongo_database(state: &AppState, connection_id: &str) -> Option { - let configs = state.configs.read().await; - configs.get(connection_id).and_then(|config| config.effective_database().map(str::to_string)) -} - -fn fallback_mongo_database(error: &str, fallback_database: Option) -> Result, String> { - fallback_database.map(|database| vec![database]).ok_or_else(|| error.to_string()) -} - -fn mongo_list_databases_unauthorized(error: &str) -> bool { - let lower = error.to_lowercase(); - lower.contains("not authorized") && lower.contains("listdatabases") -} - -use crate::db::vector_driver::CollectionInfo; - pub async fn mongo_list_collections_core( state: &AppState, connection_id: &str, database: &str, ) -> Result, String> { - ensure_document_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { - PoolKind::MongoDb(client) => { - let names = sort_names(mongo_driver::list_collections(client, database).await?); - Ok(names.into_iter().map(|n| CollectionInfo { name: n.clone(), id: n, dimension: None }).collect()) - } - PoolKind::Elasticsearch(client) => { - let names = sort_names(elasticsearch_driver::list_indices(client).await?); - Ok(names.into_iter().map(|n| CollectionInfo { name: n.clone(), id: n, dimension: None }).collect()) - } - PoolKind::VectorDb(client) => vector_driver::list_collections_with_db(&client, database).await, - PoolKind::Agent(client) => { - let mut client = client.lock().await; - let names = sort_names(client.mongo_list_collections(database).await?); - Ok(names.into_iter().map(|n| CollectionInfo { name: n.clone(), id: n, dimension: None }).collect()) - } - _ => Err("Not a MongoDB/Elasticsearch/vector connection".to_string()), - } + crate::document_ops::list_collections_core(state, connection_id, database).await } pub async fn mongo_create_database_core(state: &AppState, connection_id: &str, database: &str) -> Result<(), String> { @@ -140,54 +70,6 @@ pub async fn mongo_server_version_core( } } -#[allow(clippy::too_many_arguments)] -pub async fn document_find_documents_core( - state: &AppState, - connection_id: &str, - database: &str, - collection: &str, - skip: u64, - limit: i64, - filter: Option<&str>, - projection: Option<&str>, - sort: Option<&str>, -) -> Result { - ensure_document_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { - PoolKind::MongoDb(client) => { - mongo_driver::find_documents(client, database, collection, skip, limit, filter, projection, sort).await - } - PoolKind::Elasticsearch(client) => { - let client = client.clone(); - drop(connections); - elasticsearch_driver::find_documents(&client, collection, skip, limit, filter, sort).await - } - PoolKind::VectorDb(client) => { - let client = client.clone(); - drop(connections); - let _ = (database, filter, sort); - vector_driver::find_documents(&client, collection, skip, limit).await - } - PoolKind::Agent(client) => { - let mut client = client.lock().await; - let mut params = serde_json::json!({ - "database": database, - "collection": collection, - "skip": skip, - "limit": limit, - "filter": filter, - "sort": sort, - }); - if let Some(projection) = projection { - params["projection"] = serde_json::json!(projection); - } - client.mongo_find_documents(params).await - } - _ => Err("Not a MongoDB/Elasticsearch/vector connection".to_string()), - } -} - #[allow(clippy::too_many_arguments)] pub async fn mongo_find_documents_core( state: &AppState, @@ -200,8 +82,18 @@ pub async fn mongo_find_documents_core( projection: Option<&str>, sort: Option<&str>, ) -> Result { - document_find_documents_core(state, connection_id, database, collection, skip, limit, filter, projection, sort) - .await + crate::document_ops::find_documents_core( + state, + connection_id, + database, + collection, + skip, + limit, + filter, + projection, + sort, + ) + .await } pub async fn mongo_aggregate_documents_core( @@ -249,28 +141,7 @@ pub async fn mongo_insert_document_core( collection: &str, doc_json: &str, ) -> Result { - ensure_document_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { - PoolKind::MongoDb(client) => mongo_driver::insert_document(client, database, collection, doc_json).await, - PoolKind::Elasticsearch(client) => { - let client = client.clone(); - drop(connections); - elasticsearch_driver::insert_document(&client, collection, doc_json).await - } - PoolKind::Agent(client) => { - let mut client = client.lock().await; - let result: serde_json::Value = client - .mongo_insert_document(serde_json::json!({ - "database": database, - "collection": collection, - "doc_json": doc_json, - })) - .await?; - Ok(result.get("inserted_id").and_then(|v| v.as_str()).unwrap_or("").to_string()) - } - _ => Err("Not a MongoDB/Elasticsearch connection".to_string()), - } + crate::document_ops::insert_document_core(state, connection_id, database, collection, doc_json).await } pub async fn mongo_insert_documents_core( @@ -296,30 +167,9 @@ pub async fn mongo_update_document_core( collection: &str, id: &str, doc_json: &str, + routing: Option<&str>, ) -> Result { - ensure_document_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { - PoolKind::MongoDb(client) => mongo_driver::update_document(client, database, collection, id, doc_json).await, - PoolKind::Elasticsearch(client) => { - let client = client.clone(); - drop(connections); - elasticsearch_driver::update_document(&client, collection, id, doc_json).await - } - PoolKind::Agent(client) => { - let mut client = client.lock().await; - let result: serde_json::Value = client - .mongo_update_document(serde_json::json!({ - "database": database, - "collection": collection, - "id": id, - "doc_json": doc_json, - })) - .await?; - Ok(result.get("modified_count").and_then(|v| v.as_u64()).unwrap_or(0)) - } - _ => Err("Not a MongoDB/Elasticsearch connection".to_string()), - } + crate::document_ops::update_document_core(state, connection_id, database, collection, id, doc_json, routing).await } pub async fn mongo_update_documents_core( @@ -348,24 +198,9 @@ pub async fn mongo_delete_document_core( database: &str, collection: &str, id: &str, + routing: Option<&str>, ) -> Result { - ensure_document_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { - PoolKind::MongoDb(client) => mongo_driver::delete_document(client, database, collection, id).await, - PoolKind::Elasticsearch(client) => { - let client = client.clone(); - drop(connections); - elasticsearch_driver::delete_document(&client, collection, id).await - } - PoolKind::Agent(client) => { - let mut client = client.lock().await; - let result: serde_json::Value = - client.mongo_delete_document(mongo_document_id_params(database, collection, id)).await?; - Ok(result.get("deleted_count").and_then(|v| v.as_u64()).unwrap_or(0)) - } - _ => Err("Not a MongoDB/Elasticsearch connection".to_string()), - } + crate::document_ops::delete_document_core(state, connection_id, database, collection, id, routing).await } pub async fn mongo_delete_documents_core( @@ -386,37 +221,3 @@ pub async fn mongo_delete_documents_core( _ => Err("Not a MongoDB connection".to_string()), } } - -#[cfg(test)] -mod tests { - use super::{fallback_mongo_database, mongo_list_databases_unauthorized, sort_names}; - - #[test] - fn sorts_names_case_insensitively() { - let sorted = sort_names(vec![ - "movies".to_string(), - "Comments".to_string(), - "users".to_string(), - "embedded_movies".to_string(), - ]); - - assert_eq!(sorted, vec!["Comments", "embedded_movies", "movies", "users"]); - } - - #[test] - fn detects_mongo_list_databases_unauthorized_errors() { - assert!(mongo_list_databases_unauthorized( - "Command failed with error 13 (Unauthorized): not authorized on admin to execute command { listDatabases: 1 }", - )); - assert!(!mongo_list_databases_unauthorized("not authorized to execute command { find: \"orders\" }")); - } - - #[test] - fn falls_back_to_configured_mongo_database() { - assert_eq!( - fallback_mongo_database("not authorized", Some("app".to_string())).unwrap(), - vec!["app".to_string()], - ); - assert_eq!(fallback_mongo_database("not authorized", None).unwrap_err(), "not authorized"); - } -} diff --git a/crates/dbx-web/src/main.rs b/crates/dbx-web/src/main.rs index d25b4f9bcd..72a27ef42b 100644 --- a/crates/dbx-web/src/main.rs +++ b/crates/dbx-web/src/main.rs @@ -444,7 +444,12 @@ async fn main() { .route("/mongo/create-database", post(routes::mongo::create_database)) .route("/mongo/drop-database", post(routes::mongo::drop_database)) .route("/mongo/drop-collection", post(routes::mongo::drop_collection)) - .route("/document-store/find-documents", post(routes::mongo::document_find_documents)) + .route("/document-store/list-databases", post(routes::document_store::list_databases)) + .route("/document-store/list-collections", post(routes::document_store::list_collections)) + .route("/document-store/find-documents", post(routes::document_store::find_documents)) + .route("/document-store/insert-document", post(routes::document_store::insert_document)) + .route("/document-store/update-document", post(routes::document_store::update_document)) + .route("/document-store/delete-document", post(routes::document_store::delete_document)) .route("/mongo/find-documents", post(routes::mongo::find_documents)) .route("/mongo/server-version", post(routes::mongo::server_version)) .route("/mongo/aggregate-documents", post(routes::mongo::aggregate_documents)) diff --git a/crates/dbx-web/src/routes/document_store.rs b/crates/dbx-web/src/routes/document_store.rs new file mode 100644 index 0000000000..177e63b767 --- /dev/null +++ b/crates/dbx-web/src/routes/document_store.rs @@ -0,0 +1,195 @@ +use std::future::Future; +use std::sync::Arc; + +use axum::extract::State; +use axum::Json; +use serde::Deserialize; + +use crate::error::AppError; +use crate::state::WebState; + +async fn run_cancellable(state: &Arc, execution_id: Option, future: F) -> Result +where + F: Future>, +{ + let registered = execution_id + .as_ref() + .filter(|id| !id.trim().is_empty()) + .map(|id| state.app.running_queries.register(id.clone())); + if let Some(query) = registered.as_ref() { + let token = query.token(); + tokio::select! { + biased; + _ = token.cancelled() => Err(AppError(dbx_core::query::canceled_error())), + result = future => result.map_err(AppError), + } + } else { + future.await.map_err(AppError) + } +} + +async fn ensure_writable( + app: &dbx_core::connection::AppState, + connection_id: &str, + action: &str, +) -> Result<(), AppError> { + if let Some(name) = dbx_core::query::connection_readonly_name(app, connection_id).await { + return Err(AppError(format!( + "Read-only mode: connection '{}' has read-only protection enabled. {} blocked.", + name, action + ))); + } + Ok(()) +} + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct DocumentListDatabasesRequest { + pub connection_id: String, +} + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct DocumentListCollectionsRequest { + pub connection_id: String, + pub database: String, +} + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct DocumentFindRequest { + pub connection_id: String, + pub database: String, + pub collection: String, + pub skip: Option, + pub limit: Option, + pub filter: Option, + pub projection: Option, + pub sort: Option, + pub execution_id: Option, +} + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct DocumentInsertRequest { + pub connection_id: String, + pub database: String, + pub collection: String, + pub doc_json: String, +} + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct DocumentUpdateRequest { + pub connection_id: String, + pub database: String, + pub collection: String, + pub id: String, + pub doc_json: String, + pub routing: Option, +} + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct DocumentDeleteRequest { + pub connection_id: String, + pub database: String, + pub collection: String, + pub id: String, + pub routing: Option, +} + +pub async fn list_databases( + State(state): State>, + Json(req): Json, +) -> Result>, AppError> { + let result = dbx_core::document_ops::list_databases_core(&state.app, &req.connection_id).await.map_err(AppError)?; + Ok(Json(result)) +} + +pub async fn list_collections( + State(state): State>, + Json(req): Json, +) -> Result>, AppError> { + let result = dbx_core::document_ops::list_collections_core(&state.app, &req.connection_id, &req.database) + .await + .map_err(AppError)?; + Ok(Json(result)) +} + +pub async fn find_documents( + State(state): State>, + Json(req): Json, +) -> Result, AppError> { + let result = run_cancellable( + &state, + req.execution_id, + dbx_core::document_ops::find_documents_core( + &state.app, + &req.connection_id, + &req.database, + &req.collection, + req.skip.unwrap_or(0), + req.limit.unwrap_or(50), + req.filter.as_deref(), + req.projection.as_deref(), + req.sort.as_deref(), + ), + ) + .await?; + Ok(Json(result)) +} + +pub async fn insert_document( + State(state): State>, + Json(req): Json, +) -> Result, AppError> { + ensure_writable(&state.app, &req.connection_id, "Insert").await?; + let result = dbx_core::document_ops::insert_document_core( + &state.app, + &req.connection_id, + &req.database, + &req.collection, + &req.doc_json, + ) + .await + .map_err(AppError)?; + Ok(Json(result)) +} + +pub async fn update_document( + State(state): State>, + Json(req): Json, +) -> Result, AppError> { + ensure_writable(&state.app, &req.connection_id, "Update").await?; + let result = dbx_core::document_ops::update_document_core( + &state.app, + &req.connection_id, + &req.database, + &req.collection, + &req.id, + &req.doc_json, + req.routing.as_deref(), + ) + .await + .map_err(AppError)?; + Ok(Json(result)) +} + +pub async fn delete_document( + State(state): State>, + Json(req): Json, +) -> Result, AppError> { + ensure_writable(&state.app, &req.connection_id, "Delete").await?; + let result = dbx_core::document_ops::delete_document_core( + &state.app, + &req.connection_id, + &req.database, + &req.collection, + &req.id, + req.routing.as_deref(), + ) + .await + .map_err(AppError)?; + Ok(Json(result)) +} diff --git a/crates/dbx-web/src/routes/mod.rs b/crates/dbx-web/src/routes/mod.rs index 7630231e29..c7a22b97bb 100644 --- a/crates/dbx-web/src/routes/mod.rs +++ b/crates/dbx-web/src/routes/mod.rs @@ -5,6 +5,7 @@ pub mod cloud_sync; pub mod connection; pub mod data_compare; pub mod database_export; +pub mod document_store; pub mod etcd; pub mod history; pub mod jdbc; diff --git a/crates/dbx-web/src/routes/mongo.rs b/crates/dbx-web/src/routes/mongo.rs index 56253abd38..5796cfb835 100644 --- a/crates/dbx-web/src/routes/mongo.rs +++ b/crates/dbx-web/src/routes/mongo.rs @@ -154,6 +154,7 @@ pub struct MongoUpdateRequest { pub collection: String, pub id: String, pub doc_json: String, + pub routing: Option, } #[derive(Deserialize)] @@ -163,6 +164,7 @@ pub struct MongoDeleteRequest { pub database: String, pub collection: String, pub id: String, + pub routing: Option, } pub async fn list_databases( @@ -177,7 +179,7 @@ pub async fn list_databases( pub async fn list_collections( State(state): State>, Json(req): Json, -) -> Result>, AppError> { +) -> Result>, AppError> { let result = dbx_core::mongo_ops::mongo_list_collections_core(&state.app, &req.connection_id, &req.database) .await .map_err(AppError)?; @@ -224,30 +226,7 @@ pub async fn find_documents( let result = run_cancellable( &state, req.execution_id.clone(), - dbx_core::mongo_ops::mongo_find_documents_core( - &state.app, - &req.connection_id, - &req.database, - &req.collection, - req.skip.unwrap_or(0), - req.limit.unwrap_or(50), - req.filter.as_deref(), - req.projection.as_deref(), - req.sort.as_deref(), - ), - ) - .await?; - Ok(Json(serde_json::to_value(result).map_err(|e| AppError(e.to_string()))?)) -} - -pub async fn document_find_documents( - State(state): State>, - Json(req): Json, -) -> Result, AppError> { - let result = run_cancellable( - &state, - req.execution_id.clone(), - dbx_core::mongo_ops::document_find_documents_core( + dbx_core::document_ops::find_documents_core( &state.app, &req.connection_id, &req.database, @@ -319,7 +298,7 @@ pub async fn insert_document( Json(req): Json, ) -> Result, AppError> { ensure_writable(&state.app, &req.connection_id, "Insert").await?; - let result = dbx_core::mongo_ops::mongo_insert_document_core( + let result = dbx_core::document_ops::insert_document_core( &state.app, &req.connection_id, &req.database, @@ -353,13 +332,14 @@ pub async fn update_document( Json(req): Json, ) -> Result, AppError> { ensure_writable(&state.app, &req.connection_id, "Update").await?; - let result = dbx_core::mongo_ops::mongo_update_document_core( + let result = dbx_core::document_ops::update_document_core( &state.app, &req.connection_id, &req.database, &req.collection, &req.id, &req.doc_json, + req.routing.as_deref(), ) .await .map_err(AppError)?; @@ -390,12 +370,13 @@ pub async fn delete_document( Json(req): Json, ) -> Result, AppError> { ensure_writable(&state.app, &req.connection_id, "Delete").await?; - let result = dbx_core::mongo_ops::mongo_delete_document_core( + let result = dbx_core::document_ops::delete_document_core( &state.app, &req.connection_id, &req.database, &req.collection, &req.id, + req.routing.as_deref(), ) .await .map_err(AppError)?; diff --git a/src-tauri/src/commands/document_cmd.rs b/src-tauri/src/commands/document_cmd.rs new file mode 100644 index 0000000000..d7e48bbe93 --- /dev/null +++ b/src-tauri/src/commands/document_cmd.rs @@ -0,0 +1,131 @@ +use std::future::Future; +use std::sync::Arc; +use tauri::State; + +use crate::commands::connection::{ensure_connection_writable, AppState}; +use dbx_core::db::mongo_driver::MongoDocumentResult; +use dbx_core::document_ops::CollectionInfo; + +async fn run_cancellable(state: &Arc, execution_id: Option, future: F) -> Result +where + F: Future>, +{ + let registered_query = + execution_id.as_ref().filter(|id| !id.trim().is_empty()).map(|id| state.running_queries.register(id.clone())); + if let Some(query) = registered_query.as_ref() { + let token = query.token(); + tokio::select! { + biased; + _ = token.cancelled() => Err(dbx_core::query::canceled_error()), + result = future => result, + } + } else { + future.await + } +} + +#[tauri::command] +pub async fn document_list_databases( + state: State<'_, Arc>, + connection_id: String, +) -> Result, String> { + dbx_core::document_ops::list_databases_core(&state, &connection_id).await +} + +#[tauri::command] +pub async fn document_list_collections( + state: State<'_, Arc>, + connection_id: String, + database: String, +) -> Result, String> { + dbx_core::document_ops::list_collections_core(&state, &connection_id, &database).await +} + +#[tauri::command] +#[allow(clippy::too_many_arguments)] +pub async fn document_find_documents( + state: State<'_, Arc>, + connection_id: String, + database: String, + collection: String, + skip: u64, + limit: i64, + filter: Option, + projection: Option, + sort: Option, + execution_id: Option, +) -> Result { + let app = state.inner().clone(); + run_cancellable( + &app, + execution_id, + dbx_core::document_ops::find_documents_core( + &app, + &connection_id, + &database, + &collection, + skip, + limit, + filter.as_deref(), + projection.as_deref(), + sort.as_deref(), + ), + ) + .await +} + +#[tauri::command] +pub async fn document_insert_document( + state: State<'_, Arc>, + connection_id: String, + database: String, + collection: String, + doc_json: String, +) -> Result { + ensure_connection_writable(&state, &connection_id, "Insert").await?; + dbx_core::document_ops::insert_document_core(&state, &connection_id, &database, &collection, &doc_json).await +} + +#[tauri::command] +pub async fn document_update_document( + state: State<'_, Arc>, + connection_id: String, + database: String, + collection: String, + id: String, + doc_json: String, + routing: Option, +) -> Result { + ensure_connection_writable(&state, &connection_id, "Update").await?; + dbx_core::document_ops::update_document_core( + &state, + &connection_id, + &database, + &collection, + &id, + &doc_json, + routing.as_deref(), + ) + .await +} + +#[tauri::command] +pub async fn document_delete_document( + state: State<'_, Arc>, + connection_id: String, + database: String, + collection: String, + id: String, + routing: Option, +) -> Result { + ensure_connection_writable(&state, &connection_id, "Delete").await?; + dbx_core::document_ops::delete_document_core( + &state, + &connection_id, + &database, + &collection, + &id, + routing.as_deref(), + ) + .await +} diff --git a/src-tauri/src/commands/mod.rs b/src-tauri/src/commands/mod.rs index 7f537e8945..ee5fd55ab9 100644 --- a/src-tauri/src/commands/mod.rs +++ b/src-tauri/src/commands/mod.rs @@ -9,6 +9,7 @@ pub mod csv_export; pub mod data_compare; pub mod database_export; pub mod deep_link; +pub mod document_cmd; pub mod etcd_cmd; pub mod external_db; pub mod external_sql; diff --git a/src-tauri/src/commands/mongo_cmd.rs b/src-tauri/src/commands/mongo_cmd.rs index 7c07f11b23..63633b8aaf 100644 --- a/src-tauri/src/commands/mongo_cmd.rs +++ b/src-tauri/src/commands/mongo_cmd.rs @@ -36,7 +36,7 @@ pub async fn mongo_list_collections( state: State<'_, Arc>, connection_id: String, database: String, -) -> Result, String> { +) -> Result, String> { dbx_core::mongo_ops::mongo_list_collections_core(&state, &connection_id, &database).await } @@ -85,54 +85,17 @@ pub async fn mongo_find_documents( sort: Option, execution_id: Option, ) -> Result { - let app = state.inner().clone(); - run_cancellable( - &app, + crate::commands::document_cmd::document_find_documents( + state, + connection_id, + database, + collection, + skip, + limit, + filter, + projection, + sort, execution_id, - dbx_core::mongo_ops::mongo_find_documents_core( - &app, - &connection_id, - &database, - &collection, - skip, - limit, - filter.as_deref(), - projection.as_deref(), - sort.as_deref(), - ), - ) - .await -} - -#[tauri::command] -#[allow(clippy::too_many_arguments)] -pub async fn document_find_documents( - state: State<'_, Arc>, - connection_id: String, - database: String, - collection: String, - skip: u64, - limit: i64, - filter: Option, - projection: Option, - sort: Option, - execution_id: Option, -) -> Result { - let app = state.inner().clone(); - run_cancellable( - &app, - execution_id, - dbx_core::mongo_ops::document_find_documents_core( - &app, - &connection_id, - &database, - &collection, - skip, - limit, - filter.as_deref(), - projection.as_deref(), - sort.as_deref(), - ), ) .await } @@ -205,8 +168,7 @@ pub async fn mongo_insert_document( collection: String, doc_json: String, ) -> Result { - ensure_connection_writable(&state, &connection_id, "Insert").await?; - dbx_core::mongo_ops::mongo_insert_document_core(&state, &connection_id, &database, &collection, &doc_json).await + crate::commands::document_cmd::document_insert_document(state, connection_id, database, collection, doc_json).await } #[tauri::command] @@ -229,10 +191,18 @@ pub async fn mongo_update_document( collection: String, id: String, doc_json: String, + routing: Option, ) -> Result { - ensure_connection_writable(&state, &connection_id, "Update").await?; - dbx_core::mongo_ops::mongo_update_document_core(&state, &connection_id, &database, &collection, &id, &doc_json) - .await + crate::commands::document_cmd::document_update_document( + state, + connection_id, + database, + collection, + id, + doc_json, + routing, + ) + .await } #[tauri::command] @@ -265,9 +235,10 @@ pub async fn mongo_delete_document( database: String, collection: String, id: String, + routing: Option, ) -> Result { - ensure_connection_writable(&state, &connection_id, "Delete").await?; - dbx_core::mongo_ops::mongo_delete_document_core(&state, &connection_id, &database, &collection, &id).await + crate::commands::document_cmd::document_delete_document(state, connection_id, database, collection, id, routing) + .await } #[tauri::command] diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index e54f385010..9d8df4a8f9 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -846,15 +846,20 @@ pub fn run() { commands::mongo_cmd::mongo_create_database, commands::mongo_cmd::mongo_drop_database, commands::mongo_cmd::mongo_drop_collection, - commands::mongo_cmd::document_find_documents, + commands::document_cmd::document_list_databases, + commands::document_cmd::document_list_collections, + commands::document_cmd::document_find_documents, commands::mongo_cmd::mongo_find_documents, commands::mongo_cmd::mongo_server_version, commands::mongo_cmd::mongo_aggregate_documents, commands::mongo_cmd::mongo_create_index, + commands::document_cmd::document_insert_document, commands::mongo_cmd::mongo_insert_document, commands::mongo_cmd::mongo_insert_documents, + commands::document_cmd::document_update_document, commands::mongo_cmd::mongo_update_document, commands::mongo_cmd::mongo_update_documents, + commands::document_cmd::document_delete_document, commands::mongo_cmd::mongo_delete_document, commands::mongo_cmd::mongo_delete_documents, #[cfg(feature = "mq-admin")]