Skip to main content

mcp_postgres/actions/
maint_ext.rs

1use serde_json::{json, Value};
2use tokio_postgres::Client;
3use crate::errors::Result as MCPResult;
4
5const MAX_IDENTIFIER_LEN: usize = 255;
6
7pub async fn vacuum(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
8    let table = params.as_ref().and_then(|p| p.get("table").and_then(|v| v.as_str()));
9    let schema = params.as_ref().and_then(|p| p.get("schema").and_then(|v| v.as_str())).unwrap_or("public");
10    let full = params.as_ref().and_then(|p| p.get("full").and_then(|v| v.as_bool())).unwrap_or(false);
11    let freeze = params.as_ref().and_then(|p| p.get("freeze").and_then(|v| v.as_bool())).unwrap_or(false);
12    let verbose = params.as_ref().and_then(|p| p.get("verbose").and_then(|v| v.as_bool())).unwrap_or(false);
13
14    let mut sql = "VACUUM".to_string();
15    let opts: Vec<&str> = match (full, freeze, verbose) {
16        (true, _, _) => vec!["FULL", if verbose { "VERBOSE" } else { "" }],
17        (_, true, _) => vec!["FREEZE", if verbose { "VERBOSE" } else { "" }],
18        (_, _, true) => vec!["VERBOSE"],
19        _ => vec![],
20    };
21    let opts_str = opts.iter().filter(|s| !s.is_empty()).copied().collect::<Vec<_>>().join(" ");
22    if !opts_str.is_empty() {
23        sql.push_str(&format!(" {}", opts_str));
24    }
25
26    if let Some(t) = table {
27        sql.push_str(&format!(" {}.{}", quote_ident(schema), quote_ident(t)));
28    }
29
30    client.execute(&sql, &[]).await?;
31    Ok(json!({ "success": true, "sql": sql }))
32}
33
34pub async fn vacuum_full(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
35    let table = params.as_ref().and_then(|p| p.get("table").and_then(|v| v.as_str()));
36    let schema = params.as_ref().and_then(|p| p.get("schema").and_then(|v| v.as_str())).unwrap_or("public");
37    let verbose = params.as_ref().and_then(|p| p.get("verbose").and_then(|v| v.as_bool())).unwrap_or(false);
38    let analyze = params.as_ref().and_then(|p| p.get("analyze").and_then(|v| v.as_bool())).unwrap_or(false);
39    let freeze = params.as_ref().and_then(|p| p.get("freeze").and_then(|v| v.as_bool())).unwrap_or(false);
40
41    let mut sql = "VACUUM FULL".to_string();
42    let mut opts = Vec::new();
43    if freeze { opts.push("FREEZE"); }
44    if verbose { opts.push("VERBOSE"); }
45    if analyze { opts.push("ANALYZE"); }
46    if !opts.is_empty() {
47        sql.push_str(&format!(" {}", opts.join(" ")));
48    }
49
50    if let Some(t) = table {
51        sql.push_str(&format!(" {}.{}", quote_ident(schema), quote_ident(t)));
52    }
53
54    client.execute(&sql, &[]).await?;
55    Ok(json!({ "success": true, "sql": sql }))
56}
57
58pub async fn reindex_database(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
59    let database = params.as_ref().and_then(|p| p.get("database").and_then(|v| v.as_str()))
60        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'database' parameter".into()))?;
61
62    if database.is_empty() || database.len() > MAX_IDENTIFIER_LEN {
63        return Err(crate::errors::MCPError::InvalidParams(format!("'database' must be 1-{MAX_IDENTIFIER_LEN} characters")));
64    }
65
66    let concurrent = params.as_ref().and_then(|p| p.get("concurrent").and_then(|v| v.as_bool())).unwrap_or(false);
67    let verbose = params.as_ref().and_then(|p| p.get("verbose").and_then(|v| v.as_bool())).unwrap_or(false);
68    let schema = params.as_ref().and_then(|p| p.get("schema").and_then(|v| v.as_str()));
69
70    let mut sql = "REINDEX".to_string();
71    if concurrent { sql.push_str(" (CONCURRENTLY)"); }
72    sql.push_str(" DATABASE ");
73    if verbose { sql.push_str("VERBOSE "); }
74    sql.push_str(&quote_ident(database));
75
76    if let Some(s) = schema
77        && !s.is_empty() && s.len() <= MAX_IDENTIFIER_LEN {
78            sql.push_str(&format!(" SCHEMA {}", quote_ident(s)));
79
80            let table = params.as_ref().and_then(|p| p.get("table").and_then(|v| v.as_str()));
81            if let Some(t) = table {
82                sql.push_str(&format!(" TABLE {}", quote_ident(t)));
83            }
84
85            let index = params.as_ref().and_then(|p| p.get("index").and_then(|v| v.as_str()));
86            if let Some(i) = index {
87                sql = format!("REINDEX{} INDEX {} {}",
88                    if concurrent { " (CONCURRENTLY)" } else { "" },
89                    if verbose { "VERBOSE " } else { "" },
90                    quote_ident(i));
91            }
92    }
93
94    client.execute(&sql, &[]).await?;
95    Ok(json!({ "success": true, "sql": sql }))
96}
97
98fn quote_ident(ident: &str) -> String {
99    crate::validation::quote_ident(ident)
100}