1use dbmcp_config::{Config, DatabaseConfig};
9use dbmcp_pii::{Analyzer, Redactor};
10use dbmcp_server::{Server, server_info};
11use rmcp::RoleServer;
12use rmcp::handler::server::router::tool::ToolRouter;
13use rmcp::handler::server::tool::ToolCallContext;
14use rmcp::model::{CallToolRequestParams, CallToolResult, ListToolsResult, PaginatedRequestParams, ServerInfo, Tool};
15use rmcp::service::RequestContext;
16use rmcp::{ErrorData, ServerHandler};
17
18use crate::connection::SqliteConnection;
19use crate::tools::{
20 DropTableTool, ExplainQueryTool, ListTablesTool, ListTriggersTool, ListViewsTool, ReadQueryTool, WriteQueryTool,
21};
22
23const DESCRIPTION: &str = "Database MCP Server for SQLite";
25
26const INSTRUCTIONS: &str = r"## Workflow
28
291. Call `listTables` to discover tables in the connected database. Pass `search` to filter by name (case-insensitive substring). Pass `detailed: true` to get columns, constraints, indexes, and triggers in the same call — this supersedes the legacy `getTableSchema` workflow.
302. Call `listViews` to discover views in the connected database.
313. Call `listTriggers` to discover triggers in the connected database.
324. Use `readQuery` for read-only SQL (SELECT).
335. Use `writeQuery` for data changes (INSERT, UPDATE, DELETE, CREATE, ALTER, DROP).
346. Use `explainQuery` to analyze query execution plans and diagnose slow queries.
357. Use `dropTable` to remove a table from the database.
36
37## Constraints
38
39- The `writeQuery` and `dropTable` tools are hidden when read-only mode is active.
40- Multi-statement queries are not supported. Send one statement per request.";
41
42#[derive(Clone)]
47pub struct SqliteHandler {
48 pub(crate) config: DatabaseConfig,
49 pub(crate) connection: SqliteConnection,
50 pub(crate) redactor: Option<Redactor>,
51 tool_router: ToolRouter<Self>,
52}
53
54impl std::fmt::Debug for SqliteHandler {
55 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
56 f.debug_struct("SqliteHandler")
57 .field("read_only", &self.config.read_only)
58 .field("redact_pii", &self.redactor.is_some())
59 .field("connection", &self.connection)
60 .finish_non_exhaustive()
61 }
62}
63
64impl SqliteHandler {
65 #[must_use]
70 pub fn new(config: &Config) -> Self {
71 Self {
72 config: config.database.clone(),
73 connection: SqliteConnection::new(&config.database),
74 redactor: config
75 .pii
76 .enabled
77 .then(|| Redactor::new(Analyzer::from_pii_config(&config.pii), config.pii.operator.into())),
78 tool_router: build_tool_router(config.database.read_only),
79 }
80 }
81}
82
83impl From<SqliteHandler> for Server {
84 fn from(handler: SqliteHandler) -> Self {
86 Self::new(handler)
87 }
88}
89
90fn build_tool_router(read_only: bool) -> ToolRouter<SqliteHandler> {
92 let mut router = ToolRouter::new()
93 .with_async_tool::<ListTablesTool>()
94 .with_async_tool::<ListViewsTool>()
95 .with_async_tool::<ListTriggersTool>()
96 .with_async_tool::<ReadQueryTool>()
97 .with_async_tool::<ExplainQueryTool>();
98
99 if !read_only {
100 router = router
101 .with_async_tool::<WriteQueryTool>()
102 .with_async_tool::<DropTableTool>();
103 }
104 router
105}
106
107impl ServerHandler for SqliteHandler {
108 fn get_info(&self) -> ServerInfo {
109 let mut info = server_info();
110 info.server_info.description = Some(DESCRIPTION.into());
111 info.instructions = Some(INSTRUCTIONS.into());
112 info
113 }
114
115 async fn call_tool(
116 &self,
117 request: CallToolRequestParams,
118 context: RequestContext<RoleServer>,
119 ) -> Result<CallToolResult, ErrorData> {
120 let tcc = ToolCallContext::new(self, request, context);
121 self.tool_router.call(tcc).await
122 }
123
124 async fn list_tools(
125 &self,
126 _request: Option<PaginatedRequestParams>,
127 _context: RequestContext<RoleServer>,
128 ) -> Result<ListToolsResult, ErrorData> {
129 Ok(ListToolsResult {
130 tools: self.tool_router.list_all(),
131 next_cursor: None,
132 meta: None,
133 })
134 }
135
136 fn get_tool(&self, name: &str) -> Option<Tool> {
137 self.tool_router.get(name).cloned()
138 }
139}
140
141#[cfg(test)]
142mod tests {
143 use super::*;
144 use dbmcp_config::DatabaseBackend;
145
146 fn handler(read_only: bool) -> SqliteHandler {
147 SqliteHandler::new(&Config {
148 database: DatabaseConfig {
149 backend: DatabaseBackend::Sqlite,
150 name: Some(":memory:".into()),
151 read_only,
152 ..DatabaseConfig::default()
153 },
154 http: None,
155 pii: dbmcp_config::PiiConfig::default(),
156 })
157 }
158
159 #[tokio::test]
160 async fn router_exposes_all_seven_tools_in_read_write_mode() {
161 let router = handler(false).tool_router;
162 for name in [
163 "listTables",
164 "listViews",
165 "listTriggers",
166 "dropTable",
167 "readQuery",
168 "writeQuery",
169 "explainQuery",
170 ] {
171 assert!(router.has_route(name), "missing tool: {name}");
172 }
173 }
174
175 #[tokio::test]
176 async fn router_excludes_get_table_schema() {
177 for read_only in [false, true] {
180 let router = handler(read_only).tool_router;
181 assert!(
182 !router.has_route("getTableSchema"),
183 "getTableSchema must be absent (read_only={read_only})"
184 );
185 }
186 }
187
188 #[tokio::test]
189 async fn router_does_not_advertise_backend_specific_tools() {
190 let router = handler(false).tool_router;
191 for absent in [
192 "listDatabases",
193 "listFunctions",
194 "listProcedures",
195 "listMaterializedViews",
196 "createDatabase",
197 "dropDatabase",
198 ] {
199 assert!(!router.has_route(absent), "SQLite must not advertise {absent}");
200 }
201 }
202
203 #[tokio::test]
204 async fn router_hides_write_tools_in_read_only_mode() {
205 let router = handler(true).tool_router;
206 assert!(router.has_route("listTables"));
207 assert!(router.has_route("listViews"));
208 assert!(router.has_route("listTriggers"));
209 assert!(router.has_route("readQuery"));
210 assert!(router.has_route("explainQuery"));
211 assert!(!router.has_route("writeQuery"));
212 assert!(!router.has_route("dropTable"));
213 }
214
215 #[tokio::test]
216 async fn list_tables_annotations() {
217 let router = handler(false).tool_router;
218 let tool = router.get("listTables").expect("listTables registered");
219
220 let annotations = tool.annotations.as_ref().expect("annotations present");
221 assert_eq!(annotations.read_only_hint, Some(true));
222 assert_eq!(annotations.destructive_hint, Some(false));
223 assert_eq!(annotations.idempotent_hint, Some(true));
224 assert_eq!(annotations.open_world_hint, Some(false));
225 }
226}