fix(mongodb): support legacy count helpers

This commit is contained in:
t8y2 2026-07-03 13:25:36 +08:00
parent c702149931
commit cabe89a19b
5 changed files with 146 additions and 31 deletions

View File

@ -88,6 +88,7 @@ export function parseMongoFindCommand(input: string): MongoFindCommand | null {
const chain = source.slice(findCloseIndex + 1).trim();
if (chain && !chain.startsWith(".")) return null;
if (findChainedMethodCallIndex(chain, "count") >= 0) return null;
const sortArg = readChainedCallArgument(chain, "sort");
let sort: string | undefined;
@ -113,7 +114,13 @@ export function parseMongoFindCommand(input: string): MongoFindCommand | null {
export function parseMongoCountDocumentsCommand(input: string): MongoCountDocumentsCommand | null {
const source = input.trim().replace(/;$/, "").trim();
const target = parseCollectionMethodTarget(source, "countDocuments");
// Accept deprecated Mongo shell count helpers for old server workflows, but
// keep DBX's internal execution mapped to the countDocuments result shape.
return parseCollectionCountCommand(source, "countDocuments") ?? parseCollectionCountCommand(source, "count") ?? parseFindCountCommand(source);
}
function parseCollectionCountCommand(source: string, method: "countDocuments" | "count"): MongoCountDocumentsCommand | null {
const target = parseCollectionMethodTarget(source, method);
if (!target) return null;
const openIndex = source.indexOf("(", target.methodCallIndex);
@ -131,6 +138,28 @@ export function parseMongoCountDocumentsCommand(input: string): MongoCountDocume
};
}
function parseFindCountCommand(source: string): MongoCountDocumentsCommand | null {
const target = parseFindTarget(source);
if (!target) return null;
const findOpenIndex = source.indexOf("(", target.findCallIndex);
const findCloseIndex = findMatchingParen(source, findOpenIndex);
if (findCloseIndex < 0) return null;
const chain = source.slice(findCloseIndex + 1).trim();
if (!hasSingleEmptyChainedCall(chain, "count")) return null;
const findArgs = splitTopLevel(source.slice(findOpenIndex + 1, findCloseIndex));
if (findArgs.length > 2 && findArgs.slice(2).some((arg) => arg.trim())) return null;
const filter = normalizeJsonArgument(findArgs[0] || "{}");
if (!filter) return null;
return {
collection: target.collection,
filter,
};
}
export function parseMongoAggregateCommand(input: string): MongoAggregateCommand | null {
const source = input.trim().replace(/;$/, "").trim();
const target = parseCollectionMethodTarget(source, "aggregate");
@ -268,18 +297,20 @@ export function parseMongoCommand(input: string): ParsedMongoCommand | null {
// Keep the more specific readers ahead of generic write parsing so the
// returned kind matches the result renderer we want to use downstream.
const parsers: Array<(source: string) => MongoCommand | null> = [
(source) => {
const find = parseMongoFindCommand(source);
return find ? { kind: "find", ...find } : null;
},
(source) => {
const version = parseMongoVersionCommand(source);
return version ?? null;
},
(source) => {
// Legacy Mongo shell uses count()/find().count(); keep accepting it
// while mapping to DBX's countDocuments-compatible result path.
const count = parseMongoCountDocumentsCommand(source);
return count ? { kind: "countDocuments", ...count } : null;
},
(source) => {
const find = parseMongoFindCommand(source);
return find ? { kind: "find", ...find } : null;
},
(source) => {
const aggregate = parseMongoAggregateCommand(source);
return aggregate ? { kind: "aggregate", ...aggregate } : null;
@ -915,6 +946,15 @@ function readChainedCallArgument(source: string, name: string): string | undefin
return undefined;
}
function hasSingleEmptyChainedCall(source: string, name: string): boolean {
const trimmed = source.trim();
const match = chainedMethodCallPattern(name).exec(trimmed);
if (!match || match.index !== 0) return false;
const openIndex = trimmed.indexOf("(", match.index);
const closeIndex = findMatchingParen(trimmed, openIndex);
return closeIndex >= 0 && !trimmed.slice(openIndex + 1, closeIndex).trim() && !trimmed.slice(closeIndex + 1).trim();
}
function findChainedMethodCallIndex(source: string, name: string): number {
return chainedMethodCallPattern(name).exec(source)?.index ?? -1;
}

View File

@ -189,6 +189,27 @@ test("parseMongoCountDocumentsCommand parses db collection countDocuments", () =
});
});
test("parseMongoCountDocumentsCommand parses legacy count helpers", () => {
assert.deepEqual(parseMongoCountDocumentsCommand("db.products.count({ active: true })"), {
collection: "products",
filter: '{ "active": true }',
});
assert.deepEqual(parseMongoCountDocumentsCommand('db.getCollection("audit.logs").count()'), {
collection: "audit.logs",
filter: "{}",
});
assert.deepEqual(parseMongoCountDocumentsCommand("db.products.find({ active: true }).count()"), {
collection: "products",
filter: '{ "active": true }',
});
assert.equal(parseMongoFindCommand("db.products.find({ active: true }).count()"), null);
assert.deepEqual(parseMongoCommand("db.products.find({ active: true }).count()")?.command, {
kind: "countDocuments",
collection: "products",
filter: '{ "active": true }',
});
});
test("parseMongoAggregateCommand parses db collection aggregate", () => {
assert.deepEqual(parseMongoAggregateCommand('db.products.aggregate([{"$match":{"active":true}},{"$count":"total"}])'), {
collection: "products",

View File

@ -780,11 +780,6 @@ export async function executeQuery(config: ConnectionConfig, sql: string, option
return convertBridgeQueryResult(result, options);
}
if (config.db_type === "mongodb") {
const find = parseMongoFindCommand(sql);
if (find) {
const result = await withTimeout(mongoFindDocuments(config, find.collection, find.skip, find.limit, find.filter, find.projection, find.sort), resolveTimeoutMs(options));
return mongoDocumentsToQueryResult(result.documents.slice(0, resolveMaxRows(options)), result.total);
}
const version = parseMongoVersionCommand(sql);
if (version) {
const result = await withTimeout(mongoServerVersion(config), resolveTimeoutMs(options));
@ -795,6 +790,11 @@ export async function executeQuery(config: ConnectionConfig, sql: string, option
const result = await withTimeout(mongoFindDocuments(config, count.collection, 0, 1, count.filter), resolveTimeoutMs(options));
return { columns: ["count"], rows: [{ count: result.total }], row_count: 1 };
}
const find = parseMongoFindCommand(sql);
if (find) {
const result = await withTimeout(mongoFindDocuments(config, find.collection, find.skip, find.limit, find.filter, find.projection, find.sort), resolveTimeoutMs(options));
return mongoDocumentsToQueryResult(result.documents.slice(0, resolveMaxRows(options)), result.total);
}
const aggregate = parseMongoAggregateCommand(sql);
if (aggregate) {
const safety = evaluateMongoAggregateSafety(aggregate, sqlSafetyFromEnv());
@ -829,7 +829,7 @@ export async function executeQuery(config: ConnectionConfig, sql: string, option
return { columns: [], rows: [], row_count: result.affectedRows };
}
throw new Error(
"Use MongoDB shell-style commands, for example: db.projects.find({}).limit(100), db.version(), db.projects.countDocuments({}), db.projects.getIndexes(), db.projects.createIndex({...}), db.projects.dropIndex(\"name\"), db.projects.dropIndexes(), db.projects.insertOne({...}), db.projects.updateOne({...}, {$set: {...}}), or db.projects.deleteOne({...})",
"Use MongoDB shell-style commands, for example: db.projects.find({}).limit(100), db.version(), db.projects.countDocuments({}), db.projects.count({}), db.projects.getIndexes(), db.projects.createIndex({...}), db.projects.dropIndex(\"name\"), db.projects.dropIndexes(), db.projects.insertOne({...}), db.projects.updateOne({...}, {$set: {...}}), or db.projects.deleteOne({...})",
);
}
if (isDirectQueryType(config.db_type)) {
@ -1206,6 +1206,7 @@ export function parseMongoFindCommand(input: string): MongoFindCommand | null {
}
const chain = source.slice(findCloseIndex + 1).trim();
if (chain && !chain.startsWith(".")) return null;
if (findChainedMethodCallIndex(chain, "count") >= 0) return null;
const sortArg = readChainedCallArgument(chain, "sort");
let sort: string | undefined;
if (sortArg !== undefined) {
@ -1226,7 +1227,17 @@ export function parseMongoVersionCommand(input: string): boolean {
export function parseMongoCountDocumentsCommand(input: string): MongoCountDocumentsCommand | null {
const source = input.trim().replace(/;$/, "").trim();
const target = parseCollectionMethodTarget(source, "countDocuments");
// Accept deprecated Mongo shell count helpers for old server workflows, but
// keep DBX's internal execution mapped to the countDocuments result shape.
return (
parseCollectionCountCommand(source, "countDocuments") ??
parseCollectionCountCommand(source, "count") ??
parseFindCountCommand(source)
);
}
function parseCollectionCountCommand(source: string, method: "countDocuments" | "count"): MongoCountDocumentsCommand | null {
const target = parseCollectionMethodTarget(source, method);
if (!target) return null;
const openIndex = source.indexOf("(", target.methodCallIndex);
const closeIndex = findMatchingParen(source, openIndex);
@ -1237,6 +1248,20 @@ export function parseMongoCountDocumentsCommand(input: string): MongoCountDocume
return filter ? { collection: target.collection, filter } : null;
}
function parseFindCountCommand(source: string): MongoCountDocumentsCommand | null {
const target = parseCollectionMethodTarget(source, "find");
if (!target) return null;
const findOpenIndex = source.indexOf("(", target.methodCallIndex);
const findCloseIndex = findMatchingParen(source, findOpenIndex);
if (findCloseIndex < 0) return null;
const chain = source.slice(findCloseIndex + 1).trim();
if (!hasSingleEmptyChainedCall(chain, "count")) return null;
const findArgs = splitTopLevel(source.slice(findOpenIndex + 1, findCloseIndex));
if (findArgs.length > 2 && findArgs.slice(2).some((arg) => arg.trim())) return null;
const filter = normalizeJsonArgument(findArgs[0] || "{}");
return filter ? { collection: target.collection, filter } : null;
}
export function parseMongoAggregateCommand(input: string): MongoAggregateCommand | null {
const source = input.trim().replace(/;$/, "").trim();
const target = parseCollectionMethodTarget(source, "aggregate");
@ -1410,6 +1435,19 @@ function readChainedCallArgument(chain: string, method: string): string | undefi
return closeIndex < 0 ? undefined : chain.slice(openIndex + 1, closeIndex);
}
function hasSingleEmptyChainedCall(chain: string, method: string): boolean {
const trimmed = chain.trim();
const match = chainedMethodCallPattern(method).exec(trimmed);
if (!match || match.index !== 0) return false;
const openIndex = trimmed.indexOf("(", match.index);
const closeIndex = findMatchingParen(trimmed, openIndex);
return (
closeIndex >= 0 &&
!trimmed.slice(openIndex + 1, closeIndex).trim() &&
!trimmed.slice(closeIndex + 1).trim()
);
}
function findChainedMethodCallIndex(source: string, method: string): number {
return chainedMethodCallPattern(method).exec(source)?.index ?? -1;
}

View File

@ -129,24 +129,6 @@ export async function describeTable(config: ConnectionConfig, table: string, sch
export async function executeQuery(config: ConnectionConfig, sql: string, options?: QueryOptions): Promise<QueryResult> {
await ensureConnected(config);
if (config.db_type === "mongodb") {
const find = parseMongoFindCommand(sql);
if (find) {
const res = await apiFetch("/api/mongo/find-documents", {
method: "POST",
body: JSON.stringify({
connectionId: config.id,
database: config.database || "",
collection: find.collection,
skip: find.skip,
limit: find.limit,
filter: find.filter,
projection: find.projection,
sort: find.sort,
}),
});
const result = (await res.json()) as { documents: unknown[]; total: number };
return mongoDocumentsToQueryResult(result.documents.slice(0, options?.maxRows ?? result.documents.length), result.total);
}
if (parseMongoVersionCommand(sql)) {
const res = await apiFetch("/api/mongo/server-version", {
method: "POST",
@ -174,6 +156,24 @@ export async function executeQuery(config: ConnectionConfig, sql: string, option
const result = (await res.json()) as { documents: unknown[]; total: number };
return { columns: ["count"], rows: [{ count: result.total }], row_count: 1 };
}
const find = parseMongoFindCommand(sql);
if (find) {
const res = await apiFetch("/api/mongo/find-documents", {
method: "POST",
body: JSON.stringify({
connectionId: config.id,
database: config.database || "",
collection: find.collection,
skip: find.skip,
limit: find.limit,
filter: find.filter,
projection: find.projection,
sort: find.sort,
}),
});
const result = (await res.json()) as { documents: unknown[]; total: number };
return mongoDocumentsToQueryResult(result.documents.slice(0, options?.maxRows ?? result.documents.length), result.total);
}
const aggregate = parseMongoAggregateCommand(sql);
if (aggregate) {
const safety = evaluateMongoAggregateSafety(aggregate, sqlSafetyFromEnv());
@ -220,7 +220,7 @@ export async function executeQuery(config: ConnectionConfig, sql: string, option
return { columns: [], rows: [], row_count: result.affectedRows };
}
throw new Error(
"Use MongoDB shell-style commands, for example: db.projects.find({}).limit(100), db.version(), db.projects.countDocuments({}), db.projects.getIndexes(), db.projects.createIndex({...}), db.projects.dropIndex(\"name\"), db.projects.dropIndexes(), db.projects.insertOne({...}), db.projects.updateOne({...}, {$set: {...}}), or db.projects.deleteOne({...})",
"Use MongoDB shell-style commands, for example: db.projects.find({}).limit(100), db.version(), db.projects.countDocuments({}), db.projects.count({}), db.projects.getIndexes(), db.projects.createIndex({...}), db.projects.dropIndex(\"name\"), db.projects.dropIndexes(), db.projects.insertOne({...}), db.projects.updateOne({...}, {$set: {...}}), or db.projects.deleteOne({...})",
);
}
const res = await apiFetch("/api/query/execute", {

View File

@ -67,6 +67,22 @@ test("parseMongoCountDocumentsCommand accepts shell-style count commands", () =>
});
});
test("parseMongoCountDocumentsCommand accepts legacy count helpers", () => {
assert.deepEqual(parseMongoCountDocumentsCommand("db.projects.count({ active: true })"), {
collection: "projects",
filter: '{ "active": true }',
});
assert.deepEqual(parseMongoCountDocumentsCommand('db.getCollection("audit.logs").count()'), {
collection: "audit.logs",
filter: "{}",
});
assert.deepEqual(parseMongoCountDocumentsCommand("db.projects.find({ active: true }).count()"), {
collection: "projects",
filter: '{ "active": true }',
});
assert.equal(parseMongoFindCommand("db.projects.find({ active: true }).count()"), null);
});
test("parseMongoAggregateCommand accepts aggregate pipelines", () => {
assert.deepEqual(parseMongoAggregateCommand('db.projects.aggregate([{"$match":{"active":true}},{"$group":{"_id":"$owner","total":{"$sum":1}}}])'), {
collection: "projects",