mcp_sql/db/
dialect.rs

1use serde_json::Value;
2use sqlx::{AnyPool, Row};
3
4use crate::db::convert::row_to_json;
5use crate::db::DbBackend;
6use crate::error::McpSqlError;
7
8/// List tables with approximate row counts.
9pub async fn list_tables(pool: &AnyPool, backend: DbBackend) -> Result<Vec<Value>, McpSqlError> {
10    let sql = match backend {
11        DbBackend::Postgres => {
12            "SELECT schemaname || '.' || tablename AS table_name, \
13                    COALESCE(n_live_tup, 0) AS row_count \
14             FROM pg_tables \
15             LEFT JOIN pg_stat_user_tables ON tablename = relname AND schemaname = pg_stat_user_tables.schemaname \
16             WHERE pg_tables.schemaname NOT IN ('pg_catalog', 'information_schema') \
17             ORDER BY table_name"
18        }
19        DbBackend::Sqlite => {
20            // Get table names first
21            let name_rows = sqlx::query(
22                "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%' ORDER BY name",
23            )
24            .fetch_all(pool)
25            .await?;
26
27            let mut results = Vec::new();
28            for row in &name_rows {
29                let name: String = row.try_get("name")?;
30                // Count rows with a timeout — fall back to 0 for very large tables
31                let count: i64 = match tokio::time::timeout(
32                    std::time::Duration::from_secs(1),
33                    sqlx::query_as::<_, (i64,)>(&format!(
34                        "SELECT COUNT(*) FROM \"{}\"",
35                        name.replace('"', "\"\"")
36                    ))
37                    .fetch_one(pool),
38                )
39                .await
40                {
41                    Ok(Ok((c,))) => c,
42                    _ => 0,
43                };
44                results.push(serde_json::json!({
45                    "table_name": name,
46                    "row_count": count,
47                }));
48            }
49            return Ok(results);
50        }
51        DbBackend::Mysql => {
52            "SELECT table_name, table_rows AS row_count \
53             FROM information_schema.tables \
54             WHERE table_schema = DATABASE() \
55             ORDER BY table_name"
56        }
57    };
58
59    let rows = sqlx::query(sql).fetch_all(pool).await?;
60    Ok(rows.iter().map(row_to_json).collect())
61}
62
63/// Describe a table's columns.
64pub async fn describe_table(
65    pool: &AnyPool,
66    backend: DbBackend,
67    table: &str,
68) -> Result<Vec<Value>, McpSqlError> {
69    match backend {
70        DbBackend::Postgres => describe_table_postgres(pool, table).await,
71        DbBackend::Sqlite => describe_table_sqlite(pool, table).await,
72        DbBackend::Mysql => describe_table_mysql(pool, table).await,
73    }
74}
75
76async fn describe_table_postgres(pool: &AnyPool, table: &str) -> Result<Vec<Value>, McpSqlError> {
77    // Handle schema.table format
78    let (schema, tbl) = if let Some((s, t)) = table.split_once('.') {
79        (s, t)
80    } else {
81        ("public", table)
82    };
83
84    let sql = "SELECT c.column_name AS name, c.data_type AS type, \
85               c.is_nullable AS nullable, c.column_default AS default_value, \
86               CASE WHEN tc.constraint_type = 'PRIMARY KEY' THEN 'YES' ELSE 'NO' END AS primary_key \
87               FROM information_schema.columns c \
88               LEFT JOIN information_schema.key_column_usage kcu \
89                 ON c.table_schema = kcu.table_schema \
90                 AND c.table_name = kcu.table_name \
91                 AND c.column_name = kcu.column_name \
92               LEFT JOIN information_schema.table_constraints tc \
93                 ON kcu.constraint_name = tc.constraint_name \
94                 AND kcu.table_schema = tc.table_schema \
95                 AND tc.constraint_type = 'PRIMARY KEY' \
96               WHERE c.table_schema = $1 AND c.table_name = $2 \
97               ORDER BY c.ordinal_position";
98
99    let rows = sqlx::query(sql)
100        .bind(schema)
101        .bind(tbl)
102        .fetch_all(pool)
103        .await?;
104
105    if rows.is_empty() {
106        return Err(McpSqlError::Other(format!("Table '{table}' not found")));
107    }
108
109    // Fetch FK info
110    let fk_sql = "SELECT kcu.column_name, ccu.table_schema || '.' || ccu.table_name || '.' || ccu.column_name AS references_col \
111                   FROM information_schema.key_column_usage kcu \
112                   JOIN information_schema.referential_constraints rc \
113                     ON kcu.constraint_name = rc.constraint_name AND kcu.constraint_schema = rc.constraint_schema \
114                   JOIN information_schema.constraint_column_usage ccu \
115                     ON rc.unique_constraint_name = ccu.constraint_name AND rc.unique_constraint_schema = ccu.constraint_schema \
116                   WHERE kcu.table_schema = $1 AND kcu.table_name = $2";
117
118    let fk_rows = sqlx::query(fk_sql)
119        .bind(schema)
120        .bind(tbl)
121        .fetch_all(pool)
122        .await
123        .unwrap_or_default();
124
125    let fk_map: std::collections::HashMap<String, String> = fk_rows
126        .iter()
127        .filter_map(|r| {
128            let col: String = r.try_get("column_name").ok()?;
129            let refs: String = r.try_get("references_col").ok()?;
130            Some((col, refs))
131        })
132        .collect();
133
134    let mut result: Vec<Value> = rows.iter().map(row_to_json).collect();
135    for col in &mut result {
136        if let Value::Object(map) = col {
137            let col_name = map.get("name").and_then(|v| v.as_str()).unwrap_or("");
138            let fk = fk_map.get(col_name).map(|s| Value::String(s.clone())).unwrap_or(Value::Null);
139            map.insert("foreign_key".to_string(), fk);
140        }
141    }
142
143    Ok(result)
144}
145
146async fn describe_table_sqlite(pool: &AnyPool, table: &str) -> Result<Vec<Value>, McpSqlError> {
147    // SQLite PRAGMA doesn't support parameterized queries, so we validate the table name
148    let safe_table = sanitize_identifier(table)?;
149    let sql = format!("PRAGMA table_info(\"{}\")", safe_table);
150    let rows = sqlx::query(&sql).fetch_all(pool).await?;
151
152    if rows.is_empty() {
153        return Err(McpSqlError::Other(format!("Table '{table}' not found")));
154    }
155
156    // Fetch FK info via PRAGMA foreign_key_list
157    let fk_sql = format!("PRAGMA foreign_key_list(\"{}\")", safe_table);
158    let fk_rows = sqlx::query(&fk_sql).fetch_all(pool).await.unwrap_or_default();
159
160    let fk_map: std::collections::HashMap<String, String> = fk_rows
161        .iter()
162        .filter_map(|r| {
163            let from: String = r.try_get("from").ok()?;
164            let ref_table: String = r.try_get("table").ok()?;
165            let ref_col: String = r.try_get("to").ok()?;
166            Some((from, format!("{ref_table}.{ref_col}")))
167        })
168        .collect();
169
170    let mut result = Vec::new();
171    for row in &rows {
172        let name: String = row.try_get("name").unwrap_or_default();
173        let col_type: String = row.try_get("type").unwrap_or_default();
174        let notnull: i32 = row.try_get("notnull").unwrap_or(0);
175        let dflt_value: Option<String> = row.try_get("dflt_value").ok();
176        let pk: i32 = row.try_get("pk").unwrap_or(0);
177        let fk = fk_map.get(&name).cloned();
178
179        result.push(serde_json::json!({
180            "name": name,
181            "type": col_type,
182            "nullable": if notnull == 0 { "YES" } else { "NO" },
183            "default_value": dflt_value,
184            "primary_key": if pk > 0 { "YES" } else { "NO" },
185            "foreign_key": fk,
186        }));
187    }
188
189    Ok(result)
190}
191
192async fn describe_table_mysql(pool: &AnyPool, table: &str) -> Result<Vec<Value>, McpSqlError> {
193    let sql = "SELECT column_name AS name, column_type AS type, \
194               is_nullable AS nullable, column_default AS default_value, \
195               CASE WHEN column_key = 'PRI' THEN 'YES' ELSE 'NO' END AS primary_key \
196               FROM information_schema.columns \
197               WHERE table_schema = DATABASE() AND table_name = ? \
198               ORDER BY ordinal_position";
199
200    let rows = sqlx::query(sql).bind(table).fetch_all(pool).await?;
201
202    if rows.is_empty() {
203        return Err(McpSqlError::Other(format!("Table '{table}' not found")));
204    }
205
206    // Fetch FK info
207    let fk_sql = "SELECT column_name, CONCAT(referenced_table_name, '.', referenced_column_name) AS references_col \
208                   FROM information_schema.key_column_usage \
209                   WHERE table_schema = DATABASE() AND table_name = ? AND referenced_table_name IS NOT NULL";
210
211    let fk_rows = sqlx::query(fk_sql).bind(table).fetch_all(pool).await.unwrap_or_default();
212
213    let fk_map: std::collections::HashMap<String, String> = fk_rows
214        .iter()
215        .filter_map(|r| {
216            let col: String = r.try_get("column_name").ok()?;
217            let refs: String = r.try_get("references_col").ok()?;
218            Some((col, refs))
219        })
220        .collect();
221
222    let mut result: Vec<Value> = rows.iter().map(row_to_json).collect();
223    for col in &mut result {
224        if let Value::Object(map) = col {
225            let col_name = map.get("name").and_then(|v| v.as_str()).unwrap_or("");
226            let fk = fk_map.get(col_name).map(|s| Value::String(s.clone())).unwrap_or(Value::Null);
227            map.insert("foreign_key".to_string(), fk);
228        }
229    }
230
231    Ok(result)
232}
233
234/// Sample N rows from a table.
235pub async fn sample_data(
236    pool: &AnyPool,
237    backend: DbBackend,
238    table: &str,
239    limit: u32,
240) -> Result<Vec<Value>, McpSqlError> {
241    let safe_table = sanitize_identifier(table)?;
242    let sql = match backend {
243        DbBackend::Postgres => format!(
244            "SELECT * FROM \"{}\" TABLESAMPLE BERNOULLI (100) LIMIT {}",
245            safe_table, limit
246        ),
247        DbBackend::Sqlite => format!("SELECT * FROM \"{}\" LIMIT {}", safe_table, limit),
248        DbBackend::Mysql => format!(
249            "SELECT * FROM `{}` ORDER BY RAND() LIMIT {}",
250            safe_table, limit
251        ),
252    };
253
254    let rows = sqlx::query(&sql).fetch_all(pool).await?;
255    Ok(rows.iter().map(row_to_json).collect())
256}
257
258/// Get the correct EXPLAIN prefix for each backend.
259pub fn explain_prefix(backend: DbBackend) -> &'static str {
260    match backend {
261        DbBackend::Postgres => "EXPLAIN (FORMAT TEXT) ",
262        DbBackend::Sqlite => "EXPLAIN QUERY PLAN ",
263        DbBackend::Mysql => "EXPLAIN ",
264    }
265}
266
267/// Returns the CREATE TABLE DDL for a given table.
268pub async fn show_create_table(
269    pool: &AnyPool,
270    backend: DbBackend,
271    table: &str,
272) -> Result<String, McpSqlError> {
273    let table = sanitize_identifier(table)?;
274
275    match backend {
276        DbBackend::Sqlite => {
277            let sql =
278                format!("SELECT sql FROM sqlite_master WHERE type='table' AND name='{table}'");
279            let row = sqlx::query(&sql)
280                .fetch_optional(pool)
281                .await?
282                .ok_or_else(|| McpSqlError::Other(format!("table '{table}' not found")))?;
283            let ddl: String = row.try_get("sql")?;
284            Ok(ddl)
285        }
286        DbBackend::Mysql => {
287            let sql = format!("SHOW CREATE TABLE `{table}`");
288            let row = sqlx::query(&sql)
289                .fetch_one(pool)
290                .await
291                .map_err(|_| McpSqlError::Other(format!("table '{table}' not found")))?;
292            // MySQL returns two columns: "Table" and "Create Table"
293            let ddl: String = row.try_get(1)?;
294            Ok(ddl)
295        }
296        DbBackend::Postgres => {
297            // PostgreSQL has no built-in SHOW CREATE TABLE.
298            // Reconstruct from information_schema.
299            let rows = sqlx::query(
300                "SELECT column_name, data_type, is_nullable, column_default \
301                 FROM information_schema.columns \
302                 WHERE table_name = $1 \
303                 ORDER BY ordinal_position",
304            )
305            .bind(&table)
306            .fetch_all(pool)
307            .await?;
308
309            if rows.is_empty() {
310                return Err(McpSqlError::Other(format!("table '{table}' not found")));
311            }
312
313            let mut ddl = format!("CREATE TABLE {table} (\n");
314            for (i, row) in rows.iter().enumerate() {
315                let name: String = row.try_get("column_name")?;
316                let dtype: String = row.try_get("data_type")?;
317                let nullable: String = row.try_get("is_nullable")?;
318                let default: Option<String> = row.try_get("column_default").ok();
319
320                ddl.push_str(&format!("    {name} {}", dtype.to_uppercase()));
321                if nullable == "NO" {
322                    ddl.push_str(" NOT NULL");
323                }
324                if let Some(def) = default {
325                    ddl.push_str(&format!(" DEFAULT {def}"));
326                }
327                if i < rows.len() - 1 {
328                    ddl.push(',');
329                }
330                ddl.push('\n');
331            }
332            ddl.push_str(");");
333            Ok(ddl)
334        }
335    }
336}
337
338/// Returns index information for a table.
339pub async fn list_indexes(
340    pool: &AnyPool,
341    backend: DbBackend,
342    table: &str,
343) -> Result<Vec<Value>, McpSqlError> {
344    let table = sanitize_identifier(table)?;
345
346    match backend {
347        DbBackend::Sqlite => {
348            let index_rows = sqlx::query(&format!(
349                "PRAGMA index_list(\"{}\")",
350                table.replace('"', "\"\"")
351            ))
352            .fetch_all(pool)
353            .await?;
354
355            let mut indexes = Vec::new();
356            for row in &index_rows {
357                let name: String = row.try_get("name")?;
358                let unique: bool = row
359                    .try_get::<i32, _>("unique")
360                    .map(|v| v == 1)
361                    .unwrap_or(false);
362
363                let col_rows = sqlx::query(&format!(
364                    "PRAGMA index_info(\"{}\")",
365                    name.replace('"', "\"\"")
366                ))
367                .fetch_all(pool)
368                .await?;
369                let columns: Vec<String> = col_rows
370                    .iter()
371                    .filter_map(|r| r.try_get::<String, _>("name").ok())
372                    .collect();
373
374                indexes.push(serde_json::json!({
375                    "index_name": name,
376                    "columns": columns,
377                    "unique": unique,
378                }));
379            }
380            Ok(indexes)
381        }
382        DbBackend::Postgres => {
383            let rows = sqlx::query(
384                "SELECT indexname, indexdef FROM pg_indexes \
385                 WHERE tablename = $1 ORDER BY indexname",
386            )
387            .bind(&table)
388            .fetch_all(pool)
389            .await?;
390
391            let mut indexes = Vec::new();
392            for row in &rows {
393                let name: String = row.try_get("indexname")?;
394                let def: String = row.try_get("indexdef")?;
395                let unique = def.to_uppercase().contains("UNIQUE");
396                indexes.push(serde_json::json!({
397                    "index_name": name,
398                    "definition": def,
399                    "unique": unique,
400                }));
401            }
402            Ok(indexes)
403        }
404        DbBackend::Mysql => {
405            let rows = sqlx::query(
406                "SELECT INDEX_NAME, COLUMN_NAME, NON_UNIQUE \
407                 FROM information_schema.STATISTICS \
408                 WHERE TABLE_NAME = ? AND TABLE_SCHEMA = DATABASE() \
409                 ORDER BY INDEX_NAME, SEQ_IN_INDEX",
410            )
411            .bind(&table)
412            .fetch_all(pool)
413            .await?;
414
415            // Group columns by index name
416            let mut index_map: std::collections::HashMap<String, (Vec<String>, bool)> =
417                std::collections::HashMap::new();
418            for row in &rows {
419                let name: String = row.try_get("INDEX_NAME")?;
420                let col: String = row.try_get("COLUMN_NAME")?;
421                let non_unique: i32 = row.try_get("NON_UNIQUE")?;
422                let entry = index_map
423                    .entry(name)
424                    .or_insert_with(|| (Vec::new(), non_unique == 0));
425                entry.0.push(col);
426            }
427
428            let indexes = index_map
429                .into_iter()
430                .map(|(name, (columns, unique))| {
431                    serde_json::json!({
432                        "index_name": name,
433                        "columns": columns,
434                        "unique": unique,
435                    })
436                })
437                .collect();
438            Ok(indexes)
439        }
440    }
441}
442
443/// Validate and sanitize a SQL identifier to prevent injection.
444fn sanitize_identifier(name: &str) -> Result<String, McpSqlError> {
445    // Allow alphanumeric, underscore, dot (for schema.table), and hyphen
446    if name
447        .chars()
448        .all(|c| c.is_alphanumeric() || c == '_' || c == '.' || c == '-')
449        && !name.is_empty()
450    {
451        Ok(name.to_string())
452    } else {
453        Err(McpSqlError::InvalidSql(format!(
454            "Invalid identifier: '{name}'"
455        )))
456    }
457}
458
459#[cfg(test)]
460mod tests {
461    use super::*;
462
463    #[test]
464    fn test_sanitize_identifier() {
465        assert!(sanitize_identifier("users").is_ok());
466        assert!(sanitize_identifier("public.users").is_ok());
467        assert!(sanitize_identifier("my_table").is_ok());
468        assert!(sanitize_identifier("my-table").is_ok());
469        assert!(sanitize_identifier("").is_err());
470        assert!(sanitize_identifier("users; DROP TABLE users").is_err());
471        assert!(sanitize_identifier("users\"").is_err());
472    }
473
474    #[test]
475    fn test_explain_prefix() {
476        assert_eq!(explain_prefix(DbBackend::Postgres), "EXPLAIN (FORMAT TEXT) ");
477        assert_eq!(explain_prefix(DbBackend::Sqlite), "EXPLAIN QUERY PLAN ");
478        assert_eq!(explain_prefix(DbBackend::Mysql), "EXPLAIN ");
479    }
480}
mcp_sql/db/dialect.rs

mcp_sql/db/
dialect.rs
