mcp_postgres/actions/
maintenance.rs

1use serde_json::{json, Value};
2use tokio_postgres::Client;
3use crate::errors::Result as MCPResult;
4
5const MAX_IDENTIFIER_LEN: usize = 255;
6
7fn validate_table_name(name: &str) -> std::result::Result<(), crate::errors::MCPError> {
8    if name.is_empty() {
9        return Err(crate::errors::MCPError::InvalidParams("'table' must not be empty".into()));
10    }
11    if name.len() > MAX_IDENTIFIER_LEN {
12        return Err(crate::errors::MCPError::InvalidParams(
13            format!("'table' exceeds maximum length of {MAX_IDENTIFIER_LEN} characters (got {})", name.len())
14        ));
15    }
16    Ok(())
17}
18
19/// 21. Vacuum analyze
20pub async fn vacuum_analyze(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
21    let table_name = params
22        .as_ref()
23        .and_then(|p| p.get("table").and_then(|v| v.as_str()).map(|s| s.to_string()));
24
25    if let Some(ref table) = table_name {
26        validate_table_name(table)?;
27    }
28
29    let sql = if let Some(ref table) = table_name {
30        format!("VACUUM ANALYZE {}", table)
31    } else {
32        "VACUUM ANALYZE".to_string()
33    };
34
35    client.execute(&sql, &[]).await?;
36
37    Ok(json!({
38        "status": "success",
39        "action": "VACUUM ANALYZE",
40        "table": table_name
41    }))
42}
43
44/// 22. Analyze table
45pub async fn analyze_table(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
46    let table_name = params
47        .as_ref()
48        .and_then(|p| p.get("table").and_then(|v| v.as_str()).map(|s| s.to_string()))
49        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'table' parameter".into()))?;
50
51    validate_table_name(&table_name)?;
52
53    client.execute(&format!("ANALYZE {}", table_name), &[]).await?;
54
55    Ok(json!({
56        "status": "success",
57        "action": "ANALYZE",
58        "table": table_name
59    }))
60}
61
62/// 23. Reindex table
63pub async fn reindex_table(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
64    let table_name = params
65        .as_ref()
66        .and_then(|p| p.get("table").and_then(|v| v.as_str()).map(|s| s.to_string()))
67        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'table' parameter".into()))?;
68
69    validate_table_name(&table_name)?;
70
71    client.execute(&format!("REINDEX TABLE {}", table_name), &[]).await?;
72
73    Ok(json!({
74        "status": "success",
75        "action": "REINDEX",
76        "table": table_name
77    }))
78}
79
80/// 24. Get pg stat statements
81pub async fn get_pg_stat_statements(client: &Client, _params: &Option<&Value>) -> MCPResult<Value> {
82    let rows = client
83        .query(
84            "SELECT query, calls, mean_time, max_time, total_time
85             FROM pg_stat_statements
86             ORDER BY total_time DESC
87             LIMIT 50",
88            &[],
89        )
90        .await;
91
92    match rows {
93        Ok(r) => {
94            let statements: Vec<Value> = r
95                .iter()
96                .map(|row| {
97                    json!({
98                        "query": row.get::<_, String>(0),
99                        "calls": row.get::<_, i64>(1),
100                        "mean_time_ms": row.get::<_, f64>(2),
101                        "max_time_ms": row.get::<_, f64>(3),
102                        "total_time_ms": row.get::<_, f64>(4),
103                    })
104                })
105                .collect();
106
107            Ok(json!({ "statements": statements }))
108        }
109        Err(_) => {
110            Ok(json!({
111                "error": "pg_stat_statements extension not installed",
112                "statements": []
113            }))
114        }
115    }
116}
117
118/// 25. Reset statistics
119pub async fn reset_statistics(client: &Client, _params: &Option<&Value>) -> MCPResult<Value> {
120    client.execute("SELECT pg_stat_reset()", &[]).await?;
121
122    Ok(json!({
123        "status": "success",
124        "action": "reset_statistics",
125        "message": "All statistics counters have been reset"
126    }))
127}
128
129/// 26. Truncate table
130pub async fn truncate_table(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
131    let table = params
132        .as_ref()
133        .and_then(|p| p.get("table").and_then(|v| v.as_str()))
134        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'table' parameter".into()))?;
135
136    validate_table_name(table)?;
137
138    let cascade = params
139        .as_ref()
140        .and_then(|p| p.get("cascade").and_then(|v| v.as_bool()))
141        .unwrap_or(false);
142
143    let restart_identity = params
144        .as_ref()
145        .and_then(|p| p.get("restart_identity").and_then(|v| v.as_bool()))
146        .unwrap_or(false);
147
148    let cascade_str = if cascade { " CASCADE" } else { "" };
149    let restart_str = if restart_identity { " RESTART IDENTITY" } else { "" };
150
151    let sql = format!("TRUNCATE TABLE {}{}{}", table, restart_str, cascade_str);
152
153    client.execute(&sql, &[]).await?;
154
155    Ok(json!({
156        "status": "success",
157        "action": "TRUNCATE",
158        "table": table,
159        "cascade": cascade,
160        "restart_identity": restart_identity,
161        "message": format!("Table '{}' has been truncated", table)
162    }))
163}
mcp_postgres/actions/maintenance.rs

mcp_postgres/actions/
maintenance.rs
