diff --git a/mcp/package.json b/mcp/package.json index cd2a520cf..af61850a6 100644 --- a/mcp/package.json +++ b/mcp/package.json @@ -1,6 +1,6 @@ { "name": "@dbx-app/mcp-server", - "version": "0.2.2", + "version": "0.3.0", "mcpName": "io.github.t8y2/dbx", "description": "MCP server for DBX \u2014 query databases from Claude Code, Cursor, and other AI agents", "type": "module", diff --git a/mcp/server.json b/mcp/server.json index e3810c917..7f290a1dc 100644 --- a/mcp/server.json +++ b/mcp/server.json @@ -6,12 +6,12 @@ "url": "https://github.com/t8y2/dbx", "source": "github" }, - "version": "0.2.1", + "version": "0.3.0", "packages": [ { "registryType": "npm", "identifier": "@dbx-app/mcp-server", - "version": "0.2.1", + "version": "0.3.0", "transport": { "type": "stdio" } diff --git a/mcp/src/database.ts b/mcp/src/database.ts index 39ed5fa43..053eeb2a9 100644 --- a/mcp/src/database.ts +++ b/mcp/src/database.ts @@ -1,5 +1,8 @@ import type { ConnectionConfig } from "./connections.js"; import { createServer, connect as netConnect, type Server, type Socket } from "node:net"; +import { readFile } from "node:fs/promises"; +import { join } from "node:path"; +import { homedir, platform } from "node:os"; export interface TableInfo { name: string; @@ -225,6 +228,91 @@ function isMysqlType(dbType: string): boolean { return dbType === "mysql" || dbType === "doris" || dbType === "starrocks"; } +function isDirectType(dbType: string): boolean { + switch (dbType) { + case "postgres": + case "redshift": + case "mysql": + case "doris": + case "starrocks": + return true; + default: + return false; + } +} + +interface BridgeQueryResult { + columns: string[]; + rows: unknown[][]; + affected_rows: number; + execution_time_ms: number; + truncated: boolean; +} + +interface BridgeTableInfo { + name: string; + table_type: string; + comment: string | null; +} + +interface BridgeColumnInfo { + name: string; + data_type: string; + is_nullable: boolean; + column_default: string | null; + is_primary_key: boolean; + comment: string | null; +} + +function bridgeAppDataDir(): string { + const home = homedir(); + switch (platform()) { + case "darwin": + return join(home, "Library", "Application Support", "com.dbx.app"); + case "win32": + return join(process.env.APPDATA || join(home, "AppData", "Roaming"), "com.dbx.app"); + default: + return join(home, ".config", "com.dbx.app"); + } +} + +async function bridgeDataRequest(path: string, body: Record): Promise { + let bridgeUrl: string; + try { + const portFile = join(bridgeAppDataDir(), "mcp-bridge-port"); + const port = (await readFile(portFile, "utf-8")).trim(); + bridgeUrl = `http://127.0.0.1:${port}`; + } catch { + throw new Error("DBX desktop app is not running. This database type requires DBX to be running for query execution."); + } + const res = await fetch(`${bridgeUrl}${path}`, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(body), + }); + if (!res.ok) { + const errBody = await res.text().catch(() => ""); + let errorMsg: string; + try { + const parsed = JSON.parse(errBody); + errorMsg = parsed.error || errBody; + } catch { + errorMsg = errBody; + } + throw new Error(errorMsg || `Bridge request failed: ${res.status}`); + } + return res.json() as Promise; +} + +function convertBridgeQueryResult(result: BridgeQueryResult): QueryResult { + const rows = result.rows.slice(0, MAX_ROWS).map((row) => { + const obj: Record = {}; + result.columns.forEach((col, i) => { obj[col] = row[i]; }); + return obj; + }); + return { columns: result.columns, rows, row_count: rows.length }; +} + function withTimeout(promise: Promise, ms: number): Promise { return new Promise((resolve, reject) => { const timer = setTimeout(() => reject(new Error(`Query timed out after ${ms}ms`)), ms); @@ -272,10 +360,26 @@ async function query(config: ConnectionConfig, sql: string, params?: unknown[]): } export async function executeQuery(config: ConnectionConfig, sql: string): Promise { - return query(config, sql); + if (isDirectType(config.db_type)) { + return query(config, sql); + } + const result = await bridgeDataRequest("/data/execute-query", { + connection_name: config.name, + database: config.database || "", + sql, + }); + return convertBridgeQueryResult(result); } export async function listTables(config: ConnectionConfig, schema?: string): Promise { + if (!isDirectType(config.db_type)) { + const tables = await bridgeDataRequest("/data/list-tables", { + connection_name: config.name, + database: config.database || "", + schema: schema || "", + }); + return tables.map((t) => ({ name: t.name, type: t.table_type || "TABLE" })); + } let result: QueryResult; if (isMysqlType(config.db_type)) { result = await query(config, `SELECT TABLE_NAME AS name, TABLE_TYPE AS type FROM information_schema.TABLES WHERE TABLE_SCHEMA = DATABASE() ORDER BY TABLE_NAME`); @@ -290,6 +394,22 @@ export async function listTables(config: ConnectionConfig, schema?: string): Pro } export async function describeTable(config: ConnectionConfig, table: string, schema?: string): Promise { + if (!isDirectType(config.db_type)) { + const columns = await bridgeDataRequest("/data/describe-table", { + connection_name: config.name, + database: config.database || "", + schema: schema || "", + table, + }); + return columns.map((c) => ({ + name: c.name, + data_type: c.data_type, + is_nullable: c.is_nullable, + column_default: c.column_default, + is_primary_key: c.is_primary_key, + comment: c.comment, + })); + } let result: QueryResult; if (isMysqlType(config.db_type)) { result = await query( diff --git a/mcp/src/index.ts b/mcp/src/index.ts index cb1362ffe..7e71d786b 100644 --- a/mcp/src/index.ts +++ b/mcp/src/index.ts @@ -61,7 +61,7 @@ function mdTable(headers: string[], rows: string[][]): string { const server = new McpServer({ name: "dbx", - version: "0.2.1", + version: "0.3.0", }); server.tool( @@ -144,7 +144,7 @@ server.tool( "Add a new database connection to DBX", { name: z.string().describe("Connection name"), - db_type: z.string().describe("Database type: postgres, mysql, sqlite, redis, duckdb, clickhouse, sqlserver, mongodb, oracle, elasticsearch"), + db_type: z.string().describe("Database type: postgres, mysql, sqlite, redis, duckdb, clickhouse, sqlserver, mongodb, oracle, elasticsearch, doris, starrocks, redshift, dameng, kingbase, vastbase, goldendb, gaussdb, h2, snowflake, trino, hive, db2, informix, neo4j, cassandra, bigquery, kylin, sundb, jdbc"), host: z.string().describe("Database host"), port: z.number().describe("Database port"), username: z.string().default("").describe("Username"), diff --git a/src-tauri/src/commands/mcp_bridge.rs b/src-tauri/src/commands/mcp_bridge.rs index e5622f668..cdc2f3de3 100644 --- a/src-tauri/src/commands/mcp_bridge.rs +++ b/src-tauri/src/commands/mcp_bridge.rs @@ -21,6 +21,22 @@ struct ExecuteQueryRequest { connection_name: String, database: Option, sql: String, + schema: Option, +} + +#[derive(Deserialize)] +struct ListTablesRequest { + connection_name: String, + database: Option, + schema: Option, +} + +#[derive(Deserialize)] +struct DescribeTableRequest { + connection_name: String, + database: Option, + schema: Option, + table: String, } #[derive(Clone, Serialize)] @@ -61,7 +77,7 @@ pub fn start(app_handle: AppHandle, state: Arc) { let app = app_handle.clone(); let st = state.clone(); tokio::spawn(async move { - let mut buf = vec![0u8; 16384]; + let mut buf = vec![0u8; 65536]; let n = match stream.read(&mut buf).await { Ok(n) if n > 0 => n, _ => return, @@ -72,6 +88,12 @@ pub fn start(app_handle: AppHandle, state: Arc) { if first_line.starts_with("POST /open-table") { handle_open_table(&app, &st, body, &mut stream).await; + } else if first_line.starts_with("POST /data/list-tables") { + handle_list_tables_data(&st, body, &mut stream).await; + } else if first_line.starts_with("POST /data/describe-table") { + handle_describe_table_data(&st, body, &mut stream).await; + } else if first_line.starts_with("POST /data/execute-query") { + handle_execute_query_data(&st, body, &mut stream).await; } else if first_line.starts_with("POST /execute-query") { handle_execute_query(&app, &st, body, &mut stream).await; } else if first_line.starts_with("POST /reload-connections") { @@ -97,6 +119,35 @@ async fn respond(stream: &mut tokio::net::TcpStream, status: &str, body: &str) { let _ = stream.write_all(resp.as_bytes()).await; } +async fn respond_json(stream: &mut tokio::net::TcpStream, data: &T) { + let body = serde_json::to_string(data).unwrap_or_else(|_| "null".to_string()); + let resp = + format!("HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{body}", body.len()); + let _ = stream.write_all(resp.as_bytes()).await; +} + +async fn respond_error(stream: &mut tokio::net::TcpStream, status: &str, message: &str) { + let body = serde_json::json!({ "error": message }).to_string(); + let resp = + format!("HTTP/1.1 {status}\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{body}", body.len()); + let _ = stream.write_all(resp.as_bytes()).await; +} + +async fn resolve_connection( + state: &Arc, + 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 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) +} + async fn handle_open_table(app: &AppHandle, state: &Arc, body: &str, stream: &mut tokio::net::TcpStream) { let req: OpenTableRequest = match serde_json::from_str(body) { Ok(r) => r, @@ -153,3 +204,73 @@ async fn handle_execute_query(app: &AppHandle, state: &Arc, body: &str let _ = app.emit("mcp-execute-query", &event); respond(stream, "200 OK", "ok").await; } + +async fn handle_list_tables_data(state: &Arc, body: &str, stream: &mut tokio::net::TcpStream) { + let req: ListTablesRequest = match serde_json::from_str(body) { + Ok(r) => r, + Err(_) => { + respond_error(stream, "400 Bad Request", "Invalid JSON").await; + return; + } + }; + let config = match resolve_connection(state, &req.connection_name).await { + Ok(c) => c, + Err(e) => { + respond_error(stream, "404 Not Found", &e).await; + return; + } + }; + let database = req.database.unwrap_or_else(|| config.database.clone().unwrap_or_default()); + let schema = req.schema.unwrap_or_default(); + match dbx_core::schema::list_tables_core(state, &config.id, &database, &schema, None, None).await { + Ok(tables) => respond_json(stream, &tables).await, + Err(e) => respond_error(stream, "500 Internal Server Error", &e).await, + } +} + +async fn handle_describe_table_data(state: &Arc, body: &str, stream: &mut tokio::net::TcpStream) { + let req: DescribeTableRequest = match serde_json::from_str(body) { + Ok(r) => r, + Err(_) => { + respond_error(stream, "400 Bad Request", "Invalid JSON").await; + return; + } + }; + let config = match resolve_connection(state, &req.connection_name).await { + Ok(c) => c, + Err(e) => { + respond_error(stream, "404 Not Found", &e).await; + return; + } + }; + let database = req.database.unwrap_or_else(|| config.database.clone().unwrap_or_default()); + let schema = req.schema.unwrap_or_default(); + match dbx_core::schema::get_columns_core(state, &config.id, &database, &schema, &req.table).await { + Ok(columns) => respond_json(stream, &columns).await, + Err(e) => respond_error(stream, "500 Internal Server Error", &e).await, + } +} + +async fn handle_execute_query_data(state: &Arc, body: &str, stream: &mut tokio::net::TcpStream) { + let req: ExecuteQueryRequest = match serde_json::from_str(body) { + Ok(r) => r, + Err(_) => { + respond_error(stream, "400 Bad Request", "Invalid JSON").await; + return; + } + }; + let config = match resolve_connection(state, &req.connection_name).await { + Ok(c) => c, + Err(e) => { + respond_error(stream, "404 Not Found", &e).await; + return; + } + }; + let database = req.database.unwrap_or_else(|| config.database.clone().unwrap_or_default()); + match dbx_core::query::execute_sql_statement(state, &config.id, &database, &req.sql, req.schema.as_deref(), None) + .await + { + Ok(result) => respond_json(stream, &result).await, + Err(e) => respond_error(stream, "500 Internal Server Error", &e).await, + } +}