mcp-postgres 4.1.0

use crate::errors::Result as MCPResult;
use crate::validation::quote_ident;
use serde_json::{Value, json};
use tokio_postgres::Client;

pub async fn sample_data(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
    let table = params
        .as_ref()
        .and_then(|p| p.get("table").and_then(|v| v.as_str()))
        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'table'".into()))?;
    let schema = params
        .as_ref()
        .and_then(|p| p.get("schema").and_then(|v| v.as_str()))
        .unwrap_or("public");
    let limit = params
        .as_ref()
        .and_then(|p| p.get("limit").and_then(|v| v.as_i64()))
        .unwrap_or(100)
        .min(10000);
    let where_clause = params
        .as_ref()
        .and_then(|p| p.get("where").and_then(|v| v.as_str()));
    let order_by = params
        .as_ref()
        .and_then(|p| p.get("order_by").and_then(|v| v.as_str()));
    let anonymize = params
        .as_ref()
        .and_then(|p| p.get("anonymize_columns"))
        .and_then(|v| v.as_array())
        .map(|arr| {
            arr.iter()
                .filter_map(|v| v.as_str().map(String::from))
                .collect::<Vec<_>>()
        });

    let qualified = format!("{}.{}", quote_ident(schema), quote_ident(table));

    let order = match order_by {
        Some(col) => quote_ident(col),
        None => "RANDOM()".to_string(),
    };

    let where_sql = where_clause
        .map(|w| format!(" WHERE {}", w))
        .unwrap_or_default();

    let sql = format!(
        "SELECT * FROM {} {} ORDER BY {} LIMIT {}",
        qualified, where_sql, order, limit,
    );

    let rows = client.query(&sql, &[]).await?;

    let has_anonymize = anonymize.as_ref().is_some_and(|a| !a.is_empty());
    let mut results = Vec::new();
    for row in &rows {
        let mut obj = serde_json::Map::new();
        for (i, col) in row.columns().iter().enumerate() {
            let name = col.name();
            let masked = has_anonymize
                && anonymize
                    .as_ref()
                    .is_some_and(|a| a.iter().any(|c| c == name));

            if masked {
                let val: Option<String> = row.try_get::<_, Option<String>>(i).ok().flatten();
                match val {
                    Some(v) if v.contains('@') => {
                        let parts: Vec<&str> = v.splitn(2, '@').collect();
                        let _ = obj.insert(name.to_string(), json!(format!("****@{}", parts[1])));
                    }
                    Some(v) if v.len() > 2 => {
                        let first = &v[..1];
                        let last = &v[v.len() - 1..];
                        let _ =
                            obj.insert(name.to_string(), json!(format!("{}...{}", first, last)));
                    }
                    Some(_v) => {
                        let _ = obj.insert(name.to_string(), json!("***"));
                    }
                    None => {
                        let _ = obj.insert(name.to_string(), Value::Null);
                    }
                }
            } else if let Ok(v) = row.try_get::<_, Value>(i) {
                obj.insert(name.to_string(), v);
            } else if let Ok(v) = row.try_get::<_, String>(i) {
                obj.insert(name.to_string(), Value::String(v));
            } else if let Ok(v) = row.try_get::<_, i64>(i) {
                obj.insert(name.to_string(), json!(v));
            } else if let Ok(v) = row.try_get::<_, f64>(i) {
                obj.insert(name.to_string(), json!(v));
            } else if let Ok(v) = row.try_get::<_, bool>(i) {
                obj.insert(name.to_string(), json!(v));
            } else if let Ok(v) = row.try_get::<_, Option<String>>(i) {
                obj.insert(
                    name.to_string(),
                    v.map(Value::String).unwrap_or(Value::Null),
                );
            }
        }
        results.push(Value::Object(obj));
    }

    Ok(json!({
        "sample": results,
        "count": results.len(),
        "randomized": order_by.is_none(),
        "anonymized_columns": anonymize.unwrap_or_default(),
    }))
}