fix(mongodb): support compass shell commands

This commit is contained in:
t8y2 2026-06-01 16:24:31 +08:00
parent a0dcee51bc
commit a87b8ece31
10 changed files with 512 additions and 10 deletions

View File

@ -214,8 +214,11 @@ export const mongoListCollections = forward("mongoListCollections");
export const mongoFindDocuments = forward("mongoFindDocuments");
export const mongoAggregateDocuments = forward("mongoAggregateDocuments");
export const mongoInsertDocument = forward("mongoInsertDocument");
export const mongoInsertDocuments = forward("mongoInsertDocuments");
export const mongoUpdateDocument = forward("mongoUpdateDocument");
export const mongoUpdateDocuments = forward("mongoUpdateDocuments");
export const mongoDeleteDocument = forward("mongoDeleteDocument");
export const mongoDeleteDocuments = forward("mongoDeleteDocuments");
// History
export const saveHistory = forward("saveHistory");

View File

@ -1304,6 +1304,15 @@ export async function mongoInsertDocument(
return post("/api/mongo/insert-document", { connectionId, database, collection, docJson });
}
export async function mongoInsertDocuments(
connectionId: string,
database: string,
collection: string,
docsJson: string,
): Promise<{ affected_rows: number }> {
return post("/api/mongo/insert-documents", { connectionId, database, collection, docsJson });
}
export async function mongoUpdateDocument(
connectionId: string,
database: string,
@ -1314,6 +1323,17 @@ export async function mongoUpdateDocument(
return post("/api/mongo/update-document", { connectionId, database, collection, id, docJson });
}
export async function mongoUpdateDocuments(
connectionId: string,
database: string,
collection: string,
filterJson: string,
updateJson: string,
many: boolean,
): Promise<{ affected_rows: number }> {
return post("/api/mongo/update-documents", { connectionId, database, collection, filterJson, updateJson, many });
}
export async function mongoDeleteDocument(
connectionId: string,
database: string,
@ -1323,6 +1343,16 @@ export async function mongoDeleteDocument(
return post("/api/mongo/delete-document", { connectionId, database, collection, id });
}
export async function mongoDeleteDocuments(
connectionId: string,
database: string,
collection: string,
filterJson: string,
many: boolean,
): Promise<{ affected_rows: number }> {
return post("/api/mongo/delete-documents", { connectionId, database, collection, filterJson, many });
}
// ---------------------------------------------------------------------------
// History
// ---------------------------------------------------------------------------

View File

@ -18,6 +18,11 @@ export interface MongoAggregateCommand {
pipeline: string;
}
export type MongoWriteCommand =
| { kind: "insert"; collection: string; docsJson: string }
| { kind: "update"; collection: string; filter: string; update: string; many: boolean }
| { kind: "delete"; collection: string; filter: string; many: boolean };
export interface MongoAggregateSafetyOptions {
allowWrites?: boolean;
allowDangerous?: boolean;
@ -107,6 +112,51 @@ export function parseMongoAggregateCommand(input: string): MongoAggregateCommand
};
}
export function parseMongoWriteCommand(input: string): MongoWriteCommand | null {
const source = input.trim().replace(/;$/, "").trim();
const insertOne = parseCollectionMethodTarget(source, "insertOne");
if (insertOne) {
const args = parseMethodArgs(source, insertOne.methodCallIndex);
if (!args || args.length !== 1) return null;
const doc = normalizeJsonArgument(args[0]);
return doc ? { kind: "insert", collection: insertOne.collection, docsJson: doc } : null;
}
const insertMany = parseCollectionMethodTarget(source, "insertMany");
if (insertMany) {
const args = parseMethodArgs(source, insertMany.methodCallIndex);
if (!args || args.length !== 1) return null;
const docs = normalizeJsonArgument(args[0]);
if (!docs) return null;
return Array.isArray(JSON.parse(docs))
? { kind: "insert", collection: insertMany.collection, docsJson: docs }
: null;
}
for (const method of ["updateOne", "updateMany"] as const) {
const target = parseCollectionMethodTarget(source, method);
if (!target) continue;
const args = parseMethodArgs(source, target.methodCallIndex);
if (!args || args.length !== 2) return null;
const filter = normalizeJsonArgument(args[0]);
const update = normalizeJsonArgument(args[1]);
if (!filter || !update) return null;
return { kind: "update", collection: target.collection, filter, update, many: method === "updateMany" };
}
for (const method of ["deleteOne", "deleteMany"] as const) {
const target = parseCollectionMethodTarget(source, method);
if (!target) continue;
const args = parseMethodArgs(source, target.methodCallIndex);
if (!args || args.length !== 1) return null;
const filter = normalizeJsonArgument(args[0]);
if (!filter) return null;
return { kind: "delete", collection: target.collection, filter, many: method === "deleteMany" };
}
return null;
}
export function mongoAggregateWriteStage(pipelineJson: string): "$out" | "$merge" | null {
try {
const pipeline = JSON.parse(pipelineJson);
@ -179,6 +229,15 @@ export function mongoCountToQueryResult(total: number, executionTimeMs: number):
};
}
export function mongoWriteToQueryResult(affectedRows: number, executionTimeMs: number): QueryResult {
return {
columns: [],
rows: [],
affected_rows: affectedRows,
execution_time_ms: Math.max(0, Math.round(executionTimeMs)),
};
}
function parseFindTarget(source: string): { collection: string; findCallIndex: number } | null {
const direct = parseCollectionMethodTarget(source, "find");
if (direct) {
@ -217,7 +276,9 @@ function parseCollectionMethodTarget(
function normalizeJsonArgument(value: string): string | null {
const trimmed = value.trim();
if (!trimmed) return "{}";
const preprocessed = trimmed.replace(/ObjectId\s*\(\s*["']([^"']+)["']\s*\)/g, '{"$oid":"$1"}');
const preprocessed = quoteUnquotedObjectKeys(
convertSingleQuotedStrings(trimmed.replace(/ObjectId\s*\(\s*["']([^"']+)["']\s*\)/g, '{"$oid":"$1"}')),
);
try {
JSON.parse(preprocessed);
return preprocessed;
@ -226,6 +287,105 @@ function normalizeJsonArgument(value: string): string | null {
}
}
function parseMethodArgs(source: string, methodCallIndex: number): string[] | null {
const openIndex = source.indexOf("(", methodCallIndex);
const closeIndex = findMatchingParen(source, openIndex);
if (closeIndex < 0 || source.slice(closeIndex + 1).trim()) return null;
return splitTopLevel(source.slice(openIndex + 1, closeIndex));
}
function convertSingleQuotedStrings(source: string): string {
let result = "";
let copiedUntil = 0;
let quote: string | null = null;
let start = 0;
let value = "";
let escaped = false;
for (let i = 0; i < source.length; i += 1) {
const char = source[i];
if (!quote) {
if (char === "'") {
quote = char;
start = i;
value = "";
escaped = false;
} else if (char === '"') {
quote = char;
}
continue;
}
if (quote === '"') {
if (escaped) escaped = false;
else if (char === "\\") escaped = true;
else if (char === '"') quote = null;
continue;
}
if (escaped) {
value += char;
escaped = false;
} else if (char === "\\") {
escaped = true;
} else if (char === "'") {
result += source.slice(copiedUntil, start) + JSON.stringify(value);
copiedUntil = i + 1;
quote = null;
} else {
value += char;
}
}
return quote === "'" ? source : result + source.slice(copiedUntil);
}
function quoteUnquotedObjectKeys(source: string): string {
let result = "";
let quote: string | null = null;
let escaped = false;
for (let i = 0; i < source.length; i += 1) {
const char = source[i];
if (quote) {
result += char;
if (escaped) escaped = false;
else if (char === "\\") escaped = true;
else if (char === quote) quote = null;
continue;
}
if (char === '"' || char === "'") {
quote = char;
result += char;
continue;
}
if (/[A-Za-z_$]/.test(char) && shouldQuoteObjectKey(source, i)) {
let end = i + 1;
while (/[\w$]/.test(source[end] || "")) end += 1;
result += `"${source.slice(i, end)}"`;
i = end - 1;
continue;
}
result += char;
}
return result;
}
function shouldQuoteObjectKey(source: string, index: number): boolean {
let before = index - 1;
while (/\s/.test(source[before] || "")) before -= 1;
if (source[before] !== "{" && source[before] !== ",") return false;
let after = index + 1;
while (/[\w$]/.test(source[after] || "")) after += 1;
while (/\s/.test(source[after] || "")) after += 1;
return source[after] === ":";
}
function readChainedIntegerArgument(source: string, name: string, fallback: number): number | null {
const raw = readChainedCallArgument(source, name);
if (raw === undefined) return fallback;

View File

@ -1141,6 +1141,16 @@ export async function mongoInsertDocument(
return invoke("mongo_insert_document", { connectionId, database, collection, docJson });
}
export async function mongoInsertDocuments(
connectionId: string,
database: string,
collection: string,
docsJson: string,
): Promise<{ affected_rows: number }> {
const affectedRows = await invoke<number>("mongo_insert_documents", { connectionId, database, collection, docsJson });
return { affected_rows: affectedRows };
}
export async function mongoUpdateDocument(
connectionId: string,
database: string,
@ -1151,6 +1161,25 @@ export async function mongoUpdateDocument(
return invoke("mongo_update_document", { connectionId, database, collection, id, docJson });
}
export async function mongoUpdateDocuments(
connectionId: string,
database: string,
collection: string,
filterJson: string,
updateJson: string,
many: boolean,
): Promise<{ affected_rows: number }> {
const affectedRows = await invoke<number>("mongo_update_documents", {
connectionId,
database,
collection,
filterJson,
updateJson,
many,
});
return { affected_rows: affectedRows };
}
export async function mongoDeleteDocument(
connectionId: string,
database: string,
@ -1160,6 +1189,23 @@ export async function mongoDeleteDocument(
return invoke("mongo_delete_document", { connectionId, database, collection, id });
}
export async function mongoDeleteDocuments(
connectionId: string,
database: string,
collection: string,
filterJson: string,
many: boolean,
): Promise<{ affected_rows: number }> {
const affectedRows = await invoke<number>("mongo_delete_documents", {
connectionId,
database,
collection,
filterJson,
many,
});
return { affected_rows: affectedRows };
}
// --- History ---
export interface HistoryEntry {
id: string;

View File

@ -13,9 +13,11 @@ import {
evaluateMongoAggregateSafety,
mongoCountToQueryResult,
mongoDocumentsToQueryResult,
mongoWriteToQueryResult,
parseMongoAggregateCommand,
parseMongoCountDocumentsCommand,
parseMongoFindCommand,
parseMongoWriteCommand,
type MongoAggregateSafetyOptions,
} from "@/lib/mongoShellCommand";
import { AGENT_DRIVER_TYPES } from "@/lib/databaseCapabilities";
@ -797,6 +799,64 @@ export const useQueryStore = defineStore("query", () => {
return;
}
const mongoWrite = conn?.db_type === "mongodb" ? parseMongoWriteCommand(sql) : null;
if (mongoWrite) {
await connStore.ensureConnected(tab.connectionId);
console.info("[DBX][executeTabSql:mongo-write:start]", {
traceId,
kind: mongoWrite.kind,
collection: mongoWrite.collection,
});
let affectedRows = 0;
if (mongoWrite.kind === "insert") {
const result = await api.mongoInsertDocuments(
tab.connectionId,
tab.database,
mongoWrite.collection,
mongoWrite.docsJson,
);
affectedRows = result.affected_rows;
} else if (mongoWrite.kind === "update") {
const result = await api.mongoUpdateDocuments(
tab.connectionId,
tab.database,
mongoWrite.collection,
mongoWrite.filter,
mongoWrite.update,
mongoWrite.many,
);
affectedRows = result.affected_rows;
} else {
const result = await api.mongoDeleteDocuments(
tab.connectionId,
tab.database,
mongoWrite.collection,
mongoWrite.filter,
mongoWrite.many,
);
affectedRows = result.affected_rows;
}
console.info("[DBX][executeTabSql:mongo-write:done]", {
traceId,
affectedRows,
elapsed: elapsed(),
});
const current = tabs.value.find((t) => t.id === id);
if (current?.executionId === executionId) {
current.results = undefined;
current.activeResultIndex = undefined;
current.result = mongoWriteToQueryResult(affectedRows, performance.now() - startedAt);
touchResult(current);
current.queryAnalysis = undefined;
current.querySourceColumns = undefined;
current.queryEditabilityReason = undefined;
current.tableMeta = undefined;
current.resultBaseSql = options?.resultBaseSql ?? sql;
current.resultSortedSql = options?.resultSortedSql;
}
return;
}
console.info("[DBX][executeTabSql:execute-multi:start]", { traceId, elapsed: elapsed() });
const clientSessionId = tab.mode === "query" ? tab.id : undefined;
const executionOptions = {

View File

@ -8,6 +8,7 @@ import {
parseMongoAggregateCommand,
parseMongoCountDocumentsCommand,
parseMongoFindCommand,
parseMongoWriteCommand,
} from "../../apps/desktop/src/lib/mongoShellCommand.ts";
test("parseMongoFindCommand parses db collection find with an empty JSON filter", () => {
@ -35,8 +36,43 @@ test("parseMongoFindCommand parses getCollection find with chained sort skip and
);
});
test("parseMongoFindCommand accepts Compass-style unquoted keys and ObjectId", () => {
const command = parseMongoFindCommand("db.products.find({_id: ObjectId('6a045a92d2971e44243771a1')}).limit(1)");
assert.ok(command);
assert.equal(command.collection, "products");
assert.equal(command.limit, 1);
assert.deepEqual(JSON.parse(command.filter), { _id: { $oid: "6a045a92d2971e44243771a1" } });
});
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);
assert.equal(command.collection, "products");
assert.equal(command.limit, 2);
assert.deepEqual(JSON.parse(command.filter), { category: "Electronics" });
assert.deepEqual(JSON.parse(command.sort || "{}"), { price: -1 });
});
test("parseMongoFindCommand rejects unsupported mongo shell commands", () => {
assert.equal(parseMongoFindCommand("db.users.insertOne({})"), null);
assert.equal(parseMongoFindCommand("db.users.drop()"), null);
});
test("parseMongoWriteCommand accepts unquoted insert and update commands", () => {
assert.deepEqual(parseMongoWriteCommand("db.products.insertOne({name: 'demo', price: 1})"), {
kind: "insert",
collection: "products",
docsJson: '{"name": "demo", "price": 1}',
});
assert.deepEqual(
parseMongoWriteCommand("db.products.updateOne({_id: ObjectId('507f1f77bcf86cd799439011')}, {$set: {stock: 3}})"),
{
kind: "update",
collection: "products",
filter: '{"_id": {"$oid":"507f1f77bcf86cd799439011"}}',
update: '{"$set": {"stock": 3}}',
many: false,
},
);
});
test("parseMongoCountDocumentsCommand parses db collection countDocuments", () => {
@ -47,10 +83,13 @@ test("parseMongoCountDocumentsCommand parses db collection countDocuments", () =
});
test("parseMongoAggregateCommand parses db collection aggregate", () => {
assert.deepEqual(parseMongoAggregateCommand('db.products.aggregate([{"$match":{"active":true}},{"$count":"total"}])'), {
collection: "products",
pipeline: '[{"$match":{"active":true}},{"$count":"total"}]',
});
assert.deepEqual(
parseMongoAggregateCommand('db.products.aggregate([{"$match":{"active":true}},{"$count":"total"}])'),
{
collection: "products",
pipeline: '[{"$match":{"active":true}},{"$count":"total"}]',
},
);
});
test("parseMongoAggregateCommand accepts an empty pipeline", () => {
@ -68,13 +107,13 @@ test("parseMongoAggregateCommand rejects non-array pipelines and extra arguments
test("parseMongoAggregateCommand normalises ObjectId arguments with either quote style", () => {
const oid = "507f1f77bcf86cd799439011";
for (const quote of ["\"", "'"]) {
for (const quote of ['"', "'"]) {
const command = parseMongoAggregateCommand(
`db.orders.aggregate([{"$match":{"_id":ObjectId(${quote}${oid}${quote})}}])`,
);
assert.ok(command, `quote=${quote} should parse`);
assert.equal(command.collection, "orders");
assert.deepEqual(JSON.parse(command.pipeline), [{ "$match": { "_id": { "$oid": oid } } }]);
assert.deepEqual(JSON.parse(command.pipeline), [{ $match: { _id: { $oid: oid } } }]);
}
});
@ -87,7 +126,10 @@ test("evaluateMongoAggregateSafety blocks write stages unless MCP write flags al
const merge = parseMongoAggregateCommand('db.products.aggregate([{"$merge":{"into":"products_copy"}}])');
assert.ok(merge);
assert.equal(mongoAggregateWriteStage(merge.pipeline), "$merge");
assert.match(evaluateMongoAggregateSafety(merge, { allowWrites: true }).reason || "", /DBX_MCP_ALLOW_DANGEROUS_SQL=1/);
assert.match(
evaluateMongoAggregateSafety(merge, { allowWrites: true }).reason || "",
/DBX_MCP_ALLOW_DANGEROUS_SQL=1/,
);
assert.equal(evaluateMongoAggregateSafety(merge, { allowWrites: true, allowDangerous: true }).allowed, true);
});

View File

@ -849,7 +849,9 @@ function readChainedIntegerArgument(chain: string, method: string, fallback: num
}
function normalizeJsonArgument(arg: string): string | null {
const value = (arg.trim() || "{}").replace(/ObjectId\s*\(\s*["']([^"']+)["']\s*\)/g, '{"$oid":"$1"}');
const value = quoteUnquotedObjectKeys(
convertSingleQuotedStrings((arg.trim() || "{}").replace(/ObjectId\s*\(\s*["']([^"']+)["']\s*\)/g, '{"$oid":"$1"}')),
);
try {
JSON.parse(value);
return value;
@ -858,6 +860,98 @@ function normalizeJsonArgument(arg: string): string | null {
}
}
function convertSingleQuotedStrings(source: string): string {
let result = "";
let copiedUntil = 0;
let quote: string | null = null;
let start = 0;
let value = "";
let escaped = false;
for (let i = 0; i < source.length; i += 1) {
const char = source[i];
if (!quote) {
if (char === "'") {
quote = char;
start = i;
value = "";
escaped = false;
} else if (char === '"') {
quote = char;
}
continue;
}
if (quote === '"') {
if (escaped) escaped = false;
else if (char === "\\") escaped = true;
else if (char === '"') quote = null;
continue;
}
if (escaped) {
value += char;
escaped = false;
} else if (char === "\\") {
escaped = true;
} else if (char === "'") {
result += source.slice(copiedUntil, start) + JSON.stringify(value);
copiedUntil = i + 1;
quote = null;
} else {
value += char;
}
}
return quote === "'" ? source : result + source.slice(copiedUntil);
}
function quoteUnquotedObjectKeys(source: string): string {
let result = "";
let quote: string | null = null;
let escaped = false;
for (let i = 0; i < source.length; i += 1) {
const char = source[i];
if (quote) {
result += char;
if (escaped) escaped = false;
else if (char === "\\") escaped = true;
else if (char === quote) quote = null;
continue;
}
if (char === '"' || char === "'") {
quote = char;
result += char;
continue;
}
if (/[A-Za-z_$]/.test(char) && shouldQuoteObjectKey(source, i)) {
let end = i + 1;
while (/[\w$]/.test(source[end] || "")) end += 1;
result += `"${source.slice(i, end)}"`;
i = end - 1;
continue;
}
result += char;
}
return result;
}
function shouldQuoteObjectKey(source: string, index: number): boolean {
let before = index - 1;
while (/\s/.test(source[before] || "")) before -= 1;
if (source[before] !== "{" && source[before] !== ",") return false;
let after = index + 1;
while (/[\w$]/.test(source[after] || "")) after += 1;
while (/\s/.test(source[after] || "")) after += 1;
return source[after] === ":";
}
function isEmptyJsonObject(json: string): boolean {
try {
const parsed = JSON.parse(json);

View File

@ -21,6 +21,24 @@ test("parseMongoFindCommand accepts shell-style find commands", () => {
});
});
test("parseMongoFindCommand accepts Compass-style unquoted keys and ObjectId", () => {
const command = parseMongoFindCommand("db.products.find({_id: ObjectId('6a045a92d2971e44243771a1')}).limit(1)");
assert.ok(command);
assert.equal(command.collection, "products");
assert.equal(command.limit, 1);
assert.deepEqual(JSON.parse(command.filter), { _id: { $oid: "6a045a92d2971e44243771a1" } });
});
test("parseMongoWriteCommand accepts unquoted update operator keys", () => {
assert.deepEqual(parseMongoWriteCommand("db.projects.updateOne({_id: ObjectId('507f1f77bcf86cd799439011')}, {$set: {name: 'next'}})"), {
kind: "update",
collection: "projects",
filter: '{"_id": {"$oid":"507f1f77bcf86cd799439011"}}',
update: '{"$set": {"name": "next"}}',
many: false,
});
});
test("parseMongoCountDocumentsCommand accepts shell-style count commands", () => {
assert.deepEqual(parseMongoCountDocumentsCommand('db.projects.countDocuments({"active":true})'), {
collection: "projects",

View File

@ -76,6 +76,17 @@ pub async fn mongo_insert_document(
dbx_core::mongo_ops::mongo_insert_document_core(&state, &connection_id, &database, &collection, &doc_json).await
}
#[tauri::command]
pub async fn mongo_insert_documents(
state: State<'_, Arc<AppState>>,
connection_id: String,
database: String,
collection: String,
docs_json: String,
) -> Result<u64, String> {
dbx_core::mongo_ops::mongo_insert_documents_core(&state, &connection_id, &database, &collection, &docs_json).await
}
#[tauri::command]
pub async fn mongo_update_document(
state: State<'_, Arc<AppState>>,
@ -89,6 +100,28 @@ pub async fn mongo_update_document(
.await
}
#[tauri::command]
pub async fn mongo_update_documents(
state: State<'_, Arc<AppState>>,
connection_id: String,
database: String,
collection: String,
filter_json: String,
update_json: String,
many: bool,
) -> Result<u64, String> {
dbx_core::mongo_ops::mongo_update_documents_core(
&state,
&connection_id,
&database,
&collection,
&filter_json,
&update_json,
many,
)
.await
}
#[tauri::command]
pub async fn mongo_delete_document(
state: State<'_, Arc<AppState>>,
@ -99,3 +132,16 @@ pub async fn mongo_delete_document(
) -> Result<u64, String> {
dbx_core::mongo_ops::mongo_delete_document_core(&state, &connection_id, &database, &collection, &id).await
}
#[tauri::command]
pub async fn mongo_delete_documents(
state: State<'_, Arc<AppState>>,
connection_id: String,
database: String,
collection: String,
filter_json: String,
many: bool,
) -> Result<u64, String> {
dbx_core::mongo_ops::mongo_delete_documents_core(&state, &connection_id, &database, &collection, &filter_json, many)
.await
}

View File

@ -480,8 +480,11 @@ pub fn run() {
commands::mongo_cmd::mongo_find_documents,
commands::mongo_cmd::mongo_aggregate_documents,
commands::mongo_cmd::mongo_insert_document,
commands::mongo_cmd::mongo_insert_documents,
commands::mongo_cmd::mongo_update_document,
commands::mongo_cmd::mongo_update_documents,
commands::mongo_cmd::mongo_delete_document,
commands::mongo_cmd::mongo_delete_documents,
commands::history::save_history,
commands::history::load_history,
commands::history::clear_history,