mcp_postgres/actions/
maintenance.rs1use 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
19pub 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
44pub 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
62pub 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
80pub 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
118pub 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}