mcp_sql/
server.rs

1use std::sync::Arc;
2use std::time::Duration;
3
4use rmcp::handler::server::router::tool::ToolRouter;
5use rmcp::handler::server::wrapper::Parameters;
6use rmcp::model::*;
7use rmcp::{schemars, tool, tool_handler, tool_router, ServerHandler};
8use serde::Deserialize;
9
10use crate::db::convert::row_to_json;
11use crate::db::dialect;
12use crate::db::{DatabaseManager, DbBackend};
13use crate::error::McpSqlError;
14
15#[derive(Clone)]
16pub struct McpSqlServer {
17    db: Arc<DatabaseManager>,
18    allow_write: bool,
19    row_limit: u32,
20    query_timeout: Duration,
21    tool_router: ToolRouter<Self>,
22}
23
24// -- Tool parameter types --
25
26#[derive(Debug, Deserialize, schemars::JsonSchema)]
27pub struct DatabaseParam {
28    #[schemars(description = "Database name (optional if only one database is connected)")]
29    #[serde(default)]
30    pub database: Option<String>,
31}
32
33#[derive(Debug, Deserialize, schemars::JsonSchema)]
34pub struct DescribeTableParams {
35    #[schemars(description = "Table name to describe (use schema.table for PostgreSQL)")]
36    pub table: String,
37
38    #[schemars(description = "Database name (optional if only one database is connected)")]
39    #[serde(default)]
40    pub database: Option<String>,
41}
42
43#[derive(Debug, Deserialize, schemars::JsonSchema)]
44pub struct SampleDataParams {
45    #[schemars(description = "Table name to sample rows from")]
46    pub table: String,
47
48    #[schemars(description = "Database name (optional if only one database is connected)")]
49    #[serde(default)]
50    pub database: Option<String>,
51
52    #[schemars(description = "Number of sample rows to return (default: 5)")]
53    #[serde(default)]
54    pub limit: Option<u32>,
55}
56
57#[derive(Debug, Deserialize, schemars::JsonSchema)]
58pub struct ShowCreateTableParams {
59    #[schemars(description = "Table name")]
60    pub table: String,
61
62    #[schemars(description = "Database name (optional if only one database is connected)")]
63    #[serde(default)]
64    pub database: Option<String>,
65}
66
67#[derive(Debug, Deserialize, schemars::JsonSchema)]
68pub struct ListIndexesParams {
69    #[schemars(description = "Table name")]
70    pub table: String,
71
72    #[schemars(description = "Database name (optional if only one database is connected)")]
73    #[serde(default)]
74    pub database: Option<String>,
75}
76
77#[derive(Debug, Deserialize, schemars::JsonSchema)]
78pub struct QueryParams {
79    #[schemars(description = "SQL query to execute")]
80    pub sql: String,
81
82    #[schemars(description = "Database name (optional if only one database is connected)")]
83    #[serde(default)]
84    pub database: Option<String>,
85}
86
87impl McpSqlServer {
88    pub fn new(db: DatabaseManager, allow_write: bool, row_limit: u32, query_timeout_secs: u64) -> Self {
89        Self {
90            db: Arc::new(db),
91            allow_write,
92            row_limit,
93            query_timeout: Duration::from_secs(query_timeout_secs),
94            tool_router: Self::tool_router(),
95        }
96    }
97
98    fn err(&self, e: McpSqlError) -> ErrorData {
99        e.to_mcp_error()
100    }
101}
102
103#[tool_router]
104impl McpSqlServer {
105    #[tool(
106        name = "list_databases",
107        description = "List all connected databases with their names and types (postgres/sqlite/mysql)"
108    )]
109    async fn list_databases(&self) -> Result<CallToolResult, ErrorData> {
110        let databases: Vec<serde_json::Value> = self
111            .db
112            .databases
113            .iter()
114            .map(|d| {
115                serde_json::json!({
116                    "name": d.name,
117                    "type": d.backend.name(),
118                    "url": d.url_redacted,
119                })
120            })
121            .collect();
122
123        let text = serde_json::to_string_pretty(&databases)
124            .unwrap_or_else(|_| "[]".to_string());
125        Ok(CallToolResult::success(vec![Content::text(text)]))
126    }
127
128    #[tool(
129        name = "list_tables",
130        description = "List all tables in a database with approximate row counts"
131    )]
132    async fn list_tables(
133        &self,
134        Parameters(params): Parameters<DatabaseParam>,
135    ) -> Result<CallToolResult, ErrorData> {
136        let entry = self.db.resolve(params.database.as_deref()).map_err(|e| self.err(e))?;
137        let tables = dialect::list_tables(&entry.pool, entry.backend)
138            .await
139            .map_err(|e| self.err(e))?;
140
141        let text = serde_json::to_string_pretty(&tables)
142            .unwrap_or_else(|_| "[]".to_string());
143        Ok(CallToolResult::success(vec![Content::text(text)]))
144    }
145
146    #[tool(
147        name = "describe_table",
148        description = "Describe a table's columns with name, type, nullable, default, and primary key info"
149    )]
150    async fn describe_table(
151        &self,
152        Parameters(params): Parameters<DescribeTableParams>,
153    ) -> Result<CallToolResult, ErrorData> {
154        let entry = self.db.resolve(params.database.as_deref()).map_err(|e| self.err(e))?;
155        let columns = dialect::describe_table(&entry.pool, entry.backend, &params.table)
156            .await
157            .map_err(|e| self.err(e))?;
158
159        let text = serde_json::to_string_pretty(&columns)
160            .unwrap_or_else(|_| "[]".to_string());
161        Ok(CallToolResult::success(vec![Content::text(text)]))
162    }
163
164    #[tool(
165        name = "query",
166        description = "Execute a SQL query and return results as JSON. Read-only by default (SELECT/WITH/SHOW/PRAGMA only). Use --allow-write flag to enable write operations."
167    )]
168    async fn query(
169        &self,
170        Parameters(params): Parameters<QueryParams>,
171    ) -> Result<CallToolResult, ErrorData> {
172        let entry = self.db.resolve(params.database.as_deref()).map_err(|e| self.err(e))?;
173        let sql = params.sql.trim();
174
175        // Read-only guard
176        if !self.allow_write {
177            check_read_only(sql).map_err(|e| self.err(e))?;
178        }
179
180        // Set transaction read only for backends that support it
181        if !self.allow_write && entry.backend != DbBackend::Sqlite {
182            let read_only_sql = match entry.backend {
183                DbBackend::Postgres => "SET TRANSACTION READ ONLY",
184                DbBackend::Mysql => "SET TRANSACTION READ ONLY",
185                DbBackend::Sqlite => unreachable!(),
186            };
187            // Best effort — some connection states may not support this
188            let _ = sqlx::query(read_only_sql).execute(&entry.pool).await;
189        }
190
191        // Inject LIMIT if not present
192        let limited_sql = inject_limit(sql, self.row_limit);
193
194        let rows = tokio::time::timeout(
195            self.query_timeout,
196            sqlx::query(&limited_sql).fetch_all(&entry.pool),
197        )
198        .await
199        .map_err(|_| self.err(McpSqlError::QueryTimeout(self.query_timeout.as_secs())))?
200        .map_err(|e| self.err(McpSqlError::Database(e)))?;
201
202        let results: Vec<serde_json::Value> = rows.iter().map(row_to_json).collect();
203        let text = serde_json::to_string_pretty(&serde_json::json!({
204            "rows": results,
205            "count": results.len(),
206        }))
207        .unwrap_or_else(|_| "{}".to_string());
208
209        Ok(CallToolResult::success(vec![Content::text(text)]))
210    }
211
212    #[tool(
213        name = "explain",
214        description = "Show the query execution plan for a SQL statement. Uses the appropriate EXPLAIN syntax for the database type."
215    )]
216    async fn explain(
217        &self,
218        Parameters(params): Parameters<QueryParams>,
219    ) -> Result<CallToolResult, ErrorData> {
220        let entry = self.db.resolve(params.database.as_deref()).map_err(|e| self.err(e))?;
221        let prefix = dialect::explain_prefix(entry.backend);
222        let explain_sql = format!("{}{}", prefix, params.sql.trim());
223
224        let rows = tokio::time::timeout(
225            self.query_timeout,
226            sqlx::query(&explain_sql).fetch_all(&entry.pool),
227        )
228        .await
229        .map_err(|_| self.err(McpSqlError::QueryTimeout(self.query_timeout.as_secs())))?
230        .map_err(|e| self.err(McpSqlError::Database(e)))?;
231
232        let results: Vec<serde_json::Value> = rows.iter().map(row_to_json).collect();
233        let text = serde_json::to_string_pretty(&results)
234            .unwrap_or_else(|_| "[]".to_string());
235
236        Ok(CallToolResult::success(vec![Content::text(text)]))
237    }
238
239    #[tool(
240        name = "sample_data",
241        description = "Return sample rows from a table as JSON. Useful for previewing table contents without writing SQL."
242    )]
243    async fn sample_data(
244        &self,
245        Parameters(params): Parameters<SampleDataParams>,
246    ) -> Result<CallToolResult, ErrorData> {
247        let entry = self.db.resolve(params.database.as_deref()).map_err(|e| self.err(e))?;
248        let limit = params.limit.unwrap_or(5);
249
250        let rows = tokio::time::timeout(
251            self.query_timeout,
252            dialect::sample_data(&entry.pool, entry.backend, &params.table, limit),
253        )
254        .await
255        .map_err(|_| self.err(McpSqlError::QueryTimeout(self.query_timeout.as_secs())))?
256        .map_err(|e| self.err(e))?;
257
258        let text = serde_json::to_string_pretty(&serde_json::json!({
259            "table": params.table,
260            "rows": rows,
261            "count": rows.len(),
262        }))
263        .unwrap_or_else(|_| "{}".to_string());
264
265        Ok(CallToolResult::success(vec![Content::text(text)]))
266    }
267
268    #[tool(
269        name = "show_create_table",
270        description = "Show the CREATE TABLE DDL statement for a table"
271    )]
272    async fn show_create_table(
273        &self,
274        Parameters(params): Parameters<ShowCreateTableParams>,
275    ) -> Result<CallToolResult, ErrorData> {
276        let entry = self.db.resolve(params.database.as_deref()).map_err(|e| self.err(e))?;
277        let ddl = dialect::show_create_table(&entry.pool, entry.backend, &params.table)
278            .await
279            .map_err(|e| self.err(e))?;
280        Ok(CallToolResult::success(vec![Content::text(ddl)]))
281    }
282
283    #[tool(
284        name = "show_schema",
285        description = "Show a Mermaid ER diagram of all tables and their relationships in the database"
286    )]
287    async fn show_schema(
288        &self,
289        Parameters(params): Parameters<DatabaseParam>,
290    ) -> Result<CallToolResult, ErrorData> {
291        let entry = self.db.resolve(params.database.as_deref()).map_err(|e| self.err(e))?;
292        let diagram = crate::schema::generate_mermaid_er(&entry.pool, entry.backend)
293            .await
294            .map_err(|e| self.err(e))?;
295        Ok(CallToolResult::success(vec![Content::text(diagram)]))
296    }
297
298    #[tool(
299        name = "list_indexes",
300        description = "List all indexes on a table with column names and uniqueness"
301    )]
302    async fn list_indexes(
303        &self,
304        Parameters(params): Parameters<ListIndexesParams>,
305    ) -> Result<CallToolResult, ErrorData> {
306        let entry = self.db.resolve(params.database.as_deref()).map_err(|e| self.err(e))?;
307        let indexes = dialect::list_indexes(&entry.pool, entry.backend, &params.table)
308            .await
309            .map_err(|e| self.err(e))?;
310        let json = serde_json::to_string_pretty(&indexes).unwrap_or_default();
311        Ok(CallToolResult::success(vec![Content::text(json)]))
312    }
313
314    #[tool(
315        name = "query_dry_run",
316        description = "Validate a SQL query without executing it. Returns the query plan and any warnings."
317    )]
318    async fn query_dry_run(
319        &self,
320        Parameters(params): Parameters<QueryParams>,
321    ) -> Result<CallToolResult, ErrorData> {
322        let entry = self.db.resolve(params.database.as_deref()).map_err(|e| self.err(e))?;
323
324        // Use EXPLAIN to validate without executing
325        let explain_sql = format!(
326            "{}{}",
327            dialect::explain_prefix(entry.backend),
328            params.sql
329        );
330
331        match sqlx::query(&explain_sql).fetch_all(&entry.pool).await {
332            Ok(rows) => {
333                let plan: Vec<serde_json::Value> = rows.iter().map(row_to_json).collect();
334                let result = serde_json::json!({
335                    "valid": true,
336                    "query_plan": plan,
337                });
338                Ok(CallToolResult::success(vec![Content::text(
339                    serde_json::to_string_pretty(&result).unwrap_or_default(),
340                )]))
341            }
342            Err(e) => {
343                let result = serde_json::json!({
344                    "valid": false,
345                    "error": e.to_string(),
346                });
347                Ok(CallToolResult::success(vec![Content::text(
348                    serde_json::to_string_pretty(&result).unwrap_or_default(),
349                )]))
350            }
351        }
352    }
353}
354
355#[tool_handler]
356impl ServerHandler for McpSqlServer {
357    fn get_info(&self) -> ServerInfo {
358        ServerInfo {
359            protocol_version: ProtocolVersion::V_2024_11_05,
360            capabilities: ServerCapabilities::builder().enable_tools().build(),
361            server_info: Implementation {
362                name: "mcp-sql".to_string(),
363                version: env!("CARGO_PKG_VERSION").to_string(),
364                ..Default::default()
365            },
366            instructions: Some(
367                "SQL database server. Use list_databases to see connected databases, \
368                 list_tables to see tables, describe_table for schema details (includes foreign keys), \
369                 show_create_table for DDL statements, show_schema for a Mermaid ER diagram, \
370                 list_indexes for index details, sample_data to preview table contents, \
371                 query to run SQL, explain for query plans, and query_dry_run to validate SQL without executing."
372                    .to_string(),
373            ),
374        }
375    }
376}
377
378/// Heuristic check: only allow SELECT, WITH, SHOW, PRAGMA, EXPLAIN.
379fn check_read_only(sql: &str) -> Result<(), McpSqlError> {
380    let upper = sql.trim_start().to_uppercase();
381    let allowed_prefixes = ["SELECT", "WITH", "SHOW", "PRAGMA", "EXPLAIN"];
382    if allowed_prefixes.iter().any(|p| upper.starts_with(p)) {
383        Ok(())
384    } else {
385        Err(McpSqlError::ReadOnly(
386            "Only SELECT/WITH/SHOW/PRAGMA/EXPLAIN queries are allowed in read-only mode. \
387             Start the server with --allow-write to enable write operations."
388                .to_string(),
389        ))
390    }
391}
392
393/// Inject a LIMIT clause if the query doesn't already have one.
394fn inject_limit(sql: &str, limit: u32) -> String {
395    let upper = sql.to_uppercase();
396    // Don't inject LIMIT for non-SELECT statements or if LIMIT already present
397    if !upper.trim_start().starts_with("SELECT") && !upper.trim_start().starts_with("WITH") {
398        return sql.to_string();
399    }
400    if upper.contains(" LIMIT ") {
401        return sql.to_string();
402    }
403    // Strip trailing semicolon if present
404    let trimmed = sql.trim_end().trim_end_matches(';');
405    format!("{trimmed} LIMIT {limit}")
406}
407
408#[cfg(test)]
409mod tests {
410    use super::*;
411
412    #[test]
413    fn test_check_read_only() {
414        assert!(check_read_only("SELECT * FROM users").is_ok());
415        assert!(check_read_only("  select * from users").is_ok());
416        assert!(check_read_only("WITH cte AS (SELECT 1) SELECT * FROM cte").is_ok());
417        assert!(check_read_only("SHOW TABLES").is_ok());
418        assert!(check_read_only("PRAGMA table_info(users)").is_ok());
419        assert!(check_read_only("EXPLAIN SELECT * FROM users").is_ok());
420
421        assert!(check_read_only("INSERT INTO users VALUES (1)").is_err());
422        assert!(check_read_only("UPDATE users SET name = 'x'").is_err());
423        assert!(check_read_only("DELETE FROM users").is_err());
424        assert!(check_read_only("DROP TABLE users").is_err());
425        assert!(check_read_only("CREATE TABLE t (id INT)").is_err());
426    }
427
428    #[test]
429    fn test_inject_limit() {
430        assert_eq!(
431            inject_limit("SELECT * FROM users", 100),
432            "SELECT * FROM users LIMIT 100"
433        );
434        assert_eq!(
435            inject_limit("SELECT * FROM users;", 100),
436            "SELECT * FROM users LIMIT 100"
437        );
438        assert_eq!(
439            inject_limit("SELECT * FROM users LIMIT 10", 100),
440            "SELECT * FROM users LIMIT 10"
441        );
442        assert_eq!(
443            inject_limit("INSERT INTO users VALUES (1)", 100),
444            "INSERT INTO users VALUES (1)"
445        );
446        assert_eq!(
447            inject_limit("WITH cte AS (SELECT 1) SELECT * FROM cte", 50),
448            "WITH cte AS (SELECT 1) SELECT * FROM cte LIMIT 50"
449        );
450    }
451}
mcp_sql/server.rs

mcp_sql/
server.rs
