sql_cli/data/
query_executor.rs

1use anyhow::Result;
2
3use crate::api_client::QueryResponse;
4use crate::csv_datasource::CsvApiClient;
5use crate::data::datatable::DataTable;
6
7/// Trait for executing SQL queries against data sources
8pub trait QueryExecutor {
9    /// Execute a SQL query and return results
10    fn execute(&self, query: &str) -> Result<QueryResponse>;
11
12    /// Check if this executor can handle the given query
13    fn can_handle(&self, query: &str) -> bool;
14
15    /// Get row count without executing a query
16    fn row_count(&self) -> usize;
17
18    /// Get column count without executing a query
19    fn column_count(&self) -> usize;
20}
21
22/// Direct `DataTable` query executor (for simple SELECT * queries)
23pub struct DataTableExecutor {
24    datatable: std::sync::Arc<DataTable>,
25    table_name: String,
26}
27
28impl DataTableExecutor {
29    #[must_use]
30    pub fn new(datatable: std::sync::Arc<DataTable>, table_name: String) -> Self {
31        Self {
32            datatable,
33            table_name,
34        }
35    }
36}
37
38impl QueryExecutor for DataTableExecutor {
39    fn execute(&self, query: &str) -> Result<QueryResponse> {
40        // For now, only handle SELECT * FROM table
41        let _upper_query = query.trim().to_uppercase();
42        if !self.can_handle(query) {
43            return Err(anyhow::anyhow!(
44                "DataTableExecutor can only handle simple SELECT * queries"
45            ));
46        }
47
48        // Return a response that references the DataTable directly
49        // In the future, this will return a DataView
50        Ok(QueryResponse {
51            data: vec![], // Empty for now - TUI will use DataTable directly
52            count: self.datatable.row_count(),
53            query: crate::api_client::QueryInfo {
54                select: vec!["*".to_string()],
55                where_clause: None,
56                order_by: None,
57            },
58            source: Some("datatable".to_string()),
59            table: Some(self.table_name.clone()),
60            cached: Some(false),
61        })
62    }
63
64    fn can_handle(&self, query: &str) -> bool {
65        let upper_query = query.trim().to_uppercase();
66        upper_query.starts_with("SELECT *")
67            && !upper_query.contains(" WHERE ")
68            && !upper_query.contains(" ORDER BY ")
69            && !upper_query.contains(" LIMIT ")
70            && !upper_query.contains(" GROUP BY ")
71    }
72
73    fn row_count(&self) -> usize {
74        self.datatable.row_count()
75    }
76
77    fn column_count(&self) -> usize {
78        self.datatable.column_count()
79    }
80}
81
82/// CSV API Client query executor (for complex queries with WHERE, ORDER BY, etc.)
83pub struct CsvClientExecutor {
84    csv_client: CsvApiClient,
85    table_name: String,
86}
87
88impl CsvClientExecutor {
89    #[must_use]
90    pub fn new(csv_client: CsvApiClient, table_name: String) -> Self {
91        Self {
92            csv_client,
93            table_name,
94        }
95    }
96
97    pub fn from_csv(path: &str, table_name: &str, case_insensitive: bool) -> Result<Self> {
98        let mut csv_client = CsvApiClient::new();
99        csv_client.set_case_insensitive(case_insensitive);
100        csv_client.load_csv(path, table_name)?;
101        Ok(Self {
102            csv_client,
103            table_name: table_name.to_string(),
104        })
105    }
106
107    pub fn from_json(path: &str, table_name: &str, case_insensitive: bool) -> Result<Self> {
108        let mut csv_client = CsvApiClient::new();
109        csv_client.set_case_insensitive(case_insensitive);
110        csv_client.load_json(path, table_name)?;
111        Ok(Self {
112            csv_client,
113            table_name: table_name.to_string(),
114        })
115    }
116}
117
118impl QueryExecutor for CsvClientExecutor {
119    fn execute(&self, query: &str) -> Result<QueryResponse> {
120        let result = self.csv_client.query_csv(query)?;
121        Ok(QueryResponse {
122            data: result.data,
123            count: result.count,
124            query: crate::api_client::QueryInfo {
125                select: result.query.select,
126                where_clause: result.query.where_clause,
127                order_by: result.query.order_by,
128            },
129            source: Some("csv_client".to_string()),
130            table: Some(self.table_name.clone()),
131            cached: Some(false),
132        })
133    }
134
135    fn can_handle(&self, _query: &str) -> bool {
136        // CSV client can handle all queries
137        true
138    }
139
140    fn row_count(&self) -> usize {
141        // This is approximate - CSV client doesn't expose exact count
142        0
143    }
144
145    fn column_count(&self) -> usize {
146        // This is approximate - CSV client doesn't expose exact count
147        0
148    }
149}
150
151/// Composite query executor that tries multiple executors in order
152pub struct CompositeQueryExecutor {
153    executors: Vec<Box<dyn QueryExecutor>>,
154}
155
156impl Default for CompositeQueryExecutor {
157    fn default() -> Self {
158        Self::new()
159    }
160}
161
162impl CompositeQueryExecutor {
163    #[must_use]
164    pub fn new() -> Self {
165        Self {
166            executors: Vec::new(),
167        }
168    }
169
170    pub fn add_executor(&mut self, executor: Box<dyn QueryExecutor>) {
171        self.executors.push(executor);
172    }
173}
174
175impl QueryExecutor for CompositeQueryExecutor {
176    fn execute(&self, query: &str) -> Result<QueryResponse> {
177        // Try each executor in order
178        for executor in &self.executors {
179            if executor.can_handle(query) {
180                return executor.execute(query);
181            }
182        }
183        Err(anyhow::anyhow!("No executor can handle query: {}", query))
184    }
185
186    fn can_handle(&self, query: &str) -> bool {
187        self.executors.iter().any(|e| e.can_handle(query))
188    }
189
190    fn row_count(&self) -> usize {
191        // Return the row count from the first executor
192        self.executors.first().map_or(0, |e| e.row_count())
193    }
194
195    fn column_count(&self) -> usize {
196        // Return the column count from the first executor
197        self.executors.first().map_or(0, |e| e.column_count())
198    }
199}