972 lines
34 KiB
Rust
972 lines
34 KiB
Rust
use serde::{Deserialize, Serialize};
|
|
use std::path::{Path, PathBuf};
|
|
use std::sync::Arc;
|
|
use tauri::{AppHandle, Emitter};
|
|
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
|
use tokio::net::TcpListener;
|
|
|
|
use super::connection::AppState;
|
|
|
|
use super::connection::ensure_connection_writable;
|
|
|
|
const BIND_ADDR: &str = "127.0.0.1:0";
|
|
const MCP_BRIDGE_PORT_FILE: &str = "mcp-bridge-port";
|
|
|
|
#[derive(Deserialize)]
|
|
struct OpenTableRequest {
|
|
connection_name: String,
|
|
connection_id: Option<String>,
|
|
database: Option<String>,
|
|
schema: Option<String>,
|
|
table: String,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct ExecuteQueryRequest {
|
|
connection_name: String,
|
|
connection_id: Option<String>,
|
|
database: Option<String>,
|
|
sql: String,
|
|
schema: Option<String>,
|
|
allow_writes: Option<bool>,
|
|
allow_dangerous: Option<bool>,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct ListTablesRequest {
|
|
connection_name: String,
|
|
connection_id: Option<String>,
|
|
database: Option<String>,
|
|
schema: Option<String>,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct DescribeTableRequest {
|
|
connection_name: String,
|
|
connection_id: Option<String>,
|
|
database: Option<String>,
|
|
schema: Option<String>,
|
|
table: String,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct MongoFindDocumentsRequest {
|
|
connection_name: String,
|
|
connection_id: Option<String>,
|
|
database: Option<String>,
|
|
collection: String,
|
|
skip: Option<u64>,
|
|
limit: Option<i64>,
|
|
filter: Option<String>,
|
|
projection: Option<String>,
|
|
sort: Option<String>,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct MongoCountDocumentsRequest {
|
|
connection_name: String,
|
|
connection_id: Option<String>,
|
|
database: Option<String>,
|
|
collection: String,
|
|
filter: Option<String>,
|
|
mode: Option<String>,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct MongoServerVersionRequest {
|
|
connection_name: String,
|
|
connection_id: Option<String>,
|
|
database: Option<String>,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct MongoCollectionStatsRequest {
|
|
connection_name: String,
|
|
connection_id: Option<String>,
|
|
database: Option<String>,
|
|
collection: String,
|
|
scale: Option<serde_json::Number>,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct MongoAggregateDocumentsRequest {
|
|
connection_name: String,
|
|
connection_id: Option<String>,
|
|
database: Option<String>,
|
|
collection: String,
|
|
pipeline_json: String,
|
|
max_rows: Option<usize>,
|
|
options_json: Option<String>,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct MongoDistinctRequest {
|
|
connection_name: String,
|
|
connection_id: Option<String>,
|
|
database: Option<String>,
|
|
collection: String,
|
|
field: String,
|
|
filter: Option<String>,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct MongoCreateIndexRequest {
|
|
connection_name: String,
|
|
connection_id: Option<String>,
|
|
database: Option<String>,
|
|
collection: String,
|
|
keys_json: String,
|
|
options_json: Option<String>,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct MongoDropIndexesRequest {
|
|
connection_name: String,
|
|
connection_id: Option<String>,
|
|
database: Option<String>,
|
|
collection: String,
|
|
indexes_json: Option<String>,
|
|
single: bool,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct MongoDropCollectionRequest {
|
|
connection_name: String,
|
|
connection_id: Option<String>,
|
|
database: Option<String>,
|
|
collection: String,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct MongoInsertDocumentsRequest {
|
|
connection_name: String,
|
|
connection_id: Option<String>,
|
|
database: Option<String>,
|
|
collection: String,
|
|
docs_json: String,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct MongoUpdateDocumentsRequest {
|
|
connection_name: String,
|
|
connection_id: Option<String>,
|
|
database: Option<String>,
|
|
collection: String,
|
|
filter_json: String,
|
|
update_json: String,
|
|
many: bool,
|
|
options_json: Option<String>,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct MongoDeleteDocumentsRequest {
|
|
connection_name: String,
|
|
connection_id: Option<String>,
|
|
database: Option<String>,
|
|
collection: String,
|
|
filter_json: String,
|
|
many: bool,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct RedisCommandRequest {
|
|
connection_name: String,
|
|
connection_id: Option<String>,
|
|
db: u32,
|
|
command: String,
|
|
skip_safety_check: Option<bool>,
|
|
}
|
|
|
|
#[derive(Clone, Serialize)]
|
|
pub struct McpOpenTableEvent {
|
|
pub connection_id: String,
|
|
pub database: String,
|
|
pub schema: Option<String>,
|
|
pub table: String,
|
|
}
|
|
|
|
#[derive(Clone, Serialize)]
|
|
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<AppState>, data_dir: PathBuf) {
|
|
tauri::async_runtime::spawn(async move {
|
|
let listener = match TcpListener::bind(BIND_ADDR).await {
|
|
Ok(l) => l,
|
|
Err(e) => {
|
|
log::warn!("MCP bridge failed to bind {BIND_ADDR}: {e}");
|
|
return;
|
|
}
|
|
};
|
|
log::info!("MCP bridge listening on {BIND_ADDR}");
|
|
let actual_port = listener.local_addr().map(|a| a.port()).unwrap_or(0);
|
|
log::info!("MCP bridge assigned port {actual_port}");
|
|
// Publish into DBX's resolved data dir so DBX_DATA_DIR and portable mode share the same discovery file.
|
|
if let Err(err) = write_port_file(&data_dir, actual_port) {
|
|
log::warn!("MCP bridge failed to write port file in {}: {err}", data_dir.display());
|
|
}
|
|
loop {
|
|
let (mut stream, _) = match listener.accept().await {
|
|
Ok(s) => s,
|
|
Err(_) => continue,
|
|
};
|
|
let app = app_handle.clone();
|
|
let st = state.clone();
|
|
tokio::spawn(async move {
|
|
let mut buf = vec![0u8; 65536];
|
|
let n = match stream.read(&mut buf).await {
|
|
Ok(n) if n > 0 => n,
|
|
_ => return,
|
|
};
|
|
let request = String::from_utf8_lossy(&buf[..n]);
|
|
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") {
|
|
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/mongo/list-collections") {
|
|
handle_mongo_list_collections_data(&st, body, &mut stream).await;
|
|
} else if first_line.starts_with("POST /data/mongo/count-documents") {
|
|
handle_mongo_count_documents_data(&st, body, &mut stream).await;
|
|
} else if first_line.starts_with("POST /data/mongo/find-documents") {
|
|
handle_mongo_find_documents_data(&st, body, &mut stream).await;
|
|
} else if first_line.starts_with("POST /data/mongo/server-version") {
|
|
handle_mongo_server_version_data(&st, body, &mut stream).await;
|
|
} else if first_line.starts_with("POST /data/mongo/collection-stats") {
|
|
handle_mongo_collection_stats_data(&st, body, &mut stream).await;
|
|
} else if first_line.starts_with("POST /data/mongo/aggregate-documents") {
|
|
handle_mongo_aggregate_documents_data(&st, body, &mut stream).await;
|
|
} else if first_line.starts_with("POST /data/mongo/distinct") {
|
|
handle_mongo_distinct_data(&st, body, &mut stream).await;
|
|
} else if first_line.starts_with("POST /data/mongo/create-index") {
|
|
handle_mongo_create_index_data(&st, body, &mut stream).await;
|
|
} else if first_line.starts_with("POST /data/mongo/drop-indexes") {
|
|
handle_mongo_drop_indexes_data(&st, body, &mut stream).await;
|
|
} else if first_line.starts_with("POST /data/mongo/drop-collection") {
|
|
handle_mongo_drop_collection_data(&st, body, &mut stream).await;
|
|
} else if first_line.starts_with("POST /data/mongo/insert-documents") {
|
|
handle_mongo_insert_documents_data(&st, body, &mut stream).await;
|
|
} else if first_line.starts_with("POST /data/mongo/update-documents") {
|
|
handle_mongo_update_documents_data(&st, body, &mut stream).await;
|
|
} else if first_line.starts_with("POST /data/mongo/delete-documents") {
|
|
handle_mongo_delete_documents_data(&st, body, &mut stream).await;
|
|
} else if first_line.starts_with("POST /data/redis/execute-command") {
|
|
handle_redis_execute_command_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") {
|
|
let _ = app.emit("mcp-reload-connections", ());
|
|
respond(&mut stream, "200 OK", "ok").await;
|
|
} else {
|
|
let _ = stream.write_all(b"HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\n\r\n").await;
|
|
}
|
|
});
|
|
}
|
|
});
|
|
}
|
|
|
|
fn write_port_file(data_dir: &Path, actual_port: u16) -> std::io::Result<PathBuf> {
|
|
std::fs::create_dir_all(data_dir)?;
|
|
let path = data_dir.join(MCP_BRIDGE_PORT_FILE);
|
|
std::fs::write(&path, actual_port.to_string())?;
|
|
Ok(path)
|
|
}
|
|
|
|
fn find_config_by_name<'a>(
|
|
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 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_id: Option<&str>,
|
|
connection_name: &str,
|
|
) -> Result<crate::models::connection::ConnectionConfig, String> {
|
|
let configs = state.storage.load_connections().await.map_err(|e| e.to_string())?;
|
|
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.clone())
|
|
}
|
|
|
|
fn check_visible_database(config: &crate::models::connection::ConnectionConfig, database: &str) -> Result<(), String> {
|
|
if let Some(ref visible) = config.visible_databases {
|
|
if !visible.is_empty() && !visible.iter().any(|v| v == database) {
|
|
return Err(format!("Database '{}' is not in the visible databases list for this connection", database));
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn resolve_mongo_database(requested: Option<String>, configured: Option<String>) -> String {
|
|
requested.filter(|database| !database.trim().is_empty()).or(configured).unwrap_or_default()
|
|
}
|
|
|
|
fn resolve_mongo_target_values(
|
|
connection_id: String,
|
|
requested_database: Option<String>,
|
|
configured_database: Option<String>,
|
|
) -> (String, String) {
|
|
(connection_id, resolve_mongo_database(requested_database, configured_database))
|
|
}
|
|
|
|
async fn resolve_mongo_target(
|
|
state: &Arc<AppState>,
|
|
connection_id: Option<&str>,
|
|
connection_name: &str,
|
|
database: Option<String>,
|
|
stream: &mut tokio::net::TcpStream,
|
|
) -> Option<(String, String)> {
|
|
let config = match resolve_connection(state, connection_id, connection_name).await {
|
|
Ok(c) => c,
|
|
Err(e) => {
|
|
respond_error(stream, "404 Not Found", &e).await;
|
|
return None;
|
|
}
|
|
};
|
|
let (connection_id, database) = resolve_mongo_target_values(config.id.clone(), database, config.database.clone());
|
|
if let Err(e) = check_visible_database(&config, &database) {
|
|
respond_error(stream, "403 Forbidden", &e).await;
|
|
return None;
|
|
}
|
|
Some((connection_id, database))
|
|
}
|
|
|
|
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,
|
|
Err(_) => {
|
|
respond(stream, "400 Bad Request", "").await;
|
|
return;
|
|
}
|
|
};
|
|
let config = match resolve_connection(state, req.connection_id.as_deref(), &req.connection_name).await {
|
|
Ok(c) => c,
|
|
Err(e) => {
|
|
respond(stream, "404 Not Found", &e).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, state: &Arc<AppState>, 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 config = match resolve_connection(state, req.connection_id.as_deref(), &req.connection_name).await {
|
|
Ok(c) => c,
|
|
Err(e) => {
|
|
respond(stream, "404 Not Found", &e).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,
|
|
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;
|
|
}
|
|
|
|
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_id.as_deref(), &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();
|
|
if let Err(e) = check_visible_database(&config, &database) {
|
|
respond_error(stream, "403 Forbidden", &e).await;
|
|
return;
|
|
}
|
|
match dbx_core::schema::list_tables_core(state, &config.id, &database, &schema, None, None, 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_id.as_deref(), &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();
|
|
if let Err(e) = check_visible_database(&config, &database) {
|
|
respond_error(stream, "403 Forbidden", &e).await;
|
|
return;
|
|
}
|
|
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_mongo_list_collections_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 Some((connection_id, database)) =
|
|
resolve_mongo_target(state, req.connection_id.as_deref(), &req.connection_name, req.database, stream).await
|
|
else {
|
|
return;
|
|
};
|
|
match dbx_core::mongo_ops::mongo_list_collections_core(state, &connection_id, &database).await {
|
|
Ok(collections) => respond_json(stream, &collections).await,
|
|
Err(e) => respond_error(stream, "500 Internal Server Error", &e).await,
|
|
}
|
|
}
|
|
|
|
async fn handle_mongo_find_documents_data(state: &Arc<AppState>, body: &str, stream: &mut tokio::net::TcpStream) {
|
|
let req: MongoFindDocumentsRequest = match serde_json::from_str(body) {
|
|
Ok(r) => r,
|
|
Err(_) => {
|
|
respond_error(stream, "400 Bad Request", "Invalid JSON").await;
|
|
return;
|
|
}
|
|
};
|
|
let Some((connection_id, database)) =
|
|
resolve_mongo_target(state, req.connection_id.as_deref(), &req.connection_name, req.database, stream).await
|
|
else {
|
|
return;
|
|
};
|
|
match dbx_core::mongo_ops::mongo_find_documents_core(
|
|
state,
|
|
&connection_id,
|
|
&database,
|
|
&req.collection,
|
|
req.skip.unwrap_or(0),
|
|
req.limit.unwrap_or(100),
|
|
req.filter.as_deref(),
|
|
req.projection.as_deref(),
|
|
req.sort.as_deref(),
|
|
)
|
|
.await
|
|
{
|
|
Ok(result) => respond_json(stream, &result).await,
|
|
Err(e) => respond_error(stream, "500 Internal Server Error", &e).await,
|
|
}
|
|
}
|
|
|
|
async fn handle_mongo_count_documents_data(state: &Arc<AppState>, body: &str, stream: &mut tokio::net::TcpStream) {
|
|
let req: MongoCountDocumentsRequest = match serde_json::from_str(body) {
|
|
Ok(r) => r,
|
|
Err(_) => {
|
|
respond_error(stream, "400 Bad Request", "Invalid JSON").await;
|
|
return;
|
|
}
|
|
};
|
|
let Some((connection_id, database)) =
|
|
resolve_mongo_target(state, req.connection_id.as_deref(), &req.connection_name, req.database, stream).await
|
|
else {
|
|
return;
|
|
};
|
|
match dbx_core::mongo_ops::mongo_count_documents_core(
|
|
state,
|
|
&connection_id,
|
|
&database,
|
|
&req.collection,
|
|
req.filter.as_deref(),
|
|
req.mode.as_deref(),
|
|
)
|
|
.await
|
|
{
|
|
Ok(total) => respond_json(stream, &total).await,
|
|
Err(e) => respond_error(stream, "500 Internal Server Error", &e).await,
|
|
}
|
|
}
|
|
|
|
async fn handle_mongo_server_version_data(state: &Arc<AppState>, body: &str, stream: &mut tokio::net::TcpStream) {
|
|
let req: MongoServerVersionRequest = match serde_json::from_str(body) {
|
|
Ok(r) => r,
|
|
Err(_) => {
|
|
respond_error(stream, "400 Bad Request", "Invalid JSON").await;
|
|
return;
|
|
}
|
|
};
|
|
let Some((connection_id, database)) =
|
|
resolve_mongo_target(state, req.connection_id.as_deref(), &req.connection_name, req.database, stream).await
|
|
else {
|
|
return;
|
|
};
|
|
match dbx_core::mongo_ops::mongo_server_version_core(state, &connection_id, &database).await {
|
|
Ok(version) => respond_json(stream, &version).await,
|
|
Err(e) => respond_error(stream, "500 Internal Server Error", &e).await,
|
|
}
|
|
}
|
|
|
|
async fn handle_mongo_collection_stats_data(state: &Arc<AppState>, body: &str, stream: &mut tokio::net::TcpStream) {
|
|
let req: MongoCollectionStatsRequest = match serde_json::from_str(body) {
|
|
Ok(r) => r,
|
|
Err(_) => {
|
|
respond_error(stream, "400 Bad Request", "Invalid JSON").await;
|
|
return;
|
|
}
|
|
};
|
|
let Some((connection_id, database)) =
|
|
resolve_mongo_target(state, req.connection_id.as_deref(), &req.connection_name, req.database, stream).await
|
|
else {
|
|
return;
|
|
};
|
|
match dbx_core::mongo_ops::mongo_collection_stats_core(state, &connection_id, &database, &req.collection, req.scale)
|
|
.await
|
|
{
|
|
Ok(result) => respond_json(stream, &result).await,
|
|
Err(e) => respond_error(stream, "500 Internal Server Error", &e).await,
|
|
}
|
|
}
|
|
|
|
async fn handle_mongo_aggregate_documents_data(state: &Arc<AppState>, body: &str, stream: &mut tokio::net::TcpStream) {
|
|
let req: MongoAggregateDocumentsRequest = match serde_json::from_str(body) {
|
|
Ok(r) => r,
|
|
Err(_) => {
|
|
respond_error(stream, "400 Bad Request", "Invalid JSON").await;
|
|
return;
|
|
}
|
|
};
|
|
let Some((connection_id, database)) =
|
|
resolve_mongo_target(state, req.connection_id.as_deref(), &req.connection_name, req.database, stream).await
|
|
else {
|
|
return;
|
|
};
|
|
match dbx_core::mongo_ops::mongo_aggregate_documents_core(
|
|
state,
|
|
&connection_id,
|
|
&database,
|
|
&req.collection,
|
|
&req.pipeline_json,
|
|
req.max_rows,
|
|
req.options_json.as_deref(),
|
|
)
|
|
.await
|
|
{
|
|
Ok(result) => respond_json(stream, &result).await,
|
|
Err(e) => respond_error(stream, "500 Internal Server Error", &e).await,
|
|
}
|
|
}
|
|
|
|
async fn handle_mongo_distinct_data(state: &Arc<AppState>, body: &str, stream: &mut tokio::net::TcpStream) {
|
|
let req: MongoDistinctRequest = match serde_json::from_str(body) {
|
|
Ok(r) => r,
|
|
Err(_) => {
|
|
respond_error(stream, "400 Bad Request", "Invalid JSON").await;
|
|
return;
|
|
}
|
|
};
|
|
let Some((connection_id, database)) =
|
|
resolve_mongo_target(state, req.connection_id.as_deref(), &req.connection_name, req.database, stream).await
|
|
else {
|
|
return;
|
|
};
|
|
match dbx_core::mongo_ops::mongo_distinct_core(
|
|
state,
|
|
&connection_id,
|
|
&database,
|
|
&req.collection,
|
|
&req.field,
|
|
req.filter.as_deref(),
|
|
)
|
|
.await
|
|
{
|
|
Ok(result) => respond_json(stream, &result).await,
|
|
Err(e) => respond_error(stream, "500 Internal Server Error", &e).await,
|
|
}
|
|
}
|
|
|
|
async fn handle_mongo_create_index_data(state: &Arc<AppState>, body: &str, stream: &mut tokio::net::TcpStream) {
|
|
let req: MongoCreateIndexRequest = match serde_json::from_str(body) {
|
|
Ok(r) => r,
|
|
Err(_) => {
|
|
respond_error(stream, "400 Bad Request", "Invalid JSON").await;
|
|
return;
|
|
}
|
|
};
|
|
let Some((connection_id, database)) =
|
|
resolve_mongo_target(state, req.connection_id.as_deref(), &req.connection_name, req.database, stream).await
|
|
else {
|
|
return;
|
|
};
|
|
if let Err(e) = ensure_connection_writable(state, &connection_id, "Create index").await {
|
|
respond_error(stream, "403 Forbidden", &e).await;
|
|
return;
|
|
}
|
|
match dbx_core::mongo_ops::mongo_create_index_core(
|
|
state,
|
|
&connection_id,
|
|
&database,
|
|
&req.collection,
|
|
&req.keys_json,
|
|
req.options_json.as_deref(),
|
|
)
|
|
.await
|
|
{
|
|
Ok(name) => respond_json(stream, &serde_json::json!({ "name": name })).await,
|
|
Err(e) => respond_error(stream, "500 Internal Server Error", &e).await,
|
|
}
|
|
}
|
|
|
|
async fn handle_mongo_drop_indexes_data(state: &Arc<AppState>, body: &str, stream: &mut tokio::net::TcpStream) {
|
|
let req: MongoDropIndexesRequest = match serde_json::from_str(body) {
|
|
Ok(r) => r,
|
|
Err(_) => {
|
|
respond_error(stream, "400 Bad Request", "Invalid JSON").await;
|
|
return;
|
|
}
|
|
};
|
|
let Some((connection_id, database)) =
|
|
resolve_mongo_target(state, req.connection_id.as_deref(), &req.connection_name, req.database, stream).await
|
|
else {
|
|
return;
|
|
};
|
|
if let Err(e) = ensure_connection_writable(state, &connection_id, "Drop indexes").await {
|
|
respond_error(stream, "403 Forbidden", &e).await;
|
|
return;
|
|
}
|
|
match dbx_core::mongo_ops::mongo_drop_indexes_core(
|
|
state,
|
|
&connection_id,
|
|
&database,
|
|
&req.collection,
|
|
req.indexes_json.as_deref(),
|
|
req.single,
|
|
)
|
|
.await
|
|
{
|
|
Ok(result) => respond_json(stream, &result).await,
|
|
Err(e) => respond_error(stream, "500 Internal Server Error", &e).await,
|
|
}
|
|
}
|
|
|
|
async fn handle_mongo_drop_collection_data(state: &Arc<AppState>, body: &str, stream: &mut tokio::net::TcpStream) {
|
|
let req: MongoDropCollectionRequest = match serde_json::from_str(body) {
|
|
Ok(r) => r,
|
|
Err(_) => {
|
|
respond_error(stream, "400 Bad Request", "Invalid JSON").await;
|
|
return;
|
|
}
|
|
};
|
|
let Some((connection_id, database)) =
|
|
resolve_mongo_target(state, req.connection_id.as_deref(), &req.connection_name, req.database, stream).await
|
|
else {
|
|
return;
|
|
};
|
|
if let Err(e) = ensure_connection_writable(state, &connection_id, "Drop collection").await {
|
|
respond_error(stream, "403 Forbidden", &e).await;
|
|
return;
|
|
}
|
|
match dbx_core::mongo_ops::mongo_drop_collection_core(state, &connection_id, &database, &req.collection).await {
|
|
Ok(()) => respond_json(stream, &serde_json::json!({ "ok": true })).await,
|
|
Err(e) => respond_error(stream, "500 Internal Server Error", &e).await,
|
|
}
|
|
}
|
|
|
|
async fn handle_mongo_insert_documents_data(state: &Arc<AppState>, body: &str, stream: &mut tokio::net::TcpStream) {
|
|
let req: MongoInsertDocumentsRequest = match serde_json::from_str(body) {
|
|
Ok(r) => r,
|
|
Err(_) => {
|
|
respond_error(stream, "400 Bad Request", "Invalid JSON").await;
|
|
return;
|
|
}
|
|
};
|
|
let Some((connection_id, database)) =
|
|
resolve_mongo_target(state, req.connection_id.as_deref(), &req.connection_name, req.database, stream).await
|
|
else {
|
|
return;
|
|
};
|
|
if let Err(e) = ensure_connection_writable(state, &connection_id, "Insert").await {
|
|
respond_error(stream, "403 Forbidden", &e).await;
|
|
return;
|
|
}
|
|
match dbx_core::mongo_ops::mongo_insert_documents_core(
|
|
state,
|
|
&connection_id,
|
|
&database,
|
|
&req.collection,
|
|
&req.docs_json,
|
|
)
|
|
.await
|
|
{
|
|
Ok(inserted) => respond_json(stream, &serde_json::json!({ "affected_rows": inserted })).await,
|
|
Err(e) => respond_error(stream, "500 Internal Server Error", &e).await,
|
|
}
|
|
}
|
|
|
|
async fn handle_mongo_update_documents_data(state: &Arc<AppState>, body: &str, stream: &mut tokio::net::TcpStream) {
|
|
let req: MongoUpdateDocumentsRequest = match serde_json::from_str(body) {
|
|
Ok(r) => r,
|
|
Err(_) => {
|
|
respond_error(stream, "400 Bad Request", "Invalid JSON").await;
|
|
return;
|
|
}
|
|
};
|
|
let Some((connection_id, database)) =
|
|
resolve_mongo_target(state, req.connection_id.as_deref(), &req.connection_name, req.database, stream).await
|
|
else {
|
|
return;
|
|
};
|
|
if let Err(e) = ensure_connection_writable(state, &connection_id, "Update").await {
|
|
respond_error(stream, "403 Forbidden", &e).await;
|
|
return;
|
|
}
|
|
match dbx_core::mongo_ops::mongo_update_documents_core(
|
|
state,
|
|
&connection_id,
|
|
&database,
|
|
&req.collection,
|
|
&req.filter_json,
|
|
&req.update_json,
|
|
req.many,
|
|
req.options_json.as_deref(),
|
|
)
|
|
.await
|
|
{
|
|
Ok(modified) => respond_json(stream, &serde_json::json!({ "affected_rows": modified })).await,
|
|
Err(e) => respond_error(stream, "500 Internal Server Error", &e).await,
|
|
}
|
|
}
|
|
|
|
async fn handle_mongo_delete_documents_data(state: &Arc<AppState>, body: &str, stream: &mut tokio::net::TcpStream) {
|
|
let req: MongoDeleteDocumentsRequest = match serde_json::from_str(body) {
|
|
Ok(r) => r,
|
|
Err(_) => {
|
|
respond_error(stream, "400 Bad Request", "Invalid JSON").await;
|
|
return;
|
|
}
|
|
};
|
|
let Some((connection_id, database)) =
|
|
resolve_mongo_target(state, req.connection_id.as_deref(), &req.connection_name, req.database, stream).await
|
|
else {
|
|
return;
|
|
};
|
|
if let Err(e) = ensure_connection_writable(state, &connection_id, "Delete").await {
|
|
respond_error(stream, "403 Forbidden", &e).await;
|
|
return;
|
|
}
|
|
match dbx_core::mongo_ops::mongo_delete_documents_core(
|
|
state,
|
|
&connection_id,
|
|
&database,
|
|
&req.collection,
|
|
&req.filter_json,
|
|
req.many,
|
|
)
|
|
.await
|
|
{
|
|
Ok(deleted) => respond_json(stream, &serde_json::json!({ "affected_rows": deleted })).await,
|
|
Err(e) => respond_error(stream, "500 Internal Server Error", &e).await,
|
|
}
|
|
}
|
|
|
|
async fn handle_redis_execute_command_data(state: &Arc<AppState>, body: &str, stream: &mut tokio::net::TcpStream) {
|
|
let req: RedisCommandRequest = 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_id.as_deref(), &req.connection_name).await {
|
|
Ok(c) => c,
|
|
Err(e) => {
|
|
respond_error(stream, "404 Not Found", &e).await;
|
|
return;
|
|
}
|
|
};
|
|
let database = req.db.to_string();
|
|
if let Err(e) = check_visible_database(&config, &database) {
|
|
respond_error(stream, "403 Forbidden", &e).await;
|
|
return;
|
|
}
|
|
if let Some(name) = dbx_core::query::connection_readonly_name(state, &config.id).await {
|
|
let cmd_name = req.command.split_whitespace().next().unwrap_or("");
|
|
if dbx_core::db::redis_driver::classify_command(cmd_name)
|
|
!= dbx_core::db::redis_driver::RedisCommandSafety::Allowed
|
|
{
|
|
respond_error(
|
|
stream,
|
|
"403 Forbidden",
|
|
&format!(
|
|
"Read-only mode: connection '{}' has read-only protection enabled. Command '{}' blocked.",
|
|
name, cmd_name
|
|
),
|
|
)
|
|
.await;
|
|
return;
|
|
}
|
|
}
|
|
match dbx_core::redis_ops::redis_execute_command_core(
|
|
state,
|
|
&config.id,
|
|
req.db,
|
|
&req.command,
|
|
req.skip_safety_check.unwrap_or(false),
|
|
)
|
|
.await
|
|
{
|
|
Ok(result) => respond_json(stream, &result).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_id.as_deref(), &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());
|
|
if let Err(e) = check_visible_database(&config, &database) {
|
|
respond_error(stream, "403 Forbidden", &e).await;
|
|
return;
|
|
}
|
|
// Read-only check: reject if the connection has read-only protection and the SQL is a write
|
|
if let Err(e) = dbx_core::query::check_read_only_for_connection(state, &config.id, &req.sql).await {
|
|
respond_error(stream, "403 Forbidden", &e).await;
|
|
return;
|
|
}
|
|
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,
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::{resolve_mongo_database, resolve_mongo_target_values, write_port_file};
|
|
|
|
#[test]
|
|
fn writes_bridge_port_file_to_resolved_data_dir() {
|
|
let root = std::env::temp_dir().join(format!(
|
|
"dbx-mcp-bridge-port-test-{}-{}",
|
|
std::process::id(),
|
|
uuid::Uuid::new_v4()
|
|
));
|
|
let default_data_dir = root.join("default-app-data");
|
|
let resolved_data_dir = root.join("resolved-data");
|
|
std::fs::create_dir_all(&default_data_dir).unwrap();
|
|
|
|
let port_file = write_port_file(&resolved_data_dir, 49152).unwrap();
|
|
|
|
assert_eq!(port_file, resolved_data_dir.join("mcp-bridge-port"));
|
|
assert_eq!(std::fs::read_to_string(port_file).unwrap(), "49152");
|
|
assert!(!default_data_dir.join("mcp-bridge-port").exists());
|
|
|
|
let _ = std::fs::remove_dir_all(root);
|
|
}
|
|
|
|
#[test]
|
|
fn mongo_database_uses_configured_default_for_missing_or_blank_request() {
|
|
let configured = Some("sample_db".to_string());
|
|
|
|
assert_eq!(resolve_mongo_database(None, configured.clone()), "sample_db");
|
|
assert_eq!(resolve_mongo_database(Some(String::new()), configured.clone()), "sample_db");
|
|
assert_eq!(resolve_mongo_database(Some(" ".to_string()), configured), "sample_db");
|
|
}
|
|
|
|
#[test]
|
|
fn mongo_database_preserves_explicit_target() {
|
|
assert_eq!(resolve_mongo_database(Some("admin".to_string()), Some("sample_db".to_string())), "admin");
|
|
}
|
|
|
|
#[test]
|
|
fn mongo_target_keeps_connection_id_separate_from_database() {
|
|
assert_eq!(
|
|
resolve_mongo_target_values(
|
|
"connection-id".to_string(),
|
|
Some("sample_db".to_string()),
|
|
Some("default_db".to_string()),
|
|
),
|
|
("connection-id".to_string(), "sample_db".to_string())
|
|
);
|
|
}
|
|
}
|