ddb_core/mcp/
tools.rs

1//! MCP Tool implementations for DDB
2
3use crate::config::{Config, Table, TableCatalog};
4use crate::engine::QueryExecutor;
5use crate::lexer::Tokenizer;
6use crate::output::{format_results, OutputFormat};
7use crate::parser::Parser as SqlParser;
8use crate::parser::Statement;
9use anyhow::{anyhow, Result};
10use serde_json::{json, Value};
11use std::sync::Arc;
12use tokio::sync::RwLock;
13use tracing::info;
14
15pub struct ToolHandler {
16    config: Arc<RwLock<Config>>,
17    catalog: Arc<RwLock<TableCatalog>>,
18}
19
20impl ToolHandler {
21    pub fn new(config: Arc<RwLock<Config>>, catalog: Arc<RwLock<TableCatalog>>) -> Self {
22        Self { config, catalog }
23    }
24
25    /// List all available tools
26    pub fn list_tools(&self) -> Vec<Value> {
27        vec![
28            json!({
29                "name": "execute_query",
30                "description": "Execute a SQL query against DDB tables or CSV files",
31                "inputSchema": {
32                    "type": "object",
33                    "properties": {
34                        "query": { "type": "string", "description": "SQL query to execute" },
35                        "output_format": { "type": "string", "enum": ["json", "yaml", "csv", "table"], "default": "json" }
36                    },
37                    "required": ["query"]
38                }
39            }),
40            json!({
41                "name": "list_tables",
42                "description": "List all configured tables",
43                "inputSchema": {
44                    "type": "object",
45                    "properties": {
46                        "database": { "type": "string", "description": "Filter by database" }
47                    }
48                }
49            }),
50            json!({
51                "name": "describe_table",
52                "description": "Get table schema information",
53                "inputSchema": {
54                    "type": "object",
55                    "properties": {
56                        "table_name": { "type": "string", "description": "Table name" }
57                    },
58                    "required": ["table_name"]
59                }
60            }),
61        ]
62    }
63
64    /// Call a tool
65    pub async fn call_tool(&self, name: &str, arguments: Option<Value>) -> Result<Value> {
66        let args = arguments.unwrap_or(json!({}));
67
68        match name {
69            "execute_query" => self.execute_query(args).await,
70            "list_tables" => self.list_tables(args).await,
71            "describe_table" => self.describe_table(args).await,
72            _ => Err(anyhow!("Unknown tool: {}", name)),
73        }
74    }
75
76    async fn execute_query(&self, args: Value) -> Result<Value> {
77        let query = args["query"].as_str().ok_or_else(|| anyhow!("Missing query"))?;
78        let output_format = args["output_format"].as_str().unwrap_or("json");
79
80        info!("Executing: {}", query);
81
82        let mut tokenizer = Tokenizer::new();
83        let tokens = tokenizer.tokenize(query)?;
84        let mut parser = SqlParser::new(tokens);
85        let statement = parser.parse()?;
86
87        match statement {
88            Statement::Select(select_stmt) => {
89                let catalog = self.catalog.read().await;
90                let table_name = select_stmt.from.as_ref().ok_or_else(|| anyhow!("No FROM clause"))?;
91                let table = catalog.get_table(table_name).ok_or_else(|| anyhow!("Table not found"))?.clone();
92
93                let executor = QueryExecutor::new();
94                let results = executor.execute_select(&select_stmt, &table)?;
95                let format = OutputFormat::from_str(output_format);
96                let formatted = format_results(&results, format)?;
97
98                Ok(json!({
99                    "success": true,
100                    "rows": results.len(),
101                    "data": formatted
102                }))
103            }
104            _ => Ok(json!({ "success": false, "error": "Only SELECT supported" }))
105        }
106    }
107
108    async fn list_tables(&self, _args: Value) -> Result<Value> {
109        let catalog = self.catalog.read().await;
110        let tables = catalog.list_tables();
111        Ok(json!({ "success": true, "tables": tables }))
112    }
113
114    async fn describe_table(&self, args: Value) -> Result<Value> {
115        let table_name = args["table_name"].as_str().ok_or_else(|| anyhow!("Missing table_name"))?;
116        let catalog = self.catalog.read().await;
117        let table = catalog.get_table(table_name).ok_or_else(|| anyhow!("Table not found"))?;
118
119        let columns: Vec<Value> = table.columns.iter().map(|col| {
120            json!({ "name": col.name, "type": format!("{:?}", col.data_type), "nullable": col.nullable })
121        }).collect();
122
123        Ok(json!({
124            "success": true,
125            "table": {
126                "name": table.name,
127                "database": table.database,
128                "file": table.data_file,
129                "columns": columns
130            }
131        }))
132    }
133}
ddb_core/mcp/tools.rs

ddb_core/mcp/
tools.rs
