mcp_postgres/actions/
user_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 create_user(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
8    let username = params
9        .as_ref()
10        .and_then(|p| p.get("username").and_then(|v| v.as_str()))
11        .ok_or_else(|| {
12            crate::errors::MCPError::InvalidParams("Missing 'username' parameter".into())
13        })?;
14
15    if username.is_empty() || username.len() > MAX_IDENTIFIER_LEN {
16        return Err(crate::errors::MCPError::InvalidParams(format!(
17            "'username' must be 1-{MAX_IDENTIFIER_LEN} characters"
18        )));
19    }
20
21    let password = params
22        .as_ref()
23        .and_then(|p| p.get("password").and_then(|v| v.as_str()));
24    let valid_until = params
25        .as_ref()
26        .and_then(|p| p.get("valid_until").and_then(|v| v.as_str()));
27    let connection_limit = params
28        .as_ref()
29        .and_then(|p| p.get("connection_limit").and_then(|v| v.as_i64()));
30    let can_login = params
31        .as_ref()
32        .and_then(|p| p.get("can_login").and_then(|v| v.as_bool()));
33
34    let mut sql = format!("CREATE USER {}", quote_ident(username));
35    if let Some(pw) = password {
36        sql.push_str(&format!(" PASSWORD '{}'", pw.replace('\'', "''")));
37    }
38    if let Some(limit) = connection_limit {
39        sql.push_str(&format!(" CONNECTION LIMIT {}", limit));
40    }
41    if let Some(login) = can_login {
42        if login {
43            sql.push_str(" LOGIN");
44        } else {
45            sql.push_str(" NOLOGIN");
46        }
47    }
48    if let Some(until) = valid_until {
49        sql.push_str(&format!(" VALID UNTIL '{}'", until.replace('\'', "''")));
50    }
51
52    client.execute(&sql, &[]).await?;
53    Ok(json!({ "success": true, "user": username }))
54}
55
56pub async fn alter_user(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
57    let username = params
58        .as_ref()
59        .and_then(|p| p.get("username").and_then(|v| v.as_str()))
60        .ok_or_else(|| {
61            crate::errors::MCPError::InvalidParams("Missing 'username' parameter".into())
62        })?;
63
64    if username.is_empty() || username.len() > MAX_IDENTIFIER_LEN {
65        return Err(crate::errors::MCPError::InvalidParams(format!(
66            "'username' must be 1-{MAX_IDENTIFIER_LEN} characters"
67        )));
68    }
69
70    let password = params
71        .as_ref()
72        .and_then(|p| p.get("password").and_then(|v| v.as_str()));
73    let valid_until = params
74        .as_ref()
75        .and_then(|p| p.get("valid_until").and_then(|v| v.as_str()));
76    let connection_limit = params
77        .as_ref()
78        .and_then(|p| p.get("connection_limit").and_then(|v| v.as_i64()));
79    let can_login = params
80        .as_ref()
81        .and_then(|p| p.get("can_login").and_then(|v| v.as_bool()));
82    let new_name = params
83        .as_ref()
84        .and_then(|p| p.get("new_name").and_then(|v| v.as_str()));
85
86    if password.is_none()
87        && valid_until.is_none()
88        && connection_limit.is_none()
89        && can_login.is_none()
90        && new_name.is_none()
91    {
92        return Err(crate::errors::MCPError::InvalidParams(
93            "No attributes specified to alter".into(),
94        ));
95    }
96
97    let mut sql = format!("ALTER USER {}", quote_ident(username));
98    if let Some(pw) = password {
99        sql.push_str(&format!(" PASSWORD '{}'", pw.replace('\'', "''")));
100    }
101    if let Some(limit) = connection_limit {
102        sql.push_str(&format!(" CONNECTION LIMIT {}", limit));
103    }
104    if let Some(login) = can_login {
105        if login {
106            sql.push_str(" LOGIN");
107        } else {
108            sql.push_str(" NOLOGIN");
109        }
110    }
111    if let Some(until) = valid_until {
112        sql.push_str(&format!(" VALID UNTIL '{}'", until.replace('\'', "''")));
113    }
114    if let Some(name) = new_name {
115        sql.push_str(&format!(" RENAME TO {}", quote_ident(name)));
116    }
117
118    client.execute(&sql, &[]).await?;
119    Ok(json!({ "success": true, "user": username }))
120}
121
122pub async fn drop_user(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
123    let username = params
124        .as_ref()
125        .and_then(|p| p.get("username").and_then(|v| v.as_str()))
126        .ok_or_else(|| {
127            crate::errors::MCPError::InvalidParams("Missing 'username' parameter".into())
128        })?;
129
130    if username.is_empty() || username.len() > MAX_IDENTIFIER_LEN {
131        return Err(crate::errors::MCPError::InvalidParams(format!(
132            "'username' must be 1-{MAX_IDENTIFIER_LEN} characters"
133        )));
134    }
135
136    let if_exists = params
137        .as_ref()
138        .and_then(|p| p.get("if_exists").and_then(|v| v.as_bool()))
139        .unwrap_or(false);
140
141    let mut sql = "DROP USER".to_string();
142    if if_exists {
143        sql.push_str(" IF EXISTS");
144    }
145    sql.push_str(&format!(" {}", quote_ident(username)));
146
147    client.execute(&sql, &[]).await?;
148    Ok(json!({ "success": true, "user": username }))
149}
150
151pub async fn create_role(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
152    let role_name = params
153        .as_ref()
154        .and_then(|p| p.get("role_name").and_then(|v| v.as_str()))
155        .ok_or_else(|| {
156            crate::errors::MCPError::InvalidParams("Missing 'role_name' parameter".into())
157        })?;
158
159    if role_name.is_empty() || role_name.len() > MAX_IDENTIFIER_LEN {
160        return Err(crate::errors::MCPError::InvalidParams(format!(
161            "'role_name' must be 1-{MAX_IDENTIFIER_LEN} characters"
162        )));
163    }
164
165    let with_login = params
166        .as_ref()
167        .and_then(|p| p.get("with_login").and_then(|v| v.as_bool()))
168        .unwrap_or(false);
169    let mut sql = format!("CREATE ROLE {}", quote_ident(role_name));
170    sql.push_str(if with_login { " LOGIN" } else { " NOLOGIN" });
171
172    client.execute(&sql, &[]).await?;
173    Ok(json!({ "success": true, "role": role_name }))
174}
175
176pub async fn alter_role(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
177    let role_name = params
178        .as_ref()
179        .and_then(|p| p.get("role_name").and_then(|v| v.as_str()))
180        .ok_or_else(|| {
181            crate::errors::MCPError::InvalidParams("Missing 'role_name' parameter".into())
182        })?;
183
184    if role_name.is_empty() || role_name.len() > MAX_IDENTIFIER_LEN {
185        return Err(crate::errors::MCPError::InvalidParams(format!(
186            "'role_name' must be 1-{MAX_IDENTIFIER_LEN} characters"
187        )));
188    }
189
190    let password = params
191        .as_ref()
192        .and_then(|p| p.get("password").and_then(|v| v.as_str()));
193    let can_login = params
194        .as_ref()
195        .and_then(|p| p.get("can_login").and_then(|v| v.as_bool()));
196    let superuser = params
197        .as_ref()
198        .and_then(|p| p.get("superuser").and_then(|v| v.as_bool()));
199    let createdb = params
200        .as_ref()
201        .and_then(|p| p.get("createdb").and_then(|v| v.as_bool()));
202    let new_name = params
203        .as_ref()
204        .and_then(|p| p.get("new_name").and_then(|v| v.as_str()));
205
206    let mut sql = format!("ALTER ROLE {}", quote_ident(role_name));
207    if let Some(pw) = password {
208        sql.push_str(&format!(" PASSWORD '{}'", pw.replace('\'', "''")));
209    }
210    if let Some(login) = can_login {
211        sql.push_str(if login { " LOGIN" } else { " NOLOGIN" });
212    }
213    if let Some(su) = superuser {
214        sql.push_str(if su { " SUPERUSER" } else { " NOSUPERUSER" });
215    }
216    if let Some(db) = createdb {
217        sql.push_str(if db { " CREATEDB" } else { " NOCREATEDB" });
218    }
219    if let Some(name) = new_name {
220        sql.push_str(&format!(" RENAME TO {}", quote_ident(name)));
221    }
222
223    client.execute(&sql, &[]).await?;
224    Ok(json!({ "success": true, "role": role_name }))
225}
226
227pub async fn drop_role(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
228    let role_name = params
229        .as_ref()
230        .and_then(|p| p.get("role_name").and_then(|v| v.as_str()))
231        .ok_or_else(|| {
232            crate::errors::MCPError::InvalidParams("Missing 'role_name' parameter".into())
233        })?;
234
235    if role_name.is_empty() || role_name.len() > MAX_IDENTIFIER_LEN {
236        return Err(crate::errors::MCPError::InvalidParams(format!(
237            "'role_name' must be 1-{MAX_IDENTIFIER_LEN} characters"
238        )));
239    }
240
241    let if_exists = params
242        .as_ref()
243        .and_then(|p| p.get("if_exists").and_then(|v| v.as_bool()))
244        .unwrap_or(false);
245
246    let mut sql = "DROP ROLE".to_string();
247    if if_exists {
248        sql.push_str(" IF EXISTS");
249    }
250    sql.push_str(&format!(" {}", quote_ident(role_name)));
251
252    client.execute(&sql, &[]).await?;
253    Ok(json!({ "success": true, "role": role_name }))
254}
255
256pub async fn grant_privileges(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
257    let privilege = params
258        .as_ref()
259        .and_then(|p| p.get("privilege").and_then(|v| v.as_str()))
260        .ok_or_else(|| {
261            crate::errors::MCPError::InvalidParams("Missing 'privilege' parameter".into())
262        })?;
263    let object_type = params
264        .as_ref()
265        .and_then(|p| p.get("object_type").and_then(|v| v.as_str()))
266        .ok_or_else(|| {
267            crate::errors::MCPError::InvalidParams("Missing 'object_type' parameter".into())
268        })?;
269    let object_name = params
270        .as_ref()
271        .and_then(|p| p.get("object_name").and_then(|v| v.as_str()))
272        .ok_or_else(|| {
273            crate::errors::MCPError::InvalidParams("Missing 'object_name' parameter".into())
274        })?;
275    let grantee = params
276        .as_ref()
277        .and_then(|p| p.get("grantee").and_then(|v| v.as_str()))
278        .ok_or_else(|| {
279            crate::errors::MCPError::InvalidParams("Missing 'grantee' parameter".into())
280        })?;
281
282    let schema = params
283        .as_ref()
284        .and_then(|p| p.get("schema").and_then(|v| v.as_str()))
285        .unwrap_or("public");
286
287    let valid_types = [
288        "table",
289        "sequence",
290        "schema",
291        "database",
292        "all_tables_in_schema",
293    ];
294    if !valid_types.contains(&object_type) {
295        return Err(crate::errors::MCPError::InvalidParams(format!(
296            "Unsupported object_type '{}'. Use: {:?}",
297            object_type, valid_types
298        )));
299    }
300
301    let sql = match object_type {
302        "all_tables_in_schema" => format!(
303            "GRANT {} ON ALL TABLES IN SCHEMA {} TO {}",
304            privilege,
305            quote_ident(schema),
306            quote_ident(grantee)
307        ),
308        _ => format!(
309            "GRANT {} ON {} {} TO {}",
310            privilege,
311            object_type.to_uppercase(),
312            quote_ident(object_name),
313            quote_ident(grantee)
314        ),
315    };
316
317    client.execute(&sql, &[]).await?;
318    Ok(json!({ "success": true, "sql": sql }))
319}
320
321pub async fn revoke_privileges(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
322    let privilege = params
323        .as_ref()
324        .and_then(|p| p.get("privilege").and_then(|v| v.as_str()))
325        .ok_or_else(|| {
326            crate::errors::MCPError::InvalidParams("Missing 'privilege' parameter".into())
327        })?;
328    let object_type = params
329        .as_ref()
330        .and_then(|p| p.get("object_type").and_then(|v| v.as_str()))
331        .ok_or_else(|| {
332            crate::errors::MCPError::InvalidParams("Missing 'object_type' parameter".into())
333        })?;
334    let object_name = params
335        .as_ref()
336        .and_then(|p| p.get("object_name").and_then(|v| v.as_str()))
337        .ok_or_else(|| {
338            crate::errors::MCPError::InvalidParams("Missing 'object_name' parameter".into())
339        })?;
340    let revokee = params
341        .as_ref()
342        .and_then(|p| p.get("revokee").and_then(|v| v.as_str()))
343        .ok_or_else(|| {
344            crate::errors::MCPError::InvalidParams("Missing 'revokee' parameter".into())
345        })?;
346
347    let schema = params
348        .as_ref()
349        .and_then(|p| p.get("schema").and_then(|v| v.as_str()))
350        .unwrap_or("public");
351
352    let valid_types = [
353        "table",
354        "sequence",
355        "schema",
356        "database",
357        "all_tables_in_schema",
358    ];
359    if !valid_types.contains(&object_type) {
360        return Err(crate::errors::MCPError::InvalidParams(format!(
361            "Unsupported object_type '{}'. Use: {:?}",
362            object_type, valid_types
363        )));
364    }
365
366    let sql = match object_type {
367        "all_tables_in_schema" => format!(
368            "REVOKE {} ON ALL TABLES IN SCHEMA {} FROM {}",
369            privilege,
370            quote_ident(schema),
371            quote_ident(revokee)
372        ),
373        _ => format!(
374            "REVOKE {} ON {} {} FROM {}",
375            privilege,
376            object_type.to_uppercase(),
377            quote_ident(object_name),
378            quote_ident(revokee)
379        ),
380    };
381
382    client.execute(&sql, &[]).await?;
383    Ok(json!({ "success": true, "sql": sql }))
384}
385
386fn quote_ident(ident: &str) -> String {
387    crate::validation::quote_ident(ident)
388}
mcp_postgres/actions/user_mgmt.rs

mcp_postgres/actions/
user_mgmt.rs
