ddb_core/mcp/
tools.rs

1//! MCP Tool implementations for DDB
2
3use crate::config::{Config, 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: 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. Supports SELECT (with JOIN, GROUP BY, HAVING), INSERT, UPDATE, DELETE, and UPSERT statements.",
31                "inputSchema": {
32                    "type": "object",
33                    "properties": {
34                        "query": { "type": "string", "description": "SQL query to execute (SELECT, INSERT, UPDATE, DELETE, UPSERT)" },
35                        "output_format": { "type": "string", "enum": ["json", "yaml", "csv", "table"], "default": "json", "description": "Output format for SELECT queries" }
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                // Use execute_select_with_catalog to support JOINs
95                let results = executor.execute_select_with_catalog(&select_stmt, &table, Some(&catalog))?;
96                let format = OutputFormat::from_str(output_format);
97                let formatted = format_results(&results, format)?;
98
99                Ok(json!({
100                    "success": true,
101                    "statement": "SELECT",
102                    "rows": results.len(),
103                    "data": formatted
104                }))
105            }
106            Statement::Insert(insert_stmt) => {
107                let catalog = self.catalog.read().await;
108                let table = catalog.get_table(&insert_stmt.table)
109                    .ok_or_else(|| anyhow!("Table '{}' not found", insert_stmt.table))?
110                    .clone();
111
112                let executor = QueryExecutor::new();
113                let rows_inserted = executor.execute_insert(&insert_stmt, &table)?;
114
115                Ok(json!({
116                    "success": true,
117                    "statement": "INSERT",
118                    "rows_inserted": rows_inserted,
119                    "message": format!("{} row(s) inserted", rows_inserted)
120                }))
121            }
122            Statement::Update(update_stmt) => {
123                let catalog = self.catalog.read().await;
124                let table = catalog.get_table(&update_stmt.table)
125                    .ok_or_else(|| anyhow!("Table '{}' not found", update_stmt.table))?
126                    .clone();
127
128                let executor = QueryExecutor::new();
129                let rows_updated = executor.execute_update(&update_stmt, &table)?;
130
131                Ok(json!({
132                    "success": true,
133                    "statement": "UPDATE",
134                    "rows_updated": rows_updated,
135                    "message": format!("{} row(s) updated", rows_updated)
136                }))
137            }
138            Statement::Delete(delete_stmt) => {
139                let catalog = self.catalog.read().await;
140                let table = catalog.get_table(&delete_stmt.table)
141                    .ok_or_else(|| anyhow!("Table '{}' not found", delete_stmt.table))?
142                    .clone();
143
144                let executor = QueryExecutor::new();
145                let rows_deleted = executor.execute_delete(&delete_stmt, &table)?;
146
147                Ok(json!({
148                    "success": true,
149                    "statement": "DELETE",
150                    "rows_deleted": rows_deleted,
151                    "message": format!("{} row(s) deleted", rows_deleted)
152                }))
153            }
154            Statement::Upsert(upsert_stmt) => {
155                let catalog = self.catalog.read().await;
156                let table = catalog.get_table(&upsert_stmt.table)
157                    .ok_or_else(|| anyhow!("Table '{}' not found", upsert_stmt.table))?
158                    .clone();
159
160                let executor = QueryExecutor::new();
161                let (rows_inserted, rows_updated) = executor.execute_upsert(&upsert_stmt, &table)?;
162
163                Ok(json!({
164                    "success": true,
165                    "statement": "UPSERT",
166                    "rows_inserted": rows_inserted,
167                    "rows_updated": rows_updated,
168                    "message": format!("{} row(s) inserted, {} row(s) updated", rows_inserted, rows_updated)
169                }))
170            }
171            _ => Ok(json!({
172                "success": false,
173                "error": "Statement type not supported in MCP server"
174            }))
175        }
176    }
177
178    async fn list_tables(&self, _args: Value) -> Result<Value> {
179        let catalog = self.catalog.read().await;
180        let tables = catalog.list_tables();
181        Ok(json!({ "success": true, "tables": tables }))
182    }
183
184    async fn describe_table(&self, args: Value) -> Result<Value> {
185        let table_name = args["table_name"].as_str().ok_or_else(|| anyhow!("Missing table_name"))?;
186        let catalog = self.catalog.read().await;
187        let table = catalog.get_table(table_name).ok_or_else(|| anyhow!("Table not found"))?;
188
189        let columns: Vec<Value> = table.columns.iter().map(|col| {
190            json!({ "name": col.name, "type": format!("{:?}", col.data_type), "nullable": col.nullable })
191        }).collect();
192
193        Ok(json!({
194            "success": true,
195            "table": {
196                "name": table.name,
197                "database": table.database,
198                "file": table.data_file,
199                "columns": columns
200            }
201        }))
202    }
203}
ddb_core/mcp/tools.rs

ddb_core/mcp/
tools.rs
