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 #[must_use]
75 pub fn new(config: &Config) -> Self {
76 Self {
77 config: config.database.clone(),
78 connection: SqliteConnection::new(&config.database),
79 redactor: Redactor::from_config(&config.pii),
80 tool_router: ToolRouter::from_specs(TOOLS, config.database.read_only, true),
81 }
82 }
83}
84
85impl From<SqliteHandler> for Server {
86 fn from(handler: SqliteHandler) -> Self {
88 Self::new(handler)
89 }
90}
91
92impl ServerHandler for SqliteHandler {
93 fn get_info(&self) -> ServerInfo {
94 let mut info = server_info();
95 info.server_info.description = Some(DESCRIPTION.into());
96 info.instructions = Some(if self.config.read_only {
97 INSTRUCTIONS_READ_ONLY.into()
98 } else {
99 INSTRUCTIONS.into()
100 });
101 info
102 }
103
104 async fn call_tool(
105 &self,
106 request: CallToolRequestParams,
107 context: RequestContext<RoleServer>,
108 ) -> Result<CallToolResult, ErrorData> {
109 let tcc = ToolCallContext::new(self, request, context);
110 self.tool_router.call(tcc).await
111 }
112
113 async fn list_tools(
114 &self,
115 _request: Option<PaginatedRequestParams>,
116 _context: RequestContext<RoleServer>,
117 ) -> Result<ListToolsResult, ErrorData> {
118 Ok(ListToolsResult {
119 tools: self.tool_router.list_all(),
120 next_cursor: None,
121 meta: None,
122 })
123 }
124
125 fn get_tool(&self, name: &str) -> Option<Tool> {
126 self.tool_router.get(name).cloned()
127 }
128}
129
130#[cfg(test)]
131mod tests {
132 use super::*;
133 use dbmcp_config::DatabaseBackend;
134
135 fn handler(read_only: bool) -> SqliteHandler {
136 SqliteHandler::new(&Config {
137 database: DatabaseConfig {
138 backend: DatabaseBackend::Sqlite,
139 name: Some(":memory:".into()),
140 read_only,
141 ..DatabaseConfig::default()
142 },
143 http: None,
144 pii: dbmcp_config::PiiConfig::default(),
145 })
146 }
147
148 #[tokio::test]
149 async fn router_exposes_all_seven_tools_in_read_write_mode() {
150 let router = handler(false).tool_router;
151 for name in [
152 "listTables",
153 "listViews",
154 "listTriggers",
155 "dropTable",
156 "readQuery",
157 "writeQuery",
158 "explainQuery",
159 ] {
160 assert!(router.has_route(name), "missing tool: {name}");
161 }
162 }
163
164 #[tokio::test]
165 async fn router_excludes_get_table_schema() {
166 for read_only in [false, true] {
169 let router = handler(read_only).tool_router;
170 assert!(
171 !router.has_route("getTableSchema"),
172 "getTableSchema must be absent (read_only={read_only})"
173 );
174 }
175 }
176
177 #[tokio::test]
178 async fn router_does_not_advertise_backend_specific_tools() {
179 let router = handler(false).tool_router;
180 for absent in [
181 "listDatabases",
182 "listFunctions",
183 "listProcedures",
184 "listMaterializedViews",
185 "createDatabase",
186 "dropDatabase",
187 ] {
188 assert!(!router.has_route(absent), "SQLite must not advertise {absent}");
189 }
190 }
191
192 #[tokio::test]
193 async fn router_hides_write_tools_in_read_only_mode() {
194 let router = handler(true).tool_router;
195 assert!(router.has_route("listTables"));
196 assert!(router.has_route("listViews"));
197 assert!(router.has_route("listTriggers"));
198 assert!(router.has_route("readQuery"));
199 assert!(router.has_route("explainQuery"));
200 assert!(!router.has_route("writeQuery"));
201 assert!(!router.has_route("dropTable"));
202 }
203
204 #[tokio::test]
205 async fn instructions_match_read_only_mode() {
206 let read_write = handler(false).get_info().instructions.expect("instructions present");
207 assert!(
208 read_write.contains("writeQuery"),
209 "read-write instructions mention writeQuery"
210 );
211
212 let read_only = handler(true).get_info().instructions.expect("instructions present");
213 for tool in ["writeQuery", "dropTable"] {
214 assert!(
215 !read_only.contains(tool),
216 "read-only instructions must not mention {tool}"
217 );
218 }
219 }
220
221 #[tokio::test]
222 async fn list_tables_annotations() {
223 let router = handler(false).tool_router;
224 let tool = router.get("listTables").expect("listTables registered");
225
226 let annotations = tool.annotations.as_ref().expect("annotations present");
227 assert_eq!(annotations.read_only_hint, Some(true));
228 assert_eq!(annotations.destructive_hint, Some(false));
229 assert_eq!(annotations.idempotent_hint, Some(true));
230 assert_eq!(annotations.open_world_hint, Some(false));
231 }
232}