Skip to main content

dbmcp_sqlite/
handler.rs

1//! `SQLite` handler: composes a [`SqliteConnection`] with the MCP tool router.
2//!
3//! All pool ownership and pool initialization logic lives in the
4//! [`SqliteConnection`]. This module exposes the MCP
5//! `ServerHandler` surface and one thin delegator method that the
6//! per-tool implementations call.
7
8use 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
25/// Backend-specific description for `SQLite`.
26const DESCRIPTION: &str = "Database MCP Server for SQLite";
27
28/// Backend-specific instructions for `SQLite` in read-write mode.
29const INSTRUCTIONS: &str = include_str!("../assets/instructions/default.md");
30
31/// Backend-specific instructions for `SQLite` in read-only mode.
32const INSTRUCTIONS_READ_ONLY: &str = include_str!("../assets/instructions/read-only.md");
33
34/// Declarative tool table: `(tool, pinned, read_only)`.
35///
36/// `SQLite` has no cross-database tools and no `database` request field;
37/// its file path is always pinned, so every tool is `pinned = true` and
38/// the handler builds the router with `pinned = true`.
39const 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/// `SQLite` file-based database handler.
50///
51/// Composes one [`SqliteConnection`] (which owns the pool and
52/// the pool initialization logic) with the per-backend MCP tool router.
53#[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    /// Creates a new `SQLite` handler.
73    ///
74    /// Constructs the [`SqliteConnection`] (which builds the
75    /// lazy pool) and the MCP tool router. No file I/O happens here.
76    /// # Errors
77    ///
78    /// Returns [`dbmcp_pii::RedactorInitError`] when PII redaction is enabled
79    /// with a NER model that fails to load (fail-closed startup).
80    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    /// Wraps a [`SqliteHandler`] in the type-erased MCP server.
92    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        // Spec 046 US4: `getTableSchema` is retired on SQLite. Both read-only and
175        // read-write catalogues must no longer advertise it.
176        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}