mcp-server-sqlite 1.0.0

//! The `execute` tool: runs a SQL query against the database and returns typed
//! results. This is the primary tool exposed by the MCP server.

use rmcp::model::{Content, IntoContents};
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};

use super::ToolError;
use crate::{mcp::McpServerSqlite, traits::SqliteServerTool};

#[derive(
    Clone,
    Copy,
    Debug,
    PartialEq,
    Eq,
    PartialOrd,
    Ord,
    Hash,
    Default,
    Serialize,
    Deserialize,
    JsonSchema,
)]
/// Marker type for the `execute` tool. Implements `SqliteServerTool` to wire up
/// the tool's schema, handler, and routing. Stateless — all context comes from
/// `McpServerSqlite`.
pub struct ExecuteTool;

impl SqliteServerTool for ExecuteTool {
    const NAME: &str = "execute";

    type Context = McpServerSqlite;
    type Error = ToolError<ExecuteError>;

    type Input = ExecuteInput;
    type Output = ExecuteOutput;

    fn handle(
        ctx: &Self::Context,
        input: Self::Input,
    ) -> Result<Self::Output, Self::Error> {
        let conn = ctx
            .connection()
            .map_err(|source| ToolError::Connection { source })?;

        let mut stmt = conn.prepare(&input.query).map_err(|source| {
            if matches!(
                source,
                rusqlite::Error::SqliteFailure(
                    rusqlite::ffi::Error {
                        code: rusqlite::ffi::ErrorCode::AuthorizationForStatementDenied,
                        ..
                    },
                    _,
                )
            ) {
                ToolError::AccessDenied {
                    message: format!(
                        "the configured access control policy denied this \
                        statement: {}",
                        input.query,
                    ),
                }
            } else {
                ToolError::Tool(ExecuteError::Prepare { source })
            }
        })?;

        let column_count = stmt.column_count();
        let column_names = (0..column_count)
            .map(|i| stmt.column_name(i).unwrap_or("?").to_owned())
            .collect::<Vec<String>>();

        let rows = stmt
            .query_map([], |row: &rusqlite::Row<'_>| {
                let columns = column_names
                    .iter()
                    .enumerate()
                    .map(|(i, name)| {
                        let value = row
                            .get::<_, rusqlite::types::Value>(i)
                            .map(Value::from)
                            .unwrap_or(Value::Null);
                        (name.clone(), value)
                    })
                    .collect();
                Ok(Row { columns })
            })
            .map_err(|source| ToolError::Tool(ExecuteError::Query { source }))?
            .collect::<Result<Vec<_>, _>>()
            .map_err(|source| {
                ToolError::Tool(ExecuteError::Query { source })
            })?;

        let rows_changed = conn.changes();

        Ok(ExecuteOutput { rows, rows_changed })
    }
}

/// The input parameters for the `execute` tool.
#[derive(
    Clone,
    Debug,
    PartialEq,
    Eq,
    PartialOrd,
    Ord,
    Hash,
    Serialize,
    Deserialize,
    schemars::JsonSchema,
)]
pub struct ExecuteInput {
    /// The SQL query to execute against the database.
    #[schemars(description = "The SQL query to execute")]
    pub query: String,
}

/// The result of executing a SQL query.
#[derive(
    Clone,
    Debug,
    PartialEq,
    PartialOrd,
    Serialize,
    Deserialize,
    schemars::JsonSchema,
)]
pub struct ExecuteOutput {
    /// The rows returned by the query. Empty for statements that do not produce
    /// output (e.g. INSERT, CREATE TABLE).
    pub rows: Vec<Row>,
    /// The number of rows changed by the statement. Zero for queries that only
    /// read data.
    pub rows_changed: u64,
}

/// A single row returned from a query, represented as a mapping of column names
/// to their typed values.
#[derive(
    Clone,
    Debug,
    PartialEq,
    PartialOrd,
    Serialize,
    Deserialize,
    schemars::JsonSchema,
)]
pub struct Row {
    /// The column-value pairs for this row. Each key is the column name and
    /// each value preserves the original SQLite type.
    pub columns: std::collections::BTreeMap<String, Value>,
}

/// A dynamically-typed SQLite value preserving the original column type.
/// Serialized as a tagged enum with `kind` and `value` fields so that consumers
/// can distinguish between types unambiguously.
#[serde_with::serde_as]
#[derive(
    Clone,
    Debug,
    PartialEq,
    PartialOrd,
    Serialize,
    Deserialize,
    schemars::JsonSchema,
)]
#[serde(tag = "kind", content = "value")]
pub enum Value {
    /// A SQL NULL.
    Null,
    /// A 64-bit signed integer.
    Integer(i64),
    /// A 64-bit IEEE 754 floating-point number.
    Real(f64),
    /// A UTF-8 string.
    Text(String),
    /// Raw binary data, hex-encoded for JSON transport.
    Blob(
        #[serde_as(as = "serde_with::hex::Hex")]
        #[schemars(with = "String")]
        Vec<u8>,
    ),
}

/// Converts a rusqlite `Value` into this crate's `Value`, preserving the type
/// tag for each SQLite storage class.
impl From<rusqlite::types::Value> for Value {
    fn from(value: rusqlite::types::Value) -> Self {
        match value {
            rusqlite::types::Value::Null => Self::Null,
            rusqlite::types::Value::Integer(n) => Self::Integer(n),
            rusqlite::types::Value::Real(f) => Self::Real(f),
            rusqlite::types::Value::Text(s) => Self::Text(s),
            rusqlite::types::Value::Blob(b) => Self::Blob(b),
        }
    }
}

/// Errors specific to the `execute` tool's query preparation and result reading
/// logic.
#[derive(Debug, thiserror::Error)]
pub enum ExecuteError {
    /// SQLite failed to prepare the statement. This usually means the SQL is
    /// syntactically invalid or references a table/column that does not exist.
    #[error("failed to prepare statement: {source}")]
    Prepare {
        /// The underlying rusqlite error.
        source: rusqlite::Error,
    },
    /// The query executed but an error occurred while reading a result row.
    #[error("failed to read query results: {source}")]
    Query {
        /// The underlying rusqlite error.
        source: rusqlite::Error,
    },
}

/// Converts the execute-specific error into MCP content by rendering the
/// display string as text.
impl IntoContents for ExecuteError {
    fn into_contents(self) -> Vec<Content> {
        vec![Content::text(self.to_string())]
    }
}