From 91c77e53c41256a8c73997e088a7a80ad9a51b07 Mon Sep 17 00:00:00 2001 From: onenewcode Date: Fri, 3 Jul 2026 21:32:29 +0800 Subject: [PATCH] feat(mongodb): support collection stats shell commands Add db.coll.stats() / dataSize() / storageSize(scale) / totalIndexSize() so these MongoDB shell commands return a result grid instead of the "Use MongoDB shell-style commands" error (issue #2529). - backend: collStats via mongo_driver.collection_stats, wired through mongo_ops, the dbx-web /mongo/collection-stats route, the tauri mongo_collection_stats command, and the MCP data bridge - node-core: parseMongoCollectionStatsCommand + result formatter and executeQuery dispatch (web build + MCP path) - desktop: mirror the parser/executor/autocomplete in mongoShellCommand.ts, queryStore.ts, api/tauri/http facade, and mongoCompletion.ts - tests for the parser, result formatter, and completion entries Close #2529 --- apps/desktop/src/lib/api.ts | 1 + apps/desktop/src/lib/http.ts | 5 + apps/desktop/src/lib/mongoCompletion.ts | 4 + apps/desktop/src/lib/mongoShellCommand.ts | 56 ++++++++++ apps/desktop/src/lib/tauri.ts | 13 +++ apps/desktop/src/stores/queryStore.ts | 20 ++++ crates/dbx-core/src/db/mongo_driver.rs | 106 +++++++++++++++++++ crates/dbx-core/src/mongo_ops.rs | 18 +++- crates/dbx-web/src/main.rs | 1 + crates/dbx-web/src/routes/mongo.rs | 29 +++++ packages/app-tests/mongoCompletion.test.ts | 9 ++ packages/app-tests/mongoShellCommand.test.ts | 63 +++++++++++ packages/node-core/src/database.ts | 71 ++++++++++++- packages/node-core/src/web-backend.ts | 18 +++- packages/node-core/tests/mongo-query.test.ts | 59 ++++++++++- src-tauri/src/commands/mcp_bridge.rs | 31 ++++++ src-tauri/src/commands/mongo_cmd.rs | 18 ++++ src-tauri/src/lib.rs | 1 + 18 files changed, 518 insertions(+), 5 deletions(-) diff --git a/apps/desktop/src/lib/api.ts b/apps/desktop/src/lib/api.ts index 782d6de01e..b1bc62846d 100644 --- a/apps/desktop/src/lib/api.ts +++ b/apps/desktop/src/lib/api.ts @@ -390,6 +390,7 @@ export const documentFindDocuments = forward("documentFindDocuments"); export const mongoFindDocuments = forward("mongoFindDocuments"); export const mongoServerVersion = forward("mongoServerVersion"); export const mongoAggregateDocuments = forward("mongoAggregateDocuments"); +export const mongoCollectionStats = forward("mongoCollectionStats"); export const mongoCreateIndex = forward("mongoCreateIndex"); export const mongoDropIndexes = forward("mongoDropIndexes"); export const documentInsertDocument = forward("documentInsertDocument"); diff --git a/apps/desktop/src/lib/http.ts b/apps/desktop/src/lib/http.ts index 7e9858b3dc..794d5b168f 100644 --- a/apps/desktop/src/lib/http.ts +++ b/apps/desktop/src/lib/http.ts @@ -64,6 +64,7 @@ import type { KvPutResponse, KvDeleteResponse, MongoDocumentResult, + MongoCollectionStatsResult, MongoGridFsBucketInfo, HistoryEntry, SqlFileRequest, @@ -1859,6 +1860,10 @@ export async function mongoAggregateDocuments(connectionId: string, database: st return post("/api/mongo/aggregate-documents", { connectionId, database, collection, pipelineJson, maxRows, executionId }); } +export async function mongoCollectionStats(connectionId: string, database: string, collection: string, scale?: number, executionId?: string): Promise { + return post("/api/mongo/collection-stats", { connectionId, database, collection, scale, executionId }); +} + export async function mongoCreateIndex(connectionId: string, database: string, collection: string, keysJson: string, optionsJson?: string): Promise<{ name: string }> { return post("/api/mongo/create-index", { connectionId, database, collection, keysJson, optionsJson }); } diff --git a/apps/desktop/src/lib/mongoCompletion.ts b/apps/desktop/src/lib/mongoCompletion.ts index ddea7d7047..211bbd262c 100644 --- a/apps/desktop/src/lib/mongoCompletion.ts +++ b/apps/desktop/src/lib/mongoCompletion.ts @@ -39,6 +39,10 @@ const COLLECTION_METHODS = [ { label: "deleteOne", detail: "Delete one matching document", apply: "deleteOne({})" }, { label: "deleteMany", detail: "Delete all matching documents", apply: "deleteMany({})" }, { label: "getIndexes", detail: "List collection indexes", apply: "getIndexes()" }, + { label: "stats", detail: "Show collection statistics", apply: "stats()" }, + { label: "dataSize", detail: "Total size of documents in bytes", apply: "dataSize()" }, + { label: "storageSize", detail: "Allocated storage size in bytes", apply: "storageSize()" }, + { label: "totalIndexSize", detail: "Total size of all indexes in bytes", apply: "totalIndexSize()" }, { label: "createIndex", detail: "Create an index", apply: "createIndex({ ${field}: 1 })" }, { label: "dropIndex", detail: "Drop one index", apply: 'dropIndex("${indexName}")' }, { label: "dropIndexes", detail: "Drop collection indexes", apply: "dropIndexes()" }, diff --git a/apps/desktop/src/lib/mongoShellCommand.ts b/apps/desktop/src/lib/mongoShellCommand.ts index 5864315093..9add566732 100644 --- a/apps/desktop/src/lib/mongoShellCommand.ts +++ b/apps/desktop/src/lib/mongoShellCommand.ts @@ -31,6 +31,14 @@ export interface MongoVersionCommand { kind: "version"; } +export type MongoCollectionStatsMetric = "stats" | "dataSize" | "storageSize" | "totalIndexSize"; + +export interface MongoCollectionStatsCommand { + collection: string; + metric: MongoCollectionStatsMetric; + scale?: number; +} + type MongoWriteKind = "insert" | "update" | "delete" | "createIndex" | "dropIndex" | "dropIndexes"; export type MongoCommand = @@ -39,6 +47,7 @@ export type MongoCommand = | ({ kind: "countDocuments" } & MongoCountDocumentsCommand) | ({ kind: "aggregate" } & MongoAggregateCommand) | ({ kind: "getIndexes" } & MongoGetIndexesCommand) + | ({ kind: "collectionStats" } & MongoCollectionStatsCommand) | ({ kind: "use" } & MongoUseCommand) | { kind: "insert"; collection: string; docsJson: string } | { kind: "update"; collection: string; filter: string; update: string; many: boolean } @@ -223,6 +232,29 @@ export function parseMongoGetIndexesCommand(input: string): MongoGetIndexesComma }; } +export function parseMongoCollectionStatsCommand(input: string): MongoCollectionStatsCommand | null { + const source = input.trim().replace(/;$/, "").trim(); + for (const metric of ["stats", "dataSize", "storageSize", "totalIndexSize"] as const) { + const target = parseCollectionMethodTarget(source, metric); + if (!target) continue; + const args = parseMethodArgs(source, target.methodCallIndex); + if (!args) return null; + const scale = parseMongoCollectionStatsScale(args); + return scale === null ? null : { collection: target.collection, metric, ...(scale === undefined ? {} : { scale }) }; + } + return null; +} + +function parseMongoCollectionStatsScale(args: string[]): number | undefined | null { + if (args.length === 1 && !args[0]?.trim()) return undefined; + if (args.length !== 1) return null; + const raw = args[0].trim(); + if (!/^[+-]?(?:\d+\.?\d*|\.\d+)(?:[eE][+-]?\d+)?$/.test(raw)) return null; + const scale = Number(raw); + if (!Number.isFinite(scale)) return null; + return scale; +} + export function parseMongoUseCommand(input: string): MongoUseCommand | null { const source = input.trim().replace(/;$/, "").trim(); const match = /^use\s+([a-zA-Z0-9_-]+)$/i.exec(source); @@ -340,6 +372,10 @@ export function parseMongoCommand(input: string): ParsedMongoCommand | null { const getIndexes = parseMongoGetIndexesCommand(source); return getIndexes ? { kind: "getIndexes", ...getIndexes } : null; }, + (source) => { + const stats = parseMongoCollectionStatsCommand(source); + return stats ? { kind: "collectionStats", ...stats } : null; + }, (source) => { const write = parseMongoWriteCommand(source); return write ?? null; @@ -529,6 +565,26 @@ export function mongoIndexesToQueryResult( }; } +export function mongoCollectionStatsToQueryResult(metric: MongoCollectionStatsMetric, stats: Record, executionTimeMs: number): QueryResult { + const execution_time_ms = Math.max(0, Math.round(executionTimeMs)); + if (metric === "stats") { + const columns = ["count", "size", "avgObjSize", "storageSize", "totalIndexSize", "nindexes"]; + return { + columns, + rows: [columns.map((column) => (column in stats ? toCellValue(stats[column]) : null))], + affected_rows: 1, + execution_time_ms, + }; + } + const sourceField = metric === "dataSize" ? "size" : metric; + return { + columns: [metric], + rows: [[sourceField in stats ? toCellValue(stats[sourceField]) : null]], + affected_rows: 1, + execution_time_ms, + }; +} + function parseFindTarget(source: string): { collection: string; findCallIndex: number } | null { const direct = parseCollectionMethodTarget(source, "find"); if (direct) { diff --git a/apps/desktop/src/lib/tauri.ts b/apps/desktop/src/lib/tauri.ts index 5981bb3997..c7b63ebb2f 100644 --- a/apps/desktop/src/lib/tauri.ts +++ b/apps/desktop/src/lib/tauri.ts @@ -1484,6 +1484,15 @@ export interface MongoDocumentResult { total: number; } +export interface MongoCollectionStatsResult { + count: unknown; + size: unknown; + avgObjSize: unknown; + storageSize: unknown; + totalIndexSize: unknown; + nindexes: unknown; +} + export interface MongoGridFsFileInfo { id: string; filename?: string; @@ -1595,6 +1604,10 @@ export async function mongoAggregateDocuments(connectionId: string, database: st return invoke("mongo_aggregate_documents", { connectionId, database, collection, pipelineJson, maxRows, executionId }); } +export async function mongoCollectionStats(connectionId: string, database: string, collection: string, scale?: number, executionId?: string): Promise { + return invoke("mongo_collection_stats", { connectionId, database, collection, scale, executionId }); +} + export async function mongoCreateIndex(connectionId: string, database: string, collection: string, keysJson: string, optionsJson?: string): Promise<{ name: string }> { return invoke("mongo_create_index", { connectionId, database, collection, keysJson, optionsJson }); } diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index edac43766e..440bb0b0a0 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -11,6 +11,7 @@ import { ACTIVE_TAB_STORAGE_KEY, OPEN_TABS_STORAGE_KEY, restoreOpenTabsState, se import { evaluateMongoAggregateSafety, evaluateMongoWriteSafety, + mongoCollectionStatsToQueryResult, mongoCountToQueryResult, mongoCreateIndexToQueryResult, mongoDocumentsToQueryResult, @@ -1996,6 +1997,25 @@ export const useQueryStore = defineStore("query", () => { }); break; } + case "collectionStats": { + console.info("[DBX][executeTabSql:mongo-collection-stats:start]", { + traceId, + collection: mongoCommand.collection, + metric: mongoCommand.metric, + database: currentDatabase, + }); + const stats = await api.mongoCollectionStats(tab.connectionId, currentDatabase, mongoCommand.collection, mongoCommand.scale, executionId); + allResults.push(markQueryResultRowsRaw(mongoCollectionStatsToQueryResult(mongoCommand.metric, stats as unknown as Record, performance.now() - commandStartedAt))); + mongoEditTarget = undefined; + console.info("[DBX][executeTabSql:mongo-collection-stats:done]", { + traceId, + collection: mongoCommand.collection, + metric: mongoCommand.metric, + database: currentDatabase, + elapsed: elapsed(), + }); + break; + } case "insert": case "update": case "delete": diff --git a/crates/dbx-core/src/db/mongo_driver.rs b/crates/dbx-core/src/db/mongo_driver.rs index 158edcc3ab..9b1bccc12c 100644 --- a/crates/dbx-core/src/db/mongo_driver.rs +++ b/crates/dbx-core/src/db/mongo_driver.rs @@ -24,6 +24,19 @@ pub struct MongoDropIndexesResult { pub affected_rows: u64, } +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub struct MongoCollectionStatsResult { + pub count: serde_json::Value, + pub size: serde_json::Value, + #[serde(rename = "avgObjSize")] + pub avg_obj_size: serde_json::Value, + #[serde(rename = "storageSize")] + pub storage_size: serde_json::Value, + #[serde(rename = "totalIndexSize")] + pub total_index_size: serde_json::Value, + pub nindexes: serde_json::Value, +} + pub async fn connect(url: &str, timeout: Duration, idle_timeout: Duration) -> Result { let url = normalize_mongo_uri_direct_connection(url); let is_multi_host = is_multi_host_mongo_uri(&url); @@ -136,6 +149,52 @@ fn server_version_from_build_info(result: &Document) -> Result { result.get_str("version").map(str::to_string).map_err(|e| format!("MongoDB server version not found: {e}")) } +pub async fn collection_stats( + client: &Client, + database: &str, + collection: &str, + scale: Option, +) -> Result { + let database = database.trim(); + let collection = collection.trim(); + if database.is_empty() { + return Err("Database name is required".to_string()); + } + if collection.is_empty() { + return Err("Collection name is required".to_string()); + } + + let result = client + .database(database) + .run_command(collection_stats_command_document(collection, scale.as_ref())) + .await + .map_err(|e| e.to_string())?; + Ok(collection_stats_result_from_document(&result)) +} + +fn collection_stats_command_document(collection: &str, scale: Option<&serde_json::Number>) -> Document { + let mut command = doc! { "collStats": collection }; + if let Some(scale) = scale { + command.insert("scale", json_value_to_bson(&serde_json::Value::Number(scale.clone()))); + } + command +} + +fn collection_stats_result_from_document(result: &Document) -> MongoCollectionStatsResult { + MongoCollectionStatsResult { + count: collection_stats_field(result, "count"), + size: collection_stats_field(result, "size"), + avg_obj_size: collection_stats_field(result, "avgObjSize"), + storage_size: collection_stats_field(result, "storageSize"), + total_index_size: collection_stats_field(result, "totalIndexSize"), + nindexes: collection_stats_field(result, "nindexes"), + } +} + +fn collection_stats_field(result: &Document, key: &str) -> serde_json::Value { + result.get(key).map(bson_to_json).unwrap_or(serde_json::Value::Null) +} + pub async fn list_databases(client: &Client) -> Result, String> { client.list_database_names().await.map_err(|e| e.to_string()) } @@ -1550,6 +1609,53 @@ mod tests { assert!(error.contains("MongoDB server version not found")); } + #[test] + fn collection_stats_result_reads_expected_fields() { + let result = collection_stats_result_from_document(&doc! { + "count": 12_i64, + "size": 4096_i64, + "avgObjSize": 341.3_f64, + "storageSize": 8192_i64, + "totalIndexSize": 2048_i64, + "nindexes": 3_i32, + }); + + assert_eq!( + result, + MongoCollectionStatsResult { + count: serde_json::json!(12), + size: serde_json::json!(4096), + avg_obj_size: serde_json::json!(341.3), + storage_size: serde_json::json!(8192), + total_index_size: serde_json::json!(2048), + nindexes: serde_json::json!(3), + } + ); + } + + #[test] + fn collection_stats_result_fills_missing_fields_with_null() { + let result = collection_stats_result_from_document(&doc! { + "count": 7_i32, + "storageSize": 512_i32, + }); + + assert_eq!(result.count, serde_json::json!(7)); + assert_eq!(result.size, serde_json::Value::Null); + assert_eq!(result.avg_obj_size, serde_json::Value::Null); + assert_eq!(result.storage_size, serde_json::json!(512)); + assert_eq!(result.total_index_size, serde_json::Value::Null); + assert_eq!(result.nindexes, serde_json::Value::Null); + } + + #[test] + fn collection_stats_command_serializes_scale() { + let command = collection_stats_command_document("users", Some(&serde_json::Number::from(1024))); + + assert_eq!(command.get_str("collStats").unwrap(), "users"); + assert!(matches!(command.get("scale"), Some(Bson::Int64(1024)))); + } + #[test] fn json_object_to_document_parses_mongo_shell_isodate_strings() { let value = serde_json::json!({ diff --git a/crates/dbx-core/src/mongo_ops.rs b/crates/dbx-core/src/mongo_ops.rs index 8815399089..2969e9a6a4 100644 --- a/crates/dbx-core/src/mongo_ops.rs +++ b/crates/dbx-core/src/mongo_ops.rs @@ -1,5 +1,5 @@ use crate::connection::{AppState, PoolKind}; -use crate::db::mongo_driver::{self, MongoDocumentResult, MongoDropIndexesResult}; +use crate::db::mongo_driver::{self, MongoCollectionStatsResult, MongoDocumentResult, MongoDropIndexesResult}; use crate::document_ops::CollectionInfo; async fn ensure_document_pool(state: &AppState, connection_id: &str) -> Result<(), String> { @@ -70,6 +70,22 @@ pub async fn mongo_server_version_core( } } +pub async fn mongo_collection_stats_core( + state: &AppState, + connection_id: &str, + database: &str, + collection: &str, + scale: Option, +) -> 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::collection_stats(client, database, collection, scale).await, + PoolKind::Agent(_) => Err("MongoDB legacy agent does not support collection stats helpers".to_string()), + _ => Err("Not a MongoDB connection".to_string()), + } +} + #[allow(clippy::too_many_arguments)] pub async fn mongo_find_documents_core( state: &AppState, diff --git a/crates/dbx-web/src/main.rs b/crates/dbx-web/src/main.rs index 522563c621..9f660f2461 100644 --- a/crates/dbx-web/src/main.rs +++ b/crates/dbx-web/src/main.rs @@ -461,6 +461,7 @@ async fn main() { .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/collection-stats", post(routes::mongo::collection_stats)) .route("/mongo/aggregate-documents", post(routes::mongo::aggregate_documents)) .route("/mongo/create-index", post(routes::mongo::create_index)) .route("/mongo/drop-indexes", post(routes::mongo::drop_indexes)) diff --git a/crates/dbx-web/src/routes/mongo.rs b/crates/dbx-web/src/routes/mongo.rs index a1654e060e..bb09cfff06 100644 --- a/crates/dbx-web/src/routes/mongo.rs +++ b/crates/dbx-web/src/routes/mongo.rs @@ -86,6 +86,16 @@ pub struct MongoServerVersionRequest { pub execution_id: Option, } +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct MongoCollectionStatsRequest { + pub connection_id: String, + pub database: String, + pub collection: String, + pub scale: Option, + pub execution_id: Option, +} + #[derive(Deserialize)] #[serde(rename_all = "camelCase")] pub struct MongoAggregateRequest { @@ -288,6 +298,25 @@ pub async fn server_version( Ok(Json(serde_json::to_value(result).map_err(|e| AppError(e.to_string()))?)) } +pub async fn collection_stats( + State(state): State>, + Json(req): Json, +) -> Result, AppError> { + let result = run_cancellable( + &state, + req.execution_id.clone(), + dbx_core::mongo_ops::mongo_collection_stats_core( + &state.app, + &req.connection_id, + &req.database, + &req.collection, + req.scale, + ), + ) + .await?; + Ok(Json(serde_json::to_value(result).map_err(|e| AppError(e.to_string()))?)) +} + pub async fn aggregate_documents( State(state): State>, Json(req): Json, diff --git a/packages/app-tests/mongoCompletion.test.ts b/packages/app-tests/mongoCompletion.test.ts index dcb6b24a31..bf2d8d35fe 100644 --- a/packages/app-tests/mongoCompletion.test.ts +++ b/packages/app-tests/mongoCompletion.test.ts @@ -41,6 +41,15 @@ test("suggests collection methods after direct and getCollection references", () assert.ok(labels('db.getCollection("users").ag').includes("aggregate")); }); +test("suggests collection stats methods after a collection reference", () => { + const methodLabels = labels("db.users."); + for (const method of ["stats", "dataSize", "storageSize", "totalIndexSize"]) { + assert.ok(methodLabels.includes(method), `expected completion to include ${method}`); + } + const item = buildMongoCompletionItems("db.users.stat", "db.users.stat".length).find((candidate) => candidate.label === "stats"); + assert.equal(item?.apply, "stats()"); +}); + test("suggests cursor methods after find result chains", () => { const allItems = buildMongoCompletionItems("db.characters.find({}).", "db.characters.find({}).".length); const prefixedItems = buildMongoCompletionItems("db.characters.find({}).li", "db.characters.find({}).li".length); diff --git a/packages/app-tests/mongoShellCommand.test.ts b/packages/app-tests/mongoShellCommand.test.ts index 088d4880d7..5952ff9d88 100644 --- a/packages/app-tests/mongoShellCommand.test.ts +++ b/packages/app-tests/mongoShellCommand.test.ts @@ -4,10 +4,12 @@ import { evaluateMongoAggregateSafety, evaluateMongoWriteSafety, mongoAggregateWriteStage, + mongoCollectionStatsToQueryResult, mongoCountToQueryResult, mongoDocumentsToQueryResult, mongoIndexesToQueryResult, parseMongoAggregateCommand, + parseMongoCollectionStatsCommand, parseMongoCommand, parseMongoCountDocumentsCommand, parseMongoFindCommand, @@ -250,6 +252,67 @@ test("parseMongoGetIndexesCommand parses collection index commands", () => { assert.equal(parseMongoGetIndexesCommand("db.web_log.getIndexes({})"), null); }); +test("parseMongoCollectionStatsCommand parses collection stats commands", () => { + assert.deepEqual(parseMongoCollectionStatsCommand("db.users.stats()"), { + collection: "users", + metric: "stats", + }); + assert.deepEqual(parseMongoCollectionStatsCommand("db.users.dataSize();"), { + collection: "users", + metric: "dataSize", + }); + assert.deepEqual(parseMongoCollectionStatsCommand("db.users.storageSize(1024)"), { + collection: "users", + metric: "storageSize", + scale: 1024, + }); + assert.deepEqual(parseMongoCollectionStatsCommand("db.users.totalIndexSize()"), { + collection: "users", + metric: "totalIndexSize", + }); + assert.deepEqual(parseMongoCollectionStatsCommand('db.getCollection("audit.logs").stats()'), { + collection: "audit.logs", + metric: "stats", + }); + // A non-numeric argument is rejected rather than silently ignored. + assert.equal(parseMongoCollectionStatsCommand('db.users.storageSize("big")'), null); +}); + +test("parseMongoCommand tags collection stats commands with the collectionStats kind", () => { + const parsed = parseMongoCommand("db.users.stats()"); + assert.ok(parsed); + assert.deepEqual(parsed.command, { kind: "collectionStats", collection: "users", metric: "stats" }); +}); + +test("mongoCollectionStatsToQueryResult formats stats and single-metric results", () => { + const stats = { + count: 12, + size: 4096, + avgObjSize: 341, + storageSize: 8192, + totalIndexSize: 2048, + nindexes: 3, + }; + assert.deepEqual(mongoCollectionStatsToQueryResult("stats", stats, 5), { + columns: ["count", "size", "avgObjSize", "storageSize", "totalIndexSize", "nindexes"], + rows: [[12, 4096, 341, 8192, 2048, 3]], + affected_rows: 1, + execution_time_ms: 5, + }); + assert.deepEqual(mongoCollectionStatsToQueryResult("dataSize", stats, 0), { + columns: ["dataSize"], + rows: [[4096]], + affected_rows: 1, + execution_time_ms: 0, + }); + assert.deepEqual(mongoCollectionStatsToQueryResult("totalIndexSize", {}, 0), { + columns: ["totalIndexSize"], + rows: [[null]], + affected_rows: 1, + execution_time_ms: 0, + }); +}); + test("splitMongoCommands keeps semicolon-separated mongo commands in order", () => { const commands = splitMongoCommands(` db.users.insertOne({ name: "A" }); diff --git a/packages/node-core/src/database.ts b/packages/node-core/src/database.ts index 19b1fc7299..8f1c2a93b6 100644 --- a/packages/node-core/src/database.ts +++ b/packages/node-core/src/database.ts @@ -807,6 +807,11 @@ export async function executeQuery(config: ConnectionConfig, sql: string, option const result = await withTimeout(mongoAggregateDocuments(config, getIndexes.collection, '[{"$indexStats":{}}]', resolveMaxRows(options)), resolveTimeoutMs(options)); return mongoDocumentsToQueryResult(result.documents.slice(0, resolveMaxRows(options)), result.total); } + const collectionStats = parseMongoCollectionStatsCommand(sql); + if (collectionStats) { + const result = await withTimeout(mongoCollectionStats(config, collectionStats.collection, collectionStats.scale), resolveTimeoutMs(options)); + return mongoCollectionStatsToQueryResult(collectionStats.metric, result); + } const write = parseMongoWriteCommand(sql); if (write) { const safety = evaluateMongoWriteSafety(write, sqlSafetyFromEnv()); @@ -829,7 +834,7 @@ export async function executeQuery(config: ConnectionConfig, sql: string, option return { columns: [], rows: [], row_count: result.affectedRows }; } throw new Error( - "Use MongoDB shell-style commands, for example: db.projects.find({}).limit(100), db.version(), db.projects.countDocuments({}), db.projects.count({}), db.projects.getIndexes(), db.projects.createIndex({...}), db.projects.dropIndex(\"name\"), db.projects.dropIndexes(), db.projects.insertOne({...}), db.projects.updateOne({...}, {$set: {...}}), or db.projects.deleteOne({...})", + "Use MongoDB shell-style commands, for example: db.projects.find({}).limit(100), db.version(), db.projects.countDocuments({}), db.projects.count({}), db.projects.getIndexes(), db.projects.dataSize(), db.projects.storageSize(1024), db.projects.totalIndexSize(), db.projects.stats(), db.projects.createIndex({...}), db.projects.dropIndex(\"name\"), db.projects.dropIndexes(), db.projects.insertOne({...}), db.projects.updateOne({...}, {$set: {...}}), or db.projects.deleteOne({...})", ); } if (isDirectQueryType(config.db_type)) { @@ -1045,6 +1050,19 @@ async function mongoServerVersion(config: ConnectionConfig): Promise { }); } +async function mongoCollectionStats( + config: ConnectionConfig, + collection: string, + scale?: number, +): Promise> { + return bridgeDataRequest>("/data/mongo/collection-stats", { + connection_name: config.name, + database: config.database || "", + collection, + scale, + }); +} + async function executeMongoWrite( config: ConnectionConfig, command: MongoWriteCommand, @@ -1109,6 +1127,26 @@ async function mongoAggregateDocuments(config: ConnectionConfig, collection: str }); } +export function mongoCollectionStatsToQueryResult( + metric: MongoCollectionStatsMetric, + stats: Record, +): QueryResult { + if (metric === "stats") { + const columns = ["count", "size", "avgObjSize", "storageSize", "totalIndexSize", "nindexes"]; + const row: Record = {}; + for (const column of columns) { + row[column] = column in stats ? toCellValue(stats[column]) : null; + } + return { columns, rows: [row], row_count: 1 }; + } + const sourceField = metric === "dataSize" ? "size" : metric; + return { + columns: [metric], + rows: [{ [metric]: sourceField in stats ? toCellValue(stats[sourceField]) : null }], + row_count: 1, + }; +} + export function mongoDocumentsToQueryResult(documents: unknown[], _total: number): QueryResult { const columns: string[] = []; for (const doc of documents) { @@ -1179,6 +1217,14 @@ interface MongoGetIndexesCommand { collection: string; } +type MongoCollectionStatsMetric = "stats" | "dataSize" | "storageSize" | "totalIndexSize"; + +interface MongoCollectionStatsCommand { + collection: string; + metric: MongoCollectionStatsMetric; + scale?: number; +} + export type MongoWriteCommand = | { kind: "insert"; collection: string; docsJson: string } | { kind: "update"; collection: string; filter: string; update: string; many: boolean } @@ -1282,6 +1328,19 @@ export function parseMongoGetIndexesCommand(input: string): MongoGetIndexesComma return { collection: target.collection }; } +export function parseMongoCollectionStatsCommand(input: string): MongoCollectionStatsCommand | null { + const source = input.trim().replace(/;$/, "").trim(); + for (const metric of ["stats", "dataSize", "storageSize", "totalIndexSize"] as const) { + const target = parseCollectionMethodTarget(source, metric); + if (!target) continue; + const args = parseMethodArgs(source, target.methodCallIndex); + if (!args) return null; + const scale = parseMongoCollectionStatsScale(args); + return scale === null ? null : { collection: target.collection, metric, ...(scale === undefined ? {} : { scale }) }; + } + return null; +} + export function mongoAggregateWriteStage(pipelineJson: string): "$out" | "$merge" | null { try { const pipeline = JSON.parse(pipelineJson); @@ -1497,6 +1556,16 @@ function parseMongoDropIndexesArgument(args: string[]): string | undefined | nul return Array.isArray(parsed) && parsed.length > 0 && parsed.every((item) => typeof item === "string") ? normalized : null; } +function parseMongoCollectionStatsScale(args: string[]): number | undefined | null { + if (args.length === 1 && !args[0]?.trim()) return undefined; + if (args.length !== 1) return null; + const raw = args[0].trim(); + if (!/^[+-]?(?:\d+\.?\d*|\.\d+)(?:[eE][+-]?\d+)?$/.test(raw)) return null; + const scale = Number(raw); + if (!Number.isFinite(scale)) return null; + return scale; +} + function convertSingleQuotedStrings(source: string): string { let result = ""; let copiedUntil = 0; diff --git a/packages/node-core/src/web-backend.ts b/packages/node-core/src/web-backend.ts index 60f848ae9f..1663f8dd04 100644 --- a/packages/node-core/src/web-backend.ts +++ b/packages/node-core/src/web-backend.ts @@ -1,6 +1,6 @@ import type { ConnectionConfig } from "./connections.js"; import type { TableInfo, ColumnInfo, QueryOptions, QueryResult } from "./database.js"; -import { collectionListToTableInfos, evaluateMongoAggregateSafety, evaluateMongoWriteSafety, inferMongoColumns, mongoDocumentsToQueryResult, parseMongoAggregateCommand, parseMongoCountDocumentsCommand, parseMongoFindCommand, parseMongoGetIndexesCommand, parseMongoVersionCommand, parseMongoWriteCommand, type CollectionInfo, type MongoWriteCommand } from "./database.js"; +import { collectionListToTableInfos, evaluateMongoAggregateSafety, evaluateMongoWriteSafety, inferMongoColumns, mongoCollectionStatsToQueryResult, mongoDocumentsToQueryResult, parseMongoAggregateCommand, parseMongoCollectionStatsCommand, parseMongoCountDocumentsCommand, parseMongoFindCommand, parseMongoGetIndexesCommand, parseMongoVersionCommand, parseMongoWriteCommand, type CollectionInfo, type MongoWriteCommand } from "./database.js"; import type { RedisCommandOptions, RedisCommandResult } from "./redis-command.js"; import { sqlSafetyFromEnv } from "./sql-safety.js"; @@ -206,6 +206,20 @@ export async function executeQuery(config: ConnectionConfig, sql: string, option const result = (await res.json()) as { documents: unknown[]; total: number }; return mongoDocumentsToQueryResult(result.documents.slice(0, options?.maxRows ?? result.documents.length), result.total); } + const collectionStats = parseMongoCollectionStatsCommand(sql); + if (collectionStats) { + const res = await apiFetch("/api/mongo/collection-stats", { + method: "POST", + body: JSON.stringify({ + connectionId: config.id, + database: config.database || "", + collection: collectionStats.collection, + scale: collectionStats.scale, + }), + }); + const result = (await res.json()) as Record; + return mongoCollectionStatsToQueryResult(collectionStats.metric, result); + } const write = parseMongoWriteCommand(sql); if (write) { const safety = evaluateMongoWriteSafety(write, sqlSafetyFromEnv()); @@ -220,7 +234,7 @@ export async function executeQuery(config: ConnectionConfig, sql: string, option return { columns: [], rows: [], row_count: result.affectedRows }; } throw new Error( - "Use MongoDB shell-style commands, for example: db.projects.find({}).limit(100), db.version(), db.projects.countDocuments({}), db.projects.count({}), db.projects.getIndexes(), db.projects.createIndex({...}), db.projects.dropIndex(\"name\"), db.projects.dropIndexes(), db.projects.insertOne({...}), db.projects.updateOne({...}, {$set: {...}}), or db.projects.deleteOne({...})", + "Use MongoDB shell-style commands, for example: db.projects.find({}).limit(100), db.version(), db.projects.countDocuments({}), db.projects.count({}), db.projects.getIndexes(), db.projects.dataSize(), db.projects.storageSize(1024), db.projects.totalIndexSize(), db.projects.stats(), db.projects.createIndex({...}), db.projects.dropIndex(\"name\"), db.projects.dropIndexes(), db.projects.insertOne({...}), db.projects.updateOne({...}, {$set: {...}}), or db.projects.deleteOne({...})", ); } const res = await apiFetch("/api/query/execute", { diff --git a/packages/node-core/tests/mongo-query.test.ts b/packages/node-core/tests/mongo-query.test.ts index 44ad8eca99..5905bc41ae 100644 --- a/packages/node-core/tests/mongo-query.test.ts +++ b/packages/node-core/tests/mongo-query.test.ts @@ -1,6 +1,6 @@ import assert from "node:assert/strict"; import { test } from "vitest"; -import { executeQuery, inferMongoColumns, mongoAggregateWriteStage, mongoDocumentsToQueryResult, parseMongoAggregateCommand, parseMongoCountDocumentsCommand, parseMongoFindCommand, parseMongoGetIndexesCommand, parseMongoVersionCommand, parseMongoWriteCommand } from "../src/database.js"; +import { executeQuery, inferMongoColumns, mongoAggregateWriteStage, mongoCollectionStatsToQueryResult, mongoDocumentsToQueryResult, parseMongoAggregateCommand, parseMongoCollectionStatsCommand, parseMongoCountDocumentsCommand, parseMongoFindCommand, parseMongoGetIndexesCommand, parseMongoVersionCommand, parseMongoWriteCommand } from "../src/database.js"; test("parseMongoFindCommand accepts shell-style find commands", () => { assert.deepEqual(parseMongoFindCommand('db.getCollection("operation_logs").find({"level":"info"}).sort({"ts":-1}).skip(5).limit(10)'), { @@ -100,6 +100,63 @@ test("parseMongoGetIndexesCommand accepts shell-style index commands", () => { assert.equal(parseMongoGetIndexesCommand("db.web_log.getIndexes({})"), null); }); +test("parseMongoCollectionStatsCommand accepts Mongo shell stats helpers", () => { + assert.deepEqual(parseMongoCollectionStatsCommand("db.users.dataSize()"), { + collection: "users", + metric: "dataSize", + }); + assert.deepEqual(parseMongoCollectionStatsCommand('db.getCollection("audit.logs").dataSize(1024)'), { + collection: "audit.logs", + metric: "dataSize", + scale: 1024, + }); + assert.deepEqual(parseMongoCollectionStatsCommand("db.users.storageSize(1024)"), { + collection: "users", + metric: "storageSize", + scale: 1024, + }); + assert.deepEqual(parseMongoCollectionStatsCommand("db.users.totalIndexSize()"), { + collection: "users", + metric: "totalIndexSize", + }); + assert.deepEqual(parseMongoCollectionStatsCommand("db.users.stats()"), { + collection: "users", + metric: "stats", + }); + assert.deepEqual(parseMongoCollectionStatsCommand("db.users.stats(1024)"), { + collection: "users", + metric: "stats", + scale: 1024, + }); +}); + +test("parseMongoCollectionStatsCommand rejects unsupported stats helper arguments", () => { + assert.equal(parseMongoCollectionStatsCommand("db.users.dataSize(1, 2)"), null); + assert.equal(parseMongoCollectionStatsCommand("db.users.storageSize({scale: 1024})"), null); + assert.equal(parseMongoCollectionStatsCommand("db.users.stats().limit(1)"), null); +}); + +test("mongoCollectionStatsToQueryResult maps dataSize helper to collStats size", () => { + assert.deepEqual(mongoCollectionStatsToQueryResult("dataSize", { size: 2048 }), { + columns: ["dataSize"], + rows: [{ dataSize: 2048 }], + row_count: 1, + }); + assert.deepEqual( + mongoCollectionStatsToQueryResult("stats", { + count: 3, + size: 128, + storageSize: 512, + totalIndexSize: 64, + }), + { + columns: ["count", "size", "avgObjSize", "storageSize", "totalIndexSize", "nindexes"], + rows: [{ count: 3, size: 128, avgObjSize: null, storageSize: 512, totalIndexSize: 64, nindexes: null }], + row_count: 1, + }, + ); +}); + test("mongoAggregateWriteStage detects write stages", () => { assert.equal(mongoAggregateWriteStage('[{"$match":{"active":true}}]'), null); assert.equal(mongoAggregateWriteStage('[{"$match":{}},{"$out":"projects_dump"}]'), "$out"); diff --git a/src-tauri/src/commands/mcp_bridge.rs b/src-tauri/src/commands/mcp_bridge.rs index 807d4204df..2f86076615 100644 --- a/src-tauri/src/commands/mcp_bridge.rs +++ b/src-tauri/src/commands/mcp_bridge.rs @@ -63,6 +63,14 @@ struct MongoServerVersionRequest { database: Option, } +#[derive(Deserialize)] +struct MongoCollectionStatsRequest { + connection_name: String, + database: Option, + collection: String, + scale: Option, +} + #[derive(Deserialize)] struct MongoAggregateDocumentsRequest { connection_name: String, @@ -187,6 +195,8 @@ pub fn start(app_handle: AppHandle, state: Arc, data_dir: PathBuf) { handle_mongo_find_documents_data(&st, body, &mut stream).await; } else if first_line.starts_with("POST /data/mongo/server-version") { handle_mongo_server_version_data(&st, body, &mut stream).await; + } else if first_line.starts_with("POST /data/mongo/collection-stats") { + handle_mongo_collection_stats_data(&st, body, &mut stream).await; } else if first_line.starts_with("POST /data/mongo/aggregate-documents") { handle_mongo_aggregate_documents_data(&st, body, &mut stream).await; } else if first_line.starts_with("POST /data/mongo/create-index") { @@ -505,6 +515,27 @@ async fn handle_mongo_server_version_data(state: &Arc, body: &str, str } } +async fn handle_mongo_collection_stats_data(state: &Arc, body: &str, stream: &mut tokio::net::TcpStream) { + let req: MongoCollectionStatsRequest = 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_name, req.database, stream).await + else { + return; + }; + match dbx_core::mongo_ops::mongo_collection_stats_core(state, &pool_key, &database, &req.collection, req.scale) + .await + { + Ok(result) => respond_json(stream, &result).await, + Err(e) => respond_error(stream, "500 Internal Server Error", &e).await, + } +} + async fn handle_mongo_aggregate_documents_data(state: &Arc, body: &str, stream: &mut tokio::net::TcpStream) { let req: MongoAggregateDocumentsRequest = 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 7d37210881..7ff560f0f4 100644 --- a/src-tauri/src/commands/mongo_cmd.rs +++ b/src-tauri/src/commands/mongo_cmd.rs @@ -122,6 +122,24 @@ pub async fn mongo_server_version( .await } +#[tauri::command] +pub async fn mongo_collection_stats( + state: State<'_, Arc>, + connection_id: String, + database: String, + collection: String, + scale: Option, + execution_id: Option, +) -> Result { + let app = state.inner().clone(); + run_cancellable( + &app, + execution_id, + dbx_core::mongo_ops::mongo_collection_stats_core(&app, &connection_id, &database, &collection, scale), + ) + .await +} + #[tauri::command] pub async fn mongo_aggregate_documents( state: State<'_, Arc>, diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 499cb6d916..595f8b6b2a 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -952,6 +952,7 @@ pub fn run() { commands::document_cmd::document_delete_gridfs_file, commands::mongo_cmd::mongo_find_documents, commands::mongo_cmd::mongo_server_version, + commands::mongo_cmd::mongo_collection_stats, commands::mongo_cmd::mongo_aggregate_documents, commands::mongo_cmd::mongo_create_index, commands::mongo_cmd::mongo_drop_indexes,