Skip to main content

mcp_postgres/actions/
db_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_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
22        .iter()
23        .map(|row| {
24            json!({
25                "name": row.get::<_, String>(0),
26                "owner": row.get::<_, String>(1),
27                "encoding": row.get::<_, String>(2),
28                "collate": row.get::<_, String>(3),
29                "ctype": row.get::<_, String>(4),
30                "description": row.get::<_, Option<String>>(5),
31                "is_template": row.get::<_, bool>(6),
32            })
33        })
34        .collect();
35
36    Ok(json!({ "databases": databases }))
37}
38
39pub async fn create_database(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
40    let name = params
41        .as_ref()
42        .and_then(|p| p.get("name").and_then(|v| v.as_str()))
43        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'name' parameter".into()))?;
44
45    if name.is_empty() || name.len() > MAX_IDENTIFIER_LEN {
46        return Err(crate::errors::MCPError::InvalidParams(format!(
47            "'name' must be 1-{MAX_IDENTIFIER_LEN} characters"
48        )));
49    }
50
51    let owner = params
52        .as_ref()
53        .and_then(|p| p.get("owner").and_then(|v| v.as_str()));
54    let encoding = params
55        .as_ref()
56        .and_then(|p| p.get("encoding").and_then(|v| v.as_str()));
57    let locale = params
58        .as_ref()
59        .and_then(|p| p.get("locale").and_then(|v| v.as_str()));
60
61    let mut sql = format!("CREATE DATABASE {}", crate::validation::quote_ident(name));
62    if let Some(o) = owner {
63        sql.push_str(&format!(" OWNER {}", crate::validation::quote_ident(o)));
64    }
65    if let Some(e) = encoding {
66        sql.push_str(&format!(" ENCODING '{}'", e.replace('\'', "''")));
67    }
68    if let Some(l) = locale {
69        sql.push_str(&format!(" LC_COLLATE '{}'", l.replace('\'', "''")));
70    }
71
72    client.execute(&sql, &[]).await?;
73    Ok(json!({ "success": true, "database": name, "sql": sql }))
74}