diff --git a/crates/dbx-core/src/mongo_shell.rs b/crates/dbx-core/src/mongo_shell.rs index 95fc2e19f5..7a8fc83010 100644 --- a/crates/dbx-core/src/mongo_shell.rs +++ b/crates/dbx-core/src/mongo_shell.rs @@ -734,6 +734,21 @@ fn trim_mongo_outer_comments(mut source: &str) -> &str { } fn parse_collection_prefix(source: &str) -> Result<(String, usize), String> { + // db["orders-2024"] reaches names that are not valid identifiers, the same way the shell does. + if source.get(..3).is_some_and(|prefix| prefix.eq_ignore_ascii_case("db[")) { + let close = matching_bracket(source, 2).ok_or("Invalid db[\"collection\"] accessor.")?; + let collection = parse_string_arg(source[3..close].trim())?; + if collection.is_empty() { + return Err("Invalid MongoDB collection name.".to_string()); + } + let end = close + 1; + let suffix = &source[end..]; + let trimmed = suffix.trim_start(); + if !trimmed.starts_with('.') { + return Err("MongoDB collection method is required.".to_string()); + } + return Ok((collection, end + suffix.len() - trimmed.len())); + } if !source.get(..3).is_some_and(|prefix| prefix.eq_ignore_ascii_case("db.")) { return Err("MongoDB command must start with db..".to_string()); } @@ -1258,6 +1273,43 @@ fn matching_paren(source: &str, open: usize) -> Option { None } +/// Offset of the `]` closing the bracket at `open`, ignoring brackets inside strings. +fn matching_bracket(source: &str, open: usize) -> Option { + let mut depth = 0; + let mut quote = None; + let mut escape = false; + let mut index = open; + while index < source.len() { + let ch = source[index..].chars().next()?; + if escape { + escape = false; + index += ch.len_utf8(); + continue; + } + if quote.is_some() { + if ch == '\\' { + escape = true; + } else if Some(ch) == quote { + quote = None; + } + index += ch.len_utf8(); + continue; + } + if ch == '\'' || ch == '"' { + quote = Some(ch); + } else if ch == '[' { + depth += 1; + } else if ch == ']' { + depth -= 1; + if depth == 0 { + return Some(index); + } + } + index += ch.len_utf8(); + } + None +} + fn split_top_level(source: &str) -> Vec { if source.trim().is_empty() { return Vec::new(); @@ -1376,6 +1428,36 @@ mod tests { assert!(parse("db.items.find({}).explain('executionStats').limit(1)").unwrap_err().contains("final")); } + #[test] + fn parses_bracket_collection_accessor() { + assert_eq!( + parse(r#"db["orders-2024"].find({a: 1})"#).unwrap(), + MongoCommand::Find { + collection: "orders-2024".to_string(), + filter: r#"{"a":1}"#.to_string(), + projection: None, + sort: None, + collation: None, + skip: 0, + limit: 100, + } + ); + + // Single quotes, a dotted name, and a chained method all behave like db.. + assert_eq!( + parse("db['audit.logs'].count()").unwrap(), + MongoCommand::Count { collection: "audit.logs".to_string(), filter: "{}".to_string(), accurate: false } + ); + assert!(matches!( + parse(r#"db["orders-2024"].updateOne({a: 1}, {$set: {b: 2}}, {upsert: true})"#).unwrap(), + MongoCommand::Update { many: false, .. } + )); + + for source in [r#"db[].find({})"#, r#"db[""].find({})"#, r#"db["x"]find({})"#, r#"db["x"]"#] { + assert!(parse(source).is_err(), "{source}"); + } + } + #[test] fn parses_get_collection_and_count() { assert_eq!( diff --git a/packages/app-tests/mongoShellCommand.test.ts b/packages/app-tests/mongoShellCommand.test.ts index dc53e7c409..f6c424a331 100644 --- a/packages/app-tests/mongoShellCommand.test.ts +++ b/packages/app-tests/mongoShellCommand.test.ts @@ -357,6 +357,29 @@ test("parseMongoFindCommand does not rewrite constructor text inside strings", ( assert.deepEqual(JSON.parse(command.filter), { label: "new Date()", note: "ObjectId()" }); }); +test('parseMongoCommand accepts db["name"] bracket collection accessors', () => { + const find = parseMongoFindCommand('db["orders-2024"].find({a: 1})'); + assert.ok(find); + assert.equal(find.collection, "orders-2024"); + assert.deepEqual(JSON.parse(find.filter), { a: 1 }); + + // Single quotes and a dotted name behave like db.. + assert.equal(parseMongoFindCommand("db['audit.logs'].find({})")?.collection, "audit.logs"); + + assert.deepEqual(parseMongoWriteCommand('db["orders-2024"].updateOne({a: 1}, {$set: {b: 2}}, {upsert: true})'), { + kind: "update", + collection: "orders-2024", + filter: '{"a": 1}', + update: '{"$set": {"b": 2}}', + options: '{"upsert": true}', + many: false, + }); + + for (const source of ["db[].find({})", 'db["x"]find({})', 'db["x"]']) { + assert.equal(parseMongoCommand(source), null, source); + } +}); + test("parseMongoFindCommand accepts single-quoted string values and unquoted sort keys", () => { const command = parseMongoFindCommand("db.products.find({category: 'Electronics'}).sort({price: -1}).limit(2)"); assert.ok(command); diff --git a/packages/mongo-shell/src/json.ts b/packages/mongo-shell/src/json.ts index 5f15498301..5f5298e86c 100644 --- a/packages/mongo-shell/src/json.ts +++ b/packages/mongo-shell/src/json.ts @@ -202,6 +202,11 @@ export function parseCollectionMethodTarget(source: string, method: string): { c if (getCollection) { return { collection: getCollection[2]!, methodCallIndex: findChainedMethodCallIndex(source, method) }; } + // db["orders-2024"] reaches names that are not valid identifiers, the same way the shell does. + const bracket = new RegExp(`^db\\s*\\[\\s*(["'])(.*?)\\1\\s*\\]\\s*\\.\\s*${escapedMethod}\\s*\\(`).exec(source); + if (bracket) { + return { collection: bracket[2]!, methodCallIndex: findChainedMethodCallIndex(source, method) }; + } return null; }