Skip to main content

mcp_postgres/actions/
ext_mgmt.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 list_extensions(client: &Client, _params: &Option<&Value>) -> MCPResult<Value> {
8    let rows = client
9        .query(
10            "SELECT e.extname, e.extversion, n.nspname AS schema,
11                    c.description, e.extrelocatable
12             FROM pg_extension e
13             JOIN pg_namespace n ON n.oid = e.extnamespace
14             LEFT JOIN pg_description c ON c.objoid = e.oid
15             ORDER BY e.extname",
16            &[],
17        )
18        .await?;
19
20    let extensions: Vec<Value> = rows.iter().map(|row| {
21        json!({
22            "name": row.get::<_, String>(0),
23            "version": row.get::<_, String>(1),
24            "schema": row.get::<_, String>(2),
25            "description": row.get::<_, Option<String>>(3),
26            "relocatable": row.get::<_, bool>(4),
27        })
28    }).collect();
29
30    Ok(json!({ "extensions": extensions }))
31}
32
33pub async fn create_extension(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
34    let name = params.as_ref().and_then(|p| p.get("name").and_then(|v| v.as_str()))
35        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'name' parameter".into()))?;
36
37    if name.is_empty() || name.len() > MAX_IDENTIFIER_LEN {
38        return Err(crate::errors::MCPError::InvalidParams(format!("'name' must be 1-{MAX_IDENTIFIER_LEN} characters")));
39    }
40
41    let schema = params.as_ref().and_then(|p| p.get("schema").and_then(|v| v.as_str()));
42    let cascade = params.as_ref().and_then(|p| p.get("cascade").and_then(|v| v.as_bool())).unwrap_or(false);
43    let version = params.as_ref().and_then(|p| p.get("version").and_then(|v| v.as_str()));
44    let if_not_exists = params.as_ref().and_then(|p| p.get("if_not_exists").and_then(|v| v.as_bool())).unwrap_or(false);
45
46    let mut sql = "CREATE EXTENSION".to_string();
47    if if_not_exists { sql.push_str(" IF NOT EXISTS"); }
48    sql.push_str(&format!(" {}", crate::validation::quote_ident(name)));
49    if let Some(s) = schema { sql.push_str(&format!(" SCHEMA {}", crate::validation::quote_ident(s))); }
50    if let Some(v) = version { sql.push_str(&format!(" VERSION '{}'", v.replace('\'', "''"))); }
51    if cascade { sql.push_str(" CASCADE"); }
52
53    client.execute(&sql, &[]).await?;
54    Ok(json!({ "success": true, "extension": name, "sql": sql }))
55}
56
57pub async fn drop_extension(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
58    let name = params.as_ref().and_then(|p| p.get("name").and_then(|v| v.as_str()))
59        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'name' parameter".into()))?;
60
61    if name.is_empty() || name.len() > MAX_IDENTIFIER_LEN {
62        return Err(crate::errors::MCPError::InvalidParams(format!("'name' must be 1-{MAX_IDENTIFIER_LEN} characters")));
63    }
64
65    let if_exists = params.as_ref().and_then(|p| p.get("if_exists").and_then(|v| v.as_bool())).unwrap_or(false);
66    let cascade = params.as_ref().and_then(|p| p.get("cascade").and_then(|v| v.as_bool())).unwrap_or(false);
67
68    let mut sql = "DROP EXTENSION".to_string();
69    if if_exists { sql.push_str(" IF EXISTS"); }
70    sql.push_str(&format!(" {}", crate::validation::quote_ident(name)));
71    if cascade { sql.push_str(" CASCADE"); }
72
73    client.execute(&sql, &[]).await?;
74    Ok(json!({ "success": true, "extension": name, "sql": sql }))
75}