Skip to content

Commit 19196ef

Browse files
achristmascarlachristmascarl
authored
bypass sqlparser (#193)
* add option to bypass query parser * refactor for bypassing
1 parent 215968f commit 19196ef

File tree

13 files changed

+138
-52
lines changed

13 files changed

+138
-52
lines changed

‎.config/rainfrog_config.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ mouse_mode = true
2020
[keybindings.Editor]
2121
"<Alt-q>" = "AbortQuery"
2222
"<F5>" = "SubmitEditorQuery"
23+
"<F7>" = "SubmitEditorQueryBypassParser"
2324
"<Alt-1>" = "FocusMenu"
2425
"<Alt-2>" = "FocusEditor"
2526
"<Alt-3>" = "FocusData"

‎src/action.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,8 @@ pub enum Action {
2626
Error(String),
2727
Help,
2828
SubmitEditorQuery,
29-
Query(Vec<String>, bool), // (query_lines, execution_confirmed)
29+
SubmitEditorQueryBypassParser,
30+
Query(Vec<String>, bool, bool), // (query_lines, execution_confirmed, bypass_parser)
3031
MenuPreview(MenuPreview, String, String), // (preview, schema, table)
3132
QueryToEditor(Vec<String>),
3233
ClearHistory,

‎src/app.rs

Lines changed: 25 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
#[cfg(not(feature = "termux"))]
22
use arboard::Clipboard;
3-
use color_eyre::eyre::Result;
3+
use color_eyre::eyre::{Result, eyre};
44
use crossterm::event::{KeyEvent, MouseEvent, MouseEventKind};
55
use ratatui::{
66
Frame,
@@ -29,8 +29,8 @@ use crate::{
2929
database::{self, Database, DbTaskResult, ExecutionType, Rows},
3030
focus::Focus,
3131
popups::{
32-
PopUp, PopUpPayload, confirm_export::ConfirmExport, confirm_query::ConfirmQuery, confirm_tx::ConfirmTx,
33-
exporting::Exporting, name_favorite::NameFavorite,
32+
PopUp, PopUpPayload, confirm_bypass::ConfirmBypass, confirm_export::ConfirmExport, confirm_query::ConfirmQuery,
33+
confirm_tx::ConfirmTx, exporting::Exporting, name_favorite::NameFavorite,
3434
},
3535
tui,
3636
ui::center,
@@ -199,13 +199,13 @@ impl App {
199199
}
200200
match database.get_query_results().await? {
201201
DbTaskResult::Finished(results) => {
202-
self.components.data.set_data_state(Some(results.results), Some(results.statement_type));
202+
self.components.data.set_data_state(Some(results.results), results.statement_type);
203203
self.state.last_query_end = Some(chrono::Utc::now());
204204
self.state.query_task_running = false;
205205
},
206206
DbTaskResult::ConfirmTx(rows_affected, statement) => {
207207
self.state.last_query_end = Some(chrono::Utc::now());
208-
self.set_popup(Box::new(ConfirmTx::new(rows_affected, statement.clone())));
208+
self.set_popup(Box::new(ConfirmTx::new(rows_affected, statement)));
209209
self.state.query_task_running = true;
210210
},
211211
DbTaskResult::Pending => {
@@ -239,7 +239,11 @@ impl App {
239239
self.set_focus(Focus::Editor);
240240
},
241241
Some(PopUpPayload::ConfirmQuery(query)) => {
242-
action_tx.send(Action::Query(vec![query], true))?;
242+
action_tx.send(Action::Query(vec![query], true, false))?;
243+
self.set_focus(Focus::Editor);
244+
},
245+
Some(PopUpPayload::ConfirmBypass(query)) => {
246+
action_tx.send(Action::Query(vec![query], true, true))?;
243247
self.set_focus(Focus::Editor);
244248
},
245249
Some(PopUpPayload::ConfirmExport(confirmed)) => {
@@ -261,7 +265,7 @@ impl App {
261265
let response = database.commit_tx().await?;
262266
self.state.last_query_end = Some(chrono::Utc::now());
263267
if let Some(results) = response {
264-
self.components.data.set_data_state(Some(results.results), Some(results.statement_type));
268+
self.components.data.set_data_state(Some(results.results), results.statement_type);
265269
self.set_focus(Focus::Editor);
266270
}
267271
},
@@ -370,30 +374,39 @@ impl App {
370374
let rows = database.load_menu().await;
371375
self.components.menu.set_table_list(Some(rows));
372376
},
373-
Action::Query(query_lines, confirmed) => 'query_action: {
377+
Action::Query(query_lines, confirmed, bypass) => 'query_action: {
374378
let query_string = query_lines.clone().join(" \n");
375379
if query_string.is_empty() {
376380
break 'query_action;
377381
}
378382
self.add_to_history(query_lines.clone());
379-
let execution_info = database::get_execution_type(query_string.clone(), *confirmed, driver);
383+
if *bypass && !confirmed {
384+
log::warn!("Bypassing parser");
385+
self.set_popup(Box::new(ConfirmBypass::new(query_string.clone())));
386+
break 'query_action;
387+
}
388+
let execution_info = match *bypass && *confirmed {
389+
true => Ok((ExecutionType::Normal, None)),
390+
false => database::get_execution_type(query_string.clone(), *confirmed, driver),
391+
};
380392
match execution_info {
381393
Ok((ExecutionType::Transaction, _)) => {
382394
self.components.data.set_loading();
383395
database.start_tx(query_string).await?;
384396
self.state.last_query_start = Some(chrono::Utc::now());
385397
self.state.last_query_end = None;
386398
},
387-
Ok((ExecutionType::Confirm, statement_type)) => {
399+
Ok((ExecutionType::Confirm, Some(statement_type))) => {
388400
self.set_popup(Box::new(ConfirmQuery::new(query_string.clone(), statement_type)));
389401
},
390402
Ok((ExecutionType::Normal, _)) => {
391403
self.components.data.set_loading();
392-
database.start_query(query_string).await?;
404+
database.start_query(query_string, *bypass).await?;
393405
self.state.last_query_start = Some(chrono::Utc::now());
394406
self.state.last_query_end = None;
395407
},
396408
Err(e) => self.components.data.set_data_state(Some(Err(e)), None),
409+
_ => self.components.data.set_data_state(Some(Err(eyre!("Missing statement type but not bypass"))), None),
397410
}
398411
},
399412
Action::AbortQuery => match database.abort_query().await {
@@ -417,7 +430,7 @@ impl App {
417430
action_tx.send(Action::QueryToEditor(vec![preview_query.clone()]))?;
418431
action_tx.send(Action::FocusEditor)?;
419432
action_tx.send(Action::FocusMenu)?;
420-
action_tx.send(Action::Query(vec![preview_query.clone()], false))?;
433+
action_tx.send(Action::Query(vec![preview_query.clone()], false, false))?;
421434
},
422435

423436
Action::RequestSaveFavorite(query_lines) => {

‎src/components/data.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -375,7 +375,7 @@ impl Component for Data<'_> {
375375
}
376376

377377
fn update(&mut self, action: Action, app_state: &AppState) -> Result<Option<Action>> {
378-
if let Action::Query(query, confirmed) = action {
378+
if let Action::Query(query, confirmed, bypass) = action {
379379
self.scrollable.reset_scroll();
380380
} else if let Action::ExportData(format) = action {
381381
let DataState::HasResults(rows) = &self.data_state else {
@@ -443,7 +443,7 @@ impl Component for Data<'_> {
443443
},
444444
DataState::StatementCompleted(statement) => {
445445
f.render_widget(
446-
Paragraph::new(format!("{} statement completed", statement_type_string(statement)))
446+
Paragraph::new(format!("{} statement completed", statement_type_string(Some(statement.clone()))))
447447
.wrap(Wrap { trim: false })
448448
.block(block),
449449
area,

‎src/components/editor.rs

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,7 @@ impl Editor<'_> {
6161
Input { key: Key::Enter, alt: true, .. } | Input { key: Key::Enter, ctrl: true, .. } => {
6262
if !app_state.query_task_running {
6363
if let Some(sender) = &self.command_tx {
64-
sender.send(Action::Query(self.textarea.lines().to_vec(), false))?;
64+
sender.send(Action::Query(self.textarea.lines().to_vec(), false, false))?;
6565
self.vim_state = Vim::new(Mode::Normal);
6666
self.vim_state.register_action_handler(self.command_tx.clone())?;
6767
self.cursor_style = Mode::Normal.cursor_style();
@@ -166,9 +166,14 @@ impl Component for Editor<'_> {
166166

167167
fn update(&mut self, action: Action, app_state: &AppState) -> Result<Option<Action>> {
168168
match action {
169+
Action::SubmitEditorQueryBypassParser => {
170+
if let Some(sender) = &self.command_tx {
171+
sender.send(Action::Query(self.textarea.lines().to_vec(), false, true))?;
172+
}
173+
},
169174
Action::SubmitEditorQuery => {
170175
if let Some(sender) = &self.command_tx {
171-
sender.send(Action::Query(self.textarea.lines().to_vec(), false))?;
176+
sender.send(Action::Query(self.textarea.lines().to_vec(), false, false))?;
172177
}
173178
},
174179
Action::QueryToEditor(lines) => {

‎src/database/mod.rs

Lines changed: 18 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ pub struct Rows {
4141
#[derive(Debug)]
4242
pub struct QueryResultsWithMetadata {
4343
pub results: Result<Rows>,
44-
pub statement_type: Statement,
44+
pub statement_type: Option<Statement>,
4545
}
4646

4747
#[derive(Debug, Clone, PartialEq)]
@@ -68,7 +68,7 @@ pub type QueryTask = JoinHandle<QueryResultsWithMetadata>;
6868

6969
pub enum DbTaskResult {
7070
Finished(QueryResultsWithMetadata),
71-
ConfirmTx(Option<u64>, Statement),
71+
ConfirmTx(Option<u64>, Option<Statement>),
7272
Pending,
7373
NoTask,
7474
}
@@ -83,7 +83,7 @@ pub trait Database {
8383

8484
/// Spawns a tokio task that runs the query. The task should
8585
/// expect to be polled via the `get_query_results()` method.
86-
async fn start_query(&mut self, query: String) -> Result<()>;
86+
async fn start_query(&mut self, query: String, bypass_parser: bool) -> Result<()>;
8787

8888
/// Aborts the tokio task running the active query or transaction.
8989
/// Some drivers also kill the process that was running the query,
@@ -144,11 +144,15 @@ fn get_first_query(query: String, driver: Driver) -> Result<(String, Statement),
144144
}
145145
}
146146

147-
pub fn get_execution_type(query: String, confirmed: bool, driver: Driver) -> Result<(ExecutionType, Statement)> {
147+
pub fn get_execution_type(
148+
query: String,
149+
confirmed: bool,
150+
driver: Driver,
151+
) -> Result<(ExecutionType, Option<Statement>)> {
148152
let first_query = get_first_query(query, driver);
149153

150154
match first_query {
151-
Ok((_, statement)) => Ok((get_default_execution_type(statement.clone(), confirmed), statement.clone())),
155+
Ok((_, statement)) => Ok((get_default_execution_type(statement.clone(), confirmed), Some(statement.clone()))),
152156
Err(e) => Err(eyre::Report::new(e)),
153157
}
154158
}
@@ -189,12 +193,15 @@ fn get_default_execution_type(statement: Statement, confirmed: bool) -> Executio
189193
}
190194
}
191195

192-
pub fn statement_type_string(statement: &Statement) -> String {
193-
format!("{statement:?}").split('(').collect::<Vec<&str>>()[0].split('{').collect::<Vec<&str>>()[0]
194-
.split('[')
195-
.collect::<Vec<&str>>()[0]
196-
.trim()
197-
.to_string()
196+
pub fn statement_type_string(statement: Option<Statement>) -> String {
197+
match statement {
198+
Some(stmt) => format!("{stmt:?}").split('(').collect::<Vec<&str>>()[0].split('{').collect::<Vec<&str>>()[0]
199+
.split('[')
200+
.collect::<Vec<&str>>()[0]
201+
.trim()
202+
.to_string(),
203+
None => "UNKNOWN".to_string(),
204+
}
198205
}
199206

200207
pub fn vec_to_string<T: std::string::ToString>(vec: Vec<T>) -> String {

‎src/database/mysql.rs

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -46,8 +46,14 @@ impl Database for MySqlDriver<'_> {
4646

4747
// since it's possible for raw_sql to execute multiple queries in a single string,
4848
// we only execute the first one and then drop the rest.
49-
async fn start_query(&mut self, query: String) -> Result<()> {
50-
let (first_query, statement_type) = super::get_first_query(query, Driver::MySql)?;
49+
async fn start_query(&mut self, query: String, bypass_parser: bool) -> Result<()> {
50+
let (first_query, statement_type) = match bypass_parser {
51+
true => (query, None),
52+
false => {
53+
let (first, stmt) = super::get_first_query(query, Driver::MySql)?;
54+
(first, Some(stmt))
55+
},
56+
};
5157
let pool = self.pool.clone().unwrap();
5258
self.querying_conn = Some(Arc::new(Mutex::new(pool.acquire().await?)));
5359
let conn = self.querying_conn.clone().unwrap();
@@ -154,18 +160,18 @@ impl Database for MySqlDriver<'_> {
154160
(
155161
QueryResultsWithMetadata {
156162
results: Ok(Rows { headers: vec![], rows: vec![], rows_affected: Some(rows_affected) }),
157-
statement_type,
163+
statement_type: Some(statement_type),
158164
},
159165
tx,
160166
)
161167
},
162168
Ok(Either::Right(rows)) => {
163169
log::info!("{:?} rows affected", rows.rows_affected);
164-
(QueryResultsWithMetadata { results: Ok(rows), statement_type }, tx)
170+
(QueryResultsWithMetadata { results: Ok(rows), statement_type: Some(statement_type) }, tx)
165171
},
166172
Err(e) => {
167173
log::error!("{e:?}");
168-
(QueryResultsWithMetadata { results: Err(e), statement_type }, tx)
174+
(QueryResultsWithMetadata { results: Err(e), statement_type: Some(statement_type) }, tx)
169175
},
170176
}
171177
})));

‎src/database/postgresql.rs

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -50,8 +50,14 @@ impl Database for PostgresDriver<'_> {
5050

5151
// since it's possible for raw_sql to execute multiple queries in a single string,
5252
// we only execute the first one and then drop the rest.
53-
async fn start_query(&mut self, query: String) -> Result<()> {
54-
let (first_query, statement_type) = super::get_first_query(query, Driver::Postgres)?;
53+
async fn start_query(&mut self, query: String, bypass_parser: bool) -> Result<()> {
54+
let (first_query, statement_type) = match bypass_parser {
55+
true => (query, None),
56+
false => {
57+
let (first, stmt) = super::get_first_query(query, Driver::Postgres)?;
58+
(first, Some(stmt))
59+
},
60+
};
5561
let pool = self.pool.clone().unwrap();
5662
self.querying_conn = Some(Arc::new(Mutex::new(pool.acquire().await?)));
5763
let conn = self.querying_conn.clone().unwrap();
@@ -170,18 +176,18 @@ impl Database for PostgresDriver<'_> {
170176
(
171177
QueryResultsWithMetadata {
172178
results: Ok(Rows { headers: vec![], rows: vec![], rows_affected: Some(rows_affected) }),
173-
statement_type,
179+
statement_type: Some(statement_type),
174180
},
175181
tx,
176182
)
177183
},
178184
Ok(Either::Right(rows)) => {
179185
log::info!("{:?} rows affected", rows.rows_affected);
180-
(QueryResultsWithMetadata { results: Ok(rows), statement_type }, tx)
186+
(QueryResultsWithMetadata { results: Ok(rows), statement_type: Some(statement_type) }, tx)
181187
},
182188
Err(e) => {
183189
log::error!("{e:?}");
184-
(QueryResultsWithMetadata { results: Err(e), statement_type }, tx)
190+
(QueryResultsWithMetadata { results: Err(e), statement_type: Some(statement_type) }, tx)
185191
},
186192
}
187193
})));

‎src/database/sqlite.rs

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -43,8 +43,14 @@ impl Database for SqliteDriver<'_> {
4343

4444
// since it's possible for raw_sql to execute multiple queries in a single string,
4545
// we only execute the first one and then drop the rest.
46-
async fn start_query(&mut self, query: String) -> Result<()> {
47-
let (first_query, statement_type) = super::get_first_query(query, Driver::Sqlite)?;
46+
async fn start_query(&mut self, query: String, bypass_parser: bool) -> Result<()> {
47+
let (first_query, statement_type) = match bypass_parser {
48+
true => (query, None),
49+
false => {
50+
let (first, stmt) = super::get_first_query(query, Driver::Sqlite)?;
51+
(first, Some(stmt))
52+
},
53+
};
4854
let pool = self.pool.clone().unwrap();
4955
self.task = Some(SqliteTask::Query(tokio::spawn(async move {
5056
let results = query_with_pool(pool, first_query.clone()).await;
@@ -125,18 +131,18 @@ impl Database for SqliteDriver<'_> {
125131
(
126132
QueryResultsWithMetadata {
127133
results: Ok(Rows { headers: vec![], rows: vec![], rows_affected: Some(rows_affected) }),
128-
statement_type,
134+
statement_type: Some(statement_type),
129135
},
130136
tx,
131137
)
132138
},
133139
Ok(Either::Right(rows)) => {
134140
log::info!("{:?} rows affected", rows.rows_affected);
135-
(QueryResultsWithMetadata { results: Ok(rows), statement_type }, tx)
141+
(QueryResultsWithMetadata { results: Ok(rows), statement_type: Some(statement_type) }, tx)
136142
},
137143
Err(e) => {
138144
log::error!("{e:?}");
139-
(QueryResultsWithMetadata { results: Err(e), statement_type }, tx)
145+
(QueryResultsWithMetadata { results: Err(e), statement_type: Some(statement_type) }, tx)
140146
},
141147
}
142148
})));

‎src/popups/confirm_bypass.rs

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
use crossterm::event::KeyCode;
2+
3+
use super::{PopUp, PopUpPayload};
4+
5+
#[derive(Debug)]
6+
pub struct ConfirmBypass {
7+
pending_query: String,
8+
}
9+
10+
impl ConfirmBypass {
11+
pub fn new(pending_query: String) -> Self {
12+
Self { pending_query }
13+
}
14+
}
15+
16+
impl PopUp for ConfirmBypass {
17+
fn handle_key_events(
18+
&mut self,
19+
key: crossterm::event::KeyEvent,
20+
app_state: &mut crate::app::AppState,
21+
) -> color_eyre::eyre::Result<Option<PopUpPayload>> {
22+
match key.code {
23+
KeyCode::Char('Y') => Ok(Some(PopUpPayload::ConfirmBypass(self.pending_query.to_owned()))),
24+
KeyCode::Char('N') | KeyCode::Esc => Ok(Some(PopUpPayload::SetDataTable(None, None))),
25+
_ => Ok(None),
26+
}
27+
}
28+
29+
fn get_cta_text(&self, app_state: &crate::app::AppState) -> String {
30+
"Are you sure you want to bypass the query parser? The query will not be wrapped in a transaction, so it cannot be undone.".to_string()
31+
}
32+
33+
fn get_actions_text(&self, app_state: &crate::app::AppState) -> String {
34+
"[Y]es to confirm | [N]o to cancel".to_string()
35+
}
36+
}

0 commit comments

Comments
 (0)