ballista_cli/
exec.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
18//! Execution functions
19
20use std::fs::File;
21use std::io::BufReader;
22use std::io::prelude::*;
23use std::sync::Arc;
24use std::sync::atomic::AtomicBool;
25use std::time::Instant;
26
27use datafusion::common::Result;
28use datafusion::prelude::SessionContext;
29use rustyline::Editor;
30use rustyline::error::ReadlineError;
31
32use crate::{
33    command::{Command, OutputFormat},
34    helper::CliHelper,
35    print_options::PrintOptions,
36};
37
38/// run and execute SQL statements and commands from a file, against a context with the given print options
39pub async fn exec_from_lines(
40    ctx: &SessionContext,
41    reader: &mut BufReader<File>,
42    print_options: &PrintOptions,
43) {
44    let mut query = "".to_owned();
45
46    for line in reader.lines() {
47        match line {
48            Ok(line) if line.starts_with("--") => {
49                continue;
50            }
51            Ok(line) => {
52                let line = line.trim_end();
53                query.push_str(line);
54                if line.ends_with(';') {
55                    match exec_and_print(ctx, print_options, query).await {
56                        Ok(_) => {}
57                        Err(err) => println!("{err:?}"),
58                    }
59
60                    #[allow(clippy::assigning_clones)]
61                    {
62                        query = "".to_owned();
63                    }
64                } else {
65                    query.push('\n');
66                }
67            }
68            _ => {
69                break;
70            }
71        }
72    }
73
74    // run the left over query if the last statement doesn't contain ‘;’
75    if !query.is_empty() {
76        match exec_and_print(ctx, print_options, query).await {
77            Ok(_) => {}
78            Err(err) => println!("{err:?}"),
79        }
80    }
81}
82
83pub async fn exec_from_files(
84    files: Vec<String>,
85    ctx: &SessionContext,
86    print_options: &PrintOptions,
87) {
88    let files = files
89        .into_iter()
90        .map(|file_path| File::open(file_path).unwrap())
91        .collect::<Vec<_>>();
92    for file in files {
93        let mut reader = BufReader::new(file);
94        exec_from_lines(ctx, &mut reader, print_options).await;
95    }
96}
97
98/// run and execute SQL statements and commands against a context with the given print options
99pub async fn exec_from_repl(
100    ctx: &SessionContext,
101    print_options: &mut PrintOptions,
102    tui_mode: Arc<AtomicBool>,
103) {
104    let mut rl = Editor::new().expect("created editor");
105    rl.set_helper(Some(CliHelper::new(
106        &ctx.task_ctx().session_config().options().sql_parser.dialect,
107        print_options.color,
108    )));
109    rl.load_history(".history").ok();
110
111    let mut print_options = print_options.clone();
112
113    loop {
114        match rl.readline("❯ ") {
115            Ok(line) if line.starts_with('\\') => {
116                rl.add_history_entry(line.trim_end()).unwrap();
117                let command = line.split_whitespace().collect::<Vec<_>>().join(" ");
118                if let Ok(cmd) = &command[1..].parse::<Command>() {
119                    match cmd {
120                        Command::Quit => break,
121                        Command::OutputFormat(subcommand) => {
122                            if let Some(subcommand) = subcommand {
123                                if let Ok(command) = subcommand.parse::<OutputFormat>() {
124                                    if let Err(e) =
125                                        command.execute(&mut print_options).await
126                                    {
127                                        eprintln!("{e}")
128                                    }
129                                } else {
130                                    eprintln!(
131                                        "'\\{}' is not a valid command",
132                                        &line[1..]
133                                    );
134                                }
135                            } else {
136                                println!("Output format is {:?}.", print_options.format);
137                            }
138                        }
139                        _ => {
140                            if let Err(e) = cmd
141                                .execute(ctx, &mut print_options, tui_mode.clone())
142                                .await
143                            {
144                                eprintln!("{e}")
145                            }
146                        }
147                    }
148                } else {
149                    eprintln!("'\\{}' is not a valid command", &line[1..]);
150                }
151            }
152            Ok(line) => {
153                rl.add_history_entry(line.trim_end()).unwrap();
154                match exec_and_print(ctx, &print_options, line).await {
155                    Ok(_) => {}
156                    Err(err) => eprintln!("{err:?}"),
157                }
158            }
159            Err(ReadlineError::Interrupted) => {
160                println!("^C");
161                continue;
162            }
163            Err(ReadlineError::Eof) => {
164                println!("\\q");
165                break;
166            }
167            Err(err) => {
168                eprintln!("Unknown error happened {err:?}");
169                break;
170            }
171        }
172    }
173
174    rl.save_history(".history").ok();
175}
176
177async fn exec_and_print(
178    ctx: &SessionContext,
179    print_options: &PrintOptions,
180    sql: String,
181) -> Result<()> {
182    let now = Instant::now();
183    let df = ctx.sql(&sql).await?;
184    let df = match print_options.maxrows {
185        datafusion_cli::print_options::MaxRows::Unlimited => df,
186        datafusion_cli::print_options::MaxRows::Limited(max_rows) => {
187            df.limit(0, Some(max_rows))?
188        }
189    };
190    let schema = Arc::new(df.schema().as_arrow().clone());
191    let results = df.collect().await?;
192    let rows = &results.iter().map(|b| b.num_rows()).sum::<usize>();
193    print_options.print_batches(schema, &results, now, *rows, &Default::default())?;
194
195    Ok(())
196}
ballista_cli/exec.rs

ballista_cli/
exec.rs
