json_mcp_server/json_tools/
query.rs

1use crate::mcp::protocol::{Tool, ToolCall, ToolResult};
2use crate::mcp::server::ToolHandler;
3use async_trait::async_trait;
4use serde_json::{json, Value};
5use std::collections::HashMap;
6use std::fs;
7
8pub struct JsonQuery;
9
10impl JsonQuery {
11    pub fn new() -> Self {
12        Self
13    }
14
15    fn create_query_tool() -> Tool {
16        Tool {
17            name: "json-query".to_string(),
18            description: "Execute JSONPath queries on JSON files. Supports complex queries with filtering, projection, and various output formats.".to_string(),
19            input_schema: json!({
20                "type": "object",
21                "properties": {
22                    "file_path": {
23                        "type": "string",
24                        "description": "Path to the JSON file to query"
25                    },
26                    "query": {
27                        "type": "string",
28                        "description": "JSONPath expression to execute (e.g., '$.users[?(@.age > 25)].name')"
29                    },
30                    "format": {
31                        "type": "string",
32                        "description": "Output format: 'json' (default), 'text', or 'table'",
33                        "enum": ["json", "text", "table"],
34                        "default": "json"
35                    }
36                },
37                "required": ["file_path", "query"]
38            })
39        }
40    }
41
42    async fn handle_query(&self, args: &HashMap<String, Value>) -> anyhow::Result<ToolResult> {
43        let file_path = args.get("file_path")
44            .and_then(|v| v.as_str())
45            .ok_or_else(|| anyhow::anyhow!(
46                "file_path is required. Usage example:\n{{\n  \"file_path\": \"./data.json\",\n  \"query\": \"$.users[0].name\"\n}}"
47            ))?;
48
49        let query = args.get("query")
50            .and_then(|v| v.as_str())
51            .ok_or_else(|| anyhow::anyhow!(
52                "query is required. Usage example:\n{{\n  \"file_path\": \"./data.json\",\n  \"query\": \"$.users[0].name\"\n}}\nUse JSONPath syntax: $ (root), .property, [index], [?(@.condition)]"
53            ))?;
54
55        let format = args.get("format")
56            .and_then(|v| v.as_str())
57            .unwrap_or("json");
58
59        // Read the file
60        let content = fs::read_to_string(file_path)
61            .map_err(|e| anyhow::anyhow!("Failed to read file '{}': {}", file_path, e))?;
62
63        // Execute JSONPath query
64        let results = match jsonpath_rust::JsonPathFinder::from_str(&content, query) {
65            Ok(finder) => finder.find(),
66            Err(e) => return Ok(ToolResult::error(format!("JSONPath query error: {}", e))),
67        };
68
69        // Format output based on requested format
70        let output = match format {
71            "json" => serde_json::to_string_pretty(&results)?,
72            "text" => self.format_as_text(&results),
73            "table" => self.format_as_table(&results),
74            _ => return Ok(ToolResult::error(format!("Unknown format: {}", format))),
75        };
76
77        Ok(ToolResult::success(format!(
78            "Query results from '{}' using JSONPath '{}':\n\n{}",
79            file_path, query, output
80        )))
81    }
82
83    fn format_as_text(&self, value: &Value) -> String {
84        match value {
85            Value::Array(arr) => {
86                arr.iter()
87                    .map(|v| match v {
88                        Value::String(s) => s.clone(),
89                        _ => v.to_string().trim_matches('"').to_string(),
90                    })
91                    .collect::<Vec<_>>()
92                    .join("\n")
93            },
94            Value::String(s) => s.clone(),
95            _ => value.to_string().trim_matches('"').to_string(),
96        }
97    }
98
99    fn format_as_table(&self, value: &Value) -> String {
100        match value {
101            Value::Array(arr) => {
102                if arr.is_empty() {
103                    return "No results found".to_string();
104                }
105
106                // Try to format as table if items are objects
107                if let Some(Value::Object(first_obj)) = arr.first() {
108                    let headers: Vec<String> = first_obj.keys().cloned().collect();
109                    let mut table = vec![headers.join(" | ")];
110                    table.push(headers.iter().map(|_| "---").collect::<Vec<_>>().join(" | "));
111
112                    for item in arr {
113                        if let Value::Object(obj) = item {
114                            let row: Vec<String> = headers.iter()
115                                .map(|h| obj.get(h)
116                                    .map(|v| match v {
117                                        Value::String(s) => s.clone(),
118                                        _ => v.to_string().trim_matches('"').to_string(),
119                                    })
120                                    .unwrap_or_else(|| "".to_string()))
121                                .collect();
122                            table.push(row.join(" | "));
123                        }
124                    }
125                    table.join("\n")
126                } else {
127                    // Simple list format
128                    arr.iter()
129                        .enumerate()
130                        .map(|(i, v)| format!("{}: {}", i + 1, match v {
131                            Value::String(s) => s.clone(),
132                            _ => v.to_string().trim_matches('"').to_string(),
133                        }))
134                        .collect::<Vec<_>>()
135                        .join("\n")
136                }
137            },
138            _ => value.to_string().trim_matches('"').to_string(),
139        }
140    }
141}
142
143#[async_trait]
144impl ToolHandler for JsonQuery {
145    async fn get_tools(&self) -> anyhow::Result<Vec<Tool>> {
146        Ok(vec![Self::create_query_tool()])
147    }
148
149    async fn call_tool(&self, tool_call: ToolCall) -> anyhow::Result<ToolResult> {
150        match tool_call.name.as_str() {
151            "json-query" => self.handle_query(&tool_call.arguments).await,
152            _ => Ok(ToolResult::error(format!("Unknown tool: {}", tool_call.name))),
153        }
154    }
155}