Skip to main content

mcp_postgres/actions/
ext_mgmt.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 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
21        .iter()
22        .map(|row| {
23            json!({
24                "name": row.get::<_, String>(0),
25                "version": row.get::<_, String>(1),
26                "schema": row.get::<_, String>(2),
27                "description": row.get::<_, Option<String>>(3),
28                "relocatable": row.get::<_, bool>(4),
29            })
30        })
31        .collect();
32
33    Ok(json!({ "extensions": extensions }))
34}
35
36pub async fn create_extension(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
37    let name = params
38        .as_ref()
39        .and_then(|p| p.get("name").and_then(|v| v.as_str()))
40        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'name' parameter".into()))?;
41
42    if name.is_empty() || name.len() > MAX_IDENTIFIER_LEN {
43        return Err(crate::errors::MCPError::InvalidParams(format!(
44            "'name' must be 1-{MAX_IDENTIFIER_LEN} characters"
45        )));
46    }
47
48    let schema = params
49        .as_ref()
50        .and_then(|p| p.get("schema").and_then(|v| v.as_str()));
51    let cascade = params
52        .as_ref()
53        .and_then(|p| p.get("cascade").and_then(|v| v.as_bool()))
54        .unwrap_or(false);
55    let version = params
56        .as_ref()
57        .and_then(|p| p.get("version").and_then(|v| v.as_str()));
58    let if_not_exists = params
59        .as_ref()
60        .and_then(|p| p.get("if_not_exists").and_then(|v| v.as_bool()))
61        .unwrap_or(false);
62
63    let mut sql = "CREATE EXTENSION".to_string();
64    if if_not_exists {
65        sql.push_str(" IF NOT EXISTS");
66    }
67    sql.push_str(&format!(" {}", crate::validation::quote_ident(name)));
68    if let Some(s) = schema {
69        sql.push_str(&format!(" SCHEMA {}", crate::validation::quote_ident(s)));
70    }
71    if let Some(v) = version {
72        sql.push_str(&format!(" VERSION '{}'", v.replace('\'', "''")));
73    }
74    if cascade {
75        sql.push_str(" CASCADE");
76    }
77
78    client.execute(&sql, &[]).await?;
79    Ok(json!({ "success": true, "extension": name, "sql": sql }))
80}
81
82pub async fn drop_extension(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
83    let name = params
84        .as_ref()
85        .and_then(|p| p.get("name").and_then(|v| v.as_str()))
86        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'name' parameter".into()))?;
87
88    if name.is_empty() || name.len() > MAX_IDENTIFIER_LEN {
89        return Err(crate::errors::MCPError::InvalidParams(format!(
90            "'name' must be 1-{MAX_IDENTIFIER_LEN} characters"
91        )));
92    }
93
94    let if_exists = params
95        .as_ref()
96        .and_then(|p| p.get("if_exists").and_then(|v| v.as_bool()))
97        .unwrap_or(false);
98    let cascade = params
99        .as_ref()
100        .and_then(|p| p.get("cascade").and_then(|v| v.as_bool()))
101        .unwrap_or(false);
102
103    let mut sql = "DROP EXTENSION".to_string();
104    if if_exists {
105        sql.push_str(" IF EXISTS");
106    }
107    sql.push_str(&format!(" {}", crate::validation::quote_ident(name)));
108    if cascade {
109        sql.push_str(" CASCADE");
110    }
111
112    client.execute(&sql, &[]).await?;
113    Ok(json!({ "success": true, "extension": name, "sql": sql }))
114}