datafusion_dft/tui/state/tabs/
sql.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use core::cell::RefCell;
19
20use color_eyre::Result;
21use datafusion::arrow::array::RecordBatch;
22use datafusion::sql::sqlparser::keywords;
23use log::debug;
24use ratatui::crossterm::event::KeyEvent;
25use ratatui::style::palette::tailwind;
26use ratatui::style::{Modifier, Style};
27use ratatui::widgets::TableState;
28use tokio::task::JoinHandle;
29use tui_textarea::TextArea;
30
31use crate::config::AppConfig;
32use crate::tui::ExecutionError;
33
34pub fn get_keywords() -> Vec<String> {
35    keywords::ALL_KEYWORDS
36        .iter()
37        .map(|k| k.to_string())
38        .collect()
39}
40
41pub fn keyword_regex() -> String {
42    format!(
43        "(?i)(^|[^a-zA-Z0-9\'\"`._]*?)({})($|[^a-zA-Z0-9\'\"`._]*)",
44        get_keywords().join("|")
45    )
46}
47
48pub fn keyword_style() -> Style {
49    Style::default()
50        .bg(tailwind::BLACK)
51        .fg(tailwind::YELLOW.c100)
52        .add_modifier(Modifier::BOLD)
53}
54
55#[derive(Debug, Default, PartialEq)]
56pub enum SQLTabMode {
57    #[default]
58    Normal,
59    DDL,
60}
61
62#[derive(Debug, Default)]
63pub struct SQLTabState<'app> {
64    editor: TextArea<'app>,
65    editor_editable: bool,
66    ddl_error: bool,
67    ddl_editor: TextArea<'app>,
68    ddl_editor_editable: bool,
69    query_results_state: Option<RefCell<TableState>>,
70    result_batches: Option<Vec<RecordBatch>>,
71    current_page: Option<usize>,
72    execution_error: Option<ExecutionError>,
73    execution_task: Option<JoinHandle<Result<()>>>,
74    mode: SQLTabMode,
75}
76
77impl SQLTabState<'_> {
78    pub fn new(config: &AppConfig) -> Self {
79        let empty_text = vec!["Enter a query here.".to_string()];
80        // TODO: Enable vim mode from config?
81        let mut textarea = TextArea::new(empty_text);
82        textarea.set_style(Style::default().fg(tailwind::WHITE));
83
84        let ddl_empty_text = vec!["Write your DDL here.".to_string()];
85        let mut ddl_textarea = TextArea::new(ddl_empty_text);
86        ddl_textarea.set_style(Style::default().fg(tailwind::WHITE));
87        if config.editor.experimental_syntax_highlighting {
88            textarea.set_search_pattern(keyword_regex()).unwrap();
89            textarea.set_search_style(keyword_style());
90            ddl_textarea.set_search_pattern(keyword_regex()).unwrap();
91            ddl_textarea.set_search_style(keyword_style());
92        };
93        Self {
94            editor: textarea,
95            editor_editable: false,
96            ddl_error: false,
97            ddl_editor: ddl_textarea,
98            ddl_editor_editable: false,
99            query_results_state: None,
100            result_batches: None,
101            current_page: None,
102            execution_error: None,
103            execution_task: None,
104            mode: SQLTabMode::default(),
105        }
106    }
107
108    pub fn query_results_state(&self) -> &Option<RefCell<TableState>> {
109        &self.query_results_state
110    }
111
112    pub fn refresh_query_results_state(&mut self) {
113        self.query_results_state = Some(RefCell::new(TableState::default()));
114    }
115
116    pub fn reset_execution_results(&mut self) {
117        self.result_batches = None;
118        self.current_page = None;
119        self.execution_error = None;
120        self.refresh_query_results_state();
121    }
122
123    pub fn editor(&self) -> TextArea {
124        // TODO: Figure out how to do this without clone. Probably need logic in handler to make
125        // updates to the Widget and then pass a ref
126        self.editor.clone()
127    }
128
129    pub fn ddl_error(&self) -> bool {
130        self.ddl_error
131    }
132
133    pub fn set_ddl_error(&mut self, error: bool) {
134        self.ddl_error = error;
135    }
136
137    pub fn ddl_editor(&self) -> TextArea {
138        self.ddl_editor.clone()
139    }
140
141    pub fn active_editor_cloned(&self) -> TextArea {
142        match self.mode {
143            SQLTabMode::Normal => self.editor.clone(),
144            SQLTabMode::DDL => self.ddl_editor.clone(),
145        }
146    }
147
148    pub fn clear_placeholder(&mut self) {
149        let default = "Enter a query here.";
150        let lines = self.editor.lines();
151        let content = lines.join("");
152        if content == default {
153            self.editor
154                .move_cursor(tui_textarea::CursorMove::Jump(0, 0));
155            self.editor.delete_str(default.len());
156        }
157    }
158
159    pub fn clear_editor(&mut self, config: &AppConfig) {
160        let mut textarea = TextArea::new(vec!["".to_string()]);
161        textarea.set_style(Style::default().fg(tailwind::WHITE));
162        if config.editor.experimental_syntax_highlighting {
163            textarea.set_search_pattern(keyword_regex()).unwrap();
164            textarea.set_search_style(keyword_style());
165        };
166        self.editor = textarea;
167    }
168
169    pub fn update_editor_content(&mut self, key: KeyEvent) {
170        match self.mode {
171            SQLTabMode::Normal => self.editor.input(key),
172            SQLTabMode::DDL => self.ddl_editor.input(key),
173        };
174    }
175
176    pub fn add_ddl_to_editor(&mut self, ddl: String) {
177        debug!("Adding DDL to editor: {}", ddl);
178        self.ddl_editor.delete_line_by_end();
179        self.ddl_editor.set_yank_text(ddl);
180        self.ddl_editor.paste();
181    }
182
183    pub fn edit(&mut self) {
184        match self.mode {
185            SQLTabMode::Normal => self.editor_editable = true,
186            SQLTabMode::DDL => self.ddl_editor_editable = true,
187        };
188    }
189
190    pub fn exit_edit(&mut self) {
191        match self.mode {
192            SQLTabMode::Normal => self.editor_editable = false,
193            SQLTabMode::DDL => self.ddl_editor_editable = false,
194        };
195    }
196
197    pub fn editor_editable(&self) -> bool {
198        match self.mode {
199            SQLTabMode::Normal => self.editor_editable,
200            SQLTabMode::DDL => self.ddl_editor_editable,
201        }
202    }
203
204    pub fn editable(&self) -> bool {
205        self.editor_editable || self.ddl_editor_editable
206    }
207
208    // TODO: Create Editor struct and move this there
209    pub fn next_word(&mut self) {
210        match self.mode {
211            SQLTabMode::Normal => self
212                .editor
213                .move_cursor(tui_textarea::CursorMove::WordForward),
214            SQLTabMode::DDL => self
215                .ddl_editor
216                .move_cursor(tui_textarea::CursorMove::WordForward),
217        }
218    }
219
220    // TODO: Create Editor struct and move this there
221    pub fn previous_word(&mut self) {
222        match self.mode {
223            SQLTabMode::Normal => self.editor.move_cursor(tui_textarea::CursorMove::WordBack),
224            SQLTabMode::DDL => self
225                .ddl_editor
226                .move_cursor(tui_textarea::CursorMove::WordBack),
227        }
228    }
229
230    pub fn delete_word(&mut self) {
231        match self.mode {
232            SQLTabMode::Normal => self.editor.delete_word(),
233            SQLTabMode::DDL => self.ddl_editor.delete_word(),
234        };
235    }
236
237    pub fn add_batch(&mut self, batch: RecordBatch) {
238        if let Some(batches) = self.result_batches.as_mut() {
239            batches.push(batch);
240        } else {
241            self.result_batches = Some(vec![batch]);
242        }
243    }
244
245    pub fn current_batch(&self) -> Option<&RecordBatch> {
246        match (self.current_page, self.result_batches.as_ref()) {
247            (Some(page), Some(batches)) => batches.get(page),
248            _ => None,
249        }
250    }
251
252    pub fn batches_count(&self) -> usize {
253        if let Some(batches) = &self.result_batches {
254            batches.len()
255        } else {
256            0
257        }
258    }
259
260    pub fn execution_error(&self) -> &Option<ExecutionError> {
261        &self.execution_error
262    }
263
264    pub fn set_execution_error(&mut self, error: ExecutionError) {
265        self.execution_error = Some(error);
266    }
267
268    pub fn current_page(&self) -> Option<usize> {
269        self.current_page
270    }
271
272    pub fn next_page(&mut self) {
273        if let Some(page) = self.current_page {
274            self.current_page = Some(page + 1);
275        } else {
276            self.current_page = Some(0);
277        }
278    }
279
280    pub fn previous_page(&mut self) {
281        if let Some(page) = self.current_page {
282            if page > 0 {
283                self.current_page = Some(page - 1);
284            }
285        }
286    }
287
288    pub fn execution_task(&mut self) -> &mut Option<JoinHandle<Result<()>>> {
289        &mut self.execution_task
290    }
291
292    pub fn set_execution_task(&mut self, task: JoinHandle<Result<()>>) {
293        self.execution_task = Some(task);
294    }
295
296    pub fn mode(&self) -> &SQLTabMode {
297        &self.mode
298    }
299
300    pub fn set_mode(&mut self, mode: SQLTabMode) {
301        self.mode = mode
302    }
303
304    /// Returns the SQL to be executed.  If no text is selected it returns the entire buffer else
305    /// it returns the current selection. For DDL it returns the entire buffer.
306    pub fn sql(&self) -> String {
307        match self.mode {
308            SQLTabMode::Normal => {
309                if let Some(((start_row, start_col), (end_row, end_col))) =
310                    self.editor.selection_range()
311                {
312                    if start_row == end_row {
313                        let line = &self.editor.lines()[start_row];
314                        line.chars()
315                            .skip(start_col)
316                            .take(end_col - start_col)
317                            .collect()
318                    } else {
319                        let lines: Vec<String> = self
320                            .editor
321                            .lines()
322                            .iter()
323                            .enumerate()
324                            .map(|(i, line)| {
325                                let selected_chars: Vec<char> = if i == start_row {
326                                    line.chars().skip(start_col).collect()
327                                } else if i == end_row {
328                                    line.chars().take(end_col).collect()
329                                } else {
330                                    line.chars().collect()
331                                };
332                                selected_chars.into_iter().collect()
333                            })
334                            .collect();
335                        lines.join("\n")
336                    }
337                } else {
338                    self.editor.lines().join("\n")
339                }
340            }
341            SQLTabMode::DDL => self.ddl_editor.lines().join("\n"),
342        }
343    }
344}
datafusion_dft/tui/state/tabs/sql.rs

datafusion_dft/tui/state/tabs/
sql.rs
