fix(mongodb): accept .toArray()/.pretty() after find/aggregate

crates/dbx-core/src/mongo_shell.rs

@@ -445,6 +445,7 @@ pub fn parse(input: &str) -> Result<MongoCommand, String> {
                         verbosity: parse_explain_verbosity(&call_args)?,
                     });
                 }
+                _ if is_noop_cursor_call(&name, &call_args) => {}
                 _ => return Err(format!("Unsupported MongoDB find() chain: {name}()")),
             }
         }
@@ -503,9 +504,16 @@ pub fn parse(input: &str) -> Result<MongoCommand, String> {
     }
 
     if let Some((args, tail)) = method_call(source, prefix_end, "aggregate") {
-        if !tail.is_empty() || !(1..=2).contains(&args.len()) {
+        if !(1..=2).contains(&args.len()) {
             return Err("Invalid MongoDB aggregate() command.".to_string());
         }
+        for (name, call_args) in chained_calls(&tail)? {
+            if !is_noop_cursor_call(&name, &call_args) {
+                return Err(format!(
+                    "Unsupported MongoDB aggregate() chain: {name}(). Use pipeline stages such as $sort and $limit instead."
+                ));
+            }
+        }
         let pipeline = normalized_json(&args[0])?;
         if !parse_json_value(&pipeline).is_some_and(|value| value.is_array()) {
             return Err("MongoDB aggregate() requires a pipeline array.".to_string());
@@ -817,6 +825,12 @@ fn database_method_call(source: &str, method: &str) -> Option<(Vec<String>, Stri
     Some((split_top_level(&source[open + 1..close]), source[close + 1..].trim().to_string()))
 }
 
+/// Cursor methods that change nothing here: results are always materialised, so
+/// the `.toArray()` that mongosh and Compass append can simply be dropped.
+fn is_noop_cursor_call(name: &str, args: &[String]) -> bool {
+    matches!(name, "toArray" | "pretty") && args.is_empty()
+}
+
 fn chained_calls(chain: &str) -> Result<Vec<(String, Vec<String>)>, String> {
     let mut rest = chain.trim();
     let mut calls = Vec::new();
@@ -1754,6 +1768,43 @@ mod tests {
         assert_eq!(filter, r#"{"u":{"$uuid":"3B241101-E2BB-4255-8CAF-4136C566A962"}}"#);
     }
 
+    #[test]
+    fn drops_noop_cursor_methods_after_find_and_aggregate() {
+        // mongosh and Compass append .toArray(); results are always materialised
+        // here, so it changes nothing and must not be an error.
+        let expected = MongoCommand::Aggregate {
+            collection: "orders".to_string(),
+            pipeline: r#"[{"$match":{"a":1}}]"#.to_string(),
+            options: None,
+        };
+        for source in [
+            "db.orders.aggregate([{$match: {a: 1}}]).toArray()",
+            "db.orders.aggregate([{$match: {a: 1}}]).pretty()",
+            "db.orders.aggregate([{$match: {a: 1}}]).toArray().pretty()",
+            "db.orders.aggregate([{$match: {a: 1}}])\n  .toArray()",
+        ] {
+            assert_eq!(parse(source).unwrap(), expected, "{source}");
+        }
+
+        let with_chain = parse("db.orders.find({a: 1}).sort({b: 1}).limit(5).toArray()").unwrap();
+        assert!(matches!(with_chain, MongoCommand::Find { limit: 5, sort: Some(_), .. }), "{with_chain:?}");
+        assert!(matches!(parse("db.orders.find({a: 1}).pretty()").unwrap(), MongoCommand::Find { .. }));
+    }
+
+    #[test]
+    fn still_rejects_real_cursor_methods_after_aggregate() {
+        for source in [
+            "db.orders.aggregate([]).limit(5)",
+            "db.orders.aggregate([]).sort({a: 1})",
+            "db.orders.aggregate([]).toArray().limit(5)",
+            "db.orders.aggregate([]).toArray(1)",
+        ] {
+            let error = parse(source).unwrap_err();
+            assert!(error.contains("aggregate() chain"), "{source} => {error}");
+        }
+        assert!(parse("db.orders.find({}).toArray(1)").unwrap_err().contains("find() chain"));
+    }
+
     #[test]
     fn parses_show_databases_aliases_as_read_only() {
         for source in [

packages/app-tests/mongoShellCommand.test.ts

@@ -936,6 +936,31 @@ test("parseMongoAggregateCommand parses db collection aggregate", () => {
   });
 });
 
+test("parseMongoAggregateCommand drops the no-op cursor methods mongosh and Compass append", () => {
+  const expected = { collection: "orders", pipeline: '[{"$match": {"a": 1}}]' };
+  for (const source of [
+    "db.orders.aggregate([{$match: {a: 1}}]).toArray()",
+    "db.orders.aggregate([{$match: {a: 1}}]).pretty()",
+    "db.orders.aggregate([{$match: {a: 1}}]).toArray().pretty()",
+    "db.orders.aggregate([{$match: {a: 1}}])\n  .toArray()",
+    "db.orders.aggregate([{$match: {a: 1}}]).toArray();",
+  ]) {
+    assert.deepEqual(parseMongoAggregateCommand(source), expected, source);
+  }
+  assert.deepEqual(parseMongoAggregateCommand("db.orders.aggregate([], {allowDiskUse: true}).toArray()"), {
+    collection: "orders",
+    pipeline: "[]",
+    options: '{"allowDiskUse": true}',
+  });
+});
+
+test("parseMongoAggregateCommand still rejects cursor methods that would change the result", () => {
+  for (const source of ["db.orders.aggregate([]).limit(5)", "db.orders.aggregate([]).sort({a: 1})", "db.orders.aggregate([]).toArray().limit(5)", "db.orders.aggregate([]).toArray(1)"]) {
+    assert.equal(parseMongoAggregateCommand(source), null, source);
+    assert.match(describeMongoCommandParseFailure(source), /Chaining .* is not supported/, source);
+  }
+});
+
 test("parseMongoAggregateCommand accepts an empty pipeline", () => {
   assert.deepEqual(parseMongoAggregateCommand("db.products.aggregate([])"), {
     collection: "products",

packages/mongo-shell/src/aggregate.ts

@@ -3,15 +3,7 @@
  * JSON preprocessing lives in ./json.ts — do not re-copy helpers here.
  */
 
-import {
-  findMatchingParen,
-  hasUnclosedMongoDelimiters,
-  normalizeJsonArgument,
-  parseCollectionMethodTarget,
-  parseMongoObjectArgument,
-  splitTopLevel,
-  trimMongoOuterComments,
-} from "./json.js";
+import { findMatchingParen, hasUnclosedMongoDelimiters, normalizeJsonArgument, parseCollectionMethodTarget, parseMongoObjectArgument, splitTopLevel, trimMongoOuterComments } from "./json.js";
 
 export interface MongoAggregateCommand {
   collection: string;
@@ -28,21 +20,19 @@ export const MONGO_SHELL_COMMAND_HINT =
   "db.collection.getIndexes(), db.collection.createIndex({...}), db.createUser({...}), db.runCommand({...}), " +
   "or db.collection.insertOne({...}).";
 
-const PIPELINE_MUST_BE_ARRAY =
-  "MongoDB aggregate pipeline must be a JSON array (for example [{ $match: {} }]).";
-const OPTIONS_MUST_BE_OBJECT =
-  "MongoDB aggregate options must be a JSON object (for example { explain: true }).";
-const UNCLOSED_DELIMITERS =
-  "MongoDB command has unclosed parentheses, brackets, braces, or strings.";
-const UNSUPPORTED_CHAINING =
-  "Unsupported MongoDB aggregate form. Use db.collection.aggregate(pipeline) or " +
-  "db.collection.aggregate(pipeline, options). Chaining (for example .limit()) is not supported.";
-const EXPECTS_PIPELINE_OR_OPTIONS =
-  "MongoDB aggregate expects aggregate(pipeline) or aggregate(pipeline, options).";
+const PIPELINE_MUST_BE_ARRAY = "MongoDB aggregate pipeline must be a JSON array (for example [{ $match: {} }]).";
+const OPTIONS_MUST_BE_OBJECT = "MongoDB aggregate options must be a JSON object (for example { explain: true }).";
+const UNCLOSED_DELIMITERS = "MongoDB command has unclosed parentheses, brackets, braces, or strings.";
+const UNSUPPORTED_CHAINING = "Unsupported MongoDB aggregate form. Use db.collection.aggregate(pipeline) or " + "db.collection.aggregate(pipeline, options). Chaining (for example .limit()) is not supported.";
 
-type AggregateParseResult =
-  | { ok: true; command: MongoAggregateCommand }
-  | { ok: false; reason: string };
+/**
+ * Cursor methods that change nothing here: DBX always materialises results, so
+ * the `.toArray()` that mongosh and Compass append can simply be dropped.
+ */
+const NOOP_CURSOR_CHAIN = /^(?:\s*\.\s*(?:toArray|pretty)\s*\(\s*\))*\s*$/;
+const EXPECTS_PIPELINE_OR_OPTIONS = "MongoDB aggregate expects aggregate(pipeline) or aggregate(pipeline, options).";
+
+type AggregateParseResult = { ok: true; command: MongoAggregateCommand } | { ok: false; reason: string };
 
 export function parseMongoAggregateCommand(input: string): MongoAggregateCommand | null {
   const source = input.trim().replace(/;$/, "").trim();
@@ -70,7 +60,7 @@ function tryParseMongoAggregateCommand(source: string): AggregateParseResult | n
   const openIndex = source.indexOf("(", target.methodCallIndex);
   const closeIndex = findMatchingParen(source, openIndex);
   if (closeIndex < 0) return { ok: false, reason: UNCLOSED_DELIMITERS };
-  if (source.slice(closeIndex + 1).trim()) return { ok: false, reason: UNSUPPORTED_CHAINING };
+  if (!NOOP_CURSOR_CHAIN.test(source.slice(closeIndex + 1))) return { ok: false, reason: UNSUPPORTED_CHAINING };
 
   const args = splitTopLevel(source.slice(openIndex + 1, closeIndex));
   if (args.length < 1 || args.length > 2) return { ok: false, reason: EXPECTS_PIPELINE_OR_OPTIONS };