diff --git a/mcp/src/index.ts b/mcp/src/index.ts index 6156bb6c1..734ae6da9 100644 --- a/mcp/src/index.ts +++ b/mcp/src/index.ts @@ -134,21 +134,38 @@ server.tool( schema: z.string().optional().describe("Schema name"), }, async ({ connection_name, table, database, schema }) => { - try { - const bridgeUrl = await getBridgeUrl(); - const res = await fetch(`${bridgeUrl}/open-table`, { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ connection_name, table, database, schema }), - }); - if (res.ok) return text(`Opened ${table} in DBX`); - return text(`Failed: ${await res.text()}`); - } catch { - return text("DBX is not running. Please start DBX first."); - } + return bridgeRequest("/open-table", { connection_name, table, database, schema }, `Opened ${table} in DBX`); }, ); +server.tool( + "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"), + sql: z.string().describe("SQL query to execute"), + database: z.string().optional().describe("Database name"), + }, + async ({ connection_name, sql, database }) => { + return bridgeRequest("/execute-query", { connection_name, sql, database }, "Query sent to DBX"); + }, +); + +async function bridgeRequest(path: string, body: Record, successMsg: string) { + try { + const bridgeUrl = await getBridgeUrl(); + const res = await fetch(`${bridgeUrl}${path}`, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(body), + }); + if (res.ok) return text(successMsg); + return text(`Failed: ${await res.text()}`); + } catch { + return text("DBX is not running. Please start DBX first."); + } +} + async function main() { const transport = new StdioServerTransport(); await server.connect(transport); diff --git a/src-tauri/src/commands/mcp_bridge.rs b/src-tauri/src/commands/mcp_bridge.rs index 5e77267c0..20c8f0857 100644 --- a/src-tauri/src/commands/mcp_bridge.rs +++ b/src-tauri/src/commands/mcp_bridge.rs @@ -17,6 +17,13 @@ struct OpenTableRequest { table: String, } +#[derive(Deserialize)] +struct ExecuteQueryRequest { + connection_name: String, + database: Option, + sql: String, +} + #[derive(Clone, Serialize)] pub struct McpOpenTableEvent { pub connection_id: String, @@ -25,6 +32,13 @@ pub struct McpOpenTableEvent { pub table: String, } +#[derive(Clone, Serialize)] +pub struct McpExecuteQueryEvent { + pub connection_id: String, + pub database: String, + pub sql: String, +} + pub fn start(app_handle: AppHandle, _state: Arc) { tauri::async_runtime::spawn(async move { let listener = match TcpListener::bind(BIND_ADDR).await { @@ -47,7 +61,7 @@ pub fn start(app_handle: AppHandle, _state: Arc) { }; let app = app_handle.clone(); tokio::spawn(async move { - let mut buf = vec![0u8; 8192]; + let mut buf = vec![0u8; 16384]; let n = match stream.read(&mut buf).await { Ok(n) if n > 0 => n, _ => return, @@ -56,43 +70,13 @@ pub fn start(app_handle: AppHandle, _state: Arc) { let body = request.split("\r\n\r\n").nth(1).unwrap_or(""); let first_line = request.lines().next().unwrap_or(""); - if !first_line.starts_with("POST /open-table") { + if first_line.starts_with("POST /open-table") { + handle_open_table(&app, body, &mut stream).await; + } else if first_line.starts_with("POST /execute-query") { + handle_execute_query(&app, body, &mut stream).await; + } else { let _ = stream.write_all(b"HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\n\r\n").await; - return; } - - let req: OpenTableRequest = match serde_json::from_str(body) { - Ok(r) => r, - Err(_) => { - let _ = stream.write_all(b"HTTP/1.1 400 Bad Request\r\nContent-Length: 0\r\n\r\n").await; - return; - } - }; - - let configs = match load_configs_from_disk(&app) { - Ok(c) => c, - Err(_) => { - let _ = stream.write_all(b"HTTP/1.1 500 Internal Server Error\r\nContent-Length: 0\r\n\r\n").await; - return; - } - }; - - let found = configs.iter().find(|c| c.name.eq_ignore_ascii_case(&req.connection_name)); - let Some(config) = found else { - let resp = b"HTTP/1.1 404 Not Found\r\nContent-Length: 20\r\n\r\nConnection not found"; - let _ = stream.write_all(resp).await; - return; - }; - - let event = McpOpenTableEvent { - connection_id: config.id.clone(), - database: req.database.unwrap_or_else(|| config.database.clone().unwrap_or_default()), - schema: req.schema, - table: req.table, - }; - - let _ = app.emit("mcp-open-table", &event); - let _ = stream.write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\nok").await; }); } }); @@ -104,3 +88,55 @@ fn load_configs_from_disk(app: &AppHandle) -> Result(configs: &'a [crate::models::connection::ConnectionConfig], name: &str) -> Option<&'a crate::models::connection::ConnectionConfig> { + configs.iter().find(|c| c.name.eq_ignore_ascii_case(name)) +} + +async fn respond(stream: &mut tokio::net::TcpStream, status: &str, body: &str) { + let resp = format!("HTTP/1.1 {status}\r\nContent-Length: {}\r\n\r\n{body}", body.len()); + let _ = stream.write_all(resp.as_bytes()).await; +} + +async fn handle_open_table(app: &AppHandle, body: &str, stream: &mut tokio::net::TcpStream) { + let req: OpenTableRequest = match serde_json::from_str(body) { + Ok(r) => r, + Err(_) => { respond(stream, "400 Bad Request", "").await; return; } + }; + let configs = match load_configs_from_disk(app) { + Ok(c) => c, + Err(_) => { respond(stream, "500 Internal Server Error", "").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()), + schema: req.schema, + table: req.table, + }; + let _ = app.emit("mcp-open-table", &event); + respond(stream, "200 OK", "ok").await; +} + +async fn handle_execute_query(app: &AppHandle, body: &str, stream: &mut tokio::net::TcpStream) { + let req: ExecuteQueryRequest = match serde_json::from_str(body) { + Ok(r) => r, + Err(_) => { respond(stream, "400 Bad Request", "").await; return; } + }; + let configs = match load_configs_from_disk(app) { + Ok(c) => c, + Err(_) => { respond(stream, "500 Internal Server Error", "").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()), + sql: req.sql, + }; + let _ = app.emit("mcp-execute-query", &event); + respond(stream, "200 OK", "ok").await; +} diff --git a/src/App.vue b/src/App.vue index 6c5481195..e8774bb51 100644 --- a/src/App.vue +++ b/src/App.vue @@ -836,6 +836,10 @@ onMounted(() => { import("@tauri-apps/api/event").then(({ listen }) => { listen<{ connection_id: string; database: string; schema?: string; table: string }>("mcp-open-table", async (event) => { const { connection_id, database, schema, table } = event.payload; + + if (!connectionStore.connections.length) { + await connectionStore.initFromDisk(); + } const config = connectionStore.getConfig(connection_id); if (!config) return; connectionStore.activeConnectionId = connection_id; @@ -848,6 +852,22 @@ onMounted(() => { } else { openLineageTarget({ connectionId: connection_id, database, schema, tableName: table }); } + + getCurrentWindow().setFocus().catch(() => {}); + }); + listen<{ connection_id: string; database: string; sql: string }>("mcp-execute-query", async (event) => { + const { connection_id, database, sql } = event.payload; + if (!connectionStore.connections.length) { + await connectionStore.initFromDisk(); + } + const config = connectionStore.getConfig(connection_id); + if (!config) return; + connectionStore.activeConnectionId = connection_id; + await connectionStore.ensureConnected(connection_id); + const tabId = queryStore.createTab(connection_id, database, undefined, "query"); + queryStore.updateSql(tabId, sql); + await queryStore.executeTabSql(tabId, sql); + getCurrentWindow().setFocus().catch(() => {}); }); }).catch(() => {}); }