mcp-server-sqlite 1.0.0

//! MCP server implementation for SQLite.

use std::sync::Arc;

use rmcp::{
    ServerHandler,
    handler::server::tool::ToolRouter,
    model::{ServerCapabilities, ServerInfo},
    tool_handler,
};

use crate::{
    access_control::AuthorizationResolver,
    tools::{
        backup_tool::BackupTool, create_fts_index_tool::CreateFtsIndexTool,
        database_info_tool::DatabaseInfoTool,
        describe_table_tool::DescribeTableTool, execute_tool::ExecuteTool,
        explain_query_tool::ExplainQueryTool,
        list_foreign_keys_tool::ListForeignKeysTool,
        list_indexes_tool::ListIndexesTool, list_tables_tool::ListTablesTool,
        list_triggers_tool::ListTriggersTool, list_views_tool::ListViewsTool,
        search_fts_tool::SearchFtsTool, vacuum_tool::VacuumTool,
    },
    traits::SqliteServerTool,
};

/// The MCP server instance that handles tool calls against a SQLite database
/// with access control.
pub struct McpServerSqlite {
    /// The connection pool for the SQLite database.
    pool: r2d2::Pool<r2d2_sqlite::SqliteConnectionManager>,
    /// The resolver that evaluates access control rules for each SQL operation.
    /// Wrapped in an `Arc` so it can be moved into the per-query authorizer
    /// closure required by rusqlite.
    authorization_resolver: Arc<AuthorizationResolver>,
    /// Optional query timeout. When `Some`, each connection's progress handler
    /// aborts operations that exceed this duration.
    query_timeout: Option<std::time::Duration>,
    /// The router that dispatches tool calls to their handlers.
    tool_router: ToolRouter<Self>,
}

impl McpServerSqlite {
    /// Creates a new MCP server backed by the given connection pool and access
    /// control resolver.
    pub fn new(
        pool: r2d2::Pool<r2d2_sqlite::SqliteConnectionManager>,
        authorization_resolver: AuthorizationResolver,
        query_timeout: Option<std::time::Duration>,
    ) -> Self {
        let tool_router = Self::tool_router();
        Self {
            pool,
            authorization_resolver: Arc::new(authorization_resolver),
            query_timeout,
            tool_router,
        }
    }

    /// Obtains a pooled connection from the database pool, installs the access
    /// control authorizer, and optionally sets up the query timeout via a
    /// progress handler.
    pub fn connection(
        &self,
    ) -> Result<
        r2d2::PooledConnection<r2d2_sqlite::SqliteConnectionManager>,
        ConnectionError,
    > {
        tracing::debug!("checking out pooled connection");
        let conn = self.pool.get().map_err(|source| {
            tracing::error!(%source, "pool checkout failed");
            ConnectionError::Pool(source)
        })?;
        let resolver = self.authorization_resolver.clone();
        conn.authorizer(Some(move |ctx: rusqlite::hooks::AuthContext<'_>| {
            resolver.authorization(ctx)
        }))
        .map_err(|source| {
            tracing::error!(%source, "failed to install authorizer");
            ConnectionError::Authorizer(source)
        })?;

        if let Some(timeout) = self.query_timeout {
            let deadline = std::time::Instant::now() + timeout;
            conn.progress_handler(
                1000,
                Some(move || std::time::Instant::now() > deadline),
            )
            .map_err(|source| {
                tracing::error!(%source, "failed to install progress handler");
                ConnectionError::ProgressHandler(source)
            })?;
        }

        Ok(conn)
    }

    pub fn tool_router() -> ToolRouter<Self> {
        ToolRouter::<Self>::new()
            .with_route((ExecuteTool::tool(), ExecuteTool::handler_func()))
            .with_route((
                ListTablesTool::tool(),
                ListTablesTool::handler_func(),
            ))
            .with_route((
                DescribeTableTool::tool(),
                DescribeTableTool::handler_func(),
            ))
            .with_route((
                ListIndexesTool::tool(),
                ListIndexesTool::handler_func(),
            ))
            .with_route((
                ListForeignKeysTool::tool(),
                ListForeignKeysTool::handler_func(),
            ))
            .with_route((ListViewsTool::tool(), ListViewsTool::handler_func()))
            .with_route((
                ListTriggersTool::tool(),
                ListTriggersTool::handler_func(),
            ))
            .with_route((
                ExplainQueryTool::tool(),
                ExplainQueryTool::handler_func(),
            ))
            .with_route((BackupTool::tool(), BackupTool::handler_func()))
            .with_route((
                CreateFtsIndexTool::tool(),
                CreateFtsIndexTool::handler_func(),
            ))
            .with_route((SearchFtsTool::tool(), SearchFtsTool::handler_func()))
            .with_route((VacuumTool::tool(), VacuumTool::handler_func()))
            .with_route((
                DatabaseInfoTool::tool(),
                DatabaseInfoTool::handler_func(),
            ))
    }
}

/// Errors that can occur when checking out and configuring a pooled database
/// connection.
#[derive(Debug, thiserror::Error)]
pub enum ConnectionError {
    /// The connection pool failed to provide a connection.
    #[error("pool checkout failed: {0}")]
    Pool(r2d2::Error),
    /// The SQLite authorizer could not be installed.
    #[error("failed to install authorizer: {0}")]
    Authorizer(rusqlite::Error),
    /// The query progress handler could not be installed.
    #[error("failed to install progress handler: {0}")]
    ProgressHandler(rusqlite::Error),
}

#[tool_handler]
impl ServerHandler for McpServerSqlite {
    fn get_info(&self) -> ServerInfo {
        let mut info = ServerInfo::default();
        info.server_info.name = env!("CARGO_PKG_NAME").into();
        info.server_info.version = env!("CARGO_PKG_VERSION").into();
        info.instructions = Some(INSTRUCTIONS.into());
        info.capabilities =
            ServerCapabilities::builder().enable_tools().build();
        info
    }
}

const INSTRUCTIONS: &str = "\
You are connected to an MCP SQLite server with fine-grained access control.

Available tools:
- execute: Run any SQL query (SELECT, INSERT, UPDATE, DELETE, DDL)
- list_tables: Show all tables with their CREATE TABLE schemas
- describe_table: Show columns, types, constraints, and defaults for a table
- list_indexes: Show indexes with columns, uniqueness, and partial predicates
- list_foreign_keys: Show foreign key constraints for a table
- list_views: Show all views with their defining SQL
- list_triggers: Show triggers with event, timing, and SQL
- explain_query: Show the EXPLAIN QUERY PLAN for a SQL statement
- create_fts_index: Create a full-text search index on a table
- search_fts: Search a full-text index with BM25 ranking and snippets
- backup: Back up the database to a file
- vacuum: Reclaim unused space and defragment the database
- database_info: Show database metadata (size, page count, journal mode, etc.)

Start by calling list_tables to discover the schema, then use describe_table \
for details on specific tables. Use explain_query to optimize slow queries. \
All operations respect the server's access control policy.";