use database_mcp_config::DatabaseConfig;
use database_mcp_server::{Server, server_info};
use rmcp::RoleServer;
use rmcp::handler::server::router::tool::ToolRouter;
use rmcp::handler::server::tool::ToolCallContext;
use rmcp::model::{CallToolRequestParams, CallToolResult, ListToolsResult, PaginatedRequestParams, ServerInfo, Tool};
use rmcp::service::RequestContext;
use rmcp::{ErrorData, ServerHandler};
use crate::connection::MysqlConnection;
use crate::tools::{
CreateDatabaseTool, DropDatabaseTool, DropTableTool, ExplainQueryTool, GetTableSchemaTool, ListDatabasesTool,
ListTablesTool, ReadQueryTool, WriteQueryTool,
};
const DESCRIPTION: &str = "Database MCP Server for MySQL and MariaDB";
const INSTRUCTIONS: &str = r"## Workflow
1. Call `list_databases` to discover available databases.
2. Call `list_tables` with a `database_name` to see its tables.
3. Call `get_table_schema` with `database_name` and `table_name` to inspect columns, types, and foreign keys before writing queries.
4. Use `read_query` for read-only SQL (SELECT, SHOW, DESCRIBE, USE, EXPLAIN).
5. Use `write_query` for data changes (INSERT, UPDATE, DELETE, CREATE, ALTER, DROP).
6. Use `explain_query` to analyze query execution plans and diagnose slow queries.
7. Use `create_database` to create a new database.
8. Use `drop_database` to drop an existing database.
9. Use `drop_table` to remove a table from a database.
## Constraints
- The `write_query`, `create_database`, `drop_database`, and `drop_table` tools are hidden when read-only mode is active.
- Multi-statement queries are not supported. Send one statement per request.";
#[derive(Clone)]
pub struct MysqlHandler {
pub(crate) config: DatabaseConfig,
pub(crate) connection: MysqlConnection,
tool_router: ToolRouter<Self>,
}
impl std::fmt::Debug for MysqlHandler {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("MysqlHandler")
.field("read_only", &self.config.read_only)
.field("connection", &self.connection)
.finish_non_exhaustive()
}
}
impl MysqlHandler {
#[must_use]
pub fn new(config: &DatabaseConfig) -> Self {
Self {
config: config.clone(),
connection: MysqlConnection::new(config),
tool_router: build_tool_router(config.read_only),
}
}
}
impl From<MysqlHandler> for Server {
fn from(handler: MysqlHandler) -> Self {
Self::new(handler)
}
}
fn build_tool_router(read_only: bool) -> ToolRouter<MysqlHandler> {
let mut router = ToolRouter::new()
.with_async_tool::<ListDatabasesTool>()
.with_async_tool::<ListTablesTool>()
.with_async_tool::<GetTableSchemaTool>()
.with_async_tool::<ReadQueryTool>()
.with_async_tool::<ExplainQueryTool>();
if !read_only {
router = router
.with_async_tool::<CreateDatabaseTool>()
.with_async_tool::<DropDatabaseTool>()
.with_async_tool::<DropTableTool>()
.with_async_tool::<WriteQueryTool>();
}
router
}
impl ServerHandler for MysqlHandler {
fn get_info(&self) -> ServerInfo {
let mut info = server_info();
info.server_info.description = Some(DESCRIPTION.into());
info.instructions = Some(INSTRUCTIONS.into());
info
}
async fn call_tool(
&self,
request: CallToolRequestParams,
context: RequestContext<RoleServer>,
) -> Result<CallToolResult, ErrorData> {
let tcc = ToolCallContext::new(self, request, context);
self.tool_router.call(tcc).await
}
async fn list_tools(
&self,
_request: Option<PaginatedRequestParams>,
_context: RequestContext<RoleServer>,
) -> Result<ListToolsResult, ErrorData> {
Ok(ListToolsResult {
tools: self.tool_router.list_all(),
next_cursor: None,
meta: None,
})
}
fn get_tool(&self, name: &str) -> Option<Tool> {
self.tool_router.get(name).cloned()
}
}
#[cfg(test)]
mod tests {
use super::*;
use database_mcp_config::DatabaseBackend;
fn base_config() -> DatabaseConfig {
DatabaseConfig {
backend: DatabaseBackend::Mysql,
host: "db.example.com".into(),
port: 3307,
user: "admin".into(),
password: Some("s3cret".into()),
name: Some("mydb".into()),
..DatabaseConfig::default()
}
}
fn handler(read_only: bool) -> MysqlHandler {
MysqlHandler::new(&DatabaseConfig {
read_only,
..base_config()
})
}
#[tokio::test]
async fn router_exposes_all_nine_tools_in_read_write_mode() {
let router = handler(false).tool_router;
for name in [
"list_databases",
"list_tables",
"get_table_schema",
"read_query",
"explain_query",
"create_database",
"drop_database",
"drop_table",
"write_query",
] {
assert!(router.has_route(name), "missing tool: {name}");
}
}
#[tokio::test]
async fn router_hides_write_tools_in_read_only_mode() {
let router = handler(true).tool_router;
assert!(router.has_route("list_databases"));
assert!(router.has_route("list_tables"));
assert!(router.has_route("get_table_schema"));
assert!(router.has_route("read_query"));
assert!(router.has_route("explain_query"));
assert!(!router.has_route("write_query"));
assert!(!router.has_route("create_database"));
assert!(!router.has_route("drop_database"));
assert!(!router.has_route("drop_table"));
}
}