mcp_postgres/actions/
data_tools.rs

1use serde_json::{json, Value};
2use tokio_postgres::Client;
3use crate::errors::Result as MCPResult;
4use crate::validation::quote_ident;
5
6pub async fn sample_data(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
7    let table = params.as_ref().and_then(|p| p.get("table").and_then(|v| v.as_str()))
8        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'table'".into()))?;
9    let schema = params.as_ref().and_then(|p| p.get("schema").and_then(|v| v.as_str())).unwrap_or("public");
10    let limit = params.as_ref().and_then(|p| p.get("limit").and_then(|v| v.as_i64())).unwrap_or(100).min(10000);
11    let where_clause = params.as_ref().and_then(|p| p.get("where").and_then(|v| v.as_str()));
12    let order_by = params.as_ref().and_then(|p| p.get("order_by").and_then(|v| v.as_str()));
13    let anonymize = params.as_ref().and_then(|p| p.get("anonymize_columns"))
14        .and_then(|v| v.as_array())
15        .map(|arr| arr.iter().filter_map(|v| v.as_str().map(String::from)).collect::<Vec<_>>());
16
17    let qualified = format!("{}.{}", quote_ident(schema), quote_ident(table));
18
19    let order = match order_by {
20        Some(col) => quote_ident(col),
21        None => "RANDOM()".to_string(),
22    };
23
24    let where_sql = where_clause.map(|w| format!(" WHERE {}", w)).unwrap_or_default();
25
26    let sql = format!(
27        "SELECT * FROM {} {} ORDER BY {} LIMIT {}",
28        qualified, where_sql, order, limit,
29    );
30
31    let rows = client.query(&sql, &[]).await?;
32
33    let has_anonymize = anonymize.as_ref().is_some_and(|a| !a.is_empty());
34    let mut results = Vec::new();
35    for row in &rows {
36        let mut obj = serde_json::Map::new();
37        for (i, col) in row.columns().iter().enumerate() {
38            let name = col.name();
39            let masked = has_anonymize && anonymize.as_ref().is_some_and(|a| a.iter().any(|c| c == name));
40
41            if masked {
42                let val: Option<String> = row.try_get::<_, Option<String>>(i).ok().flatten();
43                match val {
44                    Some(v) if v.contains('@') => {
45                        let parts: Vec<&str> = v.splitn(2, '@').collect();
46                        let _ = obj.insert(name.to_string(), json!(format!("****@{}", parts[1])));
47                    }
48                    Some(v) if v.len() > 2 => {
49                        let first = &v[..1];
50                        let last = &v[v.len()-1..];
51                        let _ = obj.insert(name.to_string(), json!(format!("{}...{}", first, last)));
52                    }
53                    Some(_v) => { let _ = obj.insert(name.to_string(), json!("***")); }
54                    None => { let _ = obj.insert(name.to_string(), Value::Null); }
55                }
56            } else if let Ok(v) = row.try_get::<_, Value>(i) {
57                obj.insert(name.to_string(), v);
58            } else if let Ok(v) = row.try_get::<_, String>(i) {
59                obj.insert(name.to_string(), Value::String(v));
60            } else if let Ok(v) = row.try_get::<_, i64>(i) {
61                obj.insert(name.to_string(), json!(v));
62            } else if let Ok(v) = row.try_get::<_, f64>(i) {
63                obj.insert(name.to_string(), json!(v));
64            } else if let Ok(v) = row.try_get::<_, bool>(i) {
65                obj.insert(name.to_string(), json!(v));
66            } else if let Ok(v) = row.try_get::<_, Option<String>>(i) {
67                obj.insert(name.to_string(), v.map(Value::String).unwrap_or(Value::Null));
68            }
69        }
70        results.push(Value::Object(obj));
71    }
72
73    Ok(json!({
74        "sample": results,
75        "count": results.len(),
76        "randomized": order_by.is_none(),
77        "anonymized_columns": anonymize.unwrap_or_default(),
78    }))
79}
mcp_postgres/actions/data_tools.rs

mcp_postgres/actions/
data_tools.rs
