Skip to main content

mcp_server_sqlite/
mcp.rs

1//! MCP server implementation for SQLite.
2
3use 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
27/// The MCP server instance that handles tool calls against a SQLite database
28/// with access control.
29pub struct McpServerSqlite {
30    /// The connection pool for the SQLite database.
31    pool: r2d2::Pool<r2d2_sqlite::SqliteConnectionManager>,
32    /// The resolver that evaluates access control rules for each SQL operation.
33    /// Wrapped in an `Arc` so it can be moved into the per-query authorizer
34    /// closure required by rusqlite.
35    authorization_resolver: Arc<AuthorizationResolver>,
36    /// Optional query timeout. When `Some`, each connection's progress handler
37    /// aborts operations that exceed this duration.
38    query_timeout: Option<std::time::Duration>,
39    /// The router that dispatches tool calls to their handlers.
40    tool_router: ToolRouter<Self>,
41}
42
43impl McpServerSqlite {
44    /// Creates a new MCP server backed by the given connection pool and access
45    /// control resolver.
46    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    /// Obtains a pooled connection from the database pool, installs the access
61    /// control authorizer, and optionally sets up the query timeout via a
62    /// progress handler.
63    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/// Errors that can occur when checking out and configuring a pooled database
141/// connection.
142#[derive(Debug, thiserror::Error)]
143pub enum ConnectionError {
144    /// The connection pool failed to provide a connection.
145    #[error("pool checkout failed: {0}")]
146    Pool(r2d2::Error),
147    /// The SQLite authorizer could not be installed.
148    #[error("failed to install authorizer: {0}")]
149    Authorizer(rusqlite::Error),
150    /// The query progress handler could not be installed.
151    #[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.";