1use std::sync::Arc;
4
5use rmcp::{
6 ServerHandler,
7 handler::server::tool::ToolRouter,
8 model::{ServerCapabilities, ServerInfo},
9 tool_handler,
10};
11
12use crate::{
13 access_control::AuthorizationResolver,
14 tools::{
15 backup_tool::BackupTool, create_fts_index_tool::CreateFtsIndexTool,
16 database_info_tool::DatabaseInfoTool,
17 describe_table_tool::DescribeTableTool, execute_tool::ExecuteTool,
18 explain_query_tool::ExplainQueryTool,
19 list_foreign_keys_tool::ListForeignKeysTool,
20 list_indexes_tool::ListIndexesTool, list_tables_tool::ListTablesTool,
21 list_triggers_tool::ListTriggersTool, list_views_tool::ListViewsTool,
22 search_fts_tool::SearchFtsTool, vacuum_tool::VacuumTool,
23 },
24 traits::SqliteServerTool,
25};
26
27pub struct McpServerSqlite {
30 pool: r2d2::Pool<r2d2_sqlite::SqliteConnectionManager>,
32 authorization_resolver: Arc<AuthorizationResolver>,
36 query_timeout: Option<std::time::Duration>,
39 tool_router: ToolRouter<Self>,
41}
42
43impl McpServerSqlite {
44 pub fn new(
47 pool: r2d2::Pool<r2d2_sqlite::SqliteConnectionManager>,
48 authorization_resolver: AuthorizationResolver,
49 query_timeout: Option<std::time::Duration>,
50 ) -> Self {
51 let tool_router = Self::tool_router();
52 Self {
53 pool,
54 authorization_resolver: Arc::new(authorization_resolver),
55 query_timeout,
56 tool_router,
57 }
58 }
59
60 pub fn connection(
64 &self,
65 ) -> Result<
66 r2d2::PooledConnection<r2d2_sqlite::SqliteConnectionManager>,
67 ConnectionError,
68 > {
69 tracing::debug!("checking out pooled connection");
70 let conn = self.pool.get().map_err(|source| {
71 tracing::error!(%source, "pool checkout failed");
72 ConnectionError::Pool(source)
73 })?;
74 let resolver = self.authorization_resolver.clone();
75 conn.authorizer(Some(move |ctx: rusqlite::hooks::AuthContext<'_>| {
76 resolver.authorization(ctx)
77 }))
78 .map_err(|source| {
79 tracing::error!(%source, "failed to install authorizer");
80 ConnectionError::Authorizer(source)
81 })?;
82
83 if let Some(timeout) = self.query_timeout {
84 let deadline = std::time::Instant::now() + timeout;
85 conn.progress_handler(
86 1000,
87 Some(move || std::time::Instant::now() > deadline),
88 )
89 .map_err(|source| {
90 tracing::error!(%source, "failed to install progress handler");
91 ConnectionError::ProgressHandler(source)
92 })?;
93 }
94
95 Ok(conn)
96 }
97
98 pub fn tool_router() -> ToolRouter<Self> {
99 ToolRouter::<Self>::new()
100 .with_route((ExecuteTool::tool(), ExecuteTool::handler_func()))
101 .with_route((
102 ListTablesTool::tool(),
103 ListTablesTool::handler_func(),
104 ))
105 .with_route((
106 DescribeTableTool::tool(),
107 DescribeTableTool::handler_func(),
108 ))
109 .with_route((
110 ListIndexesTool::tool(),
111 ListIndexesTool::handler_func(),
112 ))
113 .with_route((
114 ListForeignKeysTool::tool(),
115 ListForeignKeysTool::handler_func(),
116 ))
117 .with_route((ListViewsTool::tool(), ListViewsTool::handler_func()))
118 .with_route((
119 ListTriggersTool::tool(),
120 ListTriggersTool::handler_func(),
121 ))
122 .with_route((
123 ExplainQueryTool::tool(),
124 ExplainQueryTool::handler_func(),
125 ))
126 .with_route((BackupTool::tool(), BackupTool::handler_func()))
127 .with_route((
128 CreateFtsIndexTool::tool(),
129 CreateFtsIndexTool::handler_func(),
130 ))
131 .with_route((SearchFtsTool::tool(), SearchFtsTool::handler_func()))
132 .with_route((VacuumTool::tool(), VacuumTool::handler_func()))
133 .with_route((
134 DatabaseInfoTool::tool(),
135 DatabaseInfoTool::handler_func(),
136 ))
137 }
138}
139
140#[derive(Debug, thiserror::Error)]
143pub enum ConnectionError {
144 #[error("pool checkout failed: {0}")]
146 Pool(r2d2::Error),
147 #[error("failed to install authorizer: {0}")]
149 Authorizer(rusqlite::Error),
150 #[error("failed to install progress handler: {0}")]
152 ProgressHandler(rusqlite::Error),
153}
154
155#[tool_handler]
156impl ServerHandler for McpServerSqlite {
157 fn get_info(&self) -> ServerInfo {
158 let mut info = ServerInfo::default();
159 info.server_info.name = env!("CARGO_PKG_NAME").into();
160 info.server_info.version = env!("CARGO_PKG_VERSION").into();
161 info.instructions = Some(INSTRUCTIONS.into());
162 info.capabilities =
163 ServerCapabilities::builder().enable_tools().build();
164 info
165 }
166}
167
168const INSTRUCTIONS: &str = "\
169You are connected to an MCP SQLite server with fine-grained access control.
170
171Available tools:
172- execute: Run any SQL query (SELECT, INSERT, UPDATE, DELETE, DDL)
173- list_tables: Show all tables with their CREATE TABLE schemas
174- describe_table: Show columns, types, constraints, and defaults for a table
175- list_indexes: Show indexes with columns, uniqueness, and partial predicates
176- list_foreign_keys: Show foreign key constraints for a table
177- list_views: Show all views with their defining SQL
178- list_triggers: Show triggers with event, timing, and SQL
179- explain_query: Show the EXPLAIN QUERY PLAN for a SQL statement
180- create_fts_index: Create a full-text search index on a table
181- search_fts: Search a full-text index with BM25 ranking and snippets
182- backup: Back up the database to a file
183- vacuum: Reclaim unused space and defragment the database
184- database_info: Show database metadata (size, page count, journal mode, etc.)
185
186Start by calling list_tables to discover the schema, then use describe_table \
187for details on specific tables. Use explain_query to optimize slow queries. \
188All operations respect the server's access control policy.";