datafusion_tui/app/
core.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 std::fs::File;
19use std::io::{BufWriter, Write};
20
21use datafusion::prelude::SessionConfig;
22use log::info;
23
24use crate::app::datafusion::context::{Context, QueryResults};
25use crate::app::editor::Editor;
26use crate::app::error::Result;
27use crate::app::handlers::key_event_handler;
28use crate::cli::args::Args;
29use crate::events::Key;
30
31const DATAFUSION_RC: &str = ".datafusion/.datafusionrc";
32
33pub struct Tabs {
34    pub titles: Vec<&'static str>,
35    pub index: usize,
36}
37
38#[derive(Debug, Copy, Eq, PartialEq, Clone)]
39pub enum TabItem {
40    Editor,
41    QueryHistory,
42    Context,
43    Logs,
44}
45
46impl Default for TabItem {
47    fn default() -> TabItem {
48        TabItem::Editor
49    }
50}
51
52impl TabItem {
53    pub(crate) fn all_values() -> Vec<TabItem> {
54        vec![
55            TabItem::Editor,
56            TabItem::QueryHistory,
57            TabItem::Context,
58            TabItem::Logs,
59        ]
60    }
61
62    pub(crate) fn list_index(&self) -> usize {
63        return TabItem::all_values()
64            .iter()
65            .position(|x| x == self)
66            .unwrap();
67    }
68
69    pub(crate) fn title_with_key(&self) -> String {
70        return format!("{} [{}]", self.title(), self.list_index() + 1);
71    }
72
73    pub(crate) fn title(&self) -> &'static str {
74        match self {
75            TabItem::Editor => "SQL Editor",
76            TabItem::QueryHistory => "Query History",
77            TabItem::Context => "Context",
78            TabItem::Logs => "Logs",
79        }
80    }
81}
82
83impl TryFrom<char> for TabItem {
84    type Error = String;
85
86    fn try_from(value: char) -> std::result::Result<Self, Self::Error> {
87        match value {
88            '1' => Ok(TabItem::Editor),
89            '2' => Ok(TabItem::QueryHistory),
90            '3' => Ok(TabItem::Context),
91            '4' => Ok(TabItem::Logs),
92            i => Err(format!(
93                "Invalid tab index: {}, valid range is [1..={}]",
94                i, 4
95            )),
96        }
97    }
98}
99
100impl From<TabItem> for usize {
101    fn from(item: TabItem) -> Self {
102        match item {
103            TabItem::Editor => 1,
104            TabItem::QueryHistory => 2,
105            TabItem::Context => 3,
106            TabItem::Logs => 4,
107        }
108    }
109}
110
111#[derive(Debug, Copy, Clone)]
112pub enum InputMode {
113    Normal,
114    Editing,
115    Rc,
116}
117
118impl Default for InputMode {
119    fn default() -> InputMode {
120        InputMode::Normal
121    }
122}
123
124/// Status that determines whether app should continue or exit
125#[derive(PartialEq)]
126pub enum AppReturn {
127    Continue,
128    Exit,
129}
130
131/// App holds the state of the application
132pub struct App {
133    /// Application tabs
134    pub tab_item: TabItem,
135    /// Current input mode
136    pub input_mode: InputMode,
137    /// SQL Editor and it's state
138    pub editor: Editor,
139    /// DataFusion `ExecutionContext`
140    pub context: Context,
141    /// Results from DataFusion query
142    pub query_results: Option<QueryResults>,
143}
144
145impl App {
146    pub async fn new(args: Args) -> App {
147        let execution_config = SessionConfig::new().with_information_schema(true);
148        let mut ctx: Context = match (args.host, args.port) {
149            (Some(ref h), Some(p)) => Context::new_remote(h, p).await.unwrap(),
150            _ => Context::new_local(&execution_config).await,
151        };
152
153        let files = args.file;
154
155        let rc = App::get_rc_files(args.rc);
156
157        if !files.is_empty() {
158            ctx.exec_files(files).await
159        } else if !rc.is_empty() {
160            info!("Executing '~/.datafusion/.datafusionrc' file");
161            ctx.exec_files(rc).await
162        }
163
164        App {
165            tab_item: Default::default(),
166            input_mode: Default::default(),
167            editor: Editor::default(),
168            context: ctx,
169            query_results: None,
170        }
171    }
172
173    fn get_rc_files(rc: Option<Vec<String>>) -> Vec<String> {
174        match rc {
175            Some(file) => file,
176            None => {
177                let mut files = Vec::new();
178                let home = dirs::home_dir();
179                if let Some(p) = home {
180                    let home_rc = p.join(DATAFUSION_RC);
181                    if home_rc.exists() {
182                        files.push(home_rc.into_os_string().into_string().unwrap());
183                    }
184                }
185                files
186            }
187        }
188    }
189
190    pub async fn reload_rc(&mut self) {
191        let rc = App::get_rc_files(None);
192        self.context.exec_files(rc).await;
193        info!("Reloaded .datafusionrc");
194    }
195
196    pub fn write_rc(&self) -> Result<()> {
197        let text = self.editor.input.combine_lines();
198        let rc = App::get_rc_files(None);
199        let file = File::create(rc[0].clone())?;
200        let mut writer = BufWriter::new(file);
201        writer.write_all(text.as_bytes())?;
202        Ok(())
203    }
204
205    pub async fn key_handler(&mut self, key: Key) -> AppReturn {
206        key_event_handler(self, key).await.unwrap()
207    }
208
209    pub fn update_on_tick(&mut self) -> AppReturn {
210        AppReturn::Continue
211    }
212}
213
214#[cfg(test)]
215mod test {
216    use super::*;
217
218    #[test]
219    fn test_tab_item_from_char() {
220        assert!(TabItem::try_from('0').is_err());
221        assert_eq!(TabItem::Editor, TabItem::try_from('1').unwrap());
222        assert_eq!(TabItem::QueryHistory, TabItem::try_from('2').unwrap());
223        assert_eq!(TabItem::Context, TabItem::try_from('3').unwrap());
224        assert_eq!(TabItem::Logs, TabItem::try_from('4').unwrap());
225        assert!(TabItem::try_from('5').is_err());
226    }
227
228    #[test]
229    fn test_tab_item_to_usize() {
230        (0_usize..TabItem::all_values().len()).for_each(|i| {
231            assert_eq!(
232                TabItem::all_values()[i],
233                TabItem::try_from(format!("{}", i + 1).chars().next().unwrap()).unwrap()
234            );
235            assert_eq!(TabItem::all_values()[i].list_index(), i);
236        });
237    }
238}
datafusion_tui/app/core.rs

datafusion_tui/app/
core.rs
