github/dbx
4
0
Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-10-20 00:17:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(mongodb): support aggregate options and parse errors

Browse Source
This commit is contained in:
onenewcode
2026-07-17 16:57:38 +08:00
committed by GitHub
parent 782615586c
commit bb613b1712
25 changed files with 952 additions and 486 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/desktop/src/lib/backend/http.ts
+2 -2
View File
@@ -2140,8 +2140,8 @@ export async function mongoServerVersion(connectionId: string, database: string,
return post("/api/mongo/server-version", { connectionId, database, executionId });
}
export async function mongoAggregateDocuments(connectionId: string, database: string, collection: string, pipelineJson: string, maxRows?: number, executionId?: string): Promise<MongoDocumentResult> {
return post("/api/mongo/aggregate-documents", { connectionId, database, collection, pipelineJson, maxRows, executionId });
export async function mongoAggregateDocuments(connectionId: string, database: string, collection: string, pipelineJson: string, maxRows?: number, optionsJson?: string, executionId?: string): Promise<MongoDocumentResult> {
return post("/api/mongo/aggregate-documents", { connectionId, database, collection, pipelineJson, maxRows, optionsJson, executionId });
}
export async function mongoDistinct(connectionId: string, database: string, collection: string, field: string, filter?: string, executionId?: string): Promise<MongoDocumentResult> {
apps/desktop/src/lib/backend/tauri.ts
+2 -2
View File
@@ -1911,8 +1911,8 @@ export async function mongoServerVersion(connectionId: string, database: string,
return invoke("mongo_server_version", { connectionId, database, executionId });
}
export async function mongoAggregateDocuments(connectionId: string, database: string, collection: string, pipelineJson: string, maxRows?: number, executionId?: string): Promise<MongoDocumentResult> {
return invoke("mongo_aggregate_documents", { connectionId, database, collection, pipelineJson, maxRows, executionId });
export async function mongoAggregateDocuments(connectionId: string, database: string, collection: string, pipelineJson: string, maxRows?: number, optionsJson?: string, executionId?: string): Promise<MongoDocumentResult> {
return invoke("mongo_aggregate_documents", { connectionId, database, collection, pipelineJson, maxRows, optionsJson, executionId });
}
export async function mongoDistinct(connectionId: string, database: string, collection: string, field: string, filter?: string, executionId?: string): Promise<MongoDocumentResult> {
apps/desktop/src/lib/mongo/mongoShellCommand.ts
+16 -265
View File
@@ -1,5 +1,21 @@
import type { QueryResult } from "@/types/database";
import { mongoDocumentIdForGrid } from "@/lib/mongo/mongoDocumentValues";
import {
chainedMethodCallPattern,
describeMongoCommandParseFailure,
findChainedMethodCallIndex,
findMatchingParen,
MONGO_SHELL_COMMAND_HINT,
normalizeJsonArgument,
parseCollectionMethodTarget,
parseMongoAggregateCommand,
quoteUnquotedObjectKeys,
splitTopLevel,
type MongoAggregateCommand,
} from "@dbx-app/mongo-shell";
export type { MongoAggregateCommand };
export { describeMongoCommandParseFailure, MONGO_SHELL_COMMAND_HINT, parseMongoAggregateCommand, quoteUnquotedObjectKeys };
export interface MongoFindCommand {
collection: string;
@@ -23,11 +39,6 @@ export interface MongoCountDocumentsCommand {
mode: "accurate" | "legacy";
}
export interface MongoAggregateCommand {
collection: string;
pipeline: string;
}
export interface MongoGetIndexesCommand {
collection: string;
}
@@ -308,31 +319,6 @@ function parseFindCountCommand(source: string): MongoCountDocumentsCommand | nul
};
}
export function parseMongoAggregateCommand(input: string): MongoAggregateCommand | null {
const source = input.trim().replace(/;$/, "").trim();
const target = parseCollectionMethodTarget(source, "aggregate");
if (!target) return null;
const openIndex = source.indexOf("(", target.methodCallIndex);
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) return null;
const pipeline = normalizeJsonArgument(args[0]);
if (!pipeline) return null;
try {
if (!Array.isArray(JSON.parse(pipeline))) return null;
} catch {
return null;
}
return {
collection: target.collection,
pipeline,
};
}
export function parseMongoDistinctCommand(input: string): MongoDistinctCommand | null {
const source = input.trim().replace(/;$/, "").trim();
const target = parseCollectionMethodTarget(source, "distinct");
@@ -786,82 +772,6 @@ function parseFindTarget(source: string): { collection: string; findCallIndex: n
return null;
}
function parseCollectionMethodTarget(source: string, method: string): { collection: string; methodCallIndex: number } | null {
const escapedMethod = escapeRegExp(method);
const direct = new RegExp(`^db\\s*\\.\\s*([A-Za-z_$][\\w$]*)\\s*\\.\\s*${escapedMethod}\\s*\\(`).exec(source);
if (direct) {
return {
collection: direct[1],
methodCallIndex: findChainedMethodCallIndex(source, method),
};
}
const getCollection = new RegExp(`^db\\s*\\.\\s*getCollection\\s*\\(\\s*(["'])(.*?)\\1\\s*\\)\\s*\\.\\s*${escapedMethod}\\s*\\(`).exec(source);
if (getCollection) {
return {
collection: getCollection[2],
methodCallIndex: findChainedMethodCallIndex(source, method),
};
}
return null;
}
function normalizeJsonArgument(value: string): string | null {
const trimmed = value.trim();
if (!trimmed) return "{}";
// Rewrite mongo shell constructors that are not valid JSON into the extended
// JSON the backend understands (mongo_driver::json_value_to_bson): ObjectId(x)
// -> {"$oid":x}, NumberLong(x) -> {"$numberLong":x}, and
// ISODate(x)/new Date(x) -> {"$date":x}. Without this a
// filter such as { createdAt: { $gte: ISODate("...") } } fails JSON.parse,
// the command is left unrecognized and falls through to the SQL executor,
// which rejects it with "Use MongoDB-specific commands".
const withExtendedJson = replaceMongoShellConstructors(trimmed);
const preprocessed = quoteUnquotedObjectKeys(convertSingleQuotedStrings(withExtendedJson));
try {
JSON.parse(preprocessed);
return preprocessed;
} catch {
return null;
}
}
function replaceMongoShellConstructors(source: string): string {
const constructor = /^(ObjectId|NumberLong|ISODate)\s*\(\s*["']([^"']+)["']\s*\)|^(ObjectId|NumberLong)\s*\(\s*(-?\d+)\s*\)|^(?:new\s+Date)\s*\(\s*["']([^"']+)["']\s*\)/;
let result = "";
let index = 0;
while (index < source.length) {
const quote = source[index];
if (quote === '"' || quote === "'") {
const start = index++;
while (index < source.length) {
if (source[index] === "\\") index += 2;
else if (source[index] === quote) {
index++;
break;
} else index++;
}
result += source.slice(start, index);
continue;
}
const match = source.slice(index).match(constructor);
if (!match) {
result += source[index++];
continue;
}
if (match[1]) {
result += match[1] === "ObjectId" ? `{"$oid":"${match[2]}"}` : match[1] === "NumberLong" ? `{"$numberLong":"${match[2]}"}` : `{"$date":"${match[2]}"}`;
} else if (match[3]) {
result += match[3] === "NumberLong" ? `{"$numberLong":"${match[4]}"}` : `{"$oid":"${match[4]}"}`;
} else {
result += `{"$date":"${match[5]}"}`;
}
index += match[0].length;
}
return result;
}
function parseMethodArgs(source: string, methodCallIndex: number): string[] | null {
const openIndex = source.indexOf("(", methodCallIndex);
const closeIndex = findMatchingParen(source, openIndex);
@@ -1169,52 +1079,6 @@ function trimMongoOuterCommentRange(source: string, from: number, to: number): M
};
}
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 parseMongoDropIndexArgument(args: string[]): string | null {
if (args.length !== 1 || !args[0]?.trim()) return null;
const normalized = normalizeJsonArgument(args[0]);
@@ -1235,52 +1099,6 @@ function parseMongoDropIndexesArgument(args: string[]): string | undefined | nul
return Array.isArray(parsed) && parsed.length > 0 && parsed.every((item) => typeof item === "string") ? normalized : null;
}
export 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;
@@ -1325,73 +1143,6 @@ function hasSingleEmptyChainedCall(source: string, name: string): boolean {
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;
}
function chainedMethodCallPattern(name: string): RegExp {
return new RegExp(`\\.\\s*${escapeRegExp(name)}\\s*\\(`, "g");
}
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 escapeRegExp(value: string): string {
return value.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
}
function parseNormalizedJson(json: string): unknown {
try {
return JSON.parse(json);
apps/desktop/src/stores/queryStore.ts
+8 -1
View File
@@ -17,6 +17,7 @@ import {
mongoDistinctToQueryResult,
mongoCreateIndexToQueryResult,
mongoDocumentsToQueryResult,
describeMongoCommandParseFailure,
mongoDroppedIndexesToQueryResult,
mongoIndexesToQueryResult,
mongoUseToQueryResult,
@@ -2908,6 +2909,12 @@ export const useQueryStore = defineStore("query", () => {
return;
}
if (conn?.db_type === "mongodb" && mongoCommands.length === 0 && sql.trim()) {
// Avoid falling through to the SQL executor, which only returns the generic
// "Use MongoDB-specific commands" rejection and hides parse/syntax details.
throw new Error(describeMongoCommandParseFailure(sql));
}
if (mongoCommands.length > 0) {
queryExecutionLog("info", "mongo:start", { traceId, commandCount: mongoCommands.length, sqlLength: sql.length });
@@ -2996,7 +3003,7 @@ export const useQueryStore = defineStore("query", () => {
}
queryExecutionLog("info", "mongo-aggregate:start", { traceId, collection: mongoCommand.collection, database: currentDatabase });
const aggregateMaxRows = normalizeResultPageSize(pageLimit ?? options?.pagination?.limit ?? settingsStore.editorSettings.pageSize);
const result = await api.mongoAggregateDocuments(tab.connectionId, currentDatabase, mongoCommand.collection, mongoCommand.pipeline, aggregateMaxRows, executionId);
const result = await api.mongoAggregateDocuments(tab.connectionId, currentDatabase, mongoCommand.collection, mongoCommand.pipeline, aggregateMaxRows, mongoCommand.options, executionId);
allResults.push(markQueryResultRowsRaw(annotateMongoResult(mongoDocumentsToQueryResult(result.documents, performance.now() - commandStartedAt, result.total))));
mongoEditTarget = undefined;
queryExecutionLog("info", "mongo-aggregate:done", {
apps/desktop/vite.config.ts
+2
View File
@@ -67,6 +67,8 @@ export default defineConfig(async () => ({
resolve: {
alias: {
"@": path.resolve(__dirname, "./src"),
// Prefer package source during app dev so shell parse changes need no rebuild.
"@dbx-app/mongo-shell": path.resolve(__dirname, "../../packages/mongo-shell/src/index.ts"),
},
},
clearScreen: false,
crates/dbx-core/src/db/mongo_driver.rs
+167 -4
View File
@@ -1,7 +1,7 @@
use mongodb::{
bson::{doc, oid::ObjectId, Bson, DateTime, Document},
options::{ClientOptions, GridFsBucketOptions, IndexOptions, UpdateModifications},
Client, Database, IndexModel,
Client, Cursor, Database, IndexModel,
};
use serde::{Deserialize, Serialize};
@@ -734,22 +734,86 @@ pub async fn find_documents_extended_json(
Ok(MongoDocumentResult { extended_documents: Some(documents.clone()), documents, raw_documents: None, total })
}
/// Run `db.collection.aggregate(pipeline, options)`.
///
/// One execution model for every non-explain aggregate (empty or free-form options):
/// build the server [`aggregate`](https://www.mongodb.com/docs/manual/reference/command/aggregate/)
/// command and open it with [`Database::run_cursor_command`] so session-scoped getMore/killCursors
/// stay with the driver cursor. Options are forwarded as-is (`allowDiskUse`, `cursor.batchSize`,
/// `maxTimeMS`, collation, hint, comment, let, …).
///
/// `explain: true` is the only special case: the server returns a plan document (no cursor), so
/// that path uses plain `run_command`.
///
/// Note: this intentionally uses the command/cursor path rather than `Collection::aggregate`, so
/// default read/write concern inheritance matches other free-form shell options (explicit options
/// on the command document only).
pub async fn aggregate_documents(
client: &Client,
database: &str,
collection: &str,
pipeline_json: &str,
max_rows: Option<usize>,
options_json: Option<&str>,
) -> Result<MongoDocumentResult, String> {
let json: serde_json::Value =
serde_json::from_str(pipeline_json).map_err(|e| format!("Invalid pipeline JSON: {e}"))?;
let pipeline_values = json.as_array().ok_or_else(|| "Aggregate pipeline must be a JSON array".to_string())?;
let pipeline = pipeline_values
let pipeline_docs = pipeline_values
.iter()
.map(|value| json_object_to_document(value).map_err(|e| format!("Invalid pipeline stage: {e}")))
.collect::<Result<Vec<Document>, String>>()?;
let col = client.database(database).collection::<Document>(collection);
let mut cursor = col.aggregate(pipeline).await.map_err(|e| e.to_string())?;
let options = parse_aggregate_options_document(options_json)?;
let (command, explain) = build_aggregate_command(collection, pipeline_docs, options)?;
let db = client.database(database);
if explain {
let result = db.run_command(command).await.map_err(|e| e.to_string())?;
let document = bson_to_json(&Bson::Document(result.clone()));
let extended = Bson::Document(result).into_relaxed_extjson();
return Ok(MongoDocumentResult {
documents: vec![document],
raw_documents: None,
extended_documents: Some(vec![extended]),
total: 1,
});
}
// Driver cursor owns the implicit session across aggregate + getMore + killCursors.
let mut cursor = db.run_cursor_command(command).await.map_err(|e| e.to_string())?;
drain_document_cursor(&mut cursor, max_rows).await
}
/// Build the server aggregate command document, preserving free-form options.
/// Returns `(command, explain)` so callers branch only on the explain flag.
fn build_aggregate_command(
collection: &str,
pipeline: Vec<Document>,
options: Document,
) -> Result<(Document, bool), String> {
let explain = aggregate_options_explain(&options)?;
let pipeline_bson: Vec<Bson> = pipeline.into_iter().map(Bson::Document).collect();
let mut command = doc! {
"aggregate": collection,
"pipeline": pipeline_bson,
};
for (key, value) in options {
command.insert(key, value);
}
// Server requires `cursor` unless `explain` is true.
if !explain && !command.contains_key("cursor") {
command.insert("cursor", doc! {});
}
Ok((command, explain))
}
/// Drain a driver cursor up to `max_rows`, peeking one extra document so callers can detect
/// truncation: when more rows exist, `total == max_rows + 1` while `documents.len() == max_rows`.
async fn drain_document_cursor(
cursor: &mut Cursor<Document>,
max_rows: Option<usize>,
) -> Result<MongoDocumentResult, String> {
let max_rows = max_rows.unwrap_or(100);
let fetch_limit = max_rows.saturating_add(1);
let mut documents = Vec::new();
@@ -767,6 +831,28 @@ pub async fn aggregate_documents(
Ok(MongoDocumentResult { documents, raw_documents: None, extended_documents: Some(extended_documents), total })
}
fn parse_aggregate_options_document(options_json: Option<&str>) -> Result<Document, String> {
let Some(raw) = options_json.map(str::trim).filter(|value| !value.is_empty()) else {
return Ok(Document::new());
};
let value: serde_json::Value =
serde_json::from_str(raw).map_err(|e| format!("Invalid aggregate options JSON: {e}"))?;
if !value.is_object() {
return Err("Aggregate options must be a JSON object".to_string());
}
// Free-form object conversion keeps every official aggregate option (nested
// docs, hint strings, maxTimeMS numbers, comments, let vars, …) intact.
json_object_to_document(&value).map_err(|e| format!("Invalid aggregate options: {e}"))
}
fn aggregate_options_explain(options: &Document) -> Result<bool, String> {
match options.get("explain") {
None => Ok(false),
Some(Bson::Boolean(flag)) => Ok(*flag),
Some(_) => Err("aggregate options.explain must be a boolean (for example { explain: true })".to_string()),
}
}
/// Distinct values of a field, matching the mongo shell's `db.coll.distinct(field, filter)`:
/// array fields contribute their elements rather than the whole array, and the server may
/// answer from an index with a DISTINCT_SCAN. Values are returned in the `documents` slot as
@@ -1676,6 +1762,83 @@ fn expand_object_id_string_array(items: &[serde_json::Value]) -> Bson {
mod tests {
use super::*;
#[test]
fn parse_aggregate_options_document_keeps_official_fields() {
let doc = parse_aggregate_options_document(Some(
r#"{
"explain": true,
"allowDiskUse": true,
"cursor": { "batchSize": 25 },
"maxTimeMS": 1000,
"bypassDocumentValidation": true,
"collation": { "locale": "en" },
"hint": "status_1",
"comment": "agg-test",
"let": { "year": 2024 },
"readConcern": { "level": "local" }
}"#,
))
.unwrap();
assert_eq!(doc.get_bool("explain").ok(), Some(true));
assert_eq!(doc.get_bool("allowDiskUse").ok(), Some(true));
assert_eq!(doc.get_i64("maxTimeMS").ok(), Some(1000));
assert_eq!(doc.get_str("hint").ok(), Some("status_1"));
assert_eq!(doc.get_str("comment").ok(), Some("agg-test"));
assert_eq!(doc.get_document("cursor").ok().and_then(|c| c.get_i64("batchSize").ok()), Some(25));
assert_eq!(doc.get_document("collation").ok().and_then(|c| c.get_str("locale").ok()), Some("en"));
assert_eq!(doc.get_document("let").ok().and_then(|c| c.get_i64("year").ok()), Some(2024));
assert!(aggregate_options_explain(&doc).unwrap());
let no_explain = parse_aggregate_options_document(Some(r#"{"allowDiskUse":false}"#)).unwrap();
assert!(!aggregate_options_explain(&no_explain).unwrap());
let err = aggregate_options_explain(
&parse_aggregate_options_document(Some(r#"{"explain":"executionStats"}"#)).unwrap(),
)
.unwrap_err();
assert!(err.contains("boolean"), "{err}");
}
#[test]
fn build_aggregate_command_adds_cursor_unless_explain() {
let pipeline = vec![doc! { "$match": { "active": true } }];
let (with_cursor, explain_flag) = build_aggregate_command(
"products",
pipeline.clone(),
parse_aggregate_options_document(Some(r#"{"allowDiskUse":true,"cursor":{"batchSize":1}}"#)).unwrap(),
)
.unwrap();
assert!(!explain_flag);
assert_eq!(with_cursor.get_str("aggregate").ok(), Some("products"));
assert_eq!(with_cursor.get_bool("allowDiskUse").ok(), Some(true));
assert_eq!(with_cursor.get_document("cursor").ok().and_then(|c| c.get_i64("batchSize").ok()), Some(1));
let (default_cursor, _) = build_aggregate_command(
"products",
pipeline.clone(),
parse_aggregate_options_document(Some(r#"{"comment":"agg"}"#)).unwrap(),
)
.unwrap();
assert!(default_cursor.get_document("cursor").is_ok(), "non-explain aggregate requires cursor");
assert_eq!(default_cursor.get_str("comment").ok(), Some("agg"));
// Empty options still produce a cursor command (single path with run_cursor_command).
let (empty_options, empty_explain) =
build_aggregate_command("products", pipeline.clone(), Document::new()).unwrap();
assert!(!empty_explain);
assert!(empty_options.get_document("cursor").is_ok());
let (explain, is_explain) = build_aggregate_command(
"products",
pipeline,
parse_aggregate_options_document(Some(r#"{"explain":true,"allowDiskUse":true}"#)).unwrap(),
)
.unwrap();
assert!(is_explain);
assert!(explain.get("cursor").is_none(), "explain path must not inject cursor");
assert_eq!(explain.get_bool("explain").ok(), Some(true));
}
#[test]
fn update_modifications_accept_document_and_pipeline() {
let document = json_update_to_modifications(&serde_json::json!({ "$set": { "status": "done" } })).unwrap();
crates/dbx-core/src/mongo_ops.rs
+2 -1
View File
@@ -245,12 +245,13 @@ pub async fn mongo_aggregate_documents_core(
collection: &str,
pipeline_json: &str,
max_rows: Option<usize>,
options_json: Option<&str>,
) -> Result<MongoDocumentResult, String> {
ensure_document_pool(state, connection_id).await?;
let connections = state.connections.read().await;
match connections.get(connection_id).ok_or("Not found")? {
PoolKind::MongoDb(client) => {
mongo_driver::aggregate_documents(client, database, collection, pipeline_json, max_rows).await
mongo_driver::aggregate_documents(client, database, collection, pipeline_json, max_rows, options_json).await
}
PoolKind::Agent(_) => Err("MongoDB legacy agent does not support aggregate".to_string()),
_ => Err("Not a MongoDB connection".to_string()),
crates/dbx-core/src/query.rs
+7 -3
View File
@@ -34,6 +34,10 @@ use crate::sql::{split_sql_batches, split_sql_statements, starts_with_executable
pub const QUERY_TIMEOUT: Duration = Duration::from_secs(30);
pub const MAX_ROWS: usize = 10000;
pub const QUERY_CANCELED: &str = "Query canceled";
/// Fallback when a Mongo connection hits the generic SQL executor instead of the shell path.
/// Wording must match packages/mongo-shell `MONGO_SHELL_COMMAND_HINT`
/// (desktop/CLI diagnose first; this is only the Rust SQL-executor backstop).
const MONGO_SHELL_COMMAND_HINT: &str = "Use MongoDB shell-style commands, for example: db.collection.find({}).limit(100), db.collection.aggregate([]), db.collection.aggregate([], { explain: true }), db.version(), db.collection.countDocuments({}), db.collection.distinct(\"field\"), db.collection.getIndexes(), db.collection.createIndex({...}), or db.collection.insertOne({...}).";
const SQL_OMITTED_ERROR_CONTEXT: &str =
"SQL text omitted from user-facing error; enable debug SQL diagnostics for a redacted statement.";
#[cfg(feature = "duckdb-bundled")]
@@ -1487,7 +1491,7 @@ pub async fn do_execute(
result
}
PoolKind::Redis(_) => Err("Use Redis-specific commands".to_string()),
PoolKind::MongoDb(_) => Err("Use MongoDB-specific commands".to_string()),
PoolKind::MongoDb(_) => Err(MONGO_SHELL_COMMAND_HINT.to_string()),
PoolKind::MessageQueue => Err("Use Message Queue-specific commands".to_string()),
PoolKind::Nacos => Err("Use Nacos-specific commands".to_string()),
PoolKind::InfluxDb(client) => {
@@ -1712,7 +1716,7 @@ pub async fn execute_sql_statement_with_options(
// (e.g. typos) must be rejected before any pool/key creation so that
// session-scoped pools do not leak MongoDB Clients and SSH tunnels.
if connection_is_mongodb(state, connection_id).await {
return Err("Use MongoDB-specific commands".to_string());
return Err(MONGO_SHELL_COMMAND_HINT.to_string());
}
let db_type = connection_database_type(state, connection_id).await;
@@ -1934,7 +1938,7 @@ pub async fn execute_multi_core_with_options_for_client(
) -> Result<Vec<ExecuteMultiResult>, String> {
// Reject MongoDB queries that fall through to the generic executor.
if connection_is_mongodb(state, connection_id).await {
return Err("Use MongoDB-specific commands".to_string());
return Err(MONGO_SHELL_COMMAND_HINT.to_string());
}
let pool_key = if database.is_empty() {
crates/dbx-core/tests/live_mongodb_aggregate.rs
+151
View File
@@ -0,0 +1,151 @@
//! Live MongoDB aggregate regressions.
//!
//! Run with a writable server:
//! ```text
//! DBX_LIVE_MONGODB_URL='mongodb://127.0.0.1:27017' \
//! cargo test -p dbx-core --test live_mongodb_aggregate -- --ignored --nocapture
//! ```
//!
//! `aggregate_with_small_batch_size_reads_multiple_cursor_batches` is the session/getMore
//! guard: `cursor.batchSize: 1` forces multiple batches; a broken session fails or truncates.
use std::time::Duration;
use dbx_core::db::mongo_driver;
#[tokio::test]
#[ignore = "requires DBX_LIVE_MONGODB_URL pointing at a writable MongoDB database"]
async fn aggregate_with_small_batch_size_reads_multiple_cursor_batches() {
let url = std::env::var("DBX_LIVE_MONGODB_URL").expect("DBX_LIVE_MONGODB_URL");
let client = mongo_driver::connect(&url, Duration::from_secs(10), Duration::from_secs(60)).await.unwrap();
let database = "dbx_live_aggregate";
let collection = format!("batch_pages_{}", std::process::id());
let docs: Vec<serde_json::Value> =
(0..5).map(|n| serde_json::json!({ "n": n, "label": format!("row-{n}") })).collect();
mongo_driver::insert_documents(&client, database, &collection, &serde_json::to_string(&docs).unwrap())
.await
.unwrap();
// batchSize:1 forces firstBatch + multiple getMore round-trips. Those getMore calls must share
// the driver's cursor session (run_cursor_command), or pagination fails / returns partial data.
let result = mongo_driver::aggregate_documents(
&client,
database,
&collection,
r#"[{"$sort":{"n":1}},{"$project":{"_id":0,"n":1,"label":1}}]"#,
Some(10),
Some(r#"{"allowDiskUse":true,"cursor":{"batchSize":1}}"#),
)
.await
.unwrap();
assert_eq!(result.total, 5, "expected all rows across multiple cursor batches: {:?}", result.documents);
let numbers: Vec<i64> = result
.documents
.iter()
.map(|doc| doc["n"].as_i64().or_else(|| doc["n"].as_u64().map(|v| v as i64)).expect("n"))
.collect();
assert_eq!(numbers, vec![0, 1, 2, 3, 4]);
mongo_driver::drop_collection(&client, database, &collection).await.unwrap();
}
#[tokio::test]
#[ignore = "requires DBX_LIVE_MONGODB_URL pointing at a writable MongoDB database"]
async fn aggregate_with_options_respects_max_rows_across_batches() {
let url = std::env::var("DBX_LIVE_MONGODB_URL").expect("DBX_LIVE_MONGODB_URL");
let client = mongo_driver::connect(&url, Duration::from_secs(10), Duration::from_secs(60)).await.unwrap();
let database = "dbx_live_aggregate";
let collection = format!("batch_limit_{}", std::process::id());
let docs: Vec<serde_json::Value> = (0..8).map(|n| serde_json::json!({ "n": n })).collect();
mongo_driver::insert_documents(&client, database, &collection, &serde_json::to_string(&docs).unwrap())
.await
.unwrap();
let result = mongo_driver::aggregate_documents(
&client,
database,
&collection,
r#"[{"$sort":{"n":1}},{"$project":{"_id":0,"n":1}}]"#,
Some(3),
Some(r#"{"cursor":{"batchSize":1}}"#),
)
.await
.unwrap();
assert_eq!(result.documents.len(), 3);
// drain peeks one past max_rows so total > documents.len() signals more rows available.
assert_eq!(result.total, 4);
assert_eq!(
result.documents[0]["n"].as_i64().or_else(|| result.documents[0]["n"].as_u64().map(|v| v as i64)),
Some(0)
);
assert_eq!(
result.documents[2]["n"].as_i64().or_else(|| result.documents[2]["n"].as_u64().map(|v| v as i64)),
Some(2)
);
mongo_driver::drop_collection(&client, database, &collection).await.unwrap();
}
#[tokio::test]
#[ignore = "requires DBX_LIVE_MONGODB_URL pointing at a writable MongoDB database"]
async fn aggregate_without_options_uses_same_cursor_path() {
let url = std::env::var("DBX_LIVE_MONGODB_URL").expect("DBX_LIVE_MONGODB_URL");
let client = mongo_driver::connect(&url, Duration::from_secs(10), Duration::from_secs(60)).await.unwrap();
let database = "dbx_live_aggregate";
let collection = format!("no_options_{}", std::process::id());
let docs: Vec<serde_json::Value> = (0..3).map(|n| serde_json::json!({ "n": n })).collect();
mongo_driver::insert_documents(&client, database, &collection, &serde_json::to_string(&docs).unwrap())
.await
.unwrap();
// Empty options still go through run_cursor_command (not Collection::aggregate).
let result = mongo_driver::aggregate_documents(
&client,
database,
&collection,
r#"[{"$sort":{"n":1}},{"$project":{"_id":0,"n":1}}]"#,
Some(10),
None,
)
.await
.unwrap();
assert_eq!(result.total, 3);
assert_eq!(result.documents.len(), 3);
mongo_driver::drop_collection(&client, database, &collection).await.unwrap();
}
#[tokio::test]
#[ignore = "requires DBX_LIVE_MONGODB_URL pointing at a writable MongoDB database"]
async fn aggregate_explain_option_returns_plan_document() {
let url = std::env::var("DBX_LIVE_MONGODB_URL").expect("DBX_LIVE_MONGODB_URL");
let client = mongo_driver::connect(&url, Duration::from_secs(10), Duration::from_secs(60)).await.unwrap();
let database = "dbx_live_aggregate";
let collection = format!("explain_{}", std::process::id());
mongo_driver::insert_documents(&client, database, &collection, r#"[{"n":1}]"#).await.unwrap();
let result = mongo_driver::aggregate_documents(
&client,
database,
&collection,
r#"[{"$match":{"n":1}}]"#,
Some(10),
Some(r#"{"explain":true}"#),
)
.await
.unwrap();
assert_eq!(result.total, 1);
assert_eq!(result.documents.len(), 1);
// Explain plan shape varies by server version; just require a non-empty object.
assert!(result.documents[0].is_object(), "explain should return a plan object");
mongo_driver::drop_collection(&client, database, &collection).await.unwrap();
}
crates/dbx-web/src/routes/mongo.rs
+2
View File
@@ -128,6 +128,7 @@ pub struct MongoAggregateRequest {
pub pipeline_json: String,
pub max_rows: Option<usize>,
pub execution_id: Option<String>,
pub options_json: Option<String>,
}
#[derive(Deserialize)]
@@ -439,6 +440,7 @@ pub async fn aggregate_documents(
&req.collection,
&req.pipeline_json,
req.max_rows,
req.options_json.as_deref(),
),
)
.await?;
package.json
+2 -1
View File
@@ -9,7 +9,7 @@
"dev:tauri": "tauri dev",
"dev:web": "vite --config apps/desktop/vite.config.ts --port 5173 --mode web",
"dev:backend": "node scripts/dev-backend.mjs",
"build:packages": "pnpm --filter @dbx-app/node-core build && pnpm --filter @dbx-app/cli build && pnpm --filter @dbx-app/mcp-server build",
"build:packages": "pnpm --filter @dbx-app/mongo-shell build && pnpm --filter @dbx-app/node-core build && pnpm --filter @dbx-app/cli build && pnpm --filter @dbx-app/mcp-server build",
"test:packages": "pnpm --filter @dbx-app/node-core test && pnpm --filter @dbx-app/cli test && pnpm --filter @dbx-app/mcp-server test",
"pack:packages": "rm -rf /tmp/dbx-pack-check && mkdir -p /tmp/dbx-pack-check && pnpm --filter @dbx-app/node-core pack --pack-destination /tmp/dbx-pack-check && pnpm --filter @dbx-app/cli pack --pack-destination /tmp/dbx-pack-check && pnpm --filter @dbx-app/mcp-server pack --pack-destination /tmp/dbx-pack-check",
"publish:dry-run": "pnpm build:packages && pnpm pack:packages",
@@ -32,6 +32,7 @@
"prepare": "husky"
},
"dependencies": {
"@dbx-app/mongo-shell": "workspace:^",
"@babel/runtime": "^7.29.7",
"@codemirror/autocomplete": "^6.20.2",
"@codemirror/commands": "^6.10.3",
packages/app-tests/mongoShellCommand.test.ts
Vendored
+59 -2
View File
@@ -1,6 +1,7 @@
import { strict as assert } from "node:assert";
import { test } from "vitest";
import {
describeMongoCommandParseFailure,
evaluateMongoAggregateSafety,
evaluateMongoWriteSafety,
mongoAggregateWriteStage,
@@ -439,10 +440,66 @@ test("parseMongoAggregateCommand accepts an empty pipeline", () => {
});
});
test("parseMongoAggregateCommand rejects non-array pipelines and extra arguments", () => {
test("parseMongoAggregateCommand accepts official aggregate options document", () => {
assert.deepEqual(parseMongoAggregateCommand("db.products.aggregate([], {})"), {
collection: "products",
pipeline: "[]",
options: "{}",
});
const withExplain = parseMongoAggregateCommand("db.uc_user.aggregate([], {explain: true})");
assert.equal(withExplain?.collection, "uc_user");
assert.equal(withExplain?.pipeline, "[]");
assert.deepEqual(JSON.parse(withExplain?.options ?? "null"), { explain: true });
// Official aggregate options are parsed as a free-form object and forwarded to the server.
const fullOptions = parseMongoAggregateCommand(`db.products.aggregate(
[{"$match":{"active":true}}],
{
allowDiskUse: true,
cursor: { batchSize: 50 },
maxTimeMS: 1000,
collation: { locale: "en" },
hint: "status_1",
comment: "agg-test",
let: { year: 2024 }
}
)`);
assert.equal(fullOptions?.collection, "products");
assert.deepEqual(JSON.parse(fullOptions?.pipeline ?? "null"), [{ $match: { active: true } }]);
assert.deepEqual(JSON.parse(fullOptions?.options ?? "null"), {
allowDiskUse: true,
cursor: { batchSize: 50 },
maxTimeMS: 1000,
collation: { locale: "en" },
hint: "status_1",
comment: "agg-test",
let: { year: 2024 },
});
});
test("parseMongoAggregateCommand rejects non-array pipelines and invalid options", () => {
assert.equal(parseMongoAggregateCommand('db.products.aggregate({"$match":{}})'), null);
assert.equal(parseMongoAggregateCommand("db.products.aggregate([], {})"), null);
assert.equal(parseMongoAggregateCommand("db.products.aggregate([], [])"), null);
assert.equal(parseMongoAggregateCommand("db.products.aggregate([], {explain: true"), null);
assert.equal(parseMongoAggregateCommand("db.products.aggregate([]).limit(10)"), null);
assert.equal(parseMongoAggregateCommand("db.products.aggregate([], {}, true)"), null);
});
test("describeMongoCommandParseFailure reports unclosed delimiters and shell hints", () => {
const unclosed = describeMongoCommandParseFailure("db.uc_user.aggregate([], {explain: true");
assert.match(unclosed, /unclosed/i);
const chained = describeMongoCommandParseFailure("db.products.aggregate([]).limit(10)");
assert.match(chained, /chaining|not supported/i);
const nonArrayPipeline = describeMongoCommandParseFailure('db.products.aggregate({"$match":{}})');
assert.match(nonArrayPipeline, /pipeline must be a JSON array/i);
const badOptions = describeMongoCommandParseFailure("db.products.aggregate([], [])");
assert.match(badOptions, /options must be a JSON object/i);
const generic = describeMongoCommandParseFailure("SELECT 1");
assert.match(generic, /MongoDB shell-style commands/i);
});
test("parseMongoAggregateCommand normalises ObjectId arguments with either quote style", () => {
packages/mongo-shell/package.json
+29
View File
@@ -0,0 +1,29 @@
{
"name": "@dbx-app/mongo-shell",
"version": "0.1.0",
"private": true,
"description": "Pure MongoDB shell JSON/argument parsing shared by desktop and node-core",
"license": "Apache-2.0",
"type": "module",
"exports": {
".": {
"types": "./dist/index.d.ts",
"import": "./dist/index.js",
"default": "./dist/index.js"
}
},
"files": [
"dist",
"src"
],
"scripts": {
"build": "tsc",
"prepare": "tsc"
},
"devDependencies": {
"typescript": "^5.8.3"
},
"engines": {
"node": ">=22.13.0"
}
}
packages/mongo-shell/src/aggregate.ts
+93
View File
@@ -0,0 +1,93 @@
/**
* Thin `db.collection.aggregate(pipeline[, options])` parse + diagnostics.
* JSON preprocessing lives in ./json.ts — do not re-copy helpers here.
*/
import {
findMatchingParen,
hasUnclosedMongoDelimiters,
normalizeJsonArgument,
parseCollectionMethodTarget,
parseMongoObjectArgument,
splitTopLevel,
trimMongoOuterComments,
} from "./json.js";
export interface MongoAggregateCommand {
collection: string;
pipeline: string;
/** Optional second argument to `aggregate(pipeline, options)`, e.g. `{explain: true}`. */
options?: string;
}
/** Shared unsupported-command hint (Rust SQL backstop copies this wording). */
export const MONGO_SHELL_COMMAND_HINT =
"Use MongoDB shell-style commands, for example: db.collection.find({}).limit(100), " +
"db.collection.aggregate([]), db.collection.aggregate([], { explain: true }), " +
'db.version(), db.collection.countDocuments({}), db.collection.distinct("field"), ' +
"db.collection.getIndexes(), db.collection.createIndex({...}), or db.collection.insertOne({...}).";
const PIPELINE_MUST_BE_ARRAY =
"MongoDB aggregate pipeline must be a JSON array (for example [{ $match: {} }]).";
const OPTIONS_MUST_BE_OBJECT =
"MongoDB aggregate options must be a JSON object (for example { explain: true }).";
const UNCLOSED_DELIMITERS =
"MongoDB command has unclosed parentheses, brackets, braces, or strings.";
const UNSUPPORTED_CHAINING =
"Unsupported MongoDB aggregate form. Use db.collection.aggregate(pipeline) or " +
"db.collection.aggregate(pipeline, options). Chaining (for example .limit()) is not supported.";
const EXPECTS_PIPELINE_OR_OPTIONS =
"MongoDB aggregate expects aggregate(pipeline) or aggregate(pipeline, options).";
type AggregateParseResult =
| { ok: true; command: MongoAggregateCommand }
| { ok: false; reason: string };
export function parseMongoAggregateCommand(input: string): MongoAggregateCommand | null {
const source = input.trim().replace(/;$/, "").trim();
const parsed = tryParseMongoAggregateCommand(source);
return parsed?.ok ? parsed.command : null;
}
/**
* Diagnose shell parse failures so callers do not fall through to a generic SQL rejection.
* Aggregate-shaped input gets specific reasons; everything else gets {@link MONGO_SHELL_COMMAND_HINT}.
*/
export function describeMongoCommandParseFailure(input: string): string {
const source = trimMongoOuterComments(input).trim().replace(/;$/, "").trim();
if (!source) return "Empty MongoDB command.";
if (hasUnclosedMongoDelimiters(source)) return UNCLOSED_DELIMITERS;
const aggregate = tryParseMongoAggregateCommand(source);
if (aggregate && !aggregate.ok) return aggregate.reason;
return MONGO_SHELL_COMMAND_HINT;
}
function tryParseMongoAggregateCommand(source: string): AggregateParseResult | null {
const target = parseCollectionMethodTarget(source, "aggregate");
if (!target) return null;
const openIndex = source.indexOf("(", target.methodCallIndex);
const closeIndex = findMatchingParen(source, openIndex);
if (closeIndex < 0) return { ok: false, reason: UNCLOSED_DELIMITERS };
if (source.slice(closeIndex + 1).trim()) return { ok: false, reason: UNSUPPORTED_CHAINING };
const args = splitTopLevel(source.slice(openIndex + 1, closeIndex));
if (args.length < 1 || args.length > 2) return { ok: false, reason: EXPECTS_PIPELINE_OR_OPTIONS };
if (args.length === 2 && !args[1]?.trim()) return { ok: false, reason: OPTIONS_MUST_BE_OBJECT };
const pipeline = normalizeJsonArgument(args[0] ?? "");
if (!pipeline) return { ok: false, reason: PIPELINE_MUST_BE_ARRAY };
try {
if (!Array.isArray(JSON.parse(pipeline))) return { ok: false, reason: PIPELINE_MUST_BE_ARRAY };
} catch {
return { ok: false, reason: PIPELINE_MUST_BE_ARRAY };
}
if (args.length === 2) {
const options = parseMongoObjectArgument(args[1]);
if (!options) return { ok: false, reason: OPTIONS_MUST_BE_OBJECT };
return { ok: true, command: { collection: target.collection, pipeline, options } };
}
return { ok: true, command: { collection: target.collection, pipeline } };
}
packages/mongo-shell/src/index.ts
+20
View File
@@ -0,0 +1,20 @@
export {
chainedMethodCallPattern,
escapeRegExp,
findChainedMethodCallIndex,
findMatchingParen,
hasUnclosedMongoDelimiters,
normalizeJsonArgument,
parseCollectionMethodTarget,
parseMongoObjectArgument,
quoteUnquotedObjectKeys,
splitTopLevel,
trimMongoOuterComments,
} from "./json.js";
export {
describeMongoCommandParseFailure,
MONGO_SHELL_COMMAND_HINT,
parseMongoAggregateCommand,
type MongoAggregateCommand,
} from "./aggregate.js";
packages/mongo-shell/src/json.ts
+300
View File
@@ -0,0 +1,300 @@
/**
* Shared Mongo shell → JSON argument preprocessing.
* Single home for ObjectId/ISODate rewriting, key quoting, and paren/arg splitting.
*/
/** Normalize a shell argument to JSON text the backend can parse, or null if invalid. */
export function normalizeJsonArgument(value: string): string | null {
const trimmed = value.trim();
if (!trimmed) return "{}";
// Rewrite mongo shell constructors that are not valid JSON into extended JSON
// (mongo_driver::json_value_to_bson): ObjectId / NumberLong / ISODate / new Date.
const withExtendedJson = replaceMongoShellConstructors(trimmed);
const preprocessed = quoteUnquotedObjectKeys(convertSingleQuotedStrings(withExtendedJson));
try {
JSON.parse(preprocessed);
return preprocessed;
} catch {
return null;
}
}
/** Object-shaped shell arg (options documents, etc.). */
export function parseMongoObjectArgument(arg: string | undefined): string | null {
if (!arg?.trim()) return null;
const normalized = normalizeJsonArgument(arg);
if (!normalized) return null;
try {
const value = JSON.parse(normalized) as unknown;
return value !== null && typeof value === "object" && !Array.isArray(value) ? normalized : null;
} catch {
return null;
}
}
export function parseCollectionMethodTarget(
source: string,
method: string,
): { collection: string; methodCallIndex: number } | null {
const escapedMethod = escapeRegExp(method);
const direct = new RegExp(`^db\\s*\\.\\s*([A-Za-z_$][\\w$]*)\\s*\\.\\s*${escapedMethod}\\s*\\(`).exec(source);
if (direct) {
return { collection: direct[1]!, methodCallIndex: findChainedMethodCallIndex(source, method) };
}
const getCollection = new RegExp(
`^db\\s*\\.\\s*getCollection\\s*\\(\\s*(["'])(.*?)\\1\\s*\\)\\s*\\.\\s*${escapedMethod}\\s*\\(`,
).exec(source);
if (getCollection) {
return { collection: getCollection[2]!, methodCallIndex: findChainedMethodCallIndex(source, method) };
}
return null;
}
export function findChainedMethodCallIndex(source: string, method: string): number {
return chainedMethodCallPattern(method).exec(source)?.index ?? -1;
}
export function chainedMethodCallPattern(method: string): RegExp {
return new RegExp(`\\.\\s*${escapeRegExp(method)}\\s*\\(`, "g");
}
export function escapeRegExp(value: string): string {
return value.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
}
export 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;
}
export function findMatchingParen(source: string, openIndex: number): number {
if (openIndex < 0 || 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;
}
/** True when (), [], {}, or quotes are unbalanced. */
export function hasUnclosedMongoDelimiters(source: string): boolean {
const stack: string[] = [];
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;
continue;
}
if (char === "(" || char === "[" || char === "{") {
stack.push(char);
continue;
}
if (char === ")" || char === "]" || char === "}") {
const expected = char === ")" ? "(" : char === "]" ? "[" : "{";
if (stack.pop() !== expected) return true;
}
}
return quote !== null || stack.length > 0;
}
/** Strip leading line/block comments (//, --, and block comments). */
export function trimMongoOuterComments(source: string): string {
let text = source;
for (;;) {
const trimmed = text.trimStart();
if (trimmed.startsWith("//") || trimmed.startsWith("--")) {
const nl = trimmed.indexOf("\n");
text = nl < 0 ? "" : trimmed.slice(nl + 1);
continue;
}
if (trimmed.startsWith("/*")) {
const end = trimmed.indexOf("*/");
if (end < 0) return trimmed;
text = trimmed.slice(end + 2);
continue;
}
return trimmed.trimEnd();
}
}
export 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 replaceMongoShellConstructors(source: string): string {
const constructor =
/^(ObjectId|NumberLong|ISODate)\s*\(\s*["']([^"']+)["']\s*\)|^(ObjectId|NumberLong)\s*\(\s*(-?\d+)\s*\)|^(?:new\s+Date)\s*\(\s*["']([^"']+)["']\s*\)/;
let result = "";
let index = 0;
while (index < source.length) {
const quote = source[index];
if (quote === '"' || quote === "'") {
const start = index++;
while (index < source.length) {
if (source[index] === "\\") index += 2;
else if (source[index] === quote) {
index++;
break;
} else index++;
}
result += source.slice(start, index);
continue;
}
const match = source.slice(index).match(constructor);
if (!match) {
result += source[index++]!;
continue;
}
if (match[1]) {
result +=
match[1] === "ObjectId"
? `{"$oid":"${match[2]}"}`
: match[1] === "NumberLong"
? `{"$numberLong":"${match[2]}"}`
: `{"$date":"${match[2]}"}`;
} else if (match[3]) {
result += match[3] === "NumberLong" ? `{"$numberLong":"${match[4]}"}` : `{"$oid":"${match[4]}"}`;
} else {
result += `{"$date":"${match[5]}"}`;
}
index += match[0].length;
}
return result;
}
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);
}
packages/mongo-shell/tsconfig.json
+14
View File
@@ -0,0 +1,14 @@
{
"compilerOptions": {
"target": "ES2022",
"module": "NodeNext",
"moduleResolution": "nodenext",
"strict": true,
"esModuleInterop": true,
"skipLibCheck": true,
"outDir": "dist",
"rootDir": "src",
"declaration": true
},
"include": ["src"]
}
packages/node-core/package.json
+3 -2
View File
@@ -26,10 +26,11 @@
},
"scripts": {
"test": "vitest run --config vitest.config.ts",
"build": "tsc",
"prepublishOnly": "tsc"
"build": "pnpm --filter @dbx-app/mongo-shell build && tsc",
"prepublishOnly": "pnpm run build"
},
"dependencies": {
"@dbx-app/mongo-shell": "workspace:^",
"better-sqlite3": "^12.9.0",
"ioredis": "^5.11.1",
"keytar": "^7.9.0",
packages/node-core/src/database.ts
+22 -189
View File
@@ -9,6 +9,24 @@ import { sqlSafetyFromEnv } from "./sql-safety.js";
import { isDirectQueryType } from "./diagnostics.js";
import { bridgePortFilePath } from "./paths.js";
import { parseRedisCommandArgv, classifyRedisCommand, type RedisCommandOptions, type RedisCommandResult, type RedisCommandSafety } from "./redis-command.js";
import {
chainedMethodCallPattern,
describeMongoCommandParseFailure,
findChainedMethodCallIndex,
findMatchingParen,
normalizeJsonArgument,
parseCollectionMethodTarget,
parseMongoAggregateCommand,
splitTopLevel,
type MongoAggregateCommand,
} from "@dbx-app/mongo-shell";
export {
describeMongoCommandParseFailure,
parseMongoAggregateCommand,
MONGO_SHELL_COMMAND_HINT,
} from "@dbx-app/mongo-shell";
export type { MongoAggregateCommand } from "@dbx-app/mongo-shell";
export interface TableInfo {
name: string;
@@ -893,7 +911,7 @@ export async function executeQuery(config: ConnectionConfig, sql: string, option
if (aggregate) {
const safety = evaluateMongoAggregateSafety(aggregate, sqlSafetyFromEnv());
if (!safety.allowed) throw new Error(safety.reason);
const result = await withTimeout(mongoAggregateDocuments(config, aggregate.collection, aggregate.pipeline, resolveMaxRows(options)), resolveTimeoutMs(options));
const result = await withTimeout(mongoAggregateDocuments(config, aggregate.collection, aggregate.pipeline, resolveMaxRows(options), aggregate.options), resolveTimeoutMs(options));
return mongoDocumentsToQueryResult(result.documents.slice(0, resolveMaxRows(options)), result.total);
}
const distinct = parseMongoDistinctCommand(sql);
@@ -932,9 +950,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.count({}), db.projects.distinct("status"), db.projects.getIndexes(), db.projects.dataSize(), db.projects.storageSize(1024), db.projects.totalIndexSize(), db.projects.stats(), db.projects.createIndex({...}), db.projects.dropIndex("name"), db.projects.dropIndexes(), db.projects.drop(), db.projects.insertOne({...}), db.projects.updateOne({...}, {$set: {...}}), or db.projects.deleteOne({...})',
);
throw new Error(describeMongoCommandParseFailure(sql));
}
if (isDirectQueryType(config.db_type)) {
return query(config, sql, undefined, options);
@@ -1225,7 +1241,7 @@ async function executeMongoWrite(config: ConnectionConfig, command: MongoWriteCo
return { affectedRows: result.affected_rows };
}
async function mongoAggregateDocuments(config: ConnectionConfig, collection: string, pipelineJson: string, maxRows: number): Promise<MongoDocumentResult> {
async function mongoAggregateDocuments(config: ConnectionConfig, collection: string, pipelineJson: string, maxRows: number, optionsJson?: string): Promise<MongoDocumentResult> {
return bridgeDataRequest<MongoDocumentResult>("/data/mongo/aggregate-documents", {
connection_id: config.id,
connection_name: config.name,
@@ -1233,6 +1249,7 @@ async function mongoAggregateDocuments(config: ConnectionConfig, collection: str
collection,
pipeline_json: pipelineJson,
max_rows: maxRows,
...(optionsJson ? { options_json: optionsJson } : {}),
});
}
@@ -1335,11 +1352,6 @@ interface MongoCountDocumentsCommand {
mode: "accurate" | "legacy";
}
interface MongoAggregateCommand {
collection: string;
pipeline: string;
}
interface MongoDistinctCommand {
collection: string;
field: string;
@@ -1436,17 +1448,6 @@ function parseFindCountCommand(source: string): MongoCountDocumentsCommand | nul
return filter ? { collection: target.collection, filter, mode: "legacy" } : null;
}
export function parseMongoAggregateCommand(input: string): MongoAggregateCommand | null {
const source = input.trim().replace(/;$/, "").trim();
const target = parseCollectionMethodTarget(source, "aggregate");
if (!target) return null;
const args = parseMethodArgs(source, target.methodCallIndex);
if (!args || args.length !== 1) return null;
const pipeline = normalizeJsonArgument(args[0]);
if (!pipeline) return null;
return Array.isArray(JSON.parse(pipeline)) ? { collection: target.collection, pipeline } : null;
}
export function parseMongoDistinctCommand(input: string): MongoDistinctCommand | null {
const source = input.trim().replace(/;$/, "").trim();
const target = parseCollectionMethodTarget(source, "distinct");
@@ -1635,15 +1636,6 @@ export function evaluateMongoAggregateSafety(command: MongoAggregateCommand, opt
return { allowed: true };
}
function parseCollectionMethodTarget(source: string, method: string): { collection: string; methodCallIndex: number } | null {
const escapedMethod = escapeRegExp(method);
const direct = new RegExp(`^db\\s*\\.\\s*([A-Za-z_$][\\w$]*)\\s*\\.\\s*${escapedMethod}\\s*\\(`).exec(source);
if (direct) return { collection: direct[1], methodCallIndex: findChainedMethodCallIndex(source, method) };
const quoted = new RegExp(`^db\\s*\\.\\s*getCollection\\s*\\(\\s*(['"])([^'"]+)\\1\\s*\\)\\s*\\.\\s*${escapedMethod}\\s*\\(`).exec(source);
if (quoted) return { collection: quoted[2], methodCallIndex: findChainedMethodCallIndex(source, method) };
return null;
}
function parseMethodArgs(source: string, methodCallIndex: number): string[] | null {
const openIndex = source.indexOf("(", methodCallIndex);
const closeIndex = findMatchingParen(source, openIndex);
@@ -1668,18 +1660,6 @@ function hasSingleEmptyChainedCall(chain: string, method: string): boolean {
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;
}
function chainedMethodCallPattern(method: string): RegExp {
return new RegExp(`\\.\\s*${escapeRegExp(method)}\\s*\\(`, "g");
}
function escapeRegExp(value: string): string {
return value.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
}
function readChainedIntegerArgument(chain: string, method: string, fallback: number): number | null {
const arg = readChainedCallArgument(chain, method);
if (arg === undefined) return fallback;
@@ -1687,16 +1667,6 @@ function readChainedIntegerArgument(chain: string, method: string, fallback: num
return Number(arg.trim());
}
function normalizeJsonArgument(arg: string): string | null {
const value = quoteUnquotedObjectKeys(convertSingleQuotedStrings((arg.trim() || "{}").replace(/ObjectId\s*\(\s*["']([^"']+)["']\s*\)/g, '{"$oid":"$1"}')));
try {
JSON.parse(value);
return value;
} catch {
return null;
}
}
function parseMongoDropIndexArgument(args: string[]): string | null {
if (args.length !== 1 || !args[0]?.trim()) return null;
const normalized = normalizeJsonArgument(args[0]);
@@ -1727,98 +1697,6 @@ function parseMongoCollectionStatsScale(args: string[]): number | undefined | nu
return scale;
}
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 parseNormalizedJson(json: string): unknown {
try {
return JSON.parse(json);
@@ -1848,51 +1726,6 @@ function mongoDropIndexesRequiresDangerous(command: MongoWriteCommand): boolean
return Array.isArray(parsed) && parsed.length > 1;
}
function splitTopLevel(source: string): string[] {
const parts: string[] = [];
let depth = 0;
let start = 0;
let quote: string | null = null;
for (let i = 0; i < source.length; i += 1) {
const ch = source[i];
if (quote) {
if (ch === "\\" && i + 1 < source.length) i += 1;
else if (ch === quote) quote = null;
continue;
}
if (ch === "'" || ch === '"') quote = ch;
else if (ch === "{" || ch === "[" || ch === "(") depth += 1;
else if (ch === "}" || ch === "]" || ch === ")") depth -= 1;
else if (ch === "," && depth === 0) {
parts.push(source.slice(start, i));
start = i + 1;
}
}
parts.push(source.slice(start));
return parts;
}
function findMatchingParen(source: string, openIndex: number): number {
if (openIndex < 0 || source[openIndex] !== "(") return -1;
let depth = 0;
let quote: string | null = null;
for (let i = openIndex; i < source.length; i += 1) {
const ch = source[i];
if (quote) {
if (ch === "\\" && i + 1 < source.length) i += 1;
else if (ch === quote) quote = null;
continue;
}
if (ch === "'" || ch === '"') quote = ch;
else if (ch === "(") depth += 1;
else if (ch === ")") {
depth -= 1;
if (depth === 0) return i;
}
}
return -1;
}
function isRecord(value: unknown): value is Record<string, unknown> {
return typeof value === "object" && value !== null && !Array.isArray(value);
}
packages/node-core/src/web-backend.ts
+2 -1
View File
@@ -264,6 +264,7 @@ export async function executeQuery(config: ConnectionConfig, sql: string, option
collection: aggregate.collection,
pipelineJson: aggregate.pipeline,
maxRows: options?.maxRows ?? 100,
...(aggregate.options ? { optionsJson: aggregate.options } : {}),
}),
});
const result = (await res.json()) as { documents: unknown[]; total: number };
@@ -327,7 +328,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.count({}), db.projects.distinct("status"), db.projects.getIndexes(), db.projects.dataSize(), db.projects.storageSize(1024), db.projects.totalIndexSize(), db.projects.stats(), db.projects.createIndex({...}), db.projects.dropIndex("name"), db.projects.dropIndexes(), db.projects.drop(), db.projects.insertOne({...}), db.projects.updateOne({...}, {$set: {...}}), or db.projects.deleteOne({...})',
'Use MongoDB shell-style commands, for example: db.projects.find({}).limit(100), db.projects.aggregate([]), db.projects.aggregate([], {explain:true}), db.version(), db.projects.countDocuments({}), db.projects.count({}), db.projects.distinct("status"), db.projects.getIndexes(), db.projects.dataSize(), db.projects.storageSize(1024), db.projects.totalIndexSize(), db.projects.stats(), db.projects.createIndex({...}), db.projects.dropIndex("name"), db.projects.dropIndexes(), db.projects.drop(), db.projects.insertOne({...}), db.projects.updateOne({...}, {$set: {...}}), or db.projects.deleteOne({...})',
);
}
const res = await apiFetch("/api/query/execute", {
packages/node-core/tests/mongo-query.test.ts
Vendored
+32 -13
View File
@@ -6,6 +6,7 @@ import {
mongoAggregateWriteStage,
mongoCollectionStatsToQueryResult,
mongoDocumentsToQueryResult,
describeMongoCommandParseFailure,
parseMongoAggregateCommand,
parseMongoCollectionStatsCommand,
parseMongoCountDocumentsCommand,
@@ -74,19 +75,14 @@ test("parseMongoWriteCommand accepts unquoted update operator keys", () => {
});
test("parseMongoWriteCommand accepts updateMany arrayFilters options", () => {
assert.deepEqual(
parseMongoWriteCommand(
'db.orders.updateMany({status: "open"}, {$set: {"items.$[item].status": "done"}}, {arrayFilters: [{"item.id": 7}]})',
),
{
kind: "update",
collection: "orders",
filter: '{"status": "open"}',
update: '{"$set": {"items.$[item].status": "done"}}',
options: '{"arrayFilters": [{"item.id": 7}]}',
many: true,
},
);
assert.deepEqual(parseMongoWriteCommand('db.orders.updateMany({status: "open"}, {$set: {"items.$[item].status": "done"}}, {arrayFilters: [{"item.id": 7}]})'), {
kind: "update",
collection: "orders",
filter: '{"status": "open"}',
update: '{"$set": {"items.$[item].status": "done"}}',
options: '{"arrayFilters": [{"item.id": 7}]}',
many: true,
});
});
test("parseMongoCountDocumentsCommand accepts shell-style count commands", () => {
@@ -123,6 +119,29 @@ test("parseMongoAggregateCommand accepts aggregate pipelines", () => {
});
});
test("parseMongoAggregateCommand accepts options including explain", () => {
const withExplain = parseMongoAggregateCommand("db.uc_user.aggregate([], {explain: true})");
assert.equal(withExplain?.collection, "uc_user");
assert.equal(withExplain?.pipeline, "[]");
assert.deepEqual(JSON.parse(withExplain?.options ?? "null"), { explain: true });
assert.equal(parseMongoAggregateCommand("db.uc_user.aggregate([], {explain: true"), null);
assert.equal(parseMongoAggregateCommand("db.products.aggregate([], [])"), null);
assert.equal(parseMongoAggregateCommand("db.products.aggregate([]).limit(10)"), null);
assert.deepEqual(parseMongoAggregateCommand("db.products.aggregate([], {})"), {
collection: "products",
pipeline: "[]",
options: "{}",
});
});
test("describeMongoCommandParseFailure reports aggregate-specific issues", () => {
assert.match(describeMongoCommandParseFailure("db.uc_user.aggregate([], {explain: true"), /unclosed/i);
assert.match(describeMongoCommandParseFailure("db.products.aggregate([]).limit(10)"), /chaining|not supported/i);
assert.match(describeMongoCommandParseFailure('db.products.aggregate({"$match":{}})'), /pipeline must be a JSON array/i);
assert.match(describeMongoCommandParseFailure("db.products.aggregate([], [])"), /options must be a JSON object/i);
assert.match(describeMongoCommandParseFailure("SELECT 1"), /MongoDB shell-style commands/i);
});
test("parseMongoGetIndexesCommand accepts shell-style index commands", () => {
assert.deepEqual(parseMongoGetIndexesCommand("db.web_log.getIndexes();"), {
collection: "web_log",
pnpm-lock.yaml
Generated
+12
View File
@@ -50,6 +50,9 @@ importers:
'@codemirror/view':
specifier: ^6.43.0
version: 6.43.0
'@dbx-app/mongo-shell':
specifier: workspace:^
version: link:packages/mongo-shell
'@lezer/highlight':
specifier: ^1.2.3
version: 1.2.3
@@ -267,8 +270,17 @@ importers:
specifier: ^5.8.3
version: 5.9.3
packages/mongo-shell:
devDependencies:
typescript:
specifier: ^5.8.3
version: 5.9.3
packages/node-core:
dependencies:
'@dbx-app/mongo-shell':
specifier: workspace:^
version: link:../mongo-shell
better-sqlite3:
specifier: ^12.9.0
version: 12.10.0
src-tauri/src/commands/mcp_bridge.rs
+2
View File
@@ -96,6 +96,7 @@ struct MongoAggregateDocumentsRequest {
collection: String,
pipeline_json: String,
max_rows: Option<usize>,
options_json: Option<String>,
}
#[derive(Deserialize)]
@@ -664,6 +665,7 @@ async fn handle_mongo_aggregate_documents_data(state: &Arc<AppState>, body: &str
&req.collection,
&req.pipeline_json,
req.max_rows,
req.options_json.as_deref(),
)
.await
{
src-tauri/src/commands/mongo_cmd.rs
+2
View File
@@ -203,6 +203,7 @@ pub async fn mongo_aggregate_documents(
collection: String,
pipeline_json: String,
max_rows: Option<usize>,
options_json: Option<String>,
execution_id: Option<String>,
) -> Result<MongoDocumentResult, String> {
let app = state.inner().clone();
@@ -216,6 +217,7 @@ pub async fn mongo_aggregate_documents(
&collection,
&pipeline_json,
max_rows,
options_json.as_deref(),
),
)
.await
vitest.config.ts
+1
View File
@@ -7,6 +7,7 @@ export default defineConfig({
resolve: {
alias: {
"@": path.resolve(__dirname, "apps/desktop/src"),
"@dbx-app/mongo-shell": path.resolve(__dirname, "packages/mongo-shell/src/index.ts"),
},
},
test: {