mcp_postgres/actions/
maintenance.rs

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

mcp_postgres/actions/
maintenance.rs
