feat(mongodb): support estimatedDocumentCount, db.stats, db.serverStatus

apps/desktop/src/lib/mongo/mongoCompletion.ts

@@ -87,6 +87,7 @@ const COLLECTION_METHODS = [
   { label: "aggregate", detail: "Run an aggregation pipeline", apply: "aggregate([])" },
   { label: "countDocuments", detail: "Count matching documents", apply: "countDocuments({})" },
   { label: "count", detail: "Count matching documents (legacy helper)", apply: "count({})" },
+  { label: "estimatedDocumentCount", detail: "Estimate the document count from collection metadata", apply: "estimatedDocumentCount()" },
   { label: "distinct", detail: "List the distinct values of a field", apply: 'distinct("${field}")' },
   { label: "insertOne", detail: "Insert one document", apply: "insertOne({})" },
   { label: "insertMany", detail: "Insert multiple documents", apply: "insertMany([{}])" },
@@ -113,6 +114,7 @@ const COLLECTION_METHOD_BOOST: Record<(typeof COLLECTION_METHODS)[number]["label
   findOne: 230,
   aggregate: 220,
   countDocuments: 210,
+  estimatedDocumentCount: 205,
   distinct: 200,
   insertOne: 180,
   insertMany: 170,
@@ -139,6 +141,8 @@ const COLLECTION_METHOD_BOOST: Record<(typeof COLLECTION_METHODS)[number]["label
 const DATABASE_METHODS = [
   { label: "getCollection", detail: "Reference a collection by name", apply: 'getCollection("${}")' },
   { label: "version", detail: "Show the MongoDB server version", apply: "version()" },
+  { label: "stats", detail: "Show database statistics", apply: "stats()" },
+  { label: "serverStatus", detail: "Show server status", apply: "serverStatus()" },
 ] as const;
 
 const CURSOR_METHODS = [
@@ -159,6 +163,8 @@ const ROOT_SNIPPETS = [
   { label: "db.getCollection", detail: "Reference a collection by name", apply: 'db.getCollection("${}")' },
   { label: "use", detail: "Switch the active database", apply: "use ${database}" },
   { label: "db.version", detail: "Show the MongoDB server version", apply: "db.version()" },
+  { label: "db.stats", detail: "Show database statistics", apply: "db.stats()" },
+  { label: "db.serverStatus", detail: "Show server status", apply: "db.serverStatus()" },
 ] as const;
 
 const ROOT_SNIPPET_BOOST: Record<(typeof ROOT_SNIPPETS)[number]["label"], number> = {
@@ -167,6 +173,8 @@ const ROOT_SNIPPET_BOOST: Record<(typeof ROOT_SNIPPETS)[number]["label"], number
   "db.getCollection": 330,
   use: 320,
   "db.version": 310,
+  "db.stats": 305,
+  "db.serverStatus": 300,
 };
 
 /** Role of each positional argument, by collection helper. Drives cursor classification. */

apps/desktop/src/lib/mongo/mongoShellCommand.ts

@@ -40,6 +40,7 @@ const COLLECTION_METHOD_SHAPES: Record<string, MongoMethodShape> = {
   findOne: { expects: "an optional filter, an optional projection, and optional options", roles: ["filter", "projection", "options"] },
   count: { expects: "an optional filter", roles: ["filter"] },
   countDocuments: { expects: "an optional filter", roles: ["filter"] },
+  estimatedDocumentCount: { expects: "no arguments", roles: [] },
   distinct: { expects: "a field name and an optional filter", roles: ["field", "filter"] },
   insert: { expects: "one document or an array of documents", roles: ["document"] },
   insertOne: { expects: "one document", roles: ["document"] },
@@ -69,6 +70,7 @@ const SUPPORTED_COLLECTION_METHODS = [
   "aggregate",
   "count",
   "countDocuments",
+  "estimatedDocumentCount",
   "distinct",
   "insertOne",
   "insertMany",
@@ -90,8 +92,6 @@ const SUPPORTED_COLLECTION_METHODS = [
 /** Database-level methods with a supported equivalent worth pointing at. */
 const DATABASE_METHOD_HINTS: Record<string, string> = {
   getSiblingDB: "switch databases with `use <database>` and then run the command against db.<collection>",
-  stats: "use db.runCommand({ dbStats: 1 })",
-  serverStatus: "use db.runCommand({ serverStatus: 1 })",
   adminCommand: "use db.runCommand({ ... })",
   getCollectionNames: "collections are listed in the sidebar",
   createCollection: 'collections are created on first insert, or use db.runCommand({ create: "name" })',
@@ -99,11 +99,13 @@ const DATABASE_METHOD_HINTS: Record<string, string> = {
 
 const DATABASE_METHOD_SHAPES: Record<string, MongoMethodShape> = {
   version: { expects: "no arguments", roles: [] },
+  stats: { expects: "no arguments", roles: [] },
+  serverStatus: { expects: "no arguments", roles: [] },
   createUser: { expects: "a user document and optional write concern", roles: ["user", "writeConcern"] },
   runCommand: { expects: "one command document", roles: ["command"] },
 };
 
-const SUPPORTED_DATABASE_METHODS = ["version", "createUser", "runCommand", "getCollection"];
+const SUPPORTED_DATABASE_METHODS = ["version", "stats", "serverStatus", "createUser", "runCommand", "getCollection"];
 
 const SUPPORTED_VALUE_CONSTRUCTORS = ["ObjectId", "ISODate", "new Date", "NumberLong", "NumberInt", "NumberDecimal", "UUID", "BinData", "Timestamp", "MinKey", "MaxKey"];
 
@@ -524,7 +526,23 @@ export function applyMongoFindSort(input: string, column: string, direction: "as
 
 export function parseMongoCountDocumentsCommand(input: string): MongoCountDocumentsCommand | null {
   const source = input.trim().replace(/;$/, "").trim();
-  return parseCollectionCountCommand(source, "countDocuments") ?? parseCollectionCountCommand(source, "count") ?? parseFindCountCommand(source);
+  return parseCollectionCountCommand(source, "countDocuments") ?? parseCollectionCountCommand(source, "count") ?? parseEstimatedDocumentCountCommand(source) ?? parseFindCountCommand(source);
+}
+
+/**
+ * estimatedDocumentCount() takes no filter and is metadata-backed, which is exactly
+ * the legacy count() fast path the driver already uses for a filterless count.
+ */
+function parseEstimatedDocumentCountCommand(source: string): MongoCountDocumentsCommand | null {
+  const target = parseCollectionMethodTarget(source, "estimatedDocumentCount");
+  if (!target) return null;
+
+  const openIndex = source.indexOf("(", target.methodCallIndex);
+  const closeIndex = findMatchingParen(source, openIndex);
+  if (closeIndex < 0 || source.slice(closeIndex + 1).trim()) return null;
+  if (source.slice(openIndex + 1, closeIndex).trim()) return null;
+
+  return { collection: target.collection, filter: "{}", mode: "legacy" };
 }
 
 function parseCollectionCountCommand(source: string, method: "countDocuments" | "count"): MongoCountDocumentsCommand | null {
@@ -674,8 +692,16 @@ export function parseMongoCreateUserCommand(input: string): MongoCreateUserComma
   };
 }
 
+/** The shell's shorthand for the matching runCommand, so they share its execution path. */
+const DATABASE_STATUS_COMMANDS: Record<string, string> = { stats: "dbStats", serverStatus: "serverStatus" };
+
 export function parseMongoRunCommand(input: string): MongoRunCommand | null {
   const source = input.trim().replace(/;$/, "").trim();
+  for (const [method, command] of Object.entries(DATABASE_STATUS_COMMANDS)) {
+    if (new RegExp(`^db\\s*\\.\\s*${method}\\s*\\(\\s*\\)$`, "i").test(source)) {
+      return { commandJson: JSON.stringify({ [command]: 1 }) };
+    }
+  }
   const match = /^db\s*\.\s*runCommand\s*\(/i.exec(source);
   if (!match) return null;
   const openIndex = source.indexOf("(", match.index);

crates/dbx-core/src/mongo_shell.rs

@@ -371,6 +371,16 @@ pub fn parse(input: &str) -> Result<MongoCommand, String> {
     if let Some(database) = parse_use_database(source) {
         return Ok(MongoCommand::Use { database });
     }
+    // `db.stats()` and `db.serverStatus()` are the shell's shorthand for the
+    // matching runCommand, so they execute through the same supported path.
+    for (method, command) in [("stats", "dbStats"), ("serverStatus", "serverStatus")] {
+        if let Some((args, tail)) = database_method_call(source, method) {
+            if !tail.is_empty() || !args.iter().all(|arg| arg.trim().is_empty()) {
+                return Err(format!("MongoDB db.{method}() takes no arguments."));
+            }
+            return Ok(MongoCommand::RunCommand { command_json: format!(r#"{{"{command}":1}}"#) });
+        }
+    }
     if let Some((args, tail)) = database_method_call(source, "runCommand") {
         if !tail.is_empty() || args.len() != 1 {
             return Err("MongoDB runCommand() requires exactly one command document.".to_string());
@@ -490,6 +500,18 @@ pub fn parse(input: &str) -> Result<MongoCommand, String> {
         });
     }
 
+    // estimatedDocumentCount() takes no filter and is metadata-backed, which is
+    // exactly the legacy count() fast path the driver already uses.
+    if let Some((args, tail)) = method_call(source, prefix_end, "estimatedDocumentCount") {
+        if !args.iter().all(|arg| arg.trim().is_empty()) {
+            return Err("MongoDB estimatedDocumentCount() takes no filter.".to_string());
+        }
+        if !tail.is_empty() {
+            return Err("MongoDB estimatedDocumentCount() does not support chained methods.".to_string());
+        }
+        return Ok(MongoCommand::Count { collection, filter: "{}".to_string(), accurate: false });
+    }
+
     for (method, accurate) in [("countDocuments", true), ("count", false)] {
         if let Some((args, tail)) = method_call(source, prefix_end, method) {
             if !tail.is_empty() || args.len() > 1 {
@@ -1805,6 +1827,44 @@ mod tests {
         assert!(parse("db.orders.find({}).toArray(1)").unwrap_err().contains("find() chain"));
     }
 
+    #[test]
+    fn parses_estimated_document_count_as_a_metadata_backed_count() {
+        // The driver already takes the metadata fast path for a filterless legacy
+        // count(), which is exactly what estimatedDocumentCount() asks for.
+        let expected =
+            MongoCommand::Count { collection: "orders".to_string(), filter: "{}".to_string(), accurate: false };
+        for source in [
+            "db.orders.estimatedDocumentCount()",
+            r#"db["orders"].estimatedDocumentCount()"#,
+            "db.getCollection('orders').estimatedDocumentCount();",
+        ] {
+            assert_eq!(parse(source).unwrap(), expected, "{source}");
+        }
+        assert!(!parse("db.orders.estimatedDocumentCount()").unwrap().is_mutating());
+
+        assert!(parse("db.orders.estimatedDocumentCount({a: 1})").unwrap_err().contains("no filter"));
+        assert!(parse("db.orders.estimatedDocumentCount().limit(5)").unwrap_err().contains("chained"));
+    }
+
+    #[test]
+    fn parses_db_stats_and_server_status_as_run_commands() {
+        assert_eq!(
+            parse("db.stats()").unwrap(),
+            MongoCommand::RunCommand { command_json: r#"{"dbStats":1}"#.to_string() }
+        );
+        assert_eq!(
+            parse("db . serverStatus ( ) ;").unwrap(),
+            MongoCommand::RunCommand { command_json: r#"{"serverStatus":1}"#.to_string() }
+        );
+
+        for source in ["db.stats(1)", "db.serverStatus({})"] {
+            assert!(parse(source).unwrap_err().contains("takes no arguments"), "{source}");
+        }
+
+        // db.collection.stats() still parses as collection stats, not a run command.
+        assert!(matches!(parse("db.orders.stats()").unwrap(), MongoCommand::CollectionStats { .. }));
+    }
+
     #[test]
     fn parses_show_databases_aliases_as_read_only() {
         for source in [

packages/app-tests/mongoCompletion.test.ts

@@ -134,8 +134,8 @@ test("prioritizes common read helpers and keeps destructive helpers last", () =>
   const getCollectionMethodLabels = labels('db.getCollection("order-events").');
 
   assert.deepEqual(labels("").slice(0, 5), ["db.collection.find", "db.collection.aggregate", "db.getCollection", "use", "db.version"]);
-  assert.deepEqual(methodLabels.slice(0, 5), ["find", "findOne", "aggregate", "countDocuments", "distinct"]);
-  assert.deepEqual(getCollectionMethodLabels.slice(0, 5), ["find", "findOne", "aggregate", "countDocuments", "distinct"]);
+  assert.deepEqual(methodLabels.slice(0, 6), ["find", "findOne", "aggregate", "countDocuments", "estimatedDocumentCount", "distinct"]);
+  assert.deepEqual(getCollectionMethodLabels.slice(0, 6), ["find", "findOne", "aggregate", "countDocuments", "estimatedDocumentCount", "distinct"]);
   assert.deepEqual(methodLabels.slice(-3), ["dropIndex", "dropIndexes", "drop"]);
   assert.deepEqual(labels("db.users.find({})."), ["limit", "sort", "skip", "count"]);
 });
@@ -248,6 +248,14 @@ test("treats $in, $nin and $all elements as values rather than sub-filters", ()
   assert.equal(labels("db.users.find({ $or: [{ ").includes("$oid"), false);
 });
 
+test("offers the newly supported count and database commands", () => {
+  assert.ok(labels("db.users.estim", { fields }).includes("estimatedDocumentCount"));
+  const dbLevel = labels("db.");
+  assert.ok(dbLevel.includes("stats"));
+  assert.ok(dbLevel.includes("serverStatus"));
+  assert.ok(labels("").includes("db.stats"));
+});
+
 test("offers the newer shell value constructors", () => {
   const labels = buildMongoCompletionItems("db.users.find({ _id: ", "db.users.find({ _id: ".length, { fields }).map((item) => item.label);
   for (const constructor of ["NumberInt", "NumberDecimal", "UUID", "BinData", "Timestamp", "MinKey", "MaxKey"]) {
@@ -491,8 +499,9 @@ test("suggests only helpers the shell parser accepts", () => {
   assert.ok(methodLabels.includes("count"));
   assert.ok(methodLabels.includes("drop"));
   assert.ok(methodLabels.includes("distinct"));
+  assert.ok(methodLabels.includes("estimatedDocumentCount"));
   // Suggesting a helper DBX cannot run just hands the user a command that fails.
-  for (const unsupported of ["bulkWrite", "estimatedDocumentCount", "replaceOne"]) {
+  for (const unsupported of ["bulkWrite", "replaceOne", "renameCollection"]) {
     assert.equal(methodLabels.includes(unsupported), false, `${unsupported} is not executable`);
   }
   // Cursor methods are not collection methods.

packages/app-tests/mongoShellCommand.test.ts

@@ -1161,6 +1161,35 @@ test("parseMongoAggregateCommand rejects non-array pipelines and invalid options
   assert.equal(parseMongoAggregateCommand("db.products.aggregate([], {}, true)"), null);
 });
 
+test("parseMongoCountDocumentsCommand reads estimatedDocumentCount as a metadata-backed count", () => {
+  // The driver already takes the metadata fast path for a filterless legacy count().
+  const expected = { collection: "orders", filter: "{}", mode: "legacy" };
+  for (const source of ["db.orders.estimatedDocumentCount()", 'db["orders"].estimatedDocumentCount()', "db.getCollection('orders').estimatedDocumentCount();"]) {
+    assert.deepEqual(parseMongoCountDocumentsCommand(source), expected, source);
+  }
+
+  assert.equal(parseMongoCountDocumentsCommand("db.orders.estimatedDocumentCount({a: 1})"), null);
+  assert.equal(parseMongoCountDocumentsCommand("db.orders.estimatedDocumentCount().limit(5)"), null);
+  // count() and countDocuments() keep their own modes.
+  assert.equal(parseMongoCountDocumentsCommand("db.orders.countDocuments({})")?.mode, "accurate");
+  assert.equal(parseMongoCountDocumentsCommand("db.orders.count()")?.mode, "legacy");
+});
+
+test("parseMongoRunCommand reads db.stats() and db.serverStatus() as run commands", () => {
+  assert.deepEqual(parseMongoRunCommand("db.stats()"), { commandJson: '{"dbStats":1}' });
+  assert.deepEqual(parseMongoRunCommand("db . serverStatus ( ) ;"), { commandJson: '{"serverStatus":1}' });
+  assert.deepEqual(parseMongoRunCommand("DB.STATS()"), { commandJson: '{"dbStats":1}' });
+
+  for (const source of ["db.stats(1)", "db.serverStatus({})"]) {
+    assert.equal(parseMongoRunCommand(source), null, source);
+    assert.match(describeMongoCommandParseFailure(source), /takes no arguments|expects no arguments/, source);
+  }
+
+  // db.collection.stats() is still collection stats, not a run command.
+  assert.equal(parseMongoRunCommand("db.orders.stats()"), null);
+  assert.equal(parseMongoCommand("db.orders.stats()")?.command.kind, "collectionStats");
+});
+
 test("describeMongoCommandParseFailure names an unsupported value constructor and where it is", () => {
   const message = describeMongoCommandParseFailure('db.reports.updateOne({a: 1}, {$set: {t: Foo("x")}}, {upsert: true})');
   assert.match(message, /Unsupported value Foo\(\.\.\.\) in the update argument of updateOne\(\)/);
@@ -1187,7 +1216,7 @@ test("describeMongoCommandParseFailure names unsupported methods and points at a
   assert.match(describeMongoCommandParseFailure('db["my-coll"].renameCollection("x")'), /renameCollection\(\) is not supported/);
   assert.match(describeMongoCommandParseFailure('db.getCollection("my-coll").watch()'), /watch\(\) is not supported/);
 
-  assert.match(describeMongoCommandParseFailure("db.stats()"), /db\.stats\(\) is not supported; use db\.runCommand\(\{ dbStats: 1 \}\)/);
+  assert.match(describeMongoCommandParseFailure('db.createCollection("c")'), /db\.createCollection\(\) is not supported; collections are created on first insert/);
   assert.match(describeMongoCommandParseFailure('db.getSiblingDB("other").c.find({})'), /use <database>/);
   assert.match(describeMongoCommandParseFailure("db.adminCommand({ping: 1})"), /use db\.runCommand/);
   assert.match(describeMongoCommandParseFailure("show collections"), /listed in the sidebar/);