database_mcp_sqlite/tools/
read_query.rs

1//! MCP tool: `read_query`.
2
3use std::borrow::Cow;
4
5use database_mcp_server::AppError;
6use database_mcp_server::types::QueryResponse;
7use database_mcp_sql::timeout::execute_with_timeout;
8use database_mcp_sql::validation::validate_read_only_with_dialect;
9use rmcp::handler::server::router::tool::{AsyncTool, ToolBase};
10use rmcp::model::{ErrorData, ToolAnnotations};
11use serde_json::Value;
12use sqlx::sqlite::SqliteRow;
13use sqlx_to_json::RowExt;
14
15use crate::SqliteHandler;
16use crate::types::QueryRequest;
17
18/// Marker type for the `read_query` MCP tool.
19pub(crate) struct ReadQueryTool;
20
21impl ReadQueryTool {
22    const NAME: &'static str = "read_query";
23    const DESCRIPTION: &'static str = "Execute a read-only SQL query (SELECT, SHOW, DESCRIBE, USE, EXPLAIN).";
24}
25
26impl ToolBase for ReadQueryTool {
27    type Parameter = QueryRequest;
28    type Output = QueryResponse;
29    type Error = ErrorData;
30
31    fn name() -> Cow<'static, str> {
32        Self::NAME.into()
33    }
34
35    fn description() -> Option<Cow<'static, str>> {
36        Some(Self::DESCRIPTION.into())
37    }
38
39    fn annotations() -> Option<ToolAnnotations> {
40        Some(
41            ToolAnnotations::new()
42                .read_only(true)
43                .destructive(false)
44                .idempotent(true)
45                .open_world(true),
46        )
47    }
48}
49
50impl AsyncTool<SqliteHandler> for ReadQueryTool {
51    async fn invoke(handler: &SqliteHandler, params: Self::Parameter) -> Result<Self::Output, Self::Error> {
52        Ok(handler.read_query(&params).await?)
53    }
54}
55
56impl SqliteHandler {
57    /// Executes a read-only SQL query.
58    ///
59    /// Validates that the query is read-only before executing.
60    ///
61    /// # Errors
62    ///
63    /// Returns [`AppError::ReadOnlyViolation`] if the query is not
64    /// read-only, or [`AppError::Query`] if the backend reports an error.
65    pub async fn read_query(&self, request: &QueryRequest) -> Result<QueryResponse, AppError> {
66        validate_read_only_with_dialect(&request.query, &sqlparser::dialect::SQLiteDialect {})?;
67        let pool = self.pool.clone();
68        let rows: Vec<SqliteRow> = execute_with_timeout(
69            self.config.query_timeout,
70            &request.query,
71            sqlx::query(&request.query).fetch_all(&pool),
72        )
73        .await?;
74        Ok(QueryResponse {
75            rows: Value::Array(rows.iter().map(RowExt::to_json).collect()),
76        })
77    }
78}
79
80#[cfg(test)]
81mod tests {
82    use serde_json::Value;
83    use sqlx::SqlitePool;
84    use sqlx::sqlite::SqlitePoolOptions;
85    use sqlx::sqlite::SqliteRow;
86    use sqlx_to_json::RowExt;
87
88    async fn mem_pool() -> SqlitePool {
89        SqlitePoolOptions::new()
90            .max_connections(1)
91            .connect("sqlite::memory:")
92            .await
93            .expect("in-memory SQLite")
94    }
95
96    async fn query_json(pool: &SqlitePool, sql: &str) -> Value {
97        let rows: Vec<SqliteRow> = sqlx::query(sql).fetch_all(pool).await.expect("query failed");
98        Value::Array(rows.iter().map(RowExt::to_json).collect())
99    }
100
101    #[tokio::test]
102    async fn rows_to_json_array_empty_result() {
103        let pool = mem_pool().await;
104        sqlx::query("CREATE TABLE t (v INTEGER)").execute(&pool).await.unwrap();
105
106        let rows = query_json(&pool, "SELECT v FROM t").await;
107        assert_eq!(rows, Value::Array(vec![]));
108    }
109
110    #[tokio::test]
111    async fn rows_to_json_array_multiple_rows() {
112        let pool = mem_pool().await;
113        sqlx::query("CREATE TABLE t (id INTEGER, name TEXT, score REAL)")
114            .execute(&pool)
115            .await
116            .unwrap();
117        sqlx::query("INSERT INTO t VALUES (1, 'alice', 9.5), (2, 'bob', 8.0)")
118            .execute(&pool)
119            .await
120            .unwrap();
121
122        let rows = query_json(&pool, "SELECT id, name, score FROM t ORDER BY id").await;
123        assert_eq!(rows.as_array().expect("should be array").len(), 2);
124
125        assert_eq!(rows[0]["id"], Value::Number(1.into()));
126        assert_eq!(rows[0]["name"], Value::String("alice".into()));
127        assert!(rows[0]["score"].is_number());
128
129        assert_eq!(rows[1]["id"], Value::Number(2.into()));
130        assert_eq!(rows[1]["name"], Value::String("bob".into()));
131    }
132}
database_mcp_sqlite/tools/read_query.rs

database_mcp_sqlite/tools/
read_query.rs
