From a99a34b4ab0378f584ea47041bf00c6c476e652f Mon Sep 17 00:00:00 2001 From: Guoyu Su Date: Mon, 13 Jul 2026 13:21:50 +0800 Subject: [PATCH] fix(mongodb): use dedicated count for legacy collections --- .../java/com/dbx/agent/AgentProtocol.java | 2 + .../src/main/resources/agent-protocol-v1.json | 1 + .../com/dbx/agent/mongodb/MongoAgent.java | 24 + .../com/dbx/agent/mongodb/MongoAgentTest.java | 12 + apps/desktop/src/lib/backend/api.ts | 1 + apps/desktop/src/lib/backend/http.ts | 4 + apps/desktop/src/lib/backend/tauri.ts | 4 + .../src/lib/mongo/mongoShellCommand.ts | 5 +- apps/desktop/src/stores/queryStore.ts | 6 +- crates/dbx-core/assets/agent-protocol-v1.json | 1 + crates/dbx-core/src/db/agent_driver.rs | 13 +- crates/dbx-core/src/db/mongo_driver.rs | 25 + crates/dbx-core/src/mongo_ops.rs | 44 ++ crates/dbx-web/src/main.rs | 1 + crates/dbx-web/src/routes/mongo.rs | 31 ++ package.json | 1 + packages/app-tests/mongoShellCommand.test.ts | 5 + packages/app-tests/queryStore.test.ts | 52 ++ packages/node-core/src/database.ts | 22 +- packages/node-core/src/web-backend.ts | 9 +- .../tests/mongo-count-routing.test.ts | 71 +++ packages/node-core/tests/mongo-query.test.ts | 4 + packages/node-core/tests/web-backend.test.ts | 60 ++- scripts/bench/mongodb-count.mjs | 464 ++++++++++++++++++ src-tauri/src/commands/document_cmd.rs | 6 +- src-tauri/src/commands/mcp_bridge.rs | 40 ++ src-tauri/src/commands/mongo_cmd.rs | 26 + src-tauri/src/lib.rs | 1 + 28 files changed, 916 insertions(+), 19 deletions(-) create mode 100644 packages/node-core/tests/mongo-count-routing.test.ts create mode 100644 scripts/bench/mongodb-count.mjs diff --git a/agents/common/src/main/java/com/dbx/agent/AgentProtocol.java b/agents/common/src/main/java/com/dbx/agent/AgentProtocol.java index 2dc6724007..88ddab9967 100644 --- a/agents/common/src/main/java/com/dbx/agent/AgentProtocol.java +++ b/agents/common/src/main/java/com/dbx/agent/AgentProtocol.java @@ -43,6 +43,7 @@ public final class AgentProtocol { * MongoDB read path that returns documents as relaxed Extended JSON for transfer. */ public static final String MONGO_METHOD_FIND_DOCUMENTS_EXTENDED_JSON = "find_documents_extended_json"; + public static final String MONGO_METHOD_COUNT_DOCUMENTS = "count_documents"; public static final String MONGO_METHOD_SERVER_VERSION = "server_version"; public static final String MONGO_METHOD_CREATE_INDEX = "create_index"; public static final String MONGO_METHOD_DROP_INDEXES = "drop_indexes"; @@ -124,6 +125,7 @@ public final class AgentProtocol { MONGO_METHOD_LIST_COLLECTIONS, MONGO_METHOD_FIND_DOCUMENTS, MONGO_METHOD_FIND_DOCUMENTS_EXTENDED_JSON, + MONGO_METHOD_COUNT_DOCUMENTS, MONGO_METHOD_SERVER_VERSION, MONGO_METHOD_CREATE_INDEX, MONGO_METHOD_DROP_INDEXES, diff --git a/agents/common/src/main/resources/agent-protocol-v1.json b/agents/common/src/main/resources/agent-protocol-v1.json index c108483064..29d9ca4359 100644 --- a/agents/common/src/main/resources/agent-protocol-v1.json +++ b/agents/common/src/main/resources/agent-protocol-v1.json @@ -69,6 +69,7 @@ "list_collections", "find_documents", "find_documents_extended_json", + "count_documents", "server_version", "create_index", "drop_indexes", diff --git a/agents/drivers/mongodb/src/main/java/com/dbx/agent/mongodb/MongoAgent.java b/agents/drivers/mongodb/src/main/java/com/dbx/agent/mongodb/MongoAgent.java index b5b883f90c..d062788fd0 100644 --- a/agents/drivers/mongodb/src/main/java/com/dbx/agent/mongodb/MongoAgent.java +++ b/agents/drivers/mongodb/src/main/java/com/dbx/agent/mongodb/MongoAgent.java @@ -429,6 +429,29 @@ public final class MongoAgent { return result; } + private static Object countDocuments(JsonObject params) { + MongoClient c = requireClient(); + String database = params.get("database").getAsString(); + String collection = params.get("collection").getAsString(); + Document filterDoc = documentOrNull(params, "filter"); + if (filterDoc == null) { + filterDoc = new Document(); + } + + boolean accurate = !params.has("accurate") || params.get("accurate").getAsBoolean(); + if (accurate) { + return c.getDatabase(database).getCollection(collection).countDocuments(filterDoc); + } + + // MongoDB 3.4 count() needs the legacy command to avoid the slow countDocuments path. + Document result = c.getDatabase(database).runCommand(new Document("count", collection).append("query", filterDoc)); + Object n = result.get("n"); + if (n instanceof Number number) { + return number.longValue(); + } + return 0L; + } + private static Object serverVersion(JsonObject params) { MongoClient c = requireClient(); String database = defaultString(stringOrNull(params, "database"), "admin"); @@ -846,6 +869,7 @@ public final class MongoAgent { case AgentProtocol.METHOD_LIST_INDEXES -> listIndexes(params); case AgentProtocol.MONGO_METHOD_FIND_DOCUMENTS -> findDocuments(params); case AgentProtocol.MONGO_METHOD_FIND_DOCUMENTS_EXTENDED_JSON -> findDocumentsExtendedJson(params); + case AgentProtocol.MONGO_METHOD_COUNT_DOCUMENTS -> countDocuments(params); case AgentProtocol.MONGO_METHOD_SERVER_VERSION -> serverVersion(params); case AgentProtocol.MONGO_METHOD_CREATE_INDEX -> createIndex(params); case AgentProtocol.MONGO_METHOD_DROP_INDEXES -> dropIndexes(params); diff --git a/agents/drivers/mongodb/src/test/java/com/dbx/agent/mongodb/MongoAgentTest.java b/agents/drivers/mongodb/src/test/java/com/dbx/agent/mongodb/MongoAgentTest.java index a90406cf71..63965152cb 100644 --- a/agents/drivers/mongodb/src/test/java/com/dbx/agent/mongodb/MongoAgentTest.java +++ b/agents/drivers/mongodb/src/test/java/com/dbx/agent/mongodb/MongoAgentTest.java @@ -110,6 +110,18 @@ class MongoAgentTest { assertFalse(json.getAsJsonObject("error").get("message").getAsString().contains("Unknown method")); } + @Test + void countDocumentsMethodIsRecognizedOverJsonRpc() { + String response = MongoAgent.handleRequest( + "{\"jsonrpc\":\"2.0\",\"id\":15,\"method\":\"count_documents\"," + + "\"params\":{\"database\":\"app\",\"collection\":\"orders\",\"filter\":\"{}\"}}"); + + JsonObject json = JsonParser.parseString(response).getAsJsonObject(); + assertEquals(15, json.get("id").getAsInt()); + assertEquals("Not connected", json.getAsJsonObject("error").get("message").getAsString()); + assertFalse(json.getAsJsonObject("error").get("message").getAsString().contains("Unknown method")); + } + @Test void parsesOptionalDocumentParameters() { JsonObject params = new JsonObject(); diff --git a/apps/desktop/src/lib/backend/api.ts b/apps/desktop/src/lib/backend/api.ts index b5a2c032ff..f58d3215ae 100644 --- a/apps/desktop/src/lib/backend/api.ts +++ b/apps/desktop/src/lib/backend/api.ts @@ -432,6 +432,7 @@ export const mongoDropDatabase = forward("mongoDropDatabase"); export const mongoDropCollection = forward("mongoDropCollection"); export const documentFindDocuments = forward("documentFindDocuments"); export const mongoFindDocuments = forward("mongoFindDocuments"); +export const mongoCountDocuments = forward("mongoCountDocuments"); export const mongoServerVersion = forward("mongoServerVersion"); export const mongoAggregateDocuments = forward("mongoAggregateDocuments"); export const mongoCollectionStats = forward("mongoCollectionStats"); diff --git a/apps/desktop/src/lib/backend/http.ts b/apps/desktop/src/lib/backend/http.ts index cab4f0cfda..f4013b5ff4 100644 --- a/apps/desktop/src/lib/backend/http.ts +++ b/apps/desktop/src/lib/backend/http.ts @@ -2009,6 +2009,10 @@ export async function documentFindDocuments(connectionId: string, database: stri return post("/api/document-store/find-documents", { connectionId, database, collection, skip, limit, filter, projection, sort, executionId }); } +export async function mongoCountDocuments(connectionId: string, database: string, collection: string, filter?: string, mode?: "accurate" | "legacy", executionId?: string): Promise { + return post("/api/mongo/count-documents", { connectionId, database, collection, filter, mode, executionId }); +} + export async function documentListGridFsFiles(connectionId: string, database: string, bucket: string, filter?: string, sort?: string): Promise { return post("/api/document-store/list-gridfs-files", { connectionId, database, bucket, filter, sort }); } diff --git a/apps/desktop/src/lib/backend/tauri.ts b/apps/desktop/src/lib/backend/tauri.ts index f8356cf586..0622dbff4b 100644 --- a/apps/desktop/src/lib/backend/tauri.ts +++ b/apps/desktop/src/lib/backend/tauri.ts @@ -1792,6 +1792,10 @@ export async function documentFindDocuments(connectionId: string, database: stri return invoke("document_find_documents", { connectionId, database, collection, skip, limit, filter, projection, sort, executionId }); } +export async function mongoCountDocuments(connectionId: string, database: string, collection: string, filter?: string, mode?: "accurate" | "legacy", executionId?: string): Promise { + return invoke("mongo_count_documents", { connectionId, database, collection, filter, mode, executionId }); +} + export async function documentListGridFsFiles(connectionId: string, database: string, bucket: string, filter?: string, sort?: string): Promise { return invoke("document_list_gridfs_files", { connectionId, database, bucket, filter, sort }); } diff --git a/apps/desktop/src/lib/mongo/mongoShellCommand.ts b/apps/desktop/src/lib/mongo/mongoShellCommand.ts index 449411386f..b04cd6f02b 100644 --- a/apps/desktop/src/lib/mongo/mongoShellCommand.ts +++ b/apps/desktop/src/lib/mongo/mongoShellCommand.ts @@ -12,6 +12,7 @@ export interface MongoFindCommand { export interface MongoCountDocumentsCommand { collection: string; filter: string; + mode: "accurate" | "legacy"; } export interface MongoAggregateCommand { @@ -145,8 +146,6 @@ export function applyMongoFindSort(input: string, column: string, direction: "as export function parseMongoCountDocumentsCommand(input: string): MongoCountDocumentsCommand | null { const source = input.trim().replace(/;$/, "").trim(); - // Accept deprecated Mongo shell count helpers for old server workflows, but - // keep DBX's internal execution mapped to the countDocuments result shape. return parseCollectionCountCommand(source, "countDocuments") ?? parseCollectionCountCommand(source, "count") ?? parseFindCountCommand(source); } @@ -166,6 +165,7 @@ function parseCollectionCountCommand(source: string, method: "countDocuments" | return { collection: target.collection, filter, + mode: method === "countDocuments" ? "accurate" : "legacy", }; } @@ -188,6 +188,7 @@ function parseFindCountCommand(source: string): MongoCountDocumentsCommand | nul return { collection: target.collection, filter, + mode: "legacy", }; } diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index a32896f35b..6f0500d3f5 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -2645,14 +2645,14 @@ export const useQueryStore = defineStore("query", () => { } case "countDocuments": { console.info("[DBX][executeTabSql:mongo-count:start]", { traceId, collection: mongoCommand.collection, database: currentDatabase }); - const result = await api.mongoFindDocuments(tab.connectionId, currentDatabase, mongoCommand.collection, 0, 1, mongoCommand.filter, undefined, undefined, executionId); - allResults.push(markQueryResultRowsRaw(annotateMongoResult(mongoCountToQueryResult(result.total, performance.now() - commandStartedAt)))); + const total = await api.mongoCountDocuments(tab.connectionId, currentDatabase, mongoCommand.collection, mongoCommand.filter, mongoCommand.mode, executionId); + allResults.push(markQueryResultRowsRaw(annotateMongoResult(mongoCountToQueryResult(total, performance.now() - commandStartedAt)))); mongoEditTarget = undefined; console.info("[DBX][executeTabSql:mongo-count:done]", { traceId, collection: mongoCommand.collection, database: currentDatabase, - total: result.total, + total, elapsed: elapsed(), }); break; diff --git a/crates/dbx-core/assets/agent-protocol-v1.json b/crates/dbx-core/assets/agent-protocol-v1.json index c108483064..29d9ca4359 100644 --- a/crates/dbx-core/assets/agent-protocol-v1.json +++ b/crates/dbx-core/assets/agent-protocol-v1.json @@ -69,6 +69,7 @@ "list_collections", "find_documents", "find_documents_extended_json", + "count_documents", "server_version", "create_index", "drop_indexes", diff --git a/crates/dbx-core/src/db/agent_driver.rs b/crates/dbx-core/src/db/agent_driver.rs index d65dc8e032..4b3c21252f 100644 --- a/crates/dbx-core/src/db/agent_driver.rs +++ b/crates/dbx-core/src/db/agent_driver.rs @@ -252,6 +252,7 @@ pub enum MongoAgentMethod { ListCollections, FindDocuments, FindDocumentsExtendedJson, + CountDocuments, ServerVersion, CreateIndex, DropIndexes, @@ -264,11 +265,12 @@ pub enum MongoAgentMethod { } impl MongoAgentMethod { - pub const ALL: [Self; 13] = [ + pub const ALL: [Self; 14] = [ Self::ListDatabases, Self::ListCollections, Self::FindDocuments, Self::FindDocumentsExtendedJson, + Self::CountDocuments, Self::ServerVersion, Self::CreateIndex, Self::DropIndexes, @@ -286,6 +288,7 @@ impl MongoAgentMethod { Self::ListCollections => "list_collections", Self::FindDocuments => "find_documents", Self::FindDocumentsExtendedJson => "find_documents_extended_json", + Self::CountDocuments => "count_documents", Self::ServerVersion => "server_version", Self::CreateIndex => "create_index", Self::DropIndexes => "drop_indexes", @@ -999,6 +1002,13 @@ impl AgentDriverClient { self.call_mongo_method(MongoAgentMethod::FindDocumentsExtendedJson, params).await } + pub async fn mongo_count_documents( + &mut self, + params: Value, + ) -> Result { + self.call_mongo_method(MongoAgentMethod::CountDocuments, params).await + } + pub async fn mongo_server_version( &mut self, database: &str, @@ -1727,6 +1737,7 @@ mod tests { assert_eq!(MongoAgentMethod::ListCollections.as_str(), "list_collections"); assert_eq!(MongoAgentMethod::FindDocuments.as_str(), "find_documents"); assert_eq!(MongoAgentMethod::FindDocumentsExtendedJson.as_str(), "find_documents_extended_json"); + assert_eq!(MongoAgentMethod::CountDocuments.as_str(), "count_documents"); assert_eq!(MongoAgentMethod::ServerVersion.as_str(), "server_version"); assert_eq!(MongoAgentMethod::CreateIndex.as_str(), "create_index"); assert_eq!(MongoAgentMethod::DropIndexes.as_str(), "drop_indexes"); diff --git a/crates/dbx-core/src/db/mongo_driver.rs b/crates/dbx-core/src/db/mongo_driver.rs index 474bf7bbe2..36d26544a6 100644 --- a/crates/dbx-core/src/db/mongo_driver.rs +++ b/crates/dbx-core/src/db/mongo_driver.rs @@ -592,6 +592,31 @@ pub async fn find_documents( Ok(MongoDocumentResult { documents, total }) } +pub async fn count_documents( + client: &Client, + database: &str, + collection: &str, + filter: Option<&str>, + accurate: bool, +) -> Result { + let col = client.database(database).collection::(collection); + + let filter_doc: Document = match filter { + Some(f) if !f.trim().is_empty() => { + let json: serde_json::Value = serde_json::from_str(f).map_err(|e| format!("Invalid filter JSON: {e}"))?; + json_filter_to_document(&json)? + } + _ => doc! {}, + }; + + if !accurate && filter_doc.is_empty() { + // Legacy count() permits the metadata-backed fast path; countDocuments() must scan accurately. + col.estimated_document_count().await.map_err(|e| e.to_string()) + } else { + col.count_documents(filter_doc).await.map_err(|e| e.to_string()) + } +} + /// Find MongoDB documents as relaxed Extended JSON for MongoDB transfer paths. #[allow(clippy::too_many_arguments)] pub async fn find_documents_extended_json( diff --git a/crates/dbx-core/src/mongo_ops.rs b/crates/dbx-core/src/mongo_ops.rs index 28c3643308..8388998784 100644 --- a/crates/dbx-core/src/mongo_ops.rs +++ b/crates/dbx-core/src/mongo_ops.rs @@ -121,6 +121,50 @@ pub async fn mongo_find_documents_core( .await } +pub async fn mongo_count_documents_core( + state: &AppState, + connection_id: &str, + database: &str, + collection: &str, + filter: Option<&str>, + mode: Option<&str>, +) -> Result { + let accurate = mode != Some("legacy"); + 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::count_documents(client, database, collection, filter, accurate).await + } + PoolKind::Agent(client) => { + let mut client = client.lock().await; + let params = serde_json::json!({ + "database": database, + "collection": collection, + "filter": filter, + "accurate": accurate, + }); + match client.mongo_count_documents(params.clone()).await { + Ok(total) => Ok(total), + Err(error) if is_unknown_agent_method_error(&error, "count_documents") => { + let result: MongoDocumentResult = client + .mongo_find_documents(serde_json::json!({ + "database": database, + "collection": collection, + "skip": 0, + "limit": 1, + "filter": filter, + })) + .await?; + Ok(result.total) + } + Err(error) => Err(error), + } + } + _ => Err("Not a MongoDB connection".to_string()), + } +} + /// Read MongoDB documents as relaxed Extended JSON for MongoDB transfer paths. #[allow(clippy::too_many_arguments)] pub async fn mongo_find_documents_extended_json_core( diff --git a/crates/dbx-web/src/main.rs b/crates/dbx-web/src/main.rs index aad4c55c5d..ee9a491b30 100644 --- a/crates/dbx-web/src/main.rs +++ b/crates/dbx-web/src/main.rs @@ -474,6 +474,7 @@ async fn main() { .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/count-documents", post(routes::mongo::count_documents)) .route("/mongo/server-version", post(routes::mongo::server_version)) .route("/mongo/collection-stats", post(routes::mongo::collection_stats)) .route("/mongo/aggregate-documents", post(routes::mongo::aggregate_documents)) diff --git a/crates/dbx-web/src/routes/mongo.rs b/crates/dbx-web/src/routes/mongo.rs index 5dd25c0cdb..5ec057efa6 100644 --- a/crates/dbx-web/src/routes/mongo.rs +++ b/crates/dbx-web/src/routes/mongo.rs @@ -78,6 +78,17 @@ pub struct MongoFindRequest { pub execution_id: Option, } +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct MongoCountRequest { + pub connection_id: String, + pub database: String, + pub collection: String, + pub filter: Option, + pub mode: Option, + pub execution_id: Option, +} + #[derive(Deserialize)] #[serde(rename_all = "camelCase")] pub struct MongoServerVersionRequest { @@ -286,6 +297,26 @@ pub async fn find_documents( Ok(Json(serde_json::to_value(result).map_err(|e| AppError(e.to_string()))?)) } +pub async fn count_documents( + State(state): State>, + Json(req): Json, +) -> Result, AppError> { + let result = run_cancellable( + &state, + req.execution_id, + dbx_core::mongo_ops::mongo_count_documents_core( + &state.app, + &req.connection_id, + &req.database, + &req.collection, + req.filter.as_deref(), + req.mode.as_deref(), + ), + ) + .await?; + Ok(Json(result)) +} + pub async fn server_version( State(state): State>, Json(req): Json, diff --git a/package.json b/package.json index c0b377f5c4..20c9f2fd2d 100644 --- a/package.json +++ b/package.json @@ -23,6 +23,7 @@ "build:checked": "pnpm typecheck && pnpm build", "check": "node scripts/run-check.mjs", "bench:redis-key-search": "node scripts/bench/redis-key-search.mjs", + "bench:mongodb-count": "node scripts/bench/mongodb-count.mjs", "lint": "oxlint --vue-plugin apps/desktop/src", "fmt": "oxfmt \"apps/desktop/src/**/*.{ts,vue}\"", "test": "vitest run", diff --git a/packages/app-tests/mongoShellCommand.test.ts b/packages/app-tests/mongoShellCommand.test.ts index 91f116639f..3c75787285 100644 --- a/packages/app-tests/mongoShellCommand.test.ts +++ b/packages/app-tests/mongoShellCommand.test.ts @@ -252,6 +252,7 @@ test("parseMongoCountDocumentsCommand parses db collection countDocuments", () = assert.deepEqual(parseMongoCountDocumentsCommand("db.products.countDocuments({})"), { collection: "products", filter: "{}", + mode: "accurate", }); }); @@ -259,20 +260,24 @@ test("parseMongoCountDocumentsCommand parses legacy count helpers", () => { assert.deepEqual(parseMongoCountDocumentsCommand("db.products.count({ active: true })"), { collection: "products", filter: '{ "active": true }', + mode: "legacy", }); assert.deepEqual(parseMongoCountDocumentsCommand('db.getCollection("audit.logs").count()'), { collection: "audit.logs", filter: "{}", + mode: "legacy", }); assert.deepEqual(parseMongoCountDocumentsCommand("db.products.find({ active: true }).count()"), { collection: "products", filter: '{ "active": true }', + mode: "legacy", }); assert.equal(parseMongoFindCommand("db.products.find({ active: true }).count()"), null); assert.deepEqual(parseMongoCommand("db.products.find({ active: true }).count()")?.command, { kind: "countDocuments", collection: "products", filter: '{ "active": true }', + mode: "legacy", }); }); diff --git a/packages/app-tests/queryStore.test.ts b/packages/app-tests/queryStore.test.ts index 363726c5f5..62d1f8e88e 100644 --- a/packages/app-tests/queryStore.test.ts +++ b/packages/app-tests/queryStore.test.ts @@ -3000,6 +3000,58 @@ test("replacing one paginated SQL result preserves the grouped refresh SQL", asy } }); +test("mongo count execution uses the dedicated count endpoint", async () => { + const restoreStorage = installMemoryStorage(); + setActivePinia(createPinia()); + const connectionStore = useConnectionStore(); + const store = useQueryStore(); + const originalFetch = globalThis.fetch; + let countBody: any; + + connectionStore.addEphemeralConnection({ + ...conn("mongo-1"), + db_type: "mongodb", + port: 27017, + }); + + globalThis.fetch = withConnectionHealthMock(async (input, init) => { + const url = String(input); + if (url === "/api/query/prepare-pagination-plan") { + const body = JSON.parse(String(init?.body ?? "{}")); + return new Response(JSON.stringify({ sqlToExecute: body.options.sql, useAgentResultSession: false }), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + } + if (url === "/api/mongo/count-documents") { + countBody = JSON.parse(String(init?.body ?? "{}")); + return new Response(JSON.stringify(21606536), { status: 200, headers: { "Content-Type": "application/json" } }); + } + return new Response("unexpected request", { status: 500 }); + }); + + try { + const tabId = store.createTab("mongo-1", "dbx_issue_2959", "Query", "query", ""); + await store.executeTabSql(tabId, "db.large_count.count()"); + const tab = store.tabs.find((item) => item.id === tabId); + + assert.deepEqual(countBody, { + connectionId: "mongo-1", + database: "dbx_issue_2959", + collection: "large_count", + filter: "{}", + mode: "legacy", + executionId: countBody.executionId, + }); + assert.equal(typeof countBody.executionId, "string"); + assert.deepEqual(tab?.result?.columns, ["count"]); + assert.deepEqual(tab?.result?.rows, [[21606536]]); + } finally { + globalThis.fetch = originalFetch; + restoreStorage(); + } +}); + test("mongo createIndex execution uses the dedicated create-index endpoint", async () => { const restoreStorage = installMemoryStorage(); setActivePinia(createPinia()); diff --git a/packages/node-core/src/database.ts b/packages/node-core/src/database.ts index 55fac5ff24..7037f5ec66 100644 --- a/packages/node-core/src/database.ts +++ b/packages/node-core/src/database.ts @@ -881,8 +881,8 @@ export async function executeQuery(config: ConnectionConfig, sql: string, option } const count = parseMongoCountDocumentsCommand(sql); if (count) { - const result = await withTimeout(mongoFindDocuments(config, count.collection, 0, 1, count.filter), resolveTimeoutMs(options)); - return { columns: ["count"], rows: [{ count: result.total }], row_count: 1 }; + const total = await withTimeout(mongoCountDocuments(config, count.collection, count.filter, count.mode), resolveTimeoutMs(options)); + return { columns: ["count"], rows: [{ count: total }], row_count: 1 }; } const find = parseMongoFindCommand(sql); if (find) { @@ -1125,6 +1125,17 @@ async function mongoFindDocuments(config: ConnectionConfig, collection: string, }); } +async function mongoCountDocuments(config: ConnectionConfig, collection: string, filter: string, mode: MongoCountDocumentsCommand["mode"]): Promise { + return bridgeDataRequest("/data/mongo/count-documents", { + connection_id: config.id, + connection_name: config.name, + database: config.database || "", + collection, + filter, + mode, + }); +} + async function mongoServerVersion(config: ConnectionConfig): Promise { return bridgeDataRequest("/data/mongo/server-version", { connection_id: config.id, @@ -1296,6 +1307,7 @@ interface MongoFindCommand { interface MongoCountDocumentsCommand { collection: string; filter: string; + mode: "accurate" | "legacy"; } interface MongoAggregateCommand { @@ -1364,8 +1376,6 @@ export function parseMongoVersionCommand(input: string): boolean { export function parseMongoCountDocumentsCommand(input: string): MongoCountDocumentsCommand | null { const source = input.trim().replace(/;$/, "").trim(); - // Accept deprecated Mongo shell count helpers for old server workflows, but - // keep DBX's internal execution mapped to the countDocuments result shape. return parseCollectionCountCommand(source, "countDocuments") ?? parseCollectionCountCommand(source, "count") ?? parseFindCountCommand(source); } @@ -1378,7 +1388,7 @@ function parseCollectionCountCommand(source: string, method: "countDocuments" | const args = splitTopLevel(source.slice(openIndex + 1, closeIndex)); if (args.length > 1 && args.slice(1).some((arg) => arg.trim())) return null; const filter = normalizeJsonArgument(args[0] || "{}"); - return filter ? { collection: target.collection, filter } : null; + return filter ? { collection: target.collection, filter, mode: method === "countDocuments" ? "accurate" : "legacy" } : null; } function parseFindCountCommand(source: string): MongoCountDocumentsCommand | null { @@ -1392,7 +1402,7 @@ function parseFindCountCommand(source: string): MongoCountDocumentsCommand | nul const findArgs = splitTopLevel(source.slice(findOpenIndex + 1, findCloseIndex)); if (findArgs.length > 2 && findArgs.slice(2).some((arg) => arg.trim())) return null; const filter = normalizeJsonArgument(findArgs[0] || "{}"); - return filter ? { collection: target.collection, filter } : null; + return filter ? { collection: target.collection, filter, mode: "legacy" } : null; } export function parseMongoAggregateCommand(input: string): MongoAggregateCommand | null { diff --git a/packages/node-core/src/web-backend.ts b/packages/node-core/src/web-backend.ts index e3968ee584..bf1d55fb1f 100644 --- a/packages/node-core/src/web-backend.ts +++ b/packages/node-core/src/web-backend.ts @@ -219,19 +219,18 @@ export async function executeQuery(config: ConnectionConfig, sql: string, option } const count = parseMongoCountDocumentsCommand(sql); if (count) { - const res = await apiFetch("/api/mongo/find-documents", { + const res = await apiFetch("/api/mongo/count-documents", { method: "POST", body: JSON.stringify({ connectionId: config.id, database: config.database || "", collection: count.collection, - skip: 0, - limit: 1, filter: count.filter, + mode: count.mode, }), }); - const result = (await res.json()) as { documents: unknown[]; total: number }; - return { columns: ["count"], rows: [{ count: result.total }], row_count: 1 }; + const total = (await res.json()) as number; + return { columns: ["count"], rows: [{ count: total }], row_count: 1 }; } const find = parseMongoFindCommand(sql); if (find) { diff --git a/packages/node-core/tests/mongo-count-routing.test.ts b/packages/node-core/tests/mongo-count-routing.test.ts new file mode 100644 index 0000000000..2c78ee2a86 --- /dev/null +++ b/packages/node-core/tests/mongo-count-routing.test.ts @@ -0,0 +1,71 @@ +import assert from "node:assert/strict"; +import { mkdtemp, rm, writeFile } from "node:fs/promises"; +import { createServer } from "node:http"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { test } from "vitest"; + +import type { ConnectionConfig } from "../src/connections.js"; +import { executeQuery } from "../src/database.js"; + +const mongoConfig: ConnectionConfig = { + id: "mongo-bridge", + name: "mongo-bridge", + db_type: "mongodb", + host: "127.0.0.1", + port: 27017, + username: "", + password: "", + database: "app", + ssh_enabled: false, + ssl: false, +}; + +test("direct backend routes MongoDB count commands through the count bridge", async () => { + const tempDir = await mkdtemp(join(tmpdir(), "dbx-node-core-count-")); + const previousDataDir = process.env.DBX_DATA_DIR; + let requestBody: unknown; + const server = createServer((req, res) => { + assert.equal(req.url, "/data/mongo/count-documents"); + let body = ""; + req.setEncoding("utf8"); + req.on("data", (chunk) => { + body += chunk; + }); + req.on("end", () => { + requestBody = JSON.parse(body); + res.writeHead(200, { "content-type": "application/json" }); + res.end("42"); + }); + }); + + try { + await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve)); + const address = server.address(); + if (!address || typeof address === "string") throw new Error("expected TCP bridge address"); + + process.env.DBX_DATA_DIR = tempDir; + await writeFile(join(tempDir, "mcp-bridge-port"), String(address.port)); + + const result = await executeQuery(mongoConfig, "db.projects.count({ active: true })"); + + assert.deepEqual(requestBody, { + connection_id: "mongo-bridge", + connection_name: "mongo-bridge", + database: "app", + collection: "projects", + filter: '{ "active": true }', + mode: "legacy", + }); + assert.deepEqual(result, { + columns: ["count"], + rows: [{ count: 42 }], + row_count: 1, + }); + } finally { + await new Promise((resolve) => server.close(() => resolve())); + if (previousDataDir === undefined) delete process.env.DBX_DATA_DIR; + else process.env.DBX_DATA_DIR = previousDataDir; + await rm(tempDir, { recursive: true, force: true }); + } +}); diff --git a/packages/node-core/tests/mongo-query.test.ts b/packages/node-core/tests/mongo-query.test.ts index a7d0cce049..517efd2410 100644 --- a/packages/node-core/tests/mongo-query.test.ts +++ b/packages/node-core/tests/mongo-query.test.ts @@ -93,6 +93,7 @@ test("parseMongoCountDocumentsCommand accepts shell-style count commands", () => assert.deepEqual(parseMongoCountDocumentsCommand('db.projects.countDocuments({"active":true})'), { collection: "projects", filter: '{"active":true}', + mode: "accurate", }); }); @@ -100,14 +101,17 @@ test("parseMongoCountDocumentsCommand accepts legacy count helpers", () => { assert.deepEqual(parseMongoCountDocumentsCommand("db.projects.count({ active: true })"), { collection: "projects", filter: '{ "active": true }', + mode: "legacy", }); assert.deepEqual(parseMongoCountDocumentsCommand('db.getCollection("audit.logs").count()'), { collection: "audit.logs", filter: "{}", + mode: "legacy", }); assert.deepEqual(parseMongoCountDocumentsCommand("db.projects.find({ active: true }).count()"), { collection: "projects", filter: '{ "active": true }', + mode: "legacy", }); assert.equal(parseMongoFindCommand("db.projects.find({ active: true }).count()"), null); }); diff --git a/packages/node-core/tests/web-backend.test.ts b/packages/node-core/tests/web-backend.test.ts index 176ec524d8..948c73c993 100644 --- a/packages/node-core/tests/web-backend.test.ts +++ b/packages/node-core/tests/web-backend.test.ts @@ -1,7 +1,7 @@ import assert from "node:assert/strict"; import { afterEach, test } from "vitest"; -import { loadConnections, resetWebAuthForTests } from "../src/web-backend.js"; +import { executeQuery, loadConnections, resetWebAuthForTests } from "../src/web-backend.js"; const originalFetch = globalThis.fetch; const originalWebUrl = process.env.DBX_WEB_URL; @@ -123,3 +123,61 @@ test("web backend still allows DBX Web instances with password auth disabled", a assert.deepEqual(await loadConnections(), []); }); + +test("web backend routes MongoDB count commands through count-documents", async () => { + process.env.DBX_WEB_URL = "http://127.0.0.1:4224"; + delete process.env.DBX_WEB_PASSWORD; + const calls: Array<{ url: string; init?: RequestInit }> = []; + globalThis.fetch = (async (input, init) => { + const url = String(input); + calls.push({ url, init }); + if (url.endsWith("/api/auth/check")) { + return jsonResponse({ authenticated: true, required: false, setup_required: false }); + } + if (url.endsWith("/api/connection/connect")) { + return jsonResponse({ ok: true }); + } + if (url.endsWith("/api/mongo/count-documents")) { + assert.equal(init?.method, "POST"); + assert.deepEqual(JSON.parse(String(init?.body)), { + connectionId: "mongo-web", + database: "app", + collection: "projects", + filter: '{"active":true}', + mode: "accurate", + }); + return jsonResponse(42); + } + throw new Error(`unexpected request: ${url}`); + }) as typeof fetch; + + const result = await executeQuery( + { + id: "mongo-web", + name: "mongo-web", + db_type: "mongodb", + host: "127.0.0.1", + port: 27017, + username: "", + password: "", + database: "app", + ssh_enabled: false, + ssl: false, + }, + 'db.projects.countDocuments({"active":true})', + ); + + assert.deepEqual(result, { + columns: ["count"], + rows: [{ count: 42 }], + row_count: 1, + }); + assert.deepEqual( + calls.map((call) => call.url), + [ + "http://127.0.0.1:4224/api/auth/check", + "http://127.0.0.1:4224/api/connection/connect", + "http://127.0.0.1:4224/api/mongo/count-documents", + ], + ); +}); diff --git a/scripts/bench/mongodb-count.mjs b/scripts/bench/mongodb-count.mjs new file mode 100644 index 0000000000..3b4c3a42c9 --- /dev/null +++ b/scripts/bench/mongodb-count.mjs @@ -0,0 +1,464 @@ +#!/usr/bin/env node +import { spawn } from "node:child_process"; +import { copyFileSync, existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs"; +import { dirname, resolve } from "node:path"; +import { performance } from "node:perf_hooks"; + +const DEFAULTS = { + apiBase: "http://127.0.0.1:4224/api", + container: "dbx-issue-2959-mongo", + image: "mongo:3.4", + host: "127.0.0.1", + port: 12959, + database: "dbx_issue_2959", + collection: "large_count", + expectedCount: 21_606_536, + iterations: 5, + warmups: 1, + seed: false, + forceSeed: false, + seedBatchSize: 10_000, + dbxDataDir: "", + json: false, +}; + +function parseArgs(argv) { + const options = { ...DEFAULTS }; + for (const arg of argv) { + if (arg === "--") continue; + const [rawKey, rawValue] = arg.split("=", 2); + const key = rawKey.replace(/^--/, ""); + const value = rawValue ?? "true"; + switch (key) { + case "api-base": + options.apiBase = value; + break; + case "container": + options.container = value; + break; + case "image": + options.image = value; + break; + case "host": + options.host = value; + break; + case "port": + options.port = Number.parseInt(value, 10); + break; + case "database": + options.database = value; + break; + case "collection": + options.collection = value; + break; + case "expected-count": + options.expectedCount = Number.parseInt(value, 10); + break; + case "iterations": + options.iterations = Number.parseInt(value, 10); + break; + case "warmups": + options.warmups = Number.parseInt(value, 10); + break; + case "seed": + options.seed = value !== "false"; + break; + case "force-seed": + options.forceSeed = value !== "false"; + break; + case "seed-batch-size": + options.seedBatchSize = Number.parseInt(value, 10); + break; + case "dbx-data-dir": + options.dbxDataDir = value; + break; + case "json": + options.json = value !== "false"; + break; + case "help": + printHelp(); + process.exit(0); + default: + throw new Error(`Unknown option: ${rawKey}`); + } + } + return options; +} + +function printHelp() { + console.log(`MongoDB count benchmark + +Usage: + pnpm bench:mongodb-count [options] + +Options: + --api-base=http://127.0.0.1:4224/api DBX Web API base URL + --container=dbx-issue-2959-mongo MongoDB Docker container name + --image=mongo:3.4 MongoDB image used when creating the container + --port=12959 Host port for MongoDB + --database=dbx_issue_2959 Database name + --collection=large_count Collection name + --expected-count=21606536 Expected collection count + --iterations=5 Timed iterations per case + --warmups=1 Warmup iterations per case + --dbx-data-dir=/tmp/dbx-bench Copy local Mongo agent jar into this DBX data dir + --seed Seed the collection if its count does not match + --force-seed Drop and reseed even if the collection exists + --seed-batch-size=10000 Number of docs per insertMany batch + --json Print JSON only + +Before running this benchmark, start DBX Web with the same data dir, for example: + DBX_DATA_DIR=/tmp/dbx-bench DBX_DISABLE_PASSWORD=1 cargo run -p dbx-web + +The benchmark saves one temporary connection into that DBX data dir, so use an +isolated DBX_DATA_DIR instead of your normal desktop profile. +`); +} + +function assertInteger(name, value, min = 0) { + if (!Number.isFinite(value) || Math.trunc(value) !== value || value < min) { + throw new Error(`${name} must be an integer >= ${min}, got ${value}`); + } +} + +function run(command, args, options = {}) { + const startedAt = performance.now(); + const child = spawn(command, args, { + cwd: options.cwd ?? process.cwd(), + env: process.env, + stdio: options.stdio ?? ["ignore", "pipe", "pipe"], + }); + const stdout = []; + const stderr = []; + if (child.stdout) child.stdout.on("data", (chunk) => stdout.push(chunk)); + if (child.stderr) child.stderr.on("data", (chunk) => stderr.push(chunk)); + return new Promise((resolve, reject) => { + child.on("error", reject); + child.on("close", (code) => { + const result = { + code, + stdout: Buffer.concat(stdout).toString(), + stderr: Buffer.concat(stderr).toString(), + elapsedMs: performance.now() - startedAt, + }; + if (code === 0) resolve(result); + else reject(new Error(`${command} ${args.join(" ")} failed with ${code}\n${result.stderr || result.stdout}`)); + }); + }); +} + +async function dockerExec(container, args) { + return run("docker", ["exec", container, ...args]); +} + +async function ensureMongoContainer(options) { + try { + await run("docker", ["inspect", options.container]); + await run("docker", ["start", options.container]); + } catch { + await run("docker", [ + "run", + "-d", + "--name", + options.container, + "--platform", + "linux/amd64", + "-p", + `${options.host}:${options.port}:27017`, + options.image, + "--nojournal", + "--wiredTigerCacheSizeGB", + "0.5", + ]); + } + + const deadline = Date.now() + 30_000; + while (Date.now() < deadline) { + try { + const out = await dockerExec(options.container, ["mongo", "--quiet", "--eval", "db.runCommand({ ping: 1 }).ok"]); + if (out.stdout.trim().endsWith("1")) return; + } catch { + await new Promise((resolve) => setTimeout(resolve, 500)); + } + } + throw new Error("Timed out waiting for MongoDB container"); +} + +async function mongoEval(options, script) { + return dockerExec(options.container, ["mongo", options.database, "--quiet", "--eval", script]); +} + +async function collectionCount(options) { + const out = await mongoEval(options, `print(db.${options.collection}.count())`); + const line = out.stdout.trim().split(/\r?\n/).at(-1) ?? ""; + return Number.parseInt(line, 10); +} + +async function seedMongo(options) { + const current = await collectionCount(options); + if (!options.forceSeed && current === options.expectedCount) { + return { skipped: true, count: current, elapsedMs: 0 }; + } + if (!options.seed && !options.forceSeed) { + throw new Error( + `Collection count is ${current}, expected ${options.expectedCount}. Re-run with --seed to create the benchmark dataset.`, + ); + } + + const startedAt = performance.now(); + const script = ` + var collection = db.${options.collection}; + collection.drop(); + var expected = ${options.expectedCount}; + var batchSize = ${options.seedBatchSize}; + for (var start = 0; start < expected; start += batchSize) { + var docs = []; + var end = Math.min(start + batchSize, expected); + for (var i = start; i < end; i++) { + docs.push({ seq: i, bucket: i % 10, payload: "payload-" + i }); + } + collection.insertMany(docs, { ordered: false }); + } + print(collection.count()); + `; + const out = await mongoEval(options, script); + const count = Number.parseInt(out.stdout.trim().split(/\r?\n/).at(-1) ?? "", 10); + if (count !== options.expectedCount) { + throw new Error(`Seed finished with ${count} docs, expected ${options.expectedCount}`); + } + return { skipped: false, count, elapsedMs: performance.now() - startedAt }; +} + +function localMongoAgentJar() { + return resolve("agents", "drivers", "mongodb", "build", "libs", "dbx-agent-mongodb.jar"); +} + +function syncMongoAgentJar(options) { + if (!options.dbxDataDir) return null; + const source = localMongoAgentJar(); + if (!existsSync(source)) { + throw new Error(`MongoDB agent jar not found: ${source}. Run ./agents/gradlew -p agents :mongodb:shadowJar first.`); + } + const dest = resolve(options.dbxDataDir, "agents", "drivers", "mongodb", "agent.jar"); + mkdirSync(dirname(dest), { recursive: true }); + copyFileSync(source, dest); + ensureAgentState(options.dbxDataDir); + return dest; +} + +function ensureAgentState(dbxDataDir) { + const statePath = resolve(dbxDataDir, "agents", "state.json"); + const now = new Date().toISOString(); + let state = {}; + if (existsSync(statePath)) { + state = JSON.parse(readFileSync(statePath, "utf8")); + } + state.jre_versions ??= {}; + state.installed_drivers ??= {}; + state.installed_drivers.mongodb = { + version: "0.1.0-local", + installed_at: state.installed_drivers.mongodb?.installed_at ?? now, + jre: "21", + }; + state.java_runtime = { + ...(state.java_runtime ?? {}), + mode: "system", + }; + mkdirSync(dirname(statePath), { recursive: true }); + writeFileSync(statePath, `${JSON.stringify(state, null, 2)}\n`); +} + +async function postJsonText(url, body) { + const response = await fetch(url, { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify(body), + }); + const text = await response.text(); + return { ok: response.ok, status: response.status, text, json: parseJsonMaybe(text) }; +} + +function parseJsonMaybe(text) { + try { + return JSON.parse(text); + } catch { + return undefined; + } +} + +async function ensureDbxConnection(options) { + const connectionId = `bench-mongodb-count-${options.port}`; + const config = { + id: connectionId, + name: `Bench MongoDB Count ${options.port}`, + db_type: "mongodb", + driver_profile: "mongodb-legacy", + host: options.host, + port: options.port, + username: "", + password: "", + database: options.database, + ssl: false, + connect_timeout_secs: 10, + query_timeout_secs: 30, + }; + + const save = await postJsonText(`${options.apiBase}/connection/save`, { configs: [config] }); + if (!save.ok) throw new Error(`/connection/save failed with ${save.status}\n${save.text}`); + const connect = await postJsonText(`${options.apiBase}/connection/connect`, { config }); + if (!connect.ok) throw new Error(`/connection/connect failed with ${connect.status}\n${connect.text}`); + return connectionId; +} + +async function measureCase(label, iterations, warmups, fn) { + for (let index = 0; index < warmups; index += 1) { + await fn(); + } + const samples = []; + let lastValue; + let lastPayloadBytes = 0; + for (let index = 0; index < iterations; index += 1) { + const startedAt = performance.now(); + const result = await fn(); + samples.push(performance.now() - startedAt); + lastValue = result.value; + lastPayloadBytes = result.payloadBytes ?? 0; + } + samples.sort((a, b) => a - b); + const sum = samples.reduce((total, value) => total + value, 0); + return { + label, + value: lastValue, + samples, + minMs: samples[0], + p50Ms: percentile(samples, 0.5), + p90Ms: percentile(samples, 0.9), + maxMs: samples.at(-1), + avgMs: sum / samples.length, + payloadBytes: lastPayloadBytes, + }; +} + +function percentile(sortedSamples, p) { + if (sortedSamples.length === 0) return 0; + const index = Math.min(sortedSamples.length - 1, Math.ceil(sortedSamples.length * p) - 1); + return sortedSamples[index]; +} + +async function measureNativeMongoCount(options) { + const script = `var r = db.runCommand({ count: "${options.collection}", query: {} }); print(r.n);`; + const out = await mongoEval(options, script); + return { value: Number.parseInt(out.stdout.trim().split(/\r?\n/).at(-1) ?? "", 10) }; +} + +async function measureDbxFindTotal(options, connectionId) { + const response = await postJsonText(`${options.apiBase}/document-store/find-documents`, { + connectionId, + database: options.database, + collection: options.collection, + skip: 0, + limit: 1, + filter: "{}", + }); + if (!response.ok) throw new Error(`/document-store/find-documents failed with ${response.status}\n${response.text}`); + return { value: response.json?.total, payloadBytes: Buffer.byteLength(response.text) }; +} + +async function measureDbxDedicatedCount(options, connectionId) { + const response = await postJsonText(`${options.apiBase}/mongo/count-documents`, { + connectionId, + database: options.database, + collection: options.collection, + filter: "{}", + mode: "legacy", + }); + if (!response.ok) { + return { value: `unavailable (${response.status})`, payloadBytes: Buffer.byteLength(response.text) }; + } + return { value: response.json, payloadBytes: Buffer.byteLength(response.text) }; +} + +function formatMs(ms) { + return `${Math.round(ms)}ms`; +} + +function formatBytes(bytes) { + if (!bytes) return "0 B"; + const units = ["B", "KB", "MB", "GB"]; + let value = bytes; + let unit = 0; + while (value >= 1024 && unit < units.length - 1) { + value /= 1024; + unit += 1; + } + return `${value.toFixed(unit === 0 ? 0 : 1)} ${units[unit]}`; +} + +function printReport(result) { + console.log("# MongoDB count benchmark"); + console.log(""); + console.log(`- API: ${result.config.apiBase}`); + console.log(`- MongoDB: ${result.config.container} on ${result.config.host}:${result.config.port}`); + console.log(`- Collection: ${result.config.database}.${result.config.collection}`); + console.log(`- Expected count: ${result.config.expectedCount}`); + console.log(`- Iterations: ${result.config.iterations} timed, ${result.config.warmups} warmup`); + console.log(`- Seed: ${result.seed.skipped ? "reused existing dataset" : `loaded in ${formatMs(result.seed.elapsedMs)}`}`); + if (result.agentJar) console.log(`- Synced agent jar: ${result.agentJar}`); + console.log(""); + console.log("| Case | Value | Payload | Min | P50 | P90 | Max | Avg |"); + console.log("| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: |"); + for (const row of result.measurements) { + console.log( + `| ${row.label} | ${row.value} | ${formatBytes(row.payloadBytes)} | ${formatMs(row.minMs)} | ${formatMs(row.p50Ms)} | ${formatMs(row.p90Ms)} | ${formatMs(row.maxMs)} | ${formatMs(row.avgMs)} |`, + ); + } +} + +async function main() { + const options = parseArgs(process.argv.slice(2)); + assertInteger("port", options.port, 1); + assertInteger("expected-count", options.expectedCount, 1); + assertInteger("iterations", options.iterations, 1); + assertInteger("warmups", options.warmups, 0); + assertInteger("seed-batch-size", options.seedBatchSize, 1); + + if (!options.json) console.log("Preparing MongoDB count benchmark..."); + await ensureMongoContainer(options); + const seed = await seedMongo(options); + const agentJar = syncMongoAgentJar(options); + + if (!options.json) console.log("Connecting DBX Web API..."); + const connectionId = await ensureDbxConnection(options); + + const measurements = [ + await measureCase("mongo runCommand count", options.iterations, options.warmups, () => measureNativeMongoCount(options)), + await measureCase("DBX find-documents total", options.iterations, options.warmups, () => measureDbxFindTotal(options, connectionId)), + await measureCase("DBX count-documents", options.iterations, options.warmups, () => measureDbxDedicatedCount(options, connectionId)), + ]; + + const result = { + config: { + apiBase: options.apiBase, + container: options.container, + host: options.host, + port: options.port, + database: options.database, + collection: options.collection, + expectedCount: options.expectedCount, + iterations: options.iterations, + warmups: options.warmups, + }, + seed, + agentJar, + measurements, + generatedAt: new Date().toISOString(), + }; + + if (options.json) console.log(JSON.stringify(result, null, 2)); + else printReport(result); +} + +main().catch((error) => { + console.error(error instanceof Error ? error.stack || error.message : String(error)); + process.exit(1); +}); diff --git a/src-tauri/src/commands/document_cmd.rs b/src-tauri/src/commands/document_cmd.rs index fb222b8482..82930283df 100644 --- a/src-tauri/src/commands/document_cmd.rs +++ b/src-tauri/src/commands/document_cmd.rs @@ -6,7 +6,11 @@ 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 +pub(crate) async fn run_cancellable( + state: &Arc, + execution_id: Option, + future: F, +) -> Result where F: Future>, { diff --git a/src-tauri/src/commands/mcp_bridge.rs b/src-tauri/src/commands/mcp_bridge.rs index ad885308f9..b20f51ac52 100644 --- a/src-tauri/src/commands/mcp_bridge.rs +++ b/src-tauri/src/commands/mcp_bridge.rs @@ -62,6 +62,16 @@ struct MongoFindDocumentsRequest { sort: Option, } +#[derive(Deserialize)] +struct MongoCountDocumentsRequest { + connection_name: String, + connection_id: Option, + database: Option, + collection: String, + filter: Option, + mode: Option, +} + #[derive(Deserialize)] struct MongoServerVersionRequest { connection_name: String, @@ -214,6 +224,8 @@ pub fn start(app_handle: AppHandle, state: Arc, data_dir: PathBuf) { handle_describe_table_data(&st, body, &mut stream).await; } else if first_line.starts_with("POST /data/mongo/list-collections") { handle_mongo_list_collections_data(&st, body, &mut stream).await; + } else if first_line.starts_with("POST /data/mongo/count-documents") { + handle_mongo_count_documents_data(&st, body, &mut stream).await; } else if first_line.starts_with("POST /data/mongo/find-documents") { handle_mongo_find_documents_data(&st, body, &mut stream).await; } else if first_line.starts_with("POST /data/mongo/server-version") { @@ -518,6 +530,34 @@ async fn handle_mongo_find_documents_data(state: &Arc, body: &str, str } } +async fn handle_mongo_count_documents_data(state: &Arc, body: &str, stream: &mut tokio::net::TcpStream) { + let req: MongoCountDocumentsRequest = match serde_json::from_str(body) { + Ok(r) => r, + Err(_) => { + respond_error(stream, "400 Bad Request", "Invalid JSON").await; + return; + } + }; + let Some((pool_key, database, _connection_id)) = + resolve_mongo_pool_key(state, req.connection_id.as_deref(), &req.connection_name, req.database, stream).await + else { + return; + }; + match dbx_core::mongo_ops::mongo_count_documents_core( + state, + &pool_key, + &database, + &req.collection, + req.filter.as_deref(), + req.mode.as_deref(), + ) + .await + { + Ok(total) => respond_json(stream, &total).await, + Err(e) => respond_error(stream, "500 Internal Server Error", &e).await, + } +} + async fn handle_mongo_server_version_data(state: &Arc, body: &str, stream: &mut tokio::net::TcpStream) { let req: MongoServerVersionRequest = match serde_json::from_str(body) { Ok(r) => r, diff --git a/src-tauri/src/commands/mongo_cmd.rs b/src-tauri/src/commands/mongo_cmd.rs index b3f29a004a..316237fe99 100644 --- a/src-tauri/src/commands/mongo_cmd.rs +++ b/src-tauri/src/commands/mongo_cmd.rs @@ -110,6 +110,32 @@ pub async fn mongo_find_documents( .await } +#[tauri::command] +pub async fn mongo_count_documents( + state: State<'_, Arc>, + connection_id: String, + database: String, + collection: String, + filter: Option, + mode: Option, + execution_id: Option, +) -> Result { + let app = state.inner().clone(); + crate::commands::document_cmd::run_cancellable( + &app, + execution_id, + dbx_core::mongo_ops::mongo_count_documents_core( + &app, + &connection_id, + &database, + &collection, + filter.as_deref(), + mode.as_deref(), + ), + ) + .await +} + #[tauri::command] pub async fn mongo_server_version( state: State<'_, Arc>, diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 6fc6afb396..bad079ba03 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -1128,6 +1128,7 @@ pub fn run() { commands::document_cmd::document_upload_gridfs_file, commands::document_cmd::document_delete_gridfs_file, commands::mongo_cmd::mongo_find_documents, + commands::mongo_cmd::mongo_count_documents, commands::mongo_cmd::mongo_server_version, commands::mongo_cmd::mongo_collection_stats, commands::mongo_cmd::mongo_aggregate_documents,