Skip to main content

mcp_postgres/actions/
db_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_databases(client: &Client, _params: &Option<&Value>) -> MCPResult<Value> {
8    let rows = client
9        .query(
10            "SELECT d.datname, pg_catalog.pg_get_userbyid(d.datdba) AS owner,
11                    pg_catalog.pg_encoding_to_char(d.encoding) AS encoding,
12                    d.datcollate, d.datctype,
13                    pg_catalog.shobj_description(d.oid, 'pg_database') AS description,
14                    d.datistemplate
15             FROM pg_catalog.pg_database d
16             ORDER BY d.datname",
17            &[],
18        )
19        .await?;
20
21    let databases: Vec<Value> = rows.iter().map(|row| {
22        json!({
23            "name": row.get::<_, String>(0),
24            "owner": row.get::<_, String>(1),
25            "encoding": row.get::<_, String>(2),
26            "collate": row.get::<_, String>(3),
27            "ctype": row.get::<_, String>(4),
28            "description": row.get::<_, Option<String>>(5),
29            "is_template": row.get::<_, bool>(6),
30        })
31    }).collect();
32
33    Ok(json!({ "databases": databases }))
34}
35
36pub async fn create_database(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
37    let name = params.as_ref().and_then(|p| p.get("name").and_then(|v| v.as_str()))
38        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'name' parameter".into()))?;
39
40    if name.is_empty() || name.len() > MAX_IDENTIFIER_LEN {
41        return Err(crate::errors::MCPError::InvalidParams(format!("'name' must be 1-{MAX_IDENTIFIER_LEN} characters")));
42    }
43
44    let owner = params.as_ref().and_then(|p| p.get("owner").and_then(|v| v.as_str()));
45    let encoding = params.as_ref().and_then(|p| p.get("encoding").and_then(|v| v.as_str()));
46    let locale = params.as_ref().and_then(|p| p.get("locale").and_then(|v| v.as_str()));
47
48    let mut sql = format!("CREATE DATABASE {}", crate::validation::quote_ident(name));
49    if let Some(o) = owner { sql.push_str(&format!(" OWNER {}", crate::validation::quote_ident(o))); }
50    if let Some(e) = encoding { sql.push_str(&format!(" ENCODING '{}'", e.replace('\'', "''"))); }
51    if let Some(l) = locale { sql.push_str(&format!(" LC_COLLATE '{}'", l.replace('\'', "''"))); }
52
53    client.execute(&sql, &[]).await?;
54    Ok(json!({ "success": true, "database": name, "sql": sql }))
55}