From 0674433eea76d18ee289c03e17ce223126139dcc Mon Sep 17 00:00:00 2001 From: Abeautifulsnow Date: Fri, 10 Jul 2026 14:44:40 +0800 Subject: [PATCH] fix(mcp): disambiguate same-name connections --- packages/mcp-server/src/index.ts | 171 +++++++++++--- packages/mcp-server/tests/server.test.ts | 280 +++++++++++++++++++++-- packages/node-core/src/backend.ts | 4 +- packages/node-core/src/connections.ts | 17 ++ packages/node-core/src/database.ts | 16 ++ packages/node-core/src/web-backend.ts | 7 + src-tauri/src/commands/mcp_bridge.rs | 78 ++++--- 7 files changed, 483 insertions(+), 90 deletions(-) diff --git a/packages/mcp-server/src/index.ts b/packages/mcp-server/src/index.ts index ab88b1650e..f3afd5f22a 100644 --- a/packages/mcp-server/src/index.ts +++ b/packages/mcp-server/src/index.ts @@ -40,6 +40,14 @@ function withDatabase(config: ConnectionConfig, database?: string): ConnectionCo return database === undefined ? config : { ...config, database }; } +function connectionIdentity(config: ConnectionConfig): string { + return `${config.name} (${config.id}) [${config.db_type} @ ${config.host}:${config.port}]`; +} + +function labeledText(config: ConnectionConfig, body: string): ReturnType { + return text(`[${connectionIdentity(config)}]\n${body}`); +} + function formatQueryToolResult(result: QueryResult, title?: string) { const prefix = title ? `${title}\n` : ""; if (result.columns.length === 0) return text(`${prefix}Query executed. ${result.row_count} row(s) affected.`); @@ -104,11 +112,39 @@ async function loadScopedConnections(backend: Backend, scope: McpScope): Promise return connections.filter((config) => connectionMatchesScope(config, scope)); } -async function resolveConnection(backend: Backend, scope: McpScope, requestedName?: string): Promise<{ config?: ConnectionConfig; error?: ReturnType }> { +async function resolveConnection( + backend: Backend, + scope: McpScope, + requestedId?: string, + requestedName?: string, +): Promise<{ config?: ConnectionConfig; error?: ReturnType }> { + // connection_id takes priority over connection_name when both are provided. + if (requestedId?.trim()) { + const connections = await backend.loadConnections(); + const config = connections.find((c) => c.id === requestedId.trim()); + if (!config) return { error: toolError("CONNECTION_NOT_FOUND", `Connection with id "${requestedId}" not found.`) }; + // In scoped mode, verify the resolved connection is within the scope. + if (scopeEnabled(scope) && !connectionMatchesScope(config, scope)) { + return { error: toolError("CONNECTION_OUT_OF_SCOPE", `Connection "${requestedId}" is outside this DBX AI session scope.`) }; + } + return { config }; + } + if (!scopeEnabled(scope)) { if (!requestedName?.trim()) return { error: toolError("CONNECTION_NOT_FOUND", "Connection name is required.") }; - const config = await backend.findConnection(requestedName); - return config ? { config } : { error: toolError("CONNECTION_NOT_FOUND", `Connection "${requestedName}" not found.`) }; + const connections = await backend.loadConnections(); + const matching = connections.filter((c) => c.name.toLowerCase() === requestedName.trim().toLowerCase()); + if (matching.length === 0) return { error: toolError("CONNECTION_NOT_FOUND", `Connection "${requestedName}" not found.`) }; + if (matching.length > 1) { + const lines = matching.map((c) => `- ${c.id}: ${c.db_type} @ ${c.host}:${c.port}`); + return { + error: toolError( + "AMBIGUOUS_CONNECTION", + `Multiple connections found with name "${requestedName}". Please specify connection_id:\n${lines.join("\n")}`, + ), + }; + } + return { config: matching[0] }; } const [scopedConfig] = await loadScopedConnections(backend, scope); @@ -131,25 +167,27 @@ export function createDbxMcpServer(backend: Backend, options: { isWebMode?: bool server.tool("dbx_list_connections", "List all database connections configured in DBX", {}, async () => { const connections = await loadScopedConnections(backend, scope); if (connections.length === 0) return text("No connections configured in DBX."); - const rows = connections.map((c) => [c.name, c.db_type, c.host, String(c.port), c.database || ""]); - return text(mdTable(["Name", "Type", "Host", "Port", "Database"], rows)); + const rows = connections.map((c) => [c.id, c.name, c.db_type, c.host, String(c.port), c.database || ""]); + return text(mdTable(["ID", "Name", "Type", "Host", "Port", "Database"], rows)); }); server.tool( "dbx_list_tables", "List tables and views for a database connection", { + connection_id: z.string().optional().describe("Unique ID of the DBX connection (use this to disambiguate when multiple connections share the same name)"), connection_name: z.string().optional().describe("Name of the DBX connection"), database: z.string().optional().describe("Database name"), schema: z.string().optional().describe("Schema name (default: public for PostgreSQL)"), }, - async ({ connection_name, database, schema }) => { - const { config, error } = await resolveConnection(backend, scope, connection_name); + async ({ connection_id, connection_name, database, schema }) => { + const { config, error } = await resolveConnection(backend, scope, connection_id, connection_name); if (error) return error; - const tables = await backend.listTables(withDatabase(config!, database ?? scope.database), schema); + const resolvedConfig = config!; + const tables = await backend.listTables(withDatabase(resolvedConfig, database ?? scope.database), schema); if (tables.length === 0) return text("No tables found."); const rows = tables.map((t) => [t.name, t.type]); - return text(mdTable(["Table", "Type"], rows)); + return labeledText(resolvedConfig, mdTable(["Table", "Type"], rows)); }, ); @@ -157,18 +195,20 @@ export function createDbxMcpServer(backend: Backend, options: { isWebMode?: bool "dbx_describe_table", "Get column definitions for a table", { + connection_id: z.string().optional().describe("Unique ID of the DBX connection (use this to disambiguate when multiple connections share the same name)"), connection_name: z.string().optional().describe("Name of the DBX connection"), table: z.string().describe("Table name"), database: z.string().optional().describe("Database name"), schema: z.string().optional().describe("Schema name (default: public for PostgreSQL)"), }, - async ({ connection_name, table, database, schema }) => { - const { config, error } = await resolveConnection(backend, scope, connection_name); + async ({ connection_id, connection_name, table, database, schema }) => { + const { config, error } = await resolveConnection(backend, scope, connection_id, connection_name); if (error) return error; - const columns = await backend.describeTable(withDatabase(config!, database ?? scope.database), table, schema); + const resolvedConfig = config!; + const columns = await backend.describeTable(withDatabase(resolvedConfig, database ?? scope.database), table, schema); if (columns.length === 0) return text("No columns found."); const rows = columns.map((c) => [c.is_primary_key ? `${c.name} (PK)` : c.name, c.data_type, c.is_nullable ? "YES" : "NO", c.column_default ?? "", c.comment ?? ""]); - return text(mdTable(["Column", "Type", "Nullable", "Default", "Comment"], rows)); + return labeledText(resolvedConfig, mdTable(["Column", "Type", "Nullable", "Default", "Comment"], rows)); }, ); @@ -176,12 +216,13 @@ export function createDbxMcpServer(backend: Backend, options: { isWebMode?: bool "dbx_execute_query", "Execute a SQL query on a database connection (max 100 rows returned)", { + connection_id: z.string().optional().describe("Unique ID of the DBX connection (use this to disambiguate when multiple connections share the same name)"), connection_name: z.string().optional().describe("Name of the DBX connection"), database: z.string().optional().describe("Database name"), sql: z.string().describe("SQL query to execute"), }, - async ({ connection_name, database, sql }) => { - const { config, error } = await resolveConnection(backend, scope, connection_name); + async ({ connection_id, connection_name, database, sql }) => { + const { config, error } = await resolveConnection(backend, scope, connection_id, connection_name); if (error) return error; const scopedConfig = config!; if (scopedConfig.db_type === "redis") { @@ -199,8 +240,8 @@ export function createDbxMcpServer(backend: Backend, options: { isWebMode?: bool for (const statement of statements) { results.push(await backend.executeQuery(withDatabase(scopedConfig, database ?? scope.database), statement)); } - if (results.length === 1) return formatQueryToolResult(results[0]); - return text(results.map((result, index) => formatQueryToolResult(result, `Statement ${index + 1}`).content[0].text).join("\n\n")); + if (results.length === 1) return labeledText(scopedConfig, formatQueryToolResult(results[0]).content[0].text); + return labeledText(scopedConfig, results.map((result, index) => formatQueryToolResult(result, `Statement ${index + 1}`).content[0].text).join("\n\n")); } catch (e: unknown) { const msg = e instanceof Error ? e.message : String(e); return toolError("QUERY_ERROR", msg); @@ -212,12 +253,13 @@ export function createDbxMcpServer(backend: Backend, options: { isWebMode?: bool "dbx_execute_redis_command", "Execute a Redis command on a Redis connection", { + connection_id: z.string().optional().describe("Unique ID of the DBX connection (use this to disambiguate when multiple connections share the same name)"), connection_name: z.string().optional().describe("Name of the DBX Redis connection"), db: z.number().int().min(0).optional().describe("Redis logical database number (default: scoped/default database or 0)"), command: z.string().describe("Redis command to execute, for example: GET mykey, INFO, or DBSIZE"), }, - async ({ connection_name, db, command }) => { - const { config, error } = await resolveConnection(backend, scope, connection_name); + async ({ connection_id, connection_name, db, command }) => { + const { config, error } = await resolveConnection(backend, scope, connection_id, connection_name); if (error) return error; const scopedConfig = config!; if (scopedConfig.db_type !== "redis") { @@ -232,7 +274,7 @@ export function createDbxMcpServer(backend: Backend, options: { isWebMode?: bool const result = await backend.executeRedisCommand(scopedConfig, defaultRedisDb(scopedConfig, scope, db), command, { skipSafetyCheck: safety.skipSafetyCheck, }); - return formatRedisCommandToolResult(result); + return labeledText(scopedConfig, formatRedisCommandToolResult(result).content[0].text); } catch (e: unknown) { const msg = e instanceof Error ? e.message : String(e); return toolError("REDIS_COMMAND_ERROR", msg); @@ -244,22 +286,24 @@ export function createDbxMcpServer(backend: Backend, options: { isWebMode?: bool "dbx_get_schema_context", "Get compact table and column context for writing SQL", { + connection_id: z.string().optional().describe("Unique ID of the DBX connection (use this to disambiguate when multiple connections share the same name)"), connection_name: z.string().optional().describe("Name of the DBX connection"), database: z.string().optional().describe("Database name"), schema: z.string().optional().describe("Schema name (default: public for PostgreSQL)"), tables: z.array(z.string()).optional().describe("Specific table names to include"), max_tables: z.number().int().min(1).max(20).default(8).describe("Maximum number of tables to include"), }, - async ({ connection_name, database, schema, tables, max_tables }) => { - const { config, error } = await resolveConnection(backend, scope, connection_name); + async ({ connection_id, connection_name, database, schema, tables, max_tables }) => { + const { config, error } = await resolveConnection(backend, scope, connection_id, connection_name); if (error) return error; - const context = await buildSchemaContext(backend, withDatabase(config!, database ?? scope.database), { + const resolvedConfig = config!; + const context = await buildSchemaContext(backend, withDatabase(resolvedConfig, database ?? scope.database), { schema, tables, maxTables: max_tables, }); if (context.tables.length === 0) return text("No matching tables found."); - return text(formatSchemaContext(context)); + return labeledText(resolvedConfig, formatSchemaContext(context)); }, ); @@ -313,8 +357,32 @@ export function createDbxMcpServer(backend: Backend, options: { isWebMode?: bool "Remove a database connection from DBX", { connection_name: z.string().describe("Name of the connection to remove"), + connection_id: z.string().optional().describe("Unique ID of the DBX connection (use this to remove by id instead of name)"), }, - async ({ connection_name }) => { + async ({ connection_name, connection_id }) => { + if (connection_id?.trim()) { + if (backend.removeConnectionById) { + const removed = await backend.removeConnectionById(connection_id.trim()); + if (!removed) return toolError("CONNECTION_NOT_FOUND", `Connection with id "${connection_id}" not found.`); + await notifyReload(); + return text(`Connection with id "${connection_id}" removed.`); + } + // Fallback: resolve by id then remove by name + const connections = await backend.loadConnections(); + const config = connections.find((c) => c.id === connection_id.trim()); + if (!config) return toolError("CONNECTION_NOT_FOUND", `Connection with id "${connection_id}" not found.`); + const removed = await backend.removeConnection(config.name); + if (!removed) return toolError("CONNECTION_NOT_FOUND", `Connection "${config.name}" could not be removed.`); + await notifyReload(); + return text(`Connection "${config.name}" (id: ${config.id}) removed.`); + } + const allConnections = await backend.loadConnections(); + const matching = allConnections.filter((c) => c.name.toLowerCase() === connection_name.toLowerCase()); + if (matching.length === 0) return toolError("CONNECTION_NOT_FOUND", `Connection "${connection_name}" not found.`); + if (matching.length > 1) { + const lines = matching.map((c) => `- ${c.id}: ${c.db_type} @ ${c.host}:${c.port}`); + return toolError("AMBIGUOUS_CONNECTION", `Multiple connections found with name "${connection_name}". Please specify connection_id:\n${lines.join("\n")}`); + } const removed = await backend.removeConnection(connection_name); if (!removed) return toolError("CONNECTION_NOT_FOUND", `Connection "${connection_name}" not found.`); await notifyReload(); @@ -329,15 +397,31 @@ export function createDbxMcpServer(backend: Backend, options: { isWebMode?: bool "dbx_open_table", "Open a table in DBX desktop app UI. Requires DBX to be running.", { - connection_name: z.string().describe("Name of the DBX connection"), + connection_id: z.string().optional().describe("Unique ID of the DBX connection (use this to disambiguate when multiple connections share the same name)"), + connection_name: z.string().optional().describe("Name of the DBX connection"), table: z.string().describe("Table name to open"), database: z.string().optional().describe("Database name"), schema: z.string().optional().describe("Schema name"), }, - async ({ connection_name, table, database, schema }) => { - const config = await backend.findConnection(connection_name); - if (!config) return toolError("CONNECTION_NOT_FOUND", `Connection "${connection_name}" not found.`); - return bridgeRequest("/open-table", { connection_name, table, database, schema }, `Opened ${table} in DBX`); + async ({ connection_id, connection_name, table, database, schema }) => { + let config: ConnectionConfig | undefined; + if (connection_id?.trim()) { + const connections = await backend.loadConnections(); + config = connections.find((c) => c.id === connection_id.trim()); + if (!config) return toolError("CONNECTION_NOT_FOUND", `Connection with id "${connection_id}" not found.`); + } else if (connection_name?.trim()) { + const connections = await backend.loadConnections(); + const matching = connections.filter((c) => c.name.toLowerCase() === connection_name.toLowerCase()); + if (matching.length === 0) return toolError("CONNECTION_NOT_FOUND", `Connection "${connection_name}" not found.`); + if (matching.length > 1) { + const lines = matching.map((c) => `- ${c.id}: ${c.db_type} @ ${c.host}:${c.port}`); + return toolError("AMBIGUOUS_CONNECTION", `Multiple connections found with name "${connection_name}". Please specify connection_id:\n${lines.join("\n")}`); + } + config = matching[0]; + } else { + return toolError("CONNECTION_NOT_FOUND", "Either connection_id or connection_name is required."); + } + return bridgeRequest("/open-table", { connection_id: config.id, connection_name: config.name, table, database, schema }, `Opened ${table} in DBX`); }, ); @@ -345,13 +429,29 @@ export function createDbxMcpServer(backend: Backend, options: { isWebMode?: bool "dbx_execute_and_show", "Execute a SQL query in DBX desktop app UI and show results there. Requires DBX to be running.", { - connection_name: z.string().describe("Name of the DBX connection"), + connection_id: z.string().optional().describe("Unique ID of the DBX connection (use this to disambiguate when multiple connections share the same name)"), + connection_name: z.string().optional().describe("Name of the DBX connection"), sql: z.string().describe("SQL query to execute"), database: z.string().optional().describe("Database name"), }, - async ({ connection_name, sql, database }) => { - const config = await backend.findConnection(connection_name); - if (!config) return toolError("CONNECTION_NOT_FOUND", `Connection "${connection_name}" not found.`); + async ({ connection_id, connection_name, sql, database }) => { + let config: ConnectionConfig | undefined; + if (connection_id?.trim()) { + const connections = await backend.loadConnections(); + config = connections.find((c) => c.id === connection_id.trim()); + if (!config) return toolError("CONNECTION_NOT_FOUND", `Connection with id "${connection_id}" not found.`); + } else if (connection_name?.trim()) { + const connections = await backend.loadConnections(); + const matching = connections.filter((c) => c.name.toLowerCase() === connection_name.toLowerCase()); + if (matching.length === 0) return toolError("CONNECTION_NOT_FOUND", `Connection "${connection_name}" not found.`); + if (matching.length > 1) { + const lines = matching.map((c) => `- ${c.id}: ${c.db_type} @ ${c.host}:${c.port}`); + return toolError("AMBIGUOUS_CONNECTION", `Multiple connections found with name "${connection_name}". Please specify connection_id:\n${lines.join("\n")}`); + } + config = matching[0]; + } else { + return toolError("CONNECTION_NOT_FOUND", "Either connection_id or connection_name is required."); + } const safetyOptions = sqlSafetyFromEnv(); if (config?.db_type === "mongodb") { const aggregate = parseMongoAggregateCommand(sql); @@ -368,7 +468,8 @@ export function createDbxMcpServer(backend: Backend, options: { isWebMode?: bool return bridgeRequest( "/execute-query", { - connection_name, + connection_id: config!.id, + connection_name: config!.name, sql, database, allow_writes: safetyOptions.allowWrites, diff --git a/packages/mcp-server/tests/server.test.ts b/packages/mcp-server/tests/server.test.ts index 18ca73b3de..4971433670 100644 --- a/packages/mcp-server/tests/server.test.ts +++ b/packages/mcp-server/tests/server.test.ts @@ -196,7 +196,7 @@ test("redis execute query points callers to the redis command tool", async () => const redisConnection: ConnectionConfig = { ...connection, db_type: "redis", database: "0" }; const scopedBackend: Backend = { ...backend, - findConnection: async () => redisConnection, + loadConnections: async () => [redisConnection], }; const server = createDbxMcpServer(scopedBackend, { isWebMode: true }); @@ -216,7 +216,7 @@ test("redis command tool executes redis commands on the selected database", asyn let usedCommand = ""; const scopedBackend: Backend = { ...backend, - findConnection: async () => redisConnection, + loadConnections: async () => [redisConnection], executeRedisCommand: async (_config, db, command) => { usedDb = db; usedCommand = command; @@ -242,7 +242,7 @@ test("redis command tool blocks write commands in read-only MCP sessions", async const redisConnection: ConnectionConfig = { ...connection, db_type: "redis" }; const scopedBackend: Backend = { ...backend, - findConnection: async () => redisConnection, + loadConnections: async () => [redisConnection], executeRedisCommand: async () => { executed = true; return { command: "SET", safety: "write", value: "OK" }; @@ -267,7 +267,7 @@ test("redis command tool allows dangerous redis commands only when explicitly en let skipSafetyCheck = false; const scopedBackend: Backend = { ...backend, - findConnection: async () => redisConnection, + loadConnections: async () => [redisConnection], executeRedisCommand: async (_config, _db, _command, options) => { skipSafetyCheck = options?.skipSafetyCheck ?? false; return { command: "KEYS", safety: "blocked", value: ["session:1"] }; @@ -301,7 +301,7 @@ test("mongodb list tables returns collections from the selected database", async const mongoConnection: ConnectionConfig = { ...connection, db_type: "mongodb", database: "admin" }; const scopedBackend: Backend = { ...backend, - findConnection: async () => mongoConnection, + loadConnections: async () => [mongoConnection], listTables: async (config) => { usedDatabase = config.database || ""; return [{ name: "projects", type: "COLLECTION" }]; @@ -323,7 +323,7 @@ test("mongodb describe table returns inferred document fields", async () => { const mongoConnection: ConnectionConfig = { ...connection, db_type: "mongodb" }; const scopedBackend: Backend = { ...backend, - findConnection: async () => mongoConnection, + loadConnections: async () => [mongoConnection], describeTable: async () => [ { name: "_id", @@ -359,7 +359,7 @@ test("mongodb execute query formats shell-style find results", async () => { const mongoConnection: ConnectionConfig = { ...connection, db_type: "mongodb" }; const scopedBackend: Backend = { ...backend, - findConnection: async () => mongoConnection, + loadConnections: async () => [mongoConnection], executeQuery: async () => ({ columns: ["_id", "meta", "missing"], rows: [{ _id: "1", meta: { name: "demo" }, missing: null }], @@ -395,7 +395,6 @@ test("add connection accepts H2 file paths without a port", async () => { let added: Omit | undefined; const scopedBackend: Backend = { ...backend, - findConnection: async () => undefined, addConnection: async (config) => { added = config; return { id: "h2-file", ...config }; @@ -449,23 +448,23 @@ test("query exceptions include a stable MCP error code", async () => { }); test("desktop bridge failures include a stable MCP error code", async () => { - const oldHome = process.env.HOME; const dir = await mkdtemp(join(tmpdir(), "dbx-mcp-home-")); - process.env.HOME = dir; try { - const server = createDbxMcpServer(backend, { isWebMode: false }); - const result = await (server as any)._registeredTools.dbx_open_table.handler({ - connection_name: "local", - table: "users", - }); + // Use DBX_DATA_DIR (honoured cross-platform) to point bridgePortFilePath() + // at an empty temp directory so no real bridge is reachable. + await withScopedEnv({ DBX_DATA_DIR: dir }, async () => { + const server = createDbxMcpServer(backend, { isWebMode: false }); + const result = await (server as any)._registeredTools.dbx_open_table.handler({ + connection_name: "local", + table: "users", + }); - assert.equal(result.isError, true); - assert.match(result.content[0].text, /DBX_NOT_RUNNING:/); - assert.match(result.content[0].text, /DBX is not running/); + assert.equal(result.isError, true); + assert.match(result.content[0].text, /DBX_NOT_RUNNING:/); + assert.match(result.content[0].text, /DBX is not running/); + }); } finally { - if (oldHome === undefined) delete process.env.HOME; - else process.env.HOME = oldHome; await rm(dir, { recursive: true, force: true }); } }); @@ -478,7 +477,7 @@ test("mongodb execute-and-show blocks aggregate write stages before desktop brid const mongoConnection: ConnectionConfig = { ...connection, db_type: "mongodb" }; const scopedBackend: Backend = { ...backend, - findConnection: async () => mongoConnection, + loadConnections: async () => [mongoConnection], }; const server = createDbxMcpServer(scopedBackend, { isWebMode: false }); @@ -498,3 +497,242 @@ test("mongodb execute-and-show blocks aggregate write stages before desktop brid else process.env.DBX_MCP_ALLOW_DANGEROUS_SQL = oldAllowDangerous; } }); + +test("connection_id parameter resolves correctly", async () => { + const connA: ConnectionConfig = { ...connection, id: "a1b2c3", name: "shared-name", db_type: "postgres" }; + const connB: ConnectionConfig = { ...connection, id: "d4e5f6", name: "shared-name", db_type: "redis", host: "redis.local", port: 6379 }; + const scopedBackend: Backend = { + ...backend, + loadConnections: async () => [connA, connB], + }; + const server = createDbxMcpServer(scopedBackend, { isWebMode: true }); + + // Resolve by connection_id should return the correct connection + const result = await (server as any)._registeredTools.dbx_list_tables.handler({ + connection_id: "d4e5f6", + }); + + assert.match(result.content[0].text, /users/); +}); + +test("duplicate connection names return AMBIGUOUS_CONNECTION error", async () => { + const connA: ConnectionConfig = { ...connection, id: "a1b2c3", name: "shared-name", db_type: "postgres", host: "pg.local", port: 5432 }; + const connB: ConnectionConfig = { ...connection, id: "d4e5f6", name: "shared-name", db_type: "redis", host: "redis.local", port: 6379 }; + const scopedBackend: Backend = { + ...backend, + loadConnections: async () => [connA, connB], + }; + const server = createDbxMcpServer(scopedBackend, { isWebMode: true }); + + // Using connection_name with duplicates should return AMBIGUOUS_CONNECTION + const result = await (server as any)._registeredTools.dbx_list_tables.handler({ + connection_name: "shared-name", + }); + + assert.equal(result.isError, true); + assert.match(result.content[0].text, /AMBIGUOUS_CONNECTION:/); + assert.match(result.content[0].text, /a1b2c3:/); + assert.match(result.content[0].text, /d4e5f6:/); + assert.match(result.content[0].text, /postgres @ pg.local:5432/); + assert.match(result.content[0].text, /redis @ redis.local:6379/); +}); + +test("connection_id takes priority over connection_name", async () => { + const connA: ConnectionConfig = { ...connection, id: "a1b2c3", name: "shared-name", db_type: "postgres" }; + const connB: ConnectionConfig = { ...connection, id: "d4e5f6", name: "shared-name", db_type: "mysql", host: "mysql.local", port: 3306 }; + let usedConfig: ConnectionConfig | undefined; + const scopedBackend: Backend = { + ...backend, + loadConnections: async () => [connA, connB], + listTables: async (config) => { + usedConfig = config; + return [{ name: "users", type: "BASE TABLE" }]; + }, + }; + const server = createDbxMcpServer(scopedBackend, { isWebMode: true }); + + // Provide both connection_id and connection_name; connection_id should win + await (server as any)._registeredTools.dbx_list_tables.handler({ + connection_id: "d4e5f6", + connection_name: "shared-name", + }); + + assert.equal(usedConfig?.id, "d4e5f6"); +}); + +test("dbx_list_connections includes ID column", async () => { + const server = createDbxMcpServer(backend, { isWebMode: true }); + const result = await (server as any)._registeredTools.dbx_list_connections.handler({}); + + // The table header should include the ID column + assert.match(result.content[0].text, /ID.*Name.*Type.*Host.*Port.*Database/); + // The connection's ID value "1" should appear in the table + assert.match(result.content[0].text, /1\s+\|\s+local/); +}); + +test("same name and db_type with different host/port returns AMBIGUOUS_CONNECTION", async () => { + const connA: ConnectionConfig = { + ...connection, + id: "pg-prod-us", + name: "my-db", + db_type: "postgres", + host: "10.0.1.100", + port: 5432, + database: "app", + }; + const connB: ConnectionConfig = { + ...connection, + id: "pg-prod-eu", + name: "my-db", + db_type: "postgres", + host: "10.0.2.200", + port: 5432, + database: "app", + }; + const scopedBackend: Backend = { + ...backend, + loadConnections: async () => [connA, connB], + }; + const server = createDbxMcpServer(scopedBackend, { isWebMode: true }); + + // Using connection_name with duplicates (same db_type) should still return AMBIGUOUS_CONNECTION + const result = await (server as any)._registeredTools.dbx_list_tables.handler({ + connection_name: "my-db", + }); + + assert.equal(result.isError, true); + assert.match(result.content[0].text, /AMBIGUOUS_CONNECTION:/); + assert.match(result.content[0].text, /pg-prod-us: postgres @ 10\.0\.1\.100:5432/); + assert.match(result.content[0].text, /pg-prod-eu: postgres @ 10\.0\.2\.200:5432/); +}); + +test("connection_id routes to correct host among same-name same-type connections", async () => { + const connA: ConnectionConfig = { + ...connection, + id: "pg-prod-us", + name: "my-db", + db_type: "postgres", + host: "10.0.1.100", + port: 5432, + database: "app", + }; + const connB: ConnectionConfig = { + ...connection, + id: "pg-prod-eu", + name: "my-db", + db_type: "postgres", + host: "10.0.2.200", + port: 5432, + database: "app", + }; + const usedConfigs: ConnectionConfig[] = []; + const scopedBackend: Backend = { + ...backend, + loadConnections: async () => [connA, connB], + listTables: async (config) => { + usedConfigs.push(config); + return [{ name: "orders", type: "BASE TABLE" }]; + }, + executeQuery: async (config, _sql) => { + usedConfigs.push(config); + return { columns: ["cnt"], rows: [{ cnt: 42 }], row_count: 1 }; + }, + }; + const server = createDbxMcpServer(scopedBackend, { isWebMode: true }); + + // Route to US instance via connection_id + const listResult = await (server as any)._registeredTools.dbx_list_tables.handler({ + connection_id: "pg-prod-us", + }); + assert.match(listResult.content[0].text, /orders/); + assert.equal(usedConfigs[0].id, "pg-prod-us"); + assert.equal(usedConfigs[0].host, "10.0.1.100"); + + // Route to EU instance via connection_id + const queryResult = await (server as any)._registeredTools.dbx_execute_query.handler({ + connection_id: "pg-prod-eu", + database: "app", + sql: "select count(*) as cnt from orders", + }); + assert.match(queryResult.content[0].text, /42/); + assert.equal(usedConfigs[1].id, "pg-prod-eu"); + assert.equal(usedConfigs[1].host, "10.0.2.200"); +}); + +test("tool responses are prefixed with connection identity label", async () => { + const scopedBackend: Backend = { + ...backend, + loadConnections: async () => [{ ...connection, id: "conn-xyz", name: "orders-db", db_type: "postgres", host: "10.5.5.5", port: 5432 }], + listTables: async () => [{ name: "orders", type: "BASE TABLE" }], + executeQuery: async () => ({ columns: ["cnt"], rows: [{ cnt: 7 }], row_count: 1 }), + }; + const server = createDbxMcpServer(scopedBackend, { isWebMode: true }); + + const listResult = await (server as any)._registeredTools.dbx_list_tables.handler({ connection_id: "conn-xyz" }); + assert.match(listResult.content[0].text, /^\[orders-db \(conn-xyz\) \[postgres @ 10\.5\.5\.5:5432\]\]/); + + const queryResult = await (server as any)._registeredTools.dbx_execute_query.handler({ connection_id: "conn-xyz", sql: "select count(*) as cnt from orders" }); + assert.match(queryResult.content[0].text, /^\[orders-db \(conn-xyz\) \[postgres @ 10\.5\.5\.5:5432\]\]/); +}); + +test("dbx_remove_connection with duplicate names returns AMBIGUOUS_CONNECTION", async () => { + const connA: ConnectionConfig = { ...connection, id: "db-a", name: "staging", db_type: "postgres", host: "pg-a.local" }; + const connB: ConnectionConfig = { ...connection, id: "db-b", name: "staging", db_type: "mysql", host: "mysql-b.local", port: 3306 }; + let removedName: string | undefined; + const scopedBackend: Backend = { + ...backend, + loadConnections: async () => [connA, connB], + removeConnection: async (name) => { + removedName = name; + return true; + }, + }; + const server = createDbxMcpServer(scopedBackend, { isWebMode: true }); + + const result = await (server as any)._registeredTools.dbx_remove_connection.handler({ + connection_name: "staging", + }); + + assert.equal(result.isError, true); + assert.match(result.content[0].text, /AMBIGUOUS_CONNECTION:/); + assert.match(result.content[0].text, /db-a: postgres @ pg-a\.local/); + assert.match(result.content[0].text, /db-b: mysql @ mysql-b\.local/); + // removeConnection must NOT have been called — no silent deletion + assert.equal(removedName, undefined); +}); + +test("dbx_execute_query with connection_id routes correctly on bridge-backed (SSH) connections", async () => { + const connDirect: ConnectionConfig = { ...connection, id: "pg-direct", name: "shared", db_type: "postgres", host: "direct.local", ssh_enabled: false }; + const connSsh: ConnectionConfig = { ...connection, id: "pg-ssh", name: "shared", db_type: "postgres", host: "private.local", ssh_enabled: true }; + const usedConfigs: ConnectionConfig[] = []; + const scopedBackend: Backend = { + ...backend, + loadConnections: async () => [connDirect, connSsh], + executeQuery: async (config, _sql) => { + usedConfigs.push(config); + return { columns: ["result"], rows: [{ result: "ok" }], row_count: 1 }; + }, + }; + const server = createDbxMcpServer(scopedBackend, { isWebMode: true }); + + // connection_name with two same-name connections (one SSH-backed) → AMBIGUOUS_CONNECTION + const ambigResult = await (server as any)._registeredTools.dbx_execute_query.handler({ + connection_name: "shared", + sql: "select 1", + }); + assert.equal(ambigResult.isError, true); + assert.match(ambigResult.content[0].text, /AMBIGUOUS_CONNECTION:/); + assert.match(ambigResult.content[0].text, /pg-direct/); + assert.match(ambigResult.content[0].text, /pg-ssh/); + + // connection_id routes to the SSH-backed instance and passes its config through + const result = await (server as any)._registeredTools.dbx_execute_query.handler({ + connection_id: "pg-ssh", + sql: "select 1", + }); + assert.equal(result.isError, undefined); + assert.equal(usedConfigs.length, 1); + assert.equal(usedConfigs[0].id, "pg-ssh"); + assert.equal(usedConfigs[0].host, "private.local"); + assert.equal(usedConfigs[0].ssh_enabled, true); +}); diff --git a/packages/node-core/src/backend.ts b/packages/node-core/src/backend.ts index ce571acb6d..8f0be02e2b 100644 --- a/packages/node-core/src/backend.ts +++ b/packages/node-core/src/backend.ts @@ -1,4 +1,4 @@ -import { addConnection as desktopAddConnection, findConnection as desktopFindConnection, loadConnections as desktopLoadConnections, removeConnection as desktopRemoveConnection } from "./connections.js"; +import { addConnection as desktopAddConnection, findConnection as desktopFindConnection, loadConnections as desktopLoadConnections, removeConnection as desktopRemoveConnection, removeConnectionById as desktopRemoveConnectionById } from "./connections.js"; import { closeDatabaseResources as desktopCloseDatabaseResources, describeTable as desktopDescribeTable, executeQuery as desktopExecuteQuery, executeRedisCommand as desktopExecuteRedisCommand, listTables as desktopListTables } from "./database.js"; import type { ConnectionConfig } from "./connections.js"; import type { ColumnInfo, QueryOptions, QueryResult, TableInfo } from "./database.js"; @@ -9,6 +9,7 @@ export interface Backend { findConnection(name: string): Promise; addConnection(config: Omit): Promise; removeConnection(name: string): Promise; + removeConnectionById?(id: string): Promise; listTables(config: ConnectionConfig, schema?: string): Promise; describeTable(config: ConnectionConfig, table: string, schema?: string): Promise; executeQuery(config: ConnectionConfig, sql: string, options?: QueryOptions): Promise; @@ -26,6 +27,7 @@ export async function createBackend(env: NodeJS.ProcessEnv = process.env): Promi findConnection: desktopFindConnection, addConnection: desktopAddConnection, removeConnection: desktopRemoveConnection, + removeConnectionById: desktopRemoveConnectionById, listTables: desktopListTables, describeTable: desktopDescribeTable, executeQuery: desktopExecuteQuery, diff --git a/packages/node-core/src/connections.ts b/packages/node-core/src/connections.ts index 908eb9774b..fc21d21180 100644 --- a/packages/node-core/src/connections.ts +++ b/packages/node-core/src/connections.ts @@ -279,6 +279,11 @@ export async function findConnection(name: string): Promise c.name.toLowerCase() === name.toLowerCase()); } +export async function findConnectionById(id: string): Promise { + const connections = await loadConnections(); + return connections.find((c) => c.id === id); +} + export async function addConnection(config: Omit): Promise { const id = randomUUID(); const db = openDb(); @@ -355,3 +360,15 @@ export async function removeConnection(name: string): Promise { return true; } + +export async function removeConnectionById(id: string): Promise { + const db = openDb(); + const remove = db.transaction(() => { + const result = db.prepare("DELETE FROM connections WHERE id = ?").run(id); + db.prepare("DELETE FROM connection_secrets WHERE connection_id = ?").run(id); + return result.changes > 0; + }); + const deleted = remove(); + db.close(); + return deleted; +} diff --git a/packages/node-core/src/database.ts b/packages/node-core/src/database.ts index 391077cdfd..3656e6459d 100644 --- a/packages/node-core/src/database.ts +++ b/packages/node-core/src/database.ts @@ -856,6 +856,7 @@ export async function executeQuery(config: ConnectionConfig, sql: string, option if (hasActiveSshLayer(config)) { const result = await withTimeout( bridgeDataRequest("/data/execute-query", { + connection_id: config.id, connection_name: config.name, database: config.database || "", sql, @@ -927,6 +928,7 @@ export async function executeQuery(config: ConnectionConfig, sql: string, option } const result = await withTimeout( bridgeDataRequest("/data/execute-query", { + connection_id: config.id, connection_name: config.name, database: config.database || "", sql, @@ -945,6 +947,7 @@ export async function executeRedisCommand(config: ConnectionConfig, db: number, } return withTimeout( bridgeDataRequest("/data/redis/execute-command", { + connection_id: config.id, connection_name: config.name, db, command, @@ -1029,6 +1032,7 @@ function redisTextToJson(value: string): unknown { export async function listTables(config: ConnectionConfig, schema?: string): Promise { if (config.db_type === "mongodb") { const collections = await bridgeDataRequest("/data/mongo/list-collections", { + connection_id: config.id, connection_name: config.name, database: config.database || "", schema: schema || "", @@ -1041,6 +1045,7 @@ export async function listTables(config: ConnectionConfig, schema?: string): Pro } if (hasActiveSshLayer(config) || !isDirectQueryType(config.db_type)) { const tables = await bridgeDataRequest("/data/list-tables", { + connection_id: config.id, connection_name: config.name, database: config.database || "", schema: schema || "", @@ -1074,6 +1079,7 @@ export async function describeTable(config: ConnectionConfig, table: string, sch } if (hasActiveSshLayer(config) || !isDirectQueryType(config.db_type)) { const columns = await bridgeDataRequest("/data/describe-table", { + connection_id: config.id, connection_name: config.name, database: config.database || "", schema: schema || "", @@ -1099,6 +1105,7 @@ export async function describeTable(config: ConnectionConfig, table: string, sch async function mongoFindDocuments(config: ConnectionConfig, collection: string, skip: number, limit: number, filter: string, projection?: string, sort?: string): Promise { return bridgeDataRequest("/data/mongo/find-documents", { + connection_id: config.id, connection_name: config.name, database: config.database || "", collection, @@ -1112,6 +1119,7 @@ async function mongoFindDocuments(config: ConnectionConfig, collection: string, async function mongoServerVersion(config: ConnectionConfig): Promise { return bridgeDataRequest("/data/mongo/server-version", { + connection_id: config.id, connection_name: config.name, database: config.database || "", }); @@ -1119,6 +1127,7 @@ async function mongoServerVersion(config: ConnectionConfig): Promise { async function mongoCollectionStats(config: ConnectionConfig, collection: string, scale?: number): Promise> { return bridgeDataRequest>("/data/mongo/collection-stats", { + connection_id: config.id, connection_name: config.name, database: config.database || "", collection, @@ -1129,6 +1138,7 @@ async function mongoCollectionStats(config: ConnectionConfig, collection: string async function executeMongoWrite(config: ConnectionConfig, command: MongoWriteCommand): Promise<{ affectedRows: number; indexName?: string; droppedNames?: string[] }> { if (command.kind === "insert") { const result = await bridgeDataRequest<{ affected_rows: number }>("/data/mongo/insert-documents", { + connection_id: config.id, connection_name: config.name, database: config.database || "", collection: command.collection, @@ -1138,6 +1148,7 @@ async function executeMongoWrite(config: ConnectionConfig, command: MongoWriteCo } if (command.kind === "update") { const result = await bridgeDataRequest<{ affected_rows: number }>("/data/mongo/update-documents", { + connection_id: config.id, connection_name: config.name, database: config.database || "", collection: command.collection, @@ -1149,6 +1160,7 @@ async function executeMongoWrite(config: ConnectionConfig, command: MongoWriteCo } if (command.kind === "createIndex") { const result = await bridgeDataRequest<{ name: string }>("/data/mongo/create-index", { + connection_id: config.id, connection_name: config.name, database: config.database || "", collection: command.collection, @@ -1159,6 +1171,7 @@ async function executeMongoWrite(config: ConnectionConfig, command: MongoWriteCo } if (command.kind === "dropIndex" || command.kind === "dropIndexes") { const result = await bridgeDataRequest<{ dropped_names: string[]; affected_rows: number }>("/data/mongo/drop-indexes", { + connection_id: config.id, connection_name: config.name, database: config.database || "", collection: command.collection, @@ -1169,6 +1182,7 @@ async function executeMongoWrite(config: ConnectionConfig, command: MongoWriteCo } if (command.kind === "dropCollection") { await bridgeDataRequest<{ ok: boolean }>("/data/mongo/drop-collection", { + connection_id: config.id, connection_name: config.name, database: config.database || "", collection: command.collection, @@ -1176,6 +1190,7 @@ async function executeMongoWrite(config: ConnectionConfig, command: MongoWriteCo return { affectedRows: 1 }; } const result = await bridgeDataRequest<{ affected_rows: number }>("/data/mongo/delete-documents", { + connection_id: config.id, connection_name: config.name, database: config.database || "", collection: command.collection, @@ -1187,6 +1202,7 @@ async function executeMongoWrite(config: ConnectionConfig, command: MongoWriteCo async function mongoAggregateDocuments(config: ConnectionConfig, collection: string, pipelineJson: string, maxRows: number): Promise { return bridgeDataRequest("/data/mongo/aggregate-documents", { + connection_id: config.id, connection_name: config.name, database: config.database || "", collection, diff --git a/packages/node-core/src/web-backend.ts b/packages/node-core/src/web-backend.ts index 46a56eb650..9a6c775368 100644 --- a/packages/node-core/src/web-backend.ts +++ b/packages/node-core/src/web-backend.ts @@ -150,6 +150,13 @@ export async function removeConnection(name: string): Promise { return true; } +export async function removeConnectionById(id: string): Promise { + const connection = await loadConnections().then((cs) => cs.find((c) => c.id === id)); + if (!connection) return false; + await apiFetch(`/api/connection/delete?id=${encodeURIComponent(id)}`, { method: "DELETE" }); + return true; +} + async function ensureConnected(config: ConnectionConfig): Promise { await apiFetch("/api/connection/connect", { method: "POST", diff --git a/src-tauri/src/commands/mcp_bridge.rs b/src-tauri/src/commands/mcp_bridge.rs index 9ab09dc76f..651668998e 100644 --- a/src-tauri/src/commands/mcp_bridge.rs +++ b/src-tauri/src/commands/mcp_bridge.rs @@ -15,6 +15,7 @@ const MCP_BRIDGE_PORT_FILE: &str = "mcp-bridge-port"; #[derive(Deserialize)] struct OpenTableRequest { connection_name: String, + connection_id: Option, database: Option, schema: Option, table: String, @@ -23,6 +24,7 @@ struct OpenTableRequest { #[derive(Deserialize)] struct ExecuteQueryRequest { connection_name: String, + connection_id: Option, database: Option, sql: String, schema: Option, @@ -33,6 +35,7 @@ struct ExecuteQueryRequest { #[derive(Deserialize)] struct ListTablesRequest { connection_name: String, + connection_id: Option, database: Option, schema: Option, } @@ -40,6 +43,7 @@ struct ListTablesRequest { #[derive(Deserialize)] struct DescribeTableRequest { connection_name: String, + connection_id: Option, database: Option, schema: Option, table: String, @@ -48,6 +52,7 @@ struct DescribeTableRequest { #[derive(Deserialize)] struct MongoFindDocumentsRequest { connection_name: String, + connection_id: Option, database: Option, collection: String, skip: Option, @@ -60,12 +65,14 @@ struct MongoFindDocumentsRequest { #[derive(Deserialize)] struct MongoServerVersionRequest { connection_name: String, + connection_id: Option, database: Option, } #[derive(Deserialize)] struct MongoCollectionStatsRequest { connection_name: String, + connection_id: Option, database: Option, collection: String, scale: Option, @@ -74,6 +81,7 @@ struct MongoCollectionStatsRequest { #[derive(Deserialize)] struct MongoAggregateDocumentsRequest { connection_name: String, + connection_id: Option, database: Option, collection: String, pipeline_json: String, @@ -83,6 +91,7 @@ struct MongoAggregateDocumentsRequest { #[derive(Deserialize)] struct MongoCreateIndexRequest { connection_name: String, + connection_id: Option, database: Option, collection: String, keys_json: String, @@ -92,6 +101,7 @@ struct MongoCreateIndexRequest { #[derive(Deserialize)] struct MongoDropIndexesRequest { connection_name: String, + connection_id: Option, database: Option, collection: String, indexes_json: Option, @@ -101,6 +111,7 @@ struct MongoDropIndexesRequest { #[derive(Deserialize)] struct MongoDropCollectionRequest { connection_name: String, + connection_id: Option, database: Option, collection: String, } @@ -108,6 +119,7 @@ struct MongoDropCollectionRequest { #[derive(Deserialize)] struct MongoInsertDocumentsRequest { connection_name: String, + connection_id: Option, database: Option, collection: String, docs_json: String, @@ -116,6 +128,7 @@ struct MongoInsertDocumentsRequest { #[derive(Deserialize)] struct MongoUpdateDocumentsRequest { connection_name: String, + connection_id: Option, database: Option, collection: String, filter_json: String, @@ -126,6 +139,7 @@ struct MongoUpdateDocumentsRequest { #[derive(Deserialize)] struct MongoDeleteDocumentsRequest { connection_name: String, + connection_id: Option, database: Option, collection: String, filter_json: String, @@ -135,6 +149,7 @@ struct MongoDeleteDocumentsRequest { #[derive(Deserialize)] struct RedisCommandRequest { connection_name: String, + connection_id: Option, db: u32, command: String, skip_safety_check: Option, @@ -295,17 +310,21 @@ async fn respond_error(stream: &mut tokio::net::TcpStream, status: &str, message async fn resolve_connection( state: &Arc, + connection_id: Option<&str>, connection_name: &str, ) -> Result { let configs = state.storage.load_connections().await.map_err(|e| e.to_string())?; - let config = - find_config_by_name(&configs, connection_name).ok_or_else(|| "Connection not found".to_string())?.clone(); + let config = if let Some(id) = connection_id.filter(|s| !s.is_empty()) { + configs.iter().find(|c| c.id == id).ok_or_else(|| format!("Connection with id '{}' not found", id))? + } else { + find_config_by_name(&configs, connection_name).ok_or_else(|| "Connection not found".to_string())? + }; let mut state_configs = state.configs.write().await; if !state_configs.contains_key(&config.id) { state_configs.insert(config.id.clone(), config.clone()); } drop(state_configs); - Ok(config) + Ok(config.clone()) } fn check_visible_database(config: &crate::models::connection::ConnectionConfig, database: &str) -> Result<(), String> { @@ -319,11 +338,12 @@ fn check_visible_database(config: &crate::models::connection::ConnectionConfig, async fn resolve_mongo_pool_key( state: &Arc, + connection_id: Option<&str>, connection_name: &str, database: Option, stream: &mut tokio::net::TcpStream, ) -> Option<(String, String, String)> { - let config = match resolve_connection(state, connection_name).await { + let config = match resolve_connection(state, connection_id, connection_name).await { Ok(c) => c, Err(e) => { respond_error(stream, "404 Not Found", &e).await; @@ -350,17 +370,13 @@ async fn handle_open_table(app: &AppHandle, state: &Arc, body: &str, s return; } }; - let configs = match state.storage.load_connections().await { + let config = match resolve_connection(state, req.connection_id.as_deref(), &req.connection_name).await { Ok(c) => c, - Err(_) => { - respond(stream, "500 Internal Server Error", "").await; + Err(e) => { + respond(stream, "404 Not Found", &e).await; return; } }; - let Some(config) = find_config_by_name(&configs, &req.connection_name) else { - respond(stream, "404 Not Found", "Connection not found").await; - return; - }; let event = McpOpenTableEvent { connection_id: config.id.clone(), database: req.database.unwrap_or_else(|| config.database.clone().unwrap_or_default()), @@ -379,17 +395,13 @@ async fn handle_execute_query(app: &AppHandle, state: &Arc, body: &str return; } }; - let configs = match state.storage.load_connections().await { + let config = match resolve_connection(state, req.connection_id.as_deref(), &req.connection_name).await { Ok(c) => c, - Err(_) => { - respond(stream, "500 Internal Server Error", "").await; + Err(e) => { + respond(stream, "404 Not Found", &e).await; return; } }; - let Some(config) = find_config_by_name(&configs, &req.connection_name) else { - respond(stream, "404 Not Found", "Connection not found").await; - return; - }; let event = McpExecuteQueryEvent { connection_id: config.id.clone(), database: req.database.unwrap_or_else(|| config.database.clone().unwrap_or_default()), @@ -409,7 +421,7 @@ async fn handle_list_tables_data(state: &Arc, body: &str, stream: &mut return; } }; - let config = match resolve_connection(state, &req.connection_name).await { + let config = match resolve_connection(state, req.connection_id.as_deref(), &req.connection_name).await { Ok(c) => c, Err(e) => { respond_error(stream, "404 Not Found", &e).await; @@ -436,7 +448,7 @@ async fn handle_describe_table_data(state: &Arc, body: &str, stream: & return; } }; - let config = match resolve_connection(state, &req.connection_name).await { + let config = match resolve_connection(state, req.connection_id.as_deref(), &req.connection_name).await { Ok(c) => c, Err(e) => { respond_error(stream, "404 Not Found", &e).await; @@ -464,7 +476,7 @@ async fn handle_mongo_list_collections_data(state: &Arc, body: &str, s } }; let Some((pool_key, database, _connection_id)) = - resolve_mongo_pool_key(state, &req.connection_name, req.database, stream).await + resolve_mongo_pool_key(state, req.connection_id.as_deref(), &req.connection_name, req.database, stream).await else { return; }; @@ -483,7 +495,7 @@ async fn handle_mongo_find_documents_data(state: &Arc, body: &str, str } }; let Some((pool_key, database, _connection_id)) = - resolve_mongo_pool_key(state, &req.connection_name, req.database, stream).await + resolve_mongo_pool_key(state, req.connection_id.as_deref(), &req.connection_name, req.database, stream).await else { return; }; @@ -514,7 +526,7 @@ async fn handle_mongo_server_version_data(state: &Arc, body: &str, str } }; let Some((pool_key, database, _connection_id)) = - resolve_mongo_pool_key(state, &req.connection_name, req.database, stream).await + resolve_mongo_pool_key(state, req.connection_id.as_deref(), &req.connection_name, req.database, stream).await else { return; }; @@ -533,7 +545,7 @@ async fn handle_mongo_collection_stats_data(state: &Arc, body: &str, s } }; let Some((pool_key, database, _connection_id)) = - resolve_mongo_pool_key(state, &req.connection_name, req.database, stream).await + resolve_mongo_pool_key(state, req.connection_id.as_deref(), &req.connection_name, req.database, stream).await else { return; }; @@ -554,7 +566,7 @@ async fn handle_mongo_aggregate_documents_data(state: &Arc, body: &str } }; let Some((pool_key, database, _connection_id)) = - resolve_mongo_pool_key(state, &req.connection_name, req.database, stream).await + resolve_mongo_pool_key(state, req.connection_id.as_deref(), &req.connection_name, req.database, stream).await else { return; }; @@ -582,7 +594,7 @@ async fn handle_mongo_create_index_data(state: &Arc, body: &str, strea } }; let Some((pool_key, database, connection_id)) = - resolve_mongo_pool_key(state, &req.connection_name, req.database, stream).await + resolve_mongo_pool_key(state, req.connection_id.as_deref(), &req.connection_name, req.database, stream).await else { return; }; @@ -614,7 +626,7 @@ async fn handle_mongo_drop_indexes_data(state: &Arc, body: &str, strea } }; let Some((pool_key, database, connection_id)) = - resolve_mongo_pool_key(state, &req.connection_name, req.database, stream).await + resolve_mongo_pool_key(state, req.connection_id.as_deref(), &req.connection_name, req.database, stream).await else { return; }; @@ -646,7 +658,7 @@ async fn handle_mongo_drop_collection_data(state: &Arc, body: &str, st } }; let Some((pool_key, database, connection_id)) = - resolve_mongo_pool_key(state, &req.connection_name, req.database, stream).await + resolve_mongo_pool_key(state, req.connection_id.as_deref(), &req.connection_name, req.database, stream).await else { return; }; @@ -669,7 +681,7 @@ async fn handle_mongo_insert_documents_data(state: &Arc, body: &str, s } }; let Some((pool_key, database, connection_id)) = - resolve_mongo_pool_key(state, &req.connection_name, req.database, stream).await + resolve_mongo_pool_key(state, req.connection_id.as_deref(), &req.connection_name, req.database, stream).await else { return; }; @@ -694,7 +706,7 @@ async fn handle_mongo_update_documents_data(state: &Arc, body: &str, s } }; let Some((pool_key, database, connection_id)) = - resolve_mongo_pool_key(state, &req.connection_name, req.database, stream).await + resolve_mongo_pool_key(state, req.connection_id.as_deref(), &req.connection_name, req.database, stream).await else { return; }; @@ -727,7 +739,7 @@ async fn handle_mongo_delete_documents_data(state: &Arc, body: &str, s } }; let Some((pool_key, database, connection_id)) = - resolve_mongo_pool_key(state, &req.connection_name, req.database, stream).await + resolve_mongo_pool_key(state, req.connection_id.as_deref(), &req.connection_name, req.database, stream).await else { return; }; @@ -758,7 +770,7 @@ async fn handle_redis_execute_command_data(state: &Arc, body: &str, st return; } }; - let config = match resolve_connection(state, &req.connection_name).await { + let config = match resolve_connection(state, req.connection_id.as_deref(), &req.connection_name).await { Ok(c) => c, Err(e) => { respond_error(stream, "404 Not Found", &e).await; @@ -809,7 +821,7 @@ async fn handle_execute_query_data(state: &Arc, body: &str, stream: &m return; } }; - let config = match resolve_connection(state, &req.connection_name).await { + let config = match resolve_connection(state, req.connection_id.as_deref(), &req.connection_name).await { Ok(c) => c, Err(e) => { respond_error(stream, "404 Not Found", &e).await;