From 7c51454bfdab324548c2cc1930751ffb60de451f Mon Sep 17 00:00:00 2001 From: Loc Nguyen <113083148+thailoc-dev@users.noreply.github.com> Date: Fri, 18 Sep 2026 14:06:29 +0700 Subject: [PATCH] feat(mongodb): support db.createCollection() and db.dropDatabase() --- apps/desktop/src/lib/mongo/mongoCompletion.ts | 2 + .../src/lib/mongo/mongoShellCommand.ts | 30 ++++++++- crates/dbx-drivers/src/mongo_shell.rs | 66 +++++++++++++++++++ packages/app-tests/mongoCompletion.test.ts | 2 + packages/app-tests/mongoShellCommand.test.ts | 26 +++++++- 5 files changed, 123 insertions(+), 3 deletions(-) diff --git a/apps/desktop/src/lib/mongo/mongoCompletion.ts b/apps/desktop/src/lib/mongo/mongoCompletion.ts index 49e7405633..48addea356 100644 --- a/apps/desktop/src/lib/mongo/mongoCompletion.ts +++ b/apps/desktop/src/lib/mongo/mongoCompletion.ts @@ -149,6 +149,8 @@ const DATABASE_METHODS = [ { 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()" }, + { label: "createCollection", detail: "Create a collection", apply: 'createCollection("${name}")' }, + { label: "dropDatabase", detail: "Drop the current database", apply: "dropDatabase()" }, ] as const; const CURSOR_METHODS = [ diff --git a/apps/desktop/src/lib/mongo/mongoShellCommand.ts b/apps/desktop/src/lib/mongo/mongoShellCommand.ts index 0da3179251..1deffb502d 100644 --- a/apps/desktop/src/lib/mongo/mongoShellCommand.ts +++ b/apps/desktop/src/lib/mongo/mongoShellCommand.ts @@ -100,18 +100,19 @@ const DATABASE_METHOD_HINTS: Record = { getSiblingDB: "switch databases with `use ` and then run the command against db.", adminCommand: "use db.runCommand({ ... })", getCollectionNames: "collections are listed in the sidebar", - createCollection: 'collections are created on first insert, or use db.runCommand({ create: "name" })', }; const DATABASE_METHOD_SHAPES: Record = { version: { expects: "no arguments", roles: [] }, stats: { expects: "no arguments", roles: [] }, serverStatus: { expects: "no arguments", roles: [] }, + dropDatabase: { expects: "no arguments", roles: [] }, + createCollection: { expects: "a collection name and optional options", roles: ["name", "options"] }, createUser: { expects: "a user document and optional write concern", roles: ["user", "writeConcern"] }, runCommand: { expects: "one command document", roles: ["command"] }, }; -const SUPPORTED_DATABASE_METHODS = ["version", "stats", "serverStatus", "createUser", "runCommand", "getCollection"]; +const SUPPORTED_DATABASE_METHODS = ["version", "stats", "serverStatus", "createCollection", "dropDatabase", "createUser", "runCommand", "getCollection"]; const SUPPORTED_VALUE_CONSTRUCTORS = ["ObjectId", "ISODate", "new Date", "NumberLong", "NumberInt", "NumberDecimal", "UUID", "BinData", "Timestamp", "MinKey", "MaxKey"]; @@ -770,8 +771,33 @@ 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 = { stats: "dbStats", serverStatus: "serverStatus" }; +/** `db.createCollection(name[, options])` is the create command; options pass through for the server to validate. */ +function parseMongoCreateCollectionCommand(source: string): MongoRunCommand | null { + const match = /^db\s*\.\s*createCollection\s*\(/i.exec(source); + if (!match) return null; + const openIndex = source.indexOf("(", match.index); + const closeIndex = findMatchingParen(source, openIndex); + if (closeIndex < 0 || source.slice(closeIndex + 1).trim()) return null; + const args = splitTopLevel(source.slice(openIndex + 1, closeIndex)); + if (args.length < 1 || args.length > 2) return null; + const name = /^(["'])([^"'$]+)\1$/.exec(args[0]!.trim())?.[2]?.trim(); + if (!name) return null; + const command: Record = { create: name }; + if (args[1]?.trim()) { + const optionsJson = parseMongoObjectArgument(args[1]); + if (!optionsJson) return null; + const options = JSON.parse(optionsJson) as Record; + if ("create" in options) return null; + Object.assign(command, options); + } + return { commandJson: JSON.stringify(command) }; +} + export function parseMongoRunCommand(input: string): MongoRunCommand | null { const source = input.trim().replace(/;$/, "").trim(); + if (/^db\s*\.\s*dropDatabase\s*\(\s*\)$/i.test(source)) return { commandJson: JSON.stringify({ dropDatabase: 1 }) }; + const createCollection = parseMongoCreateCollectionCommand(source); + if (createCollection) return createCollection; 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 }) }; diff --git a/crates/dbx-drivers/src/mongo_shell.rs b/crates/dbx-drivers/src/mongo_shell.rs index 95876e2f6d..b4c8c6bf84 100644 --- a/crates/dbx-drivers/src/mongo_shell.rs +++ b/crates/dbx-drivers/src/mongo_shell.rs @@ -540,6 +540,37 @@ pub fn parse(input: &str) -> Result { } // `db.stats()` and `db.serverStatus()` are the shell's shorthand for the // matching runCommand, so they execute through the same supported path. + // `db.dropDatabase()` is the shell's shorthand for the dropDatabase command; running it + // through runCommand keeps it on the supported path and classed as dangerous. + if let Some((args, tail)) = database_method_call(source, "dropDatabase") { + if !tail.is_empty() || !args.iter().all(|arg| arg.trim().is_empty()) { + return Err("MongoDB db.dropDatabase() takes no arguments.".to_string()); + } + return Ok(MongoCommand::RunCommand { command_json: r#"{"dropDatabase":1}"#.to_string() }); + } + // `db.createCollection(name[, options])` is the create command; options such as + // capped/size/max/validator pass through for the server to validate. + if let Some((args, tail)) = database_method_call(source, "createCollection") { + if !tail.is_empty() || !(1..=2).contains(&args.len()) { + return Err("MongoDB createCollection() requires a collection name and optional options.".to_string()); + } + let name = parse_string_arg(&args[0])?; + if name.trim().is_empty() || name.contains('$') { + return Err("MongoDB createCollection() requires a valid collection name.".to_string()); + } + let mut command = serde_json::Map::new(); + command.insert("create".to_string(), Value::String(name)); + if let Some(options) = args.get(1).filter(|arg| !arg.trim().is_empty()) { + let Some(Value::Object(options)) = parse_json_value(&normalized_json(options)?) else { + return Err("MongoDB createCollection() options must be a document.".to_string()); + }; + if options.contains_key("create") { + return Err("MongoDB createCollection() options must not contain create.".to_string()); + } + command.extend(options); + } + return Ok(MongoCommand::RunCommand { command_json: Value::Object(command).to_string() }); + } 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()) { @@ -2273,6 +2304,41 @@ mod tests { } } + #[test] + fn parses_drop_database_and_create_collection_as_run_commands() { + let drop = parse("db . dropDatabase ( ) ;").unwrap(); + assert_eq!(drop, MongoCommand::RunCommand { command_json: r#"{"dropDatabase":1}"#.to_string() }); + assert!(drop.is_mutating() && drop.is_dangerous()); + + assert_eq!( + parse(r#"db.createCollection("events")"#).unwrap(), + MongoCommand::RunCommand { command_json: r#"{"create":"events"}"#.to_string() } + ); + let MongoCommand::RunCommand { command_json } = + parse("db.createCollection('logs', {capped: true, size: 1048576, max: 1000})").unwrap() + else { + panic!("expected a run command"); + }; + assert_eq!( + parse_json_value(&command_json).unwrap(), + serde_json::json!({ "create": "logs", "capped": true, "size": 1048576, "max": 1000 }) + ); + + for (source, expected) in [ + ("db.dropDatabase(1)", "takes no arguments"), + ("db.createCollection()", "requires a collection name"), + ("db.createCollection(1)", "must be a string"), + (r#"db.createCollection("")"#, "valid collection name"), + (r#"db.createCollection("a$b")"#, "valid collection name"), + (r#"db.createCollection("a", [])"#, "options must be a document"), + (r#"db.createCollection("a", {create: "b"})"#, "must not contain create"), + (r#"db.createCollection("a").x()"#, "collection name and optional options"), + ] { + let error = parse(source).unwrap_err(); + assert!(error.contains(expected), "{source}\n expected: {expected}\n got: {error}"); + } + } + #[test] fn parses_show_databases_aliases_as_read_only() { for source in [ diff --git a/packages/app-tests/mongoCompletion.test.ts b/packages/app-tests/mongoCompletion.test.ts index 3e59477c04..10ff0416e6 100644 --- a/packages/app-tests/mongoCompletion.test.ts +++ b/packages/app-tests/mongoCompletion.test.ts @@ -253,6 +253,8 @@ test("offers the newly supported count and database commands", () => { const dbLevel = labels("db."); assert.ok(dbLevel.includes("stats")); assert.ok(dbLevel.includes("serverStatus")); + assert.ok(dbLevel.includes("createCollection")); + assert.ok(dbLevel.includes("dropDatabase")); assert.ok(labels("").includes("db.stats")); }); diff --git a/packages/app-tests/mongoShellCommand.test.ts b/packages/app-tests/mongoShellCommand.test.ts index 542818f354..aaac83bbb9 100644 --- a/packages/app-tests/mongoShellCommand.test.ts +++ b/packages/app-tests/mongoShellCommand.test.ts @@ -1367,6 +1367,30 @@ test("parseMongoCountDocumentsCommand reads estimatedDocumentCount as a metadata assert.equal(parseMongoCountDocumentsCommand("db.orders.count()")?.mode, "legacy"); }); +test("parseMongoRunCommand reads db.dropDatabase() and db.createCollection() as run commands", () => { + assert.deepEqual(parseMongoRunCommand("db . dropDatabase ( ) ;"), { commandJson: '{"dropDatabase":1}' }); + assert.deepEqual(parseMongoRunCommand('db.createCollection("events")'), { commandJson: '{"create":"events"}' }); + assert.deepEqual(JSON.parse(parseMongoRunCommand("db.createCollection('logs', {capped: true, size: 1048576, max: 1000})")!.commandJson), { + create: "logs", + capped: true, + size: 1048576, + max: 1000, + }); + assert.equal(parseMongoCommand("db.dropDatabase()")?.command.kind, "runCommand"); + + for (const [source, expected] of [ + ["db.dropDatabase(1)", /dropDatabase\(\) expects no arguments/], + ["db.createCollection()", /createCollection\(\) expects a collection name and optional options/], + ["db.createCollection(1)", /expects a collection name/], + ['db.createCollection("a$b")', /expects a collection name/], + ['db.createCollection("a", {create: "b"})', /expects a collection name/], + ['db.createCollection("a").x()', /Unexpected text after createCollection/], + ] as const) { + assert.equal(parseMongoCommand(source), null, source); + assert.match(describeMongoCommandParseFailure(source), expected, source); + } +}); + 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}' }); @@ -1408,7 +1432,7 @@ test("describeMongoCommandParseFailure names unsupported methods and points at a assert.match(describeMongoCommandParseFailure('db["my-coll"].validate()'), /validate\(\) is not supported/); assert.match(describeMongoCommandParseFailure('db.getCollection("my-coll").watch()'), /watch\(\) is not supported/); - assert.match(describeMongoCommandParseFailure('db.createCollection("c")'), /db\.createCollection\(\) is not supported; collections are created on first insert/); + assert.match(describeMongoCommandParseFailure("db.getCollectionNames()"), /db\.getCollectionNames\(\) is not supported; collections are listed in the sidebar/); assert.match(describeMongoCommandParseFailure('db.getSiblingDB("other").c.find({})'), /use /); assert.match(describeMongoCommandParseFailure("db.adminCommand({ping: 1})"), /use db\.runCommand/); assert.match(describeMongoCommandParseFailure("show collections"), /listed in the sidebar/);