diff --git a/apps/desktop/src/App.vue b/apps/desktop/src/App.vue index ef35ddff2..817e81c3b 100644 --- a/apps/desktop/src/App.vue +++ b/apps/desktop/src/App.vue @@ -622,7 +622,7 @@ function ensureQueryTab(): string { const tab = activeTab.value; if (tab && tab.mode === "query") return tab.id; const connId = connectionStore.activeConnectionId || connectionStore.connections[0]?.id || ""; - const db = tab?.database || connectionStore.getConfig(connId)?.database || ""; + const db = tab?.connectionId === connId ? tab.database : connectionStore.getConfig(connId)?.database || ""; return queryStore.createTab(connId, db, undefined, "query"); } diff --git a/apps/desktop/src/composables/useTauriEvents.ts b/apps/desktop/src/composables/useTauriEvents.ts index 925d82130..1a013d27c 100644 --- a/apps/desktop/src/composables/useTauriEvents.ts +++ b/apps/desktop/src/composables/useTauriEvents.ts @@ -54,9 +54,15 @@ export function useTauriEvents(deps: { } }).then((unlisten) => unlistenHandles.push(unlisten)); - listen<{ connection_id: string; database: string; sql: string }>("mcp-execute-query", async (event) => { + listen<{ + connection_id: string; + database: string; + sql: string; + allow_writes?: boolean; + allow_dangerous?: boolean; + }>("mcp-execute-query", async (event) => { try { - const { connection_id, database, sql } = event.payload; + const { connection_id, database, sql, allow_writes, allow_dangerous } = event.payload; if (!connectionStore.connections.length) await connectionStore.initFromDisk(); const config = connectionStore.getConfig(connection_id); if (!config) return; @@ -64,7 +70,9 @@ export function useTauriEvents(deps: { await connectionStore.ensureConnected(connection_id); const tabId = queryStore.createTab(connection_id, database, undefined, "query"); queryStore.updateSql(tabId, sql); - await queryStore.executeTabSql(tabId, sql); + await queryStore.executeTabSql(tabId, sql, { + mongoSafety: { allowWrites: !!allow_writes, allowDangerous: !!allow_dangerous }, + }); focusCurrentWindow(); } catch (e) { console.error("[DBX] mcp-execute-query error:", e); diff --git a/apps/desktop/src/lib/http.ts b/apps/desktop/src/lib/http.ts index 1f1436f19..6ca869393 100644 --- a/apps/desktop/src/lib/http.ts +++ b/apps/desktop/src/lib/http.ts @@ -1272,8 +1272,9 @@ export async function mongoAggregateDocuments( database: string, collection: string, pipelineJson: string, + maxRows?: number, ): Promise { - return post("/api/mongo/aggregate-documents", { connectionId, database, collection, pipelineJson }); + return post("/api/mongo/aggregate-documents", { connectionId, database, collection, pipelineJson, maxRows }); } export async function mongoInsertDocument( diff --git a/apps/desktop/src/lib/mongoShellCommand.ts b/apps/desktop/src/lib/mongoShellCommand.ts index 8fa8060e8..288c7dde8 100644 --- a/apps/desktop/src/lib/mongoShellCommand.ts +++ b/apps/desktop/src/lib/mongoShellCommand.ts @@ -18,6 +18,11 @@ export interface MongoAggregateCommand { pipeline: string; } +export interface MongoAggregateSafetyOptions { + allowWrites?: boolean; + allowDangerous?: boolean; +} + const DEFAULT_LIMIT = 100; export function parseMongoFindCommand(input: string): MongoFindCommand | null { @@ -102,6 +107,42 @@ export function parseMongoAggregateCommand(input: string): MongoAggregateCommand }; } +export function mongoAggregateWriteStage(pipelineJson: string): "$out" | "$merge" | null { + try { + const pipeline = JSON.parse(pipelineJson); + if (!Array.isArray(pipeline)) return null; + for (const stage of pipeline) { + if (!isRecord(stage)) continue; + if (Object.prototype.hasOwnProperty.call(stage, "$out")) return "$out"; + if (Object.prototype.hasOwnProperty.call(stage, "$merge")) return "$merge"; + } + } catch { + return null; + } + return null; +} + +export function evaluateMongoAggregateSafety( + command: MongoAggregateCommand, + options: MongoAggregateSafetyOptions, +): { allowed: boolean; reason?: string } { + const writeStage = mongoAggregateWriteStage(command.pipeline); + if (!writeStage) return { allowed: true }; + if (!options.allowWrites) { + return { + allowed: false, + reason: `MongoDB aggregate stage "${writeStage}" writes data. Set DBX_MCP_ALLOW_WRITES=1 to allow write commands.`, + }; + } + if (!options.allowDangerous) { + return { + allowed: false, + reason: `MongoDB aggregate stage "${writeStage}" is dangerous. Set DBX_MCP_ALLOW_DANGEROUS_SQL=1 to allow it.`, + }; + } + return { allowed: true }; +} + export function mongoDocumentsToQueryResult(documents: unknown[], executionTimeMs: number, total: number): QueryResult { const columns: string[] = []; diff --git a/apps/desktop/src/lib/tauri.ts b/apps/desktop/src/lib/tauri.ts index f3e0ed905..4ccb927f2 100644 --- a/apps/desktop/src/lib/tauri.ts +++ b/apps/desktop/src/lib/tauri.ts @@ -1103,8 +1103,9 @@ export async function mongoAggregateDocuments( database: string, collection: string, pipelineJson: string, + maxRows?: number, ): Promise { - return invoke("mongo_aggregate_documents", { connectionId, database, collection, pipelineJson }); + return invoke("mongo_aggregate_documents", { connectionId, database, collection, pipelineJson, maxRows }); } export async function mongoInsertDocument( diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index f028ba1b8..618a49643 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -10,11 +10,13 @@ import { buildExplainSql, parseExplainResult } from "@/lib/explainPlan"; import { allEditableColumnsWriteable, allPrimaryKeysPresent, sourceColumnsForResult } from "@/lib/sqlAnalysis"; import { restoreOpenTabsState, serializeOpenTabs } from "@/lib/openTabsPersistence"; import { + evaluateMongoAggregateSafety, mongoCountToQueryResult, mongoDocumentsToQueryResult, parseMongoAggregateCommand, parseMongoCountDocumentsCommand, parseMongoFindCommand, + type MongoAggregateSafetyOptions, } from "@/lib/mongoShellCommand"; import { AGENT_DRIVER_TYPES } from "@/lib/databaseCapabilities"; import { editablePrimaryKeys } from "@/lib/tableEditing"; @@ -528,6 +530,7 @@ export const useQueryStore = defineStore("query", () => { resultBaseSql?: string; resultSortedSql?: string | undefined; pagination?: { limit: number; offset: number; sessionId?: string }; + mongoSafety?: MongoAggregateSafetyOptions; }, ) { const tab = tabs.value.find((t) => t.id === id); @@ -655,6 +658,10 @@ export const useQueryStore = defineStore("query", () => { const mongoAggregate = conn?.db_type === "mongodb" ? parseMongoAggregateCommand(sql) : null; if (mongoAggregate) { + if (options?.mongoSafety) { + const safety = evaluateMongoAggregateSafety(mongoAggregate, options.mongoSafety); + if (!safety.allowed) throw new Error(safety.reason); + } await connStore.ensureConnected(tab.connectionId); console.info("[DBX][executeTabSql:mongo-aggregate:start]", { traceId, collection: mongoAggregate.collection }); const result = await api.mongoAggregateDocuments( @@ -662,6 +669,7 @@ export const useQueryStore = defineStore("query", () => { tab.database, mongoAggregate.collection, mongoAggregate.pipeline, + pageLimit, ); console.info("[DBX][executeTabSql:mongo-aggregate:done]", { traceId, diff --git a/crates/dbx-core/src/db/mongo_driver.rs b/crates/dbx-core/src/db/mongo_driver.rs index bde3241b8..5f6ee8f31 100644 --- a/crates/dbx-core/src/db/mongo_driver.rs +++ b/crates/dbx-core/src/db/mongo_driver.rs @@ -83,6 +83,7 @@ pub async fn aggregate_documents( database: &str, collection: &str, pipeline_json: &str, + max_rows: Option, ) -> Result { let json: serde_json::Value = serde_json::from_str(pipeline_json).map_err(|e| format!("Invalid pipeline JSON: {e}"))?; @@ -93,12 +94,17 @@ pub async fn aggregate_documents( .collect::, String>>()?; let col = client.database(database).collection::(collection); let mut cursor = col.aggregate(pipeline).await.map_err(|e| e.to_string())?; + let max_rows = max_rows.unwrap_or(100); + let fetch_limit = max_rows.saturating_add(1); let mut documents = Vec::new(); - while cursor.advance().await.map_err(|e| e.to_string())? { + while documents.len() < fetch_limit && cursor.advance().await.map_err(|e| e.to_string())? { let doc = cursor.deserialize_current().map_err(|e| e.to_string())?; documents.push(bson_to_json(&Bson::Document(doc))); } let total = documents.len() as u64; + if documents.len() > max_rows { + documents.truncate(max_rows); + } Ok(MongoDocumentResult { documents, total }) } diff --git a/crates/dbx-core/src/mongo_ops.rs b/crates/dbx-core/src/mongo_ops.rs index 17c417afe..0b03e0cda 100644 --- a/crates/dbx-core/src/mongo_ops.rs +++ b/crates/dbx-core/src/mongo_ops.rs @@ -77,11 +77,12 @@ pub async fn mongo_aggregate_documents_core( database: &str, collection: &str, pipeline_json: &str, + max_rows: Option, ) -> Result { 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).await + mongo_driver::aggregate_documents(client, database, collection, pipeline_json, max_rows).await } PoolKind::Agent(_) => Err("MongoDB legacy agent does not support aggregate".to_string()), _ => Err("Not a MongoDB connection".to_string()), diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index 5f4cfdc1d..2c2ab3909 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -399,14 +399,7 @@ pub async fn list_tables_core( } fn collection_names_to_tables(names: Vec, table_type: &str) -> Vec { - names - .into_iter() - .map(|name| db::TableInfo { - name, - table_type: table_type.to_string(), - comment: None, - }) - .collect() + names.into_iter().map(|name| db::TableInfo { name, table_type: table_type.to_string(), comment: None }).collect() } fn filter_table_infos(tables: Vec, filter: Option<&str>, limit: Option) -> Vec { diff --git a/crates/dbx-web/src/routes/mongo.rs b/crates/dbx-web/src/routes/mongo.rs index ad87a40aa..75d7ab9a4 100644 --- a/crates/dbx-web/src/routes/mongo.rs +++ b/crates/dbx-web/src/routes/mongo.rs @@ -39,6 +39,7 @@ pub struct MongoAggregateRequest { pub database: String, pub collection: String, pub pipeline_json: String, + pub max_rows: Option, } #[derive(Deserialize)] @@ -147,6 +148,7 @@ pub async fn aggregate_documents( &req.database, &req.collection, &req.pipeline_json, + req.max_rows, ) .await .map_err(AppError)?; diff --git a/packages/app-tests/mongoShellCommand.test.ts b/packages/app-tests/mongoShellCommand.test.ts index d26aa77d4..1a5879ba6 100644 --- a/packages/app-tests/mongoShellCommand.test.ts +++ b/packages/app-tests/mongoShellCommand.test.ts @@ -1,6 +1,8 @@ import { strict as assert } from "node:assert"; import test from "node:test"; import { + evaluateMongoAggregateSafety, + mongoAggregateWriteStage, mongoCountToQueryResult, mongoDocumentsToQueryResult, parseMongoAggregateCommand, @@ -76,6 +78,19 @@ test("parseMongoAggregateCommand normalises ObjectId arguments with either quote } }); +test("evaluateMongoAggregateSafety blocks write stages unless MCP write flags allow them", () => { + const out = parseMongoAggregateCommand('db.products.aggregate([{"$out":"products_copy"}])'); + assert.ok(out); + assert.equal(mongoAggregateWriteStage(out.pipeline), "$out"); + assert.match(evaluateMongoAggregateSafety(out, {}).reason || "", /DBX_MCP_ALLOW_WRITES=1/); + + 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.equal(evaluateMongoAggregateSafety(merge, { allowWrites: true, allowDangerous: true }).allowed, true); +}); + test("mongoCountToQueryResult returns a single count row", () => { assert.deepEqual(mongoCountToQueryResult(42, 5), { columns: ["count"], diff --git a/packages/mcp-server/src/index.ts b/packages/mcp-server/src/index.ts index cc556f8f4..3685336f1 100644 --- a/packages/mcp-server/src/index.ts +++ b/packages/mcp-server/src/index.ts @@ -6,9 +6,11 @@ import { z } from "zod"; import { buildSchemaContext, createBackend, + evaluateMongoAggregateSafety, evaluateSqlSafety, formatSchemaContext, notifyReload, + parseMongoAggregateCommand, postBridge, sqlSafetyFromEnv, type Backend, @@ -216,13 +218,24 @@ export function createDbxMcpServer(backend: Backend, options: { isWebMode?: bool }, async ({ connection_name, sql, database }) => { const config = await backend.findConnection(connection_name); - if (config?.db_type !== "mongodb") { - const safety = evaluateSqlSafety(sql, sqlSafetyFromEnv()); + const safetyOptions = sqlSafetyFromEnv(); + if (config?.db_type === "mongodb") { + const aggregate = parseMongoAggregateCommand(sql); + if (aggregate) { + const safety = evaluateMongoAggregateSafety(aggregate, safetyOptions); + if (!safety.allowed) return text(`Query blocked: ${safety.reason}`); + } + } else { + const safety = evaluateSqlSafety(sql, safetyOptions); if (!safety.allowed) return text(`Query blocked: ${safety.reason}`); } - // MongoDB shell commands bypass the SQL safety evaluator; the desktop - // app's executor applies command-aware read/write gating. - return bridgeRequest("/execute-query", { connection_name, sql, database }, "Query sent to DBX"); + // MongoDB shell commands bypass the SQL safety evaluator; pass MCP + // safety flags to the desktop executor for command-aware gating. + return bridgeRequest( + "/execute-query", + { connection_name, sql, database, allow_writes: safetyOptions.allowWrites, allow_dangerous: safetyOptions.allowDangerous }, + "Query sent to DBX", + ); }, ); } diff --git a/packages/mcp-server/tests/server.test.ts b/packages/mcp-server/tests/server.test.ts index 90b7a1e1f..10ef0960e 100644 --- a/packages/mcp-server/tests/server.test.ts +++ b/packages/mcp-server/tests/server.test.ts @@ -117,3 +117,32 @@ test("mongodb execute query formats shell-style find results", async () => { assert.match(result.content[0].text, /demo/); assert.match(result.content[0].text, /1 row\(s\)/); }); + +test("mongodb execute-and-show blocks aggregate write stages before desktop bridge", async () => { + const oldAllowWrites = process.env.DBX_MCP_ALLOW_WRITES; + const oldAllowDangerous = process.env.DBX_MCP_ALLOW_DANGEROUS_SQL; + delete process.env.DBX_MCP_ALLOW_WRITES; + delete process.env.DBX_MCP_ALLOW_DANGEROUS_SQL; + const mongoConnection: ConnectionConfig = { ...connection, db_type: "mongodb" }; + const scopedBackend: Backend = { + ...backend, + findConnection: async () => mongoConnection, + }; + const server = createDbxMcpServer(scopedBackend, { isWebMode: false }); + + try { + const result = await (server as any)._registeredTools.dbx_execute_and_show.handler({ + connection_name: "local", + database: "pystrument", + sql: 'db.projects.aggregate([{"$out":"projects_dump"}])', + }); + + assert.match(result.content[0].text, /Query blocked:/); + assert.match(result.content[0].text, /DBX_MCP_ALLOW_WRITES=1/); + } finally { + if (oldAllowWrites === undefined) delete process.env.DBX_MCP_ALLOW_WRITES; + else process.env.DBX_MCP_ALLOW_WRITES = oldAllowWrites; + if (oldAllowDangerous === undefined) delete process.env.DBX_MCP_ALLOW_DANGEROUS_SQL; + else process.env.DBX_MCP_ALLOW_DANGEROUS_SQL = oldAllowDangerous; + } +}); diff --git a/packages/node-core/src/database.ts b/packages/node-core/src/database.ts index b1c044323..0ab418c5e 100644 --- a/packages/node-core/src/database.ts +++ b/packages/node-core/src/database.ts @@ -427,7 +427,10 @@ 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), resolveTimeoutMs(options)); + const result = await withTimeout( + mongoAggregateDocuments(config, aggregate.collection, aggregate.pipeline, resolveMaxRows(options)), + resolveTimeoutMs(options), + ); return mongoDocumentsToQueryResult(result.documents.slice(0, resolveMaxRows(options)), result.total); } const write = parseMongoWriteCommand(sql); @@ -599,12 +602,14 @@ async function mongoAggregateDocuments( config: ConnectionConfig, collection: string, pipelineJson: string, + maxRows: number, ): Promise { return bridgeDataRequest("/data/mongo/aggregate-documents", { connection_name: config.name, database: config.database || "", collection, pipeline_json: pipelineJson, + max_rows: maxRows, }); } diff --git a/packages/node-core/src/web-backend.ts b/packages/node-core/src/web-backend.ts index 7f48dd442..cd181bda2 100644 --- a/packages/node-core/src/web-backend.ts +++ b/packages/node-core/src/web-backend.ts @@ -182,6 +182,7 @@ export async function executeQuery(config: ConnectionConfig, sql: string, option database: config.database || "", collection: aggregate.collection, pipelineJson: aggregate.pipeline, + maxRows: options?.maxRows ?? 100, }), }); const result = (await res.json()) as { documents: unknown[]; total: number }; @@ -260,4 +261,3 @@ async function executeMongoWrite(config: ConnectionConfig, command: MongoWriteCo const result = (await res.json()) as { affected_rows: number }; return result.affected_rows; } - diff --git a/src-tauri/src/commands/connection.rs b/src-tauri/src/commands/connection.rs index 9291d1105..3f5a21ea7 100644 --- a/src-tauri/src/commands/connection.rs +++ b/src-tauri/src/commands/connection.rs @@ -143,6 +143,7 @@ mod tests { ssh_key_passphrase: String::new(), ssh_expose_lan: false, ssh_connect_timeout_secs: dbx_core::models::connection::default_ssh_connect_timeout_secs(), + connect_timeout_secs: dbx_core::models::connection::default_connect_timeout_secs(), query_timeout_secs: dbx_core::models::connection::default_query_timeout_secs(), proxy_enabled: false, proxy_type: ProxyType::Socks5, diff --git a/src-tauri/src/commands/mcp_bridge.rs b/src-tauri/src/commands/mcp_bridge.rs index c01035e2b..a7bd50ddb 100644 --- a/src-tauri/src/commands/mcp_bridge.rs +++ b/src-tauri/src/commands/mcp_bridge.rs @@ -22,6 +22,8 @@ struct ExecuteQueryRequest { database: Option, sql: String, schema: Option, + allow_writes: Option, + allow_dangerous: Option, } #[derive(Deserialize)] @@ -56,6 +58,7 @@ struct MongoAggregateDocumentsRequest { database: Option, collection: String, pipeline_json: String, + max_rows: Option, } #[derive(Deserialize)] @@ -98,6 +101,8 @@ pub struct McpExecuteQueryEvent { pub connection_id: String, pub database: String, pub sql: String, + pub allow_writes: bool, + pub allow_dangerous: bool, } pub fn start(app_handle: AppHandle, state: Arc) { @@ -282,6 +287,8 @@ async fn handle_execute_query(app: &AppHandle, state: &Arc, body: &str connection_id: config.id.clone(), database: req.database.unwrap_or_else(|| config.database.clone().unwrap_or_default()), sql: req.sql, + allow_writes: req.allow_writes.unwrap_or(false), + allow_dangerous: req.allow_dangerous.unwrap_or(false), }; let _ = app.emit("mcp-execute-query", &event); respond(stream, "200 OK", "ok").await; @@ -398,6 +405,7 @@ async fn handle_mongo_aggregate_documents_data(state: &Arc, body: &str &database, &req.collection, &req.pipeline_json, + req.max_rows, ) .await { diff --git a/src-tauri/src/commands/mongo_cmd.rs b/src-tauri/src/commands/mongo_cmd.rs index 48e28e423..65069eaeb 100644 --- a/src-tauri/src/commands/mongo_cmd.rs +++ b/src-tauri/src/commands/mongo_cmd.rs @@ -52,9 +52,17 @@ pub async fn mongo_aggregate_documents( database: String, collection: String, pipeline_json: String, + max_rows: Option, ) -> Result { - dbx_core::mongo_ops::mongo_aggregate_documents_core(&state, &connection_id, &database, &collection, &pipeline_json) - .await + dbx_core::mongo_ops::mongo_aggregate_documents_core( + &state, + &connection_id, + &database, + &collection, + &pipeline_json, + max_rows, + ) + .await } #[tauri::command] diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 4cb1ef85a..179eb1454 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -8,11 +8,11 @@ use commands::connection::AppState; use dbx_core::storage::Storage; use std::sync::Arc; use std::time::Instant; +use tauri::Manager; use tauri::{ menu::MenuBuilder, tray::{MouseButton, MouseButtonState, TrayIconBuilder, TrayIconEvent}, }; -use tauri::Manager; #[cfg(target_os = "macos")] use tauri::{Emitter, RunEvent}; #[cfg(any(windows, target_os = "linux"))]