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

mcp_postgres/actions/
user_mgmt.rs
