mcp_postgres/actions/
user_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 create_user(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
8    let username = params.as_ref().and_then(|p| p.get("username").and_then(|v| v.as_str()))
9        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'username' parameter".into()))?;
10
11    if username.is_empty() || username.len() > MAX_IDENTIFIER_LEN {
12        return Err(crate::errors::MCPError::InvalidParams(format!("'username' must be 1-{MAX_IDENTIFIER_LEN} characters")));
13    }
14
15    let password = params.as_ref().and_then(|p| p.get("password").and_then(|v| v.as_str()));
16    let valid_until = params.as_ref().and_then(|p| p.get("valid_until").and_then(|v| v.as_str()));
17    let connection_limit = params.as_ref().and_then(|p| p.get("connection_limit").and_then(|v| v.as_i64()));
18    let can_login = params.as_ref().and_then(|p| p.get("can_login").and_then(|v| v.as_bool()));
19
20    let mut sql = format!("CREATE USER {}", quote_ident(username));
21    if let Some(pw) = password {
22        sql.push_str(&format!(" PASSWORD '{}'", pw.replace('\'', "''")));
23    }
24    if let Some(limit) = connection_limit {
25        sql.push_str(&format!(" CONNECTION LIMIT {}", limit));
26    }
27    if let Some(login) = can_login {
28        if login {
29            sql.push_str(" LOGIN");
30        } else {
31            sql.push_str(" NOLOGIN");
32        }
33    }
34    if let Some(until) = valid_until {
35        sql.push_str(&format!(" VALID UNTIL '{}'", until.replace('\'', "''")));
36    }
37
38    client.execute(&sql, &[]).await?;
39    Ok(json!({ "success": true, "user": username }))
40}
41
42pub async fn alter_user(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
43    let username = params.as_ref().and_then(|p| p.get("username").and_then(|v| v.as_str()))
44        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'username' parameter".into()))?;
45
46    if username.is_empty() || username.len() > MAX_IDENTIFIER_LEN {
47        return Err(crate::errors::MCPError::InvalidParams(format!("'username' must be 1-{MAX_IDENTIFIER_LEN} characters")));
48    }
49
50    let password = params.as_ref().and_then(|p| p.get("password").and_then(|v| v.as_str()));
51    let valid_until = params.as_ref().and_then(|p| p.get("valid_until").and_then(|v| v.as_str()));
52    let connection_limit = params.as_ref().and_then(|p| p.get("connection_limit").and_then(|v| v.as_i64()));
53    let can_login = params.as_ref().and_then(|p| p.get("can_login").and_then(|v| v.as_bool()));
54    let new_name = params.as_ref().and_then(|p| p.get("new_name").and_then(|v| v.as_str()));
55
56    if password.is_none() && valid_until.is_none() && connection_limit.is_none() && can_login.is_none() && new_name.is_none() {
57        return Err(crate::errors::MCPError::InvalidParams("No attributes specified to alter".into()));
58    }
59
60    let mut sql = format!("ALTER USER {}", quote_ident(username));
61    if let Some(pw) = password {
62        sql.push_str(&format!(" PASSWORD '{}'", pw.replace('\'', "''")));
63    }
64    if let Some(limit) = connection_limit {
65        sql.push_str(&format!(" CONNECTION LIMIT {}", limit));
66    }
67    if let Some(login) = can_login {
68        if login { sql.push_str(" LOGIN"); } else { sql.push_str(" NOLOGIN"); }
69    }
70    if let Some(until) = valid_until {
71        sql.push_str(&format!(" VALID UNTIL '{}'", until.replace('\'', "''")));
72    }
73    if let Some(name) = new_name {
74        sql.push_str(&format!(" RENAME TO {}", quote_ident(name)));
75    }
76
77    client.execute(&sql, &[]).await?;
78    Ok(json!({ "success": true, "user": username }))
79}
80
81pub async fn drop_user(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
82    let username = params.as_ref().and_then(|p| p.get("username").and_then(|v| v.as_str()))
83        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'username' parameter".into()))?;
84
85    if username.is_empty() || username.len() > MAX_IDENTIFIER_LEN {
86        return Err(crate::errors::MCPError::InvalidParams(format!("'username' must be 1-{MAX_IDENTIFIER_LEN} characters")));
87    }
88
89    let if_exists = params.as_ref().and_then(|p| p.get("if_exists").and_then(|v| v.as_bool())).unwrap_or(false);
90
91    let mut sql = "DROP USER".to_string();
92    if if_exists { sql.push_str(" IF EXISTS"); }
93    sql.push_str(&format!(" {}", quote_ident(username)));
94
95    client.execute(&sql, &[]).await?;
96    Ok(json!({ "success": true, "user": username }))
97}
98
99pub async fn create_role(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
100    let role_name = params.as_ref().and_then(|p| p.get("role_name").and_then(|v| v.as_str()))
101        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'role_name' parameter".into()))?;
102
103    if role_name.is_empty() || role_name.len() > MAX_IDENTIFIER_LEN {
104        return Err(crate::errors::MCPError::InvalidParams(format!("'role_name' must be 1-{MAX_IDENTIFIER_LEN} characters")));
105    }
106
107    let with_login = params.as_ref().and_then(|p| p.get("with_login").and_then(|v| v.as_bool())).unwrap_or(false);
108    let mut sql = format!("CREATE ROLE {}", quote_ident(role_name));
109    sql.push_str(if with_login { " LOGIN" } else { " NOLOGIN" });
110
111    client.execute(&sql, &[]).await?;
112    Ok(json!({ "success": true, "role": role_name }))
113}
114
115pub async fn alter_role(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
116    let role_name = params.as_ref().and_then(|p| p.get("role_name").and_then(|v| v.as_str()))
117        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'role_name' parameter".into()))?;
118
119    if role_name.is_empty() || role_name.len() > MAX_IDENTIFIER_LEN {
120        return Err(crate::errors::MCPError::InvalidParams(format!("'role_name' must be 1-{MAX_IDENTIFIER_LEN} characters")));
121    }
122
123    let password = params.as_ref().and_then(|p| p.get("password").and_then(|v| v.as_str()));
124    let can_login = params.as_ref().and_then(|p| p.get("can_login").and_then(|v| v.as_bool()));
125    let superuser = params.as_ref().and_then(|p| p.get("superuser").and_then(|v| v.as_bool()));
126    let createdb = params.as_ref().and_then(|p| p.get("createdb").and_then(|v| v.as_bool()));
127    let new_name = params.as_ref().and_then(|p| p.get("new_name").and_then(|v| v.as_str()));
128
129    let mut sql = format!("ALTER ROLE {}", quote_ident(role_name));
130    if let Some(pw) = password {
131        sql.push_str(&format!(" PASSWORD '{}'", pw.replace('\'', "''")));
132    }
133    if let Some(login) = can_login {
134        sql.push_str(if login { " LOGIN" } else { " NOLOGIN" });
135    }
136    if let Some(su) = superuser {
137        sql.push_str(if su { " SUPERUSER" } else { " NOSUPERUSER" });
138    }
139    if let Some(db) = createdb {
140        sql.push_str(if db { " CREATEDB" } else { " NOCREATEDB" });
141    }
142    if let Some(name) = new_name {
143        sql.push_str(&format!(" RENAME TO {}", quote_ident(name)));
144    }
145
146    client.execute(&sql, &[]).await?;
147    Ok(json!({ "success": true, "role": role_name }))
148}
149
150pub async fn drop_role(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
151    let role_name = params.as_ref().and_then(|p| p.get("role_name").and_then(|v| v.as_str()))
152        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'role_name' parameter".into()))?;
153
154    if role_name.is_empty() || role_name.len() > MAX_IDENTIFIER_LEN {
155        return Err(crate::errors::MCPError::InvalidParams(format!("'role_name' must be 1-{MAX_IDENTIFIER_LEN} characters")));
156    }
157
158    let if_exists = params.as_ref().and_then(|p| p.get("if_exists").and_then(|v| v.as_bool())).unwrap_or(false);
159
160    let mut sql = "DROP ROLE".to_string();
161    if if_exists { sql.push_str(" IF EXISTS"); }
162    sql.push_str(&format!(" {}", quote_ident(role_name)));
163
164    client.execute(&sql, &[]).await?;
165    Ok(json!({ "success": true, "role": role_name }))
166}
167
168pub async fn grant_privileges(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
169    let privilege = params.as_ref().and_then(|p| p.get("privilege").and_then(|v| v.as_str()))
170        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'privilege' parameter".into()))?;
171    let object_type = params.as_ref().and_then(|p| p.get("object_type").and_then(|v| v.as_str()))
172        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'object_type' parameter".into()))?;
173    let object_name = params.as_ref().and_then(|p| p.get("object_name").and_then(|v| v.as_str()))
174        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'object_name' parameter".into()))?;
175    let grantee = params.as_ref().and_then(|p| p.get("grantee").and_then(|v| v.as_str()))
176        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'grantee' parameter".into()))?;
177
178    let schema = params.as_ref().and_then(|p| p.get("schema").and_then(|v| v.as_str())).unwrap_or("public");
179
180    let valid_types = ["table", "sequence", "schema", "database", "all_tables_in_schema"];
181    if !valid_types.contains(&object_type) {
182        return Err(crate::errors::MCPError::InvalidParams(format!("Unsupported object_type '{}'. Use: {:?}", object_type, valid_types)));
183    }
184
185    let sql = match object_type {
186        "all_tables_in_schema" => format!("GRANT {} ON ALL TABLES IN SCHEMA {} TO {}", privilege, quote_ident(schema), quote_ident(grantee)),
187        _ => format!("GRANT {} ON {} {} TO {}", privilege, object_type.to_uppercase(), quote_ident(object_name), quote_ident(grantee)),
188    };
189
190    client.execute(&sql, &[]).await?;
191    Ok(json!({ "success": true, "sql": sql }))
192}
193
194pub async fn revoke_privileges(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
195    let privilege = params.as_ref().and_then(|p| p.get("privilege").and_then(|v| v.as_str()))
196        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'privilege' parameter".into()))?;
197    let object_type = params.as_ref().and_then(|p| p.get("object_type").and_then(|v| v.as_str()))
198        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'object_type' parameter".into()))?;
199    let object_name = params.as_ref().and_then(|p| p.get("object_name").and_then(|v| v.as_str()))
200        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'object_name' parameter".into()))?;
201    let revokee = params.as_ref().and_then(|p| p.get("revokee").and_then(|v| v.as_str()))
202        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'revokee' parameter".into()))?;
203
204    let schema = params.as_ref().and_then(|p| p.get("schema").and_then(|v| v.as_str())).unwrap_or("public");
205
206    let valid_types = ["table", "sequence", "schema", "database", "all_tables_in_schema"];
207    if !valid_types.contains(&object_type) {
208        return Err(crate::errors::MCPError::InvalidParams(format!("Unsupported object_type '{}'. Use: {:?}", object_type, valid_types)));
209    }
210
211    let sql = match object_type {
212        "all_tables_in_schema" => format!("REVOKE {} ON ALL TABLES IN SCHEMA {} FROM {}", privilege, quote_ident(schema), quote_ident(revokee)),
213        _ => format!("REVOKE {} ON {} {} FROM {}", privilege, object_type.to_uppercase(), quote_ident(object_name), quote_ident(revokee)),
214    };
215
216    client.execute(&sql, &[]).await?;
217    Ok(json!({ "success": true, "sql": sql }))
218}
219
220fn quote_ident(ident: &str) -> String {
221    crate::validation::quote_ident(ident)
222}
mcp_postgres/actions/user_mgmt.rs

mcp_postgres/actions/
user_mgmt.rs
