feat(mcp): add bridge data endpoints for all database types

Add /data/list-tables, /data/describe-table, /data/execute-query
bridge endpoints so the MCP server can delegate queries for non-PG/MySQL
database types (Oracle, H2, Trino, Hive, Neo4j, Cassandra, etc.)
through the DBX desktop app. Direct PG/MySQL connections remain
unchanged. Bump MCP server version to 0.3.0.
This commit is contained in:
t8y2
2026-05-13 18:45:02 +08:00
parent a4ad5cd08b
commit 1937ee9fd3
5 changed files with 248 additions and 7 deletions
+122 -1
View File
@@ -21,6 +21,22 @@ struct ExecuteQueryRequest {
connection_name: String,
database: Option<String>,
sql: String,
schema: Option<String>,
}
#[derive(Deserialize)]
struct ListTablesRequest {
connection_name: String,
database: Option<String>,
schema: Option<String>,
}
#[derive(Deserialize)]
struct DescribeTableRequest {
connection_name: String,
database: Option<String>,
schema: Option<String>,
table: String,
}
#[derive(Clone, Serialize)]
@@ -61,7 +77,7 @@ pub fn start(app_handle: AppHandle, state: Arc<AppState>) {
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<AppState>) {
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<T: Serialize>(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<AppState>,
connection_name: &str,
) -> Result<crate::models::connection::ConnectionConfig, String> {
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<AppState>, 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<AppState>, body: &str
let _ = app.emit("mcp-execute-query", &event);
respond(stream, "200 OK", "ok").await;
}
async fn handle_list_tables_data(state: &Arc<AppState>, 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<AppState>, 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<AppState>, 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,
}
}