1use dbmcp_config::{Config, DatabaseConfig};
9use dbmcp_pii::Redactor;
10use dbmcp_server::{Server, Shutdown, ShutdownFuture, ToolRouterExt, ToolSpec, server_info};
11use rmcp::RoleServer;
12use rmcp::handler::server::router::tool::ToolRouter;
13use rmcp::handler::server::tool::ToolCallContext;
14use rmcp::model::{
15 CallToolRequestParams, CallToolResponse, ListToolsResult, PaginatedRequestParams, ServerConfig, Tool,
16};
17use rmcp::service::RequestContext;
18use rmcp::{ErrorData, ServerHandler};
19
20use crate::connection::SqliteConnection;
21use crate::tools::{
22 DropTableTool, ExplainQueryTool, ListTablesTool, ListTriggersTool, ListViewsTool, ReadQueryTool, WriteQueryTool,
23};
24
25const DESCRIPTION: &str = "Database MCP Server for SQLite";
27
28const INSTRUCTIONS: &str = include_str!("../assets/instructions/default.md");
30
31const INSTRUCTIONS_READ_ONLY: &str = include_str!("../assets/instructions/read-only.md");
33
34const TOOLS: &[ToolSpec<SqliteHandler>] = &[
40 ToolSpec::async_tool::<ListTablesTool>(true, false),
41 ToolSpec::async_tool::<ListViewsTool>(true, false),
42 ToolSpec::async_tool::<ListTriggersTool>(true, false),
43 ToolSpec::async_tool::<ReadQueryTool>(true, false),
44 ToolSpec::async_tool::<ExplainQueryTool>(true, false),
45 ToolSpec::async_tool::<WriteQueryTool>(true, true),
46 ToolSpec::async_tool::<DropTableTool>(true, true),
47];
48
49#[derive(Clone)]
54pub struct SqliteHandler {
55 pub(crate) config: DatabaseConfig,
56 pub(crate) connection: SqliteConnection,
57 pub(crate) redactor: Option<Redactor>,
58 tool_router: ToolRouter<Self>,
59}
60
61impl std::fmt::Debug for SqliteHandler {
62 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
63 f.debug_struct("SqliteHandler")
64 .field("read_only", &self.config.read_only)
65 .field("redact_pii", &self.redactor.is_some())
66 .field("connection", &self.connection)
67 .finish_non_exhaustive()
68 }
69}
70
71impl SqliteHandler {
72 pub fn new(config: &Config) -> Result<Self, dbmcp_pii::RedactorInitError> {
81 Ok(Self {
82 config: config.database.clone(),
83 connection: SqliteConnection::new(&config.database),
84 redactor: Redactor::from_config(&config.pii)?,
85 tool_router: ToolRouter::from_specs(TOOLS, config.database.read_only, true),
86 })
87 }
88}
89
90impl From<SqliteHandler> for Server {
91 fn from(handler: SqliteHandler) -> Self {
93 Self::new(handler)
94 }
95}
96
97impl Shutdown for SqliteHandler {
98 fn shutdown(&self) -> ShutdownFuture<'_> {
99 Box::pin(async move { self.connection.shutdown().await })
100 }
101}
102
103impl ServerHandler for SqliteHandler {
104 fn get_info(&self) -> ServerConfig {
105 let mut info = server_info();
106 info.server_info.description = Some(DESCRIPTION.into());
107 info.instructions = Some(if self.config.read_only {
108 INSTRUCTIONS_READ_ONLY.into()
109 } else {
110 INSTRUCTIONS.into()
111 });
112 info
113 }
114
115 async fn call_tool(
116 &self,
117 request: CallToolRequestParams,
118 context: RequestContext<RoleServer>,
119 ) -> Result<CallToolResponse, 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::with_all_items(self.tool_router.list_all()))
130 }
131
132 fn get_tool(&self, name: &str) -> Option<Tool> {
133 self.tool_router.get(name).cloned()
134 }
135}
136
137#[cfg(test)]
138mod tests {
139 use super::*;
140 use dbmcp_config::DatabaseBackend;
141
142 fn handler(read_only: bool) -> SqliteHandler {
143 SqliteHandler::new(&Config {
144 database: DatabaseConfig {
145 backend: DatabaseBackend::Sqlite,
146 name: Some(":memory:".into()),
147 read_only,
148 ..DatabaseConfig::default()
149 },
150 http: None,
151 pii: dbmcp_config::PiiConfig::default(),
152 })
153 .expect("handler builds in test")
154 }
155
156 #[tokio::test]
157 async fn router_exposes_all_seven_tools_in_read_write_mode() {
158 let router = handler(false).tool_router;
159 for name in [
160 "listTables",
161 "listViews",
162 "listTriggers",
163 "dropTable",
164 "readQuery",
165 "writeQuery",
166 "explainQuery",
167 ] {
168 assert!(router.has_route(name), "missing tool: {name}");
169 }
170 }
171
172 #[tokio::test]
173 async fn router_excludes_get_table_schema() {
174 for read_only in [false, true] {
177 let router = handler(read_only).tool_router;
178 assert!(
179 !router.has_route("getTableSchema"),
180 "getTableSchema must be absent (read_only={read_only})"
181 );
182 }
183 }
184
185 #[tokio::test]
186 async fn router_does_not_advertise_backend_specific_tools() {
187 let router = handler(false).tool_router;
188 for absent in [
189 "listDatabases",
190 "listFunctions",
191 "listProcedures",
192 "listMaterializedViews",
193 "createDatabase",
194 "dropDatabase",
195 ] {
196 assert!(!router.has_route(absent), "SQLite must not advertise {absent}");
197 }
198 }
199
200 #[tokio::test]
201 async fn router_hides_write_tools_in_read_only_mode() {
202 let router = handler(true).tool_router;
203 assert!(router.has_route("listTables"));
204 assert!(router.has_route("listViews"));
205 assert!(router.has_route("listTriggers"));
206 assert!(router.has_route("readQuery"));
207 assert!(router.has_route("explainQuery"));
208 assert!(!router.has_route("writeQuery"));
209 assert!(!router.has_route("dropTable"));
210 }
211
212 #[tokio::test]
213 async fn instructions_match_read_only_mode() {
214 let read_write = handler(false).get_info().instructions.expect("instructions present");
215 assert!(
216 read_write.contains("writeQuery"),
217 "read-write instructions mention writeQuery"
218 );
219
220 let read_only = handler(true).get_info().instructions.expect("instructions present");
221 for tool in ["writeQuery", "dropTable"] {
222 assert!(
223 !read_only.contains(tool),
224 "read-only instructions must not mention {tool}"
225 );
226 }
227 }
228
229 #[tokio::test]
230 async fn list_tables_annotations() {
231 let router = handler(false).tool_router;
232 let tool = router.get("listTables").expect("listTables registered");
233
234 let annotations = tool.annotations.as_ref().expect("annotations present");
235 assert_eq!(annotations.read_only_hint, Some(true));
236 assert_eq!(annotations.destructive_hint, Some(false));
237 assert_eq!(annotations.idempotent_hint, Some(true));
238 assert_eq!(annotations.open_world_hint, Some(false));
239 }
240}