fix(elasticsearch): pass routing for document writes

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<T, F>(state: &Arc<AppState>, execution_id: Option<String>, future: F) -> Result<T, String>
+where
+    F: Future<Output = Result<T, String>>,
+{
+    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<AppState>>,
+    connection_id: String,
+) -> Result<Vec<String>, String> {
+    dbx_core::document_ops::list_databases_core(&state, &connection_id).await
+}
+
+#[tauri::command]
+pub async fn document_list_collections(
+    state: State<'_, Arc<AppState>>,
+    connection_id: String,
+    database: String,
+) -> Result<Vec<CollectionInfo>, 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<AppState>>,
+    connection_id: String,
+    database: String,
+    collection: String,
+    skip: u64,
+    limit: i64,
+    filter: Option<String>,
+    projection: Option<String>,
+    sort: Option<String>,
+    execution_id: Option<String>,
+) -> Result<MongoDocumentResult, String> {
+    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<AppState>>,
+    connection_id: String,
+    database: String,
+    collection: String,
+    doc_json: String,
+) -> Result<String, String> {
+    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<AppState>>,
+    connection_id: String,
+    database: String,
+    collection: String,
+    id: String,
+    doc_json: String,
+    routing: Option<String>,
+) -> Result<u64, String> {
+    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<AppState>>,
+    connection_id: String,
+    database: String,
+    collection: String,
+    id: String,
+    routing: Option<String>,
+) -> Result<u64, String> {
+    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
+}

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;

src-tauri/src/commands/mongo_cmd.rs

@@ -36,7 +36,7 @@ pub async fn mongo_list_collections(
     state: State<'_, Arc<AppState>>,
     connection_id: String,
     database: String,
-) -> Result<Vec<dbx_core::db::vector_driver::CollectionInfo>, String> {
+) -> Result<Vec<dbx_core::document_ops::CollectionInfo>, 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<String>,
     execution_id: Option<String>,
 ) -> Result<MongoDocumentResult, String> {
-    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<AppState>>,
-    connection_id: String,
-    database: String,
-    collection: String,
-    skip: u64,
-    limit: i64,
-    filter: Option<String>,
-    projection: Option<String>,
-    sort: Option<String>,
-    execution_id: Option<String>,
-) -> Result<MongoDocumentResult, String> {
-    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<String, String> {
-    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<String>,
 ) -> Result<u64, String> {
-    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<String>,
 ) -> Result<u64, String> {
-    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]