mcp_postgres/actions/
data_tools.rs

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

mcp_postgres/actions/
data_tools.rs
