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