Skip to main content

mcp_postgres/actions/
maint_ext.rs

1use crate::errors::Result as MCPResult;
2use serde_json::{Value, json};
3use tokio_postgres::Client;
4
5const MAX_IDENTIFIER_LEN: usize = 255;
6
7pub async fn vacuum(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
8    let table = params
9        .as_ref()
10        .and_then(|p| p.get("table").and_then(|v| v.as_str()));
11    let schema = params
12        .as_ref()
13        .and_then(|p| p.get("schema").and_then(|v| v.as_str()))
14        .unwrap_or("public");
15    let full = params
16        .as_ref()
17        .and_then(|p| p.get("full").and_then(|v| v.as_bool()))
18        .unwrap_or(false);
19    let freeze = params
20        .as_ref()
21        .and_then(|p| p.get("freeze").and_then(|v| v.as_bool()))
22        .unwrap_or(false);
23    let verbose = params
24        .as_ref()
25        .and_then(|p| p.get("verbose").and_then(|v| v.as_bool()))
26        .unwrap_or(false);
27
28    let mut sql = "VACUUM".to_string();
29    let opts: Vec<&str> = match (full, freeze, verbose) {
30        (true, _, _) => vec!["FULL", if verbose { "VERBOSE" } else { "" }],
31        (_, true, _) => vec!["FREEZE", if verbose { "VERBOSE" } else { "" }],
32        (_, _, true) => vec!["VERBOSE"],
33        _ => vec![],
34    };
35    let opts_str = opts
36        .iter()
37        .filter(|s| !s.is_empty())
38        .copied()
39        .collect::<Vec<_>>()
40        .join(" ");
41    if !opts_str.is_empty() {
42        sql.push_str(&format!(" {}", opts_str));
43    }
44
45    if let Some(t) = table {
46        sql.push_str(&format!(" {}.{}", quote_ident(schema), quote_ident(t)));
47    }
48
49    client.execute(&sql, &[]).await?;
50    Ok(json!({ "success": true, "sql": sql }))
51}
52
53pub async fn vacuum_full(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
54    let table = params
55        .as_ref()
56        .and_then(|p| p.get("table").and_then(|v| v.as_str()));
57    let schema = params
58        .as_ref()
59        .and_then(|p| p.get("schema").and_then(|v| v.as_str()))
60        .unwrap_or("public");
61    let verbose = params
62        .as_ref()
63        .and_then(|p| p.get("verbose").and_then(|v| v.as_bool()))
64        .unwrap_or(false);
65    let analyze = params
66        .as_ref()
67        .and_then(|p| p.get("analyze").and_then(|v| v.as_bool()))
68        .unwrap_or(false);
69    let freeze = params
70        .as_ref()
71        .and_then(|p| p.get("freeze").and_then(|v| v.as_bool()))
72        .unwrap_or(false);
73
74    let mut sql = "VACUUM FULL".to_string();
75    let mut opts = Vec::new();
76    if freeze {
77        opts.push("FREEZE");
78    }
79    if verbose {
80        opts.push("VERBOSE");
81    }
82    if analyze {
83        opts.push("ANALYZE");
84    }
85    if !opts.is_empty() {
86        sql.push_str(&format!(" {}", opts.join(" ")));
87    }
88
89    if let Some(t) = table {
90        sql.push_str(&format!(" {}.{}", quote_ident(schema), quote_ident(t)));
91    }
92
93    client.execute(&sql, &[]).await?;
94    Ok(json!({ "success": true, "sql": sql }))
95}
96
97pub async fn reindex_database(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
98    let database = params
99        .as_ref()
100        .and_then(|p| p.get("database").and_then(|v| v.as_str()))
101        .ok_or_else(|| {
102            crate::errors::MCPError::InvalidParams("Missing 'database' parameter".into())
103        })?;
104
105    if database.is_empty() || database.len() > MAX_IDENTIFIER_LEN {
106        return Err(crate::errors::MCPError::InvalidParams(format!(
107            "'database' must be 1-{MAX_IDENTIFIER_LEN} characters"
108        )));
109    }
110
111    let concurrent = params
112        .as_ref()
113        .and_then(|p| p.get("concurrent").and_then(|v| v.as_bool()))
114        .unwrap_or(false);
115    let verbose = params
116        .as_ref()
117        .and_then(|p| p.get("verbose").and_then(|v| v.as_bool()))
118        .unwrap_or(false);
119    let schema = params
120        .as_ref()
121        .and_then(|p| p.get("schema").and_then(|v| v.as_str()));
122
123    let mut sql = "REINDEX".to_string();
124    if concurrent {
125        sql.push_str(" (CONCURRENTLY)");
126    }
127    sql.push_str(" DATABASE ");
128    if verbose {
129        sql.push_str("VERBOSE ");
130    }
131    sql.push_str(&quote_ident(database));
132
133    if let Some(s) = schema
134        && !s.is_empty()
135        && s.len() <= MAX_IDENTIFIER_LEN
136    {
137        sql.push_str(&format!(" SCHEMA {}", quote_ident(s)));
138
139        let table = params
140            .as_ref()
141            .and_then(|p| p.get("table").and_then(|v| v.as_str()));
142        if let Some(t) = table {
143            sql.push_str(&format!(" TABLE {}", quote_ident(t)));
144        }
145
146        let index = params
147            .as_ref()
148            .and_then(|p| p.get("index").and_then(|v| v.as_str()));
149        if let Some(i) = index {
150            sql = format!(
151                "REINDEX{} INDEX {} {}",
152                if concurrent { " (CONCURRENTLY)" } else { "" },
153                if verbose { "VERBOSE " } else { "" },
154                quote_ident(i)
155            );
156        }
157    }
158
159    client.execute(&sql, &[]).await?;
160    Ok(json!({ "success": true, "sql": sql }))
161}
162
163fn quote_ident(ident: &str) -> String {
164    crate::validation::quote_ident(ident)
165}