diff --git a/src/lib/mongoShellCommand.ts b/src/lib/mongoShellCommand.ts new file mode 100644 index 0000000000..3effba6e27 --- /dev/null +++ b/src/lib/mongoShellCommand.ts @@ -0,0 +1,197 @@ +import type { QueryResult } from "@/types/database"; + +export interface MongoFindCommand { + collection: string; + filter: string; + skip: number; + limit: number; + sort?: string; +} + +const DEFAULT_LIMIT = 100; + +export function parseMongoFindCommand(input: string): MongoFindCommand | null { + const source = input.trim().replace(/;$/, "").trim(); + 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 findArgs = splitTopLevel(source.slice(findOpenIndex + 1, findCloseIndex)); + const filter = normalizeJsonArgument(findArgs[0] || "{}"); + if (!filter) return null; + + const chain = source.slice(findCloseIndex + 1).trim(); + if (chain && !chain.startsWith(".")) return null; + + const sortArg = readChainedCallArgument(chain, "sort"); + let sort: string | undefined; + if (sortArg !== undefined) { + const parsedSort = normalizeJsonArgument(sortArg); + if (!parsedSort) return null; + sort = parsedSort; + } + + const skip = readChainedIntegerArgument(chain, "skip", 0); + const limit = readChainedIntegerArgument(chain, "limit", DEFAULT_LIMIT); + if (skip === null || limit === null) return null; + + return { + collection: target.collection, + filter, + skip, + limit, + sort, + }; +} + +export function mongoDocumentsToQueryResult(documents: unknown[], executionTimeMs: number, total: number): QueryResult { + const columns: string[] = []; + + for (const doc of documents) { + if (isRecord(doc)) { + for (const key of Object.keys(doc)) { + if (!columns.includes(key)) columns.push(key); + } + } else if (!columns.includes("value")) { + columns.push("value"); + } + } + + const rows = documents.map((doc) => { + if (isRecord(doc)) return columns.map((column) => toCellValue(doc[column])); + return columns.map((column) => (column === "value" ? toCellValue(doc) : null)); + }); + + return { + columns, + rows, + affected_rows: total, + execution_time_ms: Math.max(0, Math.round(executionTimeMs)), + truncated: total > documents.length, + }; +} + +function parseFindTarget(source: string): { collection: string; findCallIndex: number } | null { + const direct = /^db\.([A-Za-z_$][\w$]*)\.find\s*\(/.exec(source); + if (direct) { + return { collection: direct[1], findCallIndex: source.indexOf(".find", direct[0].length - ".find(".length) }; + } + + const getCollection = /^db\.getCollection\s*\(\s*(["'])(.*?)\1\s*\)\.find\s*\(/.exec(source); + if (getCollection) { + return { + collection: getCollection[2], + findCallIndex: source.indexOf(".find", getCollection[0].length - ".find(".length), + }; + } + + return null; +} + +function normalizeJsonArgument(value: string): string | null { + const trimmed = value.trim(); + if (!trimmed) return "{}"; + try { + JSON.parse(trimmed); + return trimmed; + } catch { + return null; + } +} + +function readChainedIntegerArgument(source: string, name: string, fallback: number): number | null { + const raw = readChainedCallArgument(source, name); + if (raw === undefined) return fallback; + const value = Number(raw.trim()); + if (!Number.isSafeInteger(value) || value < 0) return null; + return value; +} + +function readChainedCallArgument(source: string, name: string): string | undefined { + const call = `.${name}`; + let index = source.indexOf(call); + while (index >= 0) { + const afterName = index + call.length; + const openIndex = skipWhitespace(source, afterName); + if (source[openIndex] === "(") { + const closeIndex = findMatchingParen(source, openIndex); + if (closeIndex >= 0) return source.slice(openIndex + 1, closeIndex); + } + index = source.indexOf(call, afterName); + } + return undefined; +} + +function skipWhitespace(source: string, index: number) { + let cursor = index; + while (/\s/.test(source[cursor] || "")) cursor += 1; + return cursor; +} + +function splitTopLevel(source: string): string[] { + const parts: string[] = []; + let start = 0; + let depth = 0; + let quote: string | null = null; + let escaped = false; + + for (let i = 0; i < source.length; i += 1) { + const char = source[i]; + if (quote) { + if (escaped) escaped = false; + else if (char === "\\") escaped = true; + else if (char === quote) quote = null; + continue; + } + + if (char === '"' || char === "'") quote = char; + else if (char === "{" || char === "[" || char === "(") depth += 1; + else if (char === "}" || char === "]" || char === ")") depth -= 1; + else if (char === "," && depth === 0) { + parts.push(source.slice(start, i).trim()); + start = i + 1; + } + } + + parts.push(source.slice(start).trim()); + return parts; +} + +function findMatchingParen(source: string, openIndex: number): number { + if (source[openIndex] !== "(") return -1; + let depth = 0; + let quote: string | null = null; + let escaped = false; + + for (let i = openIndex; i < source.length; i += 1) { + const char = source[i]; + if (quote) { + if (escaped) escaped = false; + else if (char === "\\") escaped = true; + else if (char === quote) quote = null; + continue; + } + + if (char === '"' || char === "'") quote = char; + else if (char === "(") depth += 1; + else if (char === ")") { + depth -= 1; + if (depth === 0) return i; + } + } + + return -1; +} + +function isRecord(value: unknown): value is Record { + return !!value && typeof value === "object" && !Array.isArray(value); +} + +function toCellValue(value: unknown): string | number | boolean | null { + if (value === undefined || value === null) return null; + if (typeof value === "string" || typeof value === "number" || typeof value === "boolean") return value; + return JSON.stringify(value); +} diff --git a/src/stores/queryStore.ts b/src/stores/queryStore.ts index d3094ddb9e..c8c4c4ea99 100644 --- a/src/stores/queryStore.ts +++ b/src/stores/queryStore.ts @@ -8,6 +8,7 @@ import { closeAllTabsState, closeOtherTabsState } from "@/lib/tabCloseActions"; import { buildExplainSql, parseExplainResult } from "@/lib/explainPlan"; import { analyzeEditableQuery, allPrimaryKeysPresent } from "@/lib/sqlAnalysis"; import { restoreOpenTabsState, serializeOpenTabs } from "@/lib/openTabsPersistence"; +import { mongoDocumentsToQueryResult, parseMongoFindCommand } from "@/lib/mongoShellCommand"; import * as api from "@/lib/api"; import { useConnectionStore } from "@/stores/connectionStore"; import { isTauriRuntime } from "@/lib/tauriRuntime"; @@ -385,6 +386,39 @@ export const useQueryStore = defineStore("query", () => { sql, }); try { + const connStore = useConnectionStore(); + const conn = connStore.getConfig(tab.connectionId); + const mongoFind = conn?.db_type === "mongodb" ? parseMongoFindCommand(sql) : null; + if (mongoFind) { + console.info("[DBX][executeTabSql:mongo-find:start]", { traceId, collection: mongoFind.collection }); + const result = await api.mongoFindDocuments( + tab.connectionId, + tab.database, + mongoFind.collection, + mongoFind.skip, + mongoFind.limit, + mongoFind.filter, + mongoFind.sort, + ); + console.info("[DBX][executeTabSql:mongo-find:done]", { + traceId, + rowCount: result.documents.length, + total: result.total, + elapsed: elapsed(), + }); + const current = tabs.value.find((t) => t.id === id); + if (current?.executionId === executionId) { + current.results = undefined; + current.activeResultIndex = undefined; + current.result = mongoDocumentsToQueryResult(result.documents, performance.now() - startedAt, result.total); + current.queryAnalysis = 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 results = await api.executeMulti(tab.connectionId, tab.database, sql, tab.schema, executionId); console.info("[DBX][executeTabSql:execute-multi:done]", { diff --git a/tests/mongoShellCommand.test.ts b/tests/mongoShellCommand.test.ts new file mode 100644 index 0000000000..ca2d37dd0d --- /dev/null +++ b/tests/mongoShellCommand.test.ts @@ -0,0 +1,52 @@ +import { strict as assert } from "node:assert"; +import test from "node:test"; +import { mongoDocumentsToQueryResult, parseMongoFindCommand } from "../src/lib/mongoShellCommand.ts"; + +test("parseMongoFindCommand parses db collection find with an empty JSON filter", () => { + assert.deepEqual(parseMongoFindCommand("db.users.find({})"), { + collection: "users", + filter: "{}", + skip: 0, + limit: 100, + sort: undefined, + }); +}); + +test("parseMongoFindCommand parses getCollection find with chained sort skip and limit", () => { + assert.deepEqual( + parseMongoFindCommand( + 'db.getCollection("audit.logs").find({"level":"warn"}).sort({"createdAt":-1}).skip(20).limit(10)', + ), + { + collection: "audit.logs", + filter: '{"level":"warn"}', + skip: 20, + limit: 10, + sort: '{"createdAt":-1}', + }, + ); +}); + +test("parseMongoFindCommand rejects unsupported mongo shell commands", () => { + assert.equal(parseMongoFindCommand("db.users.insertOne({})"), null); +}); + +test("mongoDocumentsToQueryResult turns mongo documents into grid rows", () => { + const result = mongoDocumentsToQueryResult( + [ + { _id: "1", name: "Ada", profile: { role: "admin" } }, + { _id: "2", active: true, name: "Lin" }, + ], + 5, + 12, + ); + + assert.deepEqual(result.columns, ["_id", "name", "profile", "active"]); + assert.deepEqual(result.rows, [ + ["1", "Ada", '{"role":"admin"}', null], + ["2", "Lin", null, true], + ]); + assert.equal(result.affected_rows, 12); + assert.equal(result.execution_time_ms, 5); + assert.equal(result.truncated, true); +});