public-backup/rainfrog
1
0
Fork 0
mirror of https://github.com/achristmascarl/rainfrog.git synced 2025-12-05 19:06:12 -06:00
Code Issues Packages Projects Releases 30 Wiki Activity

bypass sqlparser (#193)

Browse Source 
* add option to bypass query parser

* refactor for bypassing
This commit is contained in:
carl
2025-07-24 17:12:04 -04:00
committed by GitHub
parent 215968fad1
commit 19196ef279
13 changed files with 138 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
.config/rainfrog_config.toml
View File

@@ -20,6 +20,7 @@ mouse_mode = true
[keybindings.Editor]
"<Alt-q>" = "AbortQuery"
"<F5>" = "SubmitEditorQuery"
"<F7>" = "SubmitEditorQueryBypassParser"
"<Alt-1>" = "FocusMenu"
"<Alt-2>" = "FocusEditor"
"<Alt-3>" = "FocusData"

3
src/action.rs
View File

@@ -26,7 +26,8 @@ pub enum Action {
Error(String),
Help,
SubmitEditorQuery,
Query(Vec<String>, bool), // (query_lines, execution_confirmed)
SubmitEditorQueryBypassParser,
Query(Vec<String>, bool, bool), // (query_lines, execution_confirmed, bypass_parser)
MenuPreview(MenuPreview, String, String), // (preview, schema, table)
QueryToEditor(Vec<String>),
ClearHistory,

37
src/app.rs
View File

@@ -1,6 +1,6 @@
#[cfg(not(feature = "termux"))]
use arboard::Clipboard;
use color_eyre::eyre::Result;
use color_eyre::eyre::{Result, eyre};
use crossterm::event::{KeyEvent, MouseEvent, MouseEventKind};
use ratatui::{
Frame,
@@ -29,8 +29,8 @@ use crate::{
database::{self, Database, DbTaskResult, ExecutionType, Rows},
focus::Focus,
popups::{
PopUp, PopUpPayload, confirm_export::ConfirmExport, confirm_query::ConfirmQuery, confirm_tx::ConfirmTx,
exporting::Exporting, name_favorite::NameFavorite,
PopUp, PopUpPayload, confirm_bypass::ConfirmBypass, confirm_export::ConfirmExport, confirm_query::ConfirmQuery,
confirm_tx::ConfirmTx, exporting::Exporting, name_favorite::NameFavorite,
},
tui,
ui::center,
@@ -199,13 +199,13 @@ impl App {
}
match database.get_query_results().await? {
DbTaskResult::Finished(results) => {
self.components.data.set_data_state(Some(results.results), Some(results.statement_type));
self.components.data.set_data_state(Some(results.results), results.statement_type);
self.state.last_query_end = Some(chrono::Utc::now());
self.state.query_task_running = false;
},
DbTaskResult::ConfirmTx(rows_affected, statement) => {
self.state.last_query_end = Some(chrono::Utc::now());
self.set_popup(Box::new(ConfirmTx::new(rows_affected, statement.clone())));
self.set_popup(Box::new(ConfirmTx::new(rows_affected, statement)));
self.state.query_task_running = true;
},
DbTaskResult::Pending => {
@@ -239,7 +239,11 @@ impl App {
self.set_focus(Focus::Editor);
},
Some(PopUpPayload::ConfirmQuery(query)) => {
action_tx.send(Action::Query(vec![query], true))?;
action_tx.send(Action::Query(vec![query], true, false))?;
self.set_focus(Focus::Editor);
},
Some(PopUpPayload::ConfirmBypass(query)) => {
action_tx.send(Action::Query(vec![query], true, true))?;
self.set_focus(Focus::Editor);
},
Some(PopUpPayload::ConfirmExport(confirmed)) => {
@@ -261,7 +265,7 @@ impl App {
let response = database.commit_tx().await?;
self.state.last_query_end = Some(chrono::Utc::now());
if let Some(results) = response {
self.components.data.set_data_state(Some(results.results), Some(results.statement_type));
self.components.data.set_data_state(Some(results.results), results.statement_type);
self.set_focus(Focus::Editor);
}
},
@@ -370,13 +374,21 @@ impl App {
let rows = database.load_menu().await;
self.components.menu.set_table_list(Some(rows));
},
Action::Query(query_lines, confirmed) => 'query_action: {
Action::Query(query_lines, confirmed, bypass) => 'query_action: {
let query_string = query_lines.clone().join(" \n");
if query_string.is_empty() {
break 'query_action;
}
self.add_to_history(query_lines.clone());
let execution_info = database::get_execution_type(query_string.clone(), *confirmed, driver);
if *bypass && !confirmed {
log::warn!("Bypassing parser");
self.set_popup(Box::new(ConfirmBypass::new(query_string.clone())));
break 'query_action;
}
let execution_info = match *bypass && *confirmed {
true => Ok((ExecutionType::Normal, None)),
false => database::get_execution_type(query_string.clone(), *confirmed, driver),
};
match execution_info {
Ok((ExecutionType::Transaction, _)) => {
self.components.data.set_loading();
@@ -384,16 +396,17 @@ impl App {
self.state.last_query_start = Some(chrono::Utc::now());
self.state.last_query_end = None;
},
Ok((ExecutionType::Confirm, statement_type)) => {
Ok((ExecutionType::Confirm, Some(statement_type))) => {
self.set_popup(Box::new(ConfirmQuery::new(query_string.clone(), statement_type)));
},
Ok((ExecutionType::Normal, _)) => {
self.components.data.set_loading();
database.start_query(query_string).await?;
database.start_query(query_string, *bypass).await?;
self.state.last_query_start = Some(chrono::Utc::now());
self.state.last_query_end = None;
},
Err(e) => self.components.data.set_data_state(Some(Err(e)), None),
_ => self.components.data.set_data_state(Some(Err(eyre!("Missing statement type but not bypass"))), None),
}
},
Action::AbortQuery => match database.abort_query().await {
@@ -417,7 +430,7 @@ impl App {
action_tx.send(Action::QueryToEditor(vec![preview_query.clone()]))?;
action_tx.send(Action::FocusEditor)?;
action_tx.send(Action::FocusMenu)?;
action_tx.send(Action::Query(vec![preview_query.clone()], false))?;
action_tx.send(Action::Query(vec![preview_query.clone()], false, false))?;
},
Action::RequestSaveFavorite(query_lines) => {

4
src/components/data.rs
View File

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

9
src/components/editor.rs
View File

@@ -61,7 +61,7 @@ impl Editor<'_> {
Input { key: Key::Enter, alt: true, .. } | Input { key: Key::Enter, ctrl: true, .. } => {
if !app_state.query_task_running {
if let Some(sender) = &self.command_tx {
sender.send(Action::Query(self.textarea.lines().to_vec(), false))?;
sender.send(Action::Query(self.textarea.lines().to_vec(), false, false))?;
self.vim_state = Vim::new(Mode::Normal);
self.vim_state.register_action_handler(self.command_tx.clone())?;
self.cursor_style = Mode::Normal.cursor_style();
@@ -166,9 +166,14 @@ impl Component for Editor<'_> {
fn update(&mut self, action: Action, app_state: &AppState) -> Result<Option<Action>> {
match action {
Action::SubmitEditorQueryBypassParser => {
if let Some(sender) = &self.command_tx {
sender.send(Action::Query(self.textarea.lines().to_vec(), false, true))?;
}
},
Action::SubmitEditorQuery => {
if let Some(sender) = &self.command_tx {
sender.send(Action::Query(self.textarea.lines().to_vec(), false))?;
sender.send(Action::Query(self.textarea.lines().to_vec(), false, false))?;
}
},
Action::QueryToEditor(lines) => {

29
src/database/mod.rs
View File

@@ -41,7 +41,7 @@ pub struct Rows {
#[derive(Debug)]
pub struct QueryResultsWithMetadata {
pub results: Result<Rows>,
pub statement_type: Statement,
pub statement_type: Option<Statement>,
}
#[derive(Debug, Clone, PartialEq)]
@@ -68,7 +68,7 @@ pub type QueryTask = JoinHandle<QueryResultsWithMetadata>;
pub enum DbTaskResult {
Finished(QueryResultsWithMetadata),
ConfirmTx(Option<u64>, Statement),
ConfirmTx(Option<u64>, Option<Statement>),
Pending,
NoTask,
}
@@ -83,7 +83,7 @@ pub trait Database {
/// Spawns a tokio task that runs the query. The task should
/// expect to be polled via the `get_query_results()` method.
async fn start_query(&mut self, query: String) -> Result<()>;
async fn start_query(&mut self, query: String, bypass_parser: bool) -> Result<()>;
/// Aborts the tokio task running the active query or transaction.
/// 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),
}
}
pub fn get_execution_type(query: String, confirmed: bool, driver: Driver) -> Result<(ExecutionType, Statement)> {
pub fn get_execution_type(
query: String,
confirmed: bool,
driver: Driver,
) -> Result<(ExecutionType, Option<Statement>)> {
let first_query = get_first_query(query, driver);
match first_query {
Ok((_, statement)) => Ok((get_default_execution_type(statement.clone(), confirmed), statement.clone())),
Ok((_, statement)) => Ok((get_default_execution_type(statement.clone(), confirmed), Some(statement.clone()))),
Err(e) => Err(eyre::Report::new(e)),
}
}
@@ -189,12 +193,15 @@ fn get_default_execution_type(statement: Statement, confirmed: bool) -> Executio
}
}
pub fn statement_type_string(statement: &Statement) -> String {
format!("{statement:?}").split('(').collect::<Vec<&str>>()[0].split('{').collect::<Vec<&str>>()[0]
.split('[')
.collect::<Vec<&str>>()[0]
.trim()
.to_string()
pub fn statement_type_string(statement: Option<Statement>) -> String {
match statement {
Some(stmt) => format!("{stmt:?}").split('(').collect::<Vec<&str>>()[0].split('{').collect::<Vec<&str>>()[0]
.split('[')
.collect::<Vec<&str>>()[0]
.trim()
.to_string(),
None => "UNKNOWN".to_string(),
}
}
pub fn vec_to_string<T: std::string::ToString>(vec: Vec<T>) -> String {

16
src/database/mysql.rs
View File

@@ -46,8 +46,14 @@ impl Database for MySqlDriver<'_> {
// since it's possible for raw_sql to execute multiple queries in a single string,
// we only execute the first one and then drop the rest.
async fn start_query(&mut self, query: String) -> Result<()> {
let (first_query, statement_type) = super::get_first_query(query, Driver::MySql)?;
async fn start_query(&mut self, query: String, bypass_parser: bool) -> Result<()> {
let (first_query, statement_type) = match bypass_parser {
true => (query, None),
false => {
let (first, stmt) = super::get_first_query(query, Driver::MySql)?;
(first, Some(stmt))
},
};
let pool = self.pool.clone().unwrap();
self.querying_conn = Some(Arc::new(Mutex::new(pool.acquire().await?)));
let conn = self.querying_conn.clone().unwrap();
@@ -154,18 +160,18 @@ impl Database for MySqlDriver<'_> {
(
QueryResultsWithMetadata {
results: Ok(Rows { headers: vec![], rows: vec![], rows_affected: Some(rows_affected) }),
statement_type,
statement_type: Some(statement_type),
},
tx,
)
},
Ok(Either::Right(rows)) => {
log::info!("{:?} rows affected", rows.rows_affected);
(QueryResultsWithMetadata { results: Ok(rows), statement_type }, tx)
(QueryResultsWithMetadata { results: Ok(rows), statement_type: Some(statement_type) }, tx)
},
Err(e) => {
log::error!("{e:?}");
(QueryResultsWithMetadata { results: Err(e), statement_type }, tx)
(QueryResultsWithMetadata { results: Err(e), statement_type: Some(statement_type) }, tx)
},
}
})));

16
src/database/postgresql.rs
View File

@@ -50,8 +50,14 @@ impl Database for PostgresDriver<'_> {
// since it's possible for raw_sql to execute multiple queries in a single string,
// we only execute the first one and then drop the rest.
async fn start_query(&mut self, query: String) -> Result<()> {
let (first_query, statement_type) = super::get_first_query(query, Driver::Postgres)?;
async fn start_query(&mut self, query: String, bypass_parser: bool) -> Result<()> {
let (first_query, statement_type) = match bypass_parser {
true => (query, None),
false => {
let (first, stmt) = super::get_first_query(query, Driver::Postgres)?;
(first, Some(stmt))
},
};
let pool = self.pool.clone().unwrap();
self.querying_conn = Some(Arc::new(Mutex::new(pool.acquire().await?)));
let conn = self.querying_conn.clone().unwrap();
@@ -170,18 +176,18 @@ impl Database for PostgresDriver<'_> {
(
QueryResultsWithMetadata {
results: Ok(Rows { headers: vec![], rows: vec![], rows_affected: Some(rows_affected) }),
statement_type,
statement_type: Some(statement_type),
},
tx,
)
},
Ok(Either::Right(rows)) => {
log::info!("{:?} rows affected", rows.rows_affected);
(QueryResultsWithMetadata { results: Ok(rows), statement_type }, tx)
(QueryResultsWithMetadata { results: Ok(rows), statement_type: Some(statement_type) }, tx)
},
Err(e) => {
log::error!("{e:?}");
(QueryResultsWithMetadata { results: Err(e), statement_type }, tx)
(QueryResultsWithMetadata { results: Err(e), statement_type: Some(statement_type) }, tx)
},
}
})));

16
src/database/sqlite.rs
View File

@@ -43,8 +43,14 @@ impl Database for SqliteDriver<'_> {
// since it's possible for raw_sql to execute multiple queries in a single string,
// we only execute the first one and then drop the rest.
async fn start_query(&mut self, query: String) -> Result<()> {
let (first_query, statement_type) = super::get_first_query(query, Driver::Sqlite)?;
async fn start_query(&mut self, query: String, bypass_parser: bool) -> Result<()> {
let (first_query, statement_type) = match bypass_parser {
true => (query, None),
false => {
let (first, stmt) = super::get_first_query(query, Driver::Sqlite)?;
(first, Some(stmt))
},
};
let pool = self.pool.clone().unwrap();
self.task = Some(SqliteTask::Query(tokio::spawn(async move {
let results = query_with_pool(pool, first_query.clone()).await;
@@ -125,18 +131,18 @@ impl Database for SqliteDriver<'_> {
(
QueryResultsWithMetadata {
results: Ok(Rows { headers: vec![], rows: vec![], rows_affected: Some(rows_affected) }),
statement_type,
statement_type: Some(statement_type),
},
tx,
)
},
Ok(Either::Right(rows)) => {
log::info!("{:?} rows affected", rows.rows_affected);
(QueryResultsWithMetadata { results: Ok(rows), statement_type }, tx)
(QueryResultsWithMetadata { results: Ok(rows), statement_type: Some(statement_type) }, tx)
},
Err(e) => {
log::error!("{e:?}");
(QueryResultsWithMetadata { results: Err(e), statement_type }, tx)
(QueryResultsWithMetadata { results: Err(e), statement_type: Some(statement_type) }, tx)
},
}
})));

36
src/popups/confirm_bypass.rs Normal file
View File

@@ -0,0 +1,36 @@
use crossterm::event::KeyCode;
use super::{PopUp, PopUpPayload};
#[derive(Debug)]
pub struct ConfirmBypass {
pending_query: String,
}
impl ConfirmBypass {
pub fn new(pending_query: String) -> Self {
Self { pending_query }
}
}
impl PopUp for ConfirmBypass {
fn handle_key_events(
&mut self,
key: crossterm::event::KeyEvent,
app_state: &mut crate::app::AppState,
) -> color_eyre::eyre::Result<Option<PopUpPayload>> {
match key.code {
KeyCode::Char('Y') => Ok(Some(PopUpPayload::ConfirmBypass(self.pending_query.to_owned()))),
KeyCode::Char('N') | KeyCode::Esc => Ok(Some(PopUpPayload::SetDataTable(None, None))),
_ => Ok(None),
}
}
fn get_cta_text(&self, app_state: &crate::app::AppState) -> String {
"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()
}
fn get_actions_text(&self, app_state: &crate::app::AppState) -> String {
"[Y]es to confirm | [N]o to cancel".to_string()
}
}

4
src/popups/confirm_query.rs
View File

@@ -34,13 +34,13 @@ impl PopUp for ConfirmQuery {
Statement::Explain { statement, .. } => {
format!(
"Are you sure you want to run an EXPLAIN ANALYZE that will run a {} statement?",
statement_type_string(&statement).to_uppercase(),
statement_type_string(Some(*statement.clone())).to_uppercase(),
)
},
_ => {
format!(
"Are you sure you want to use a {} statement?",
statement_type_string(&self.statement_type).to_uppercase()
statement_type_string(Some(self.statement_type.clone())).to_uppercase()
)
},
}

17
src/popups/confirm_tx.rs
View File

@@ -7,11 +7,11 @@ use crate::database::statement_type_string;
#[derive(Debug)]
pub struct ConfirmTx {
rows_affected: Option<u64>,
statement_type: Statement,
statement_type: Option<Statement>,
}
impl ConfirmTx {
pub fn new(rows_affected: Option<u64>, statement_type: Statement) -> Self {
pub fn new(rows_affected: Option<u64>, statement_type: Option<Statement>) -> Self {
Self { rows_affected, statement_type }
}
}
@@ -32,23 +32,26 @@ impl PopUp for ConfirmTx {
fn get_cta_text(&self, app_state: &crate::app::AppState) -> String {
let rows_affected = self.rows_affected.unwrap_or_default();
match self.statement_type.clone() {
Statement::Delete(_) | Statement::Insert(_) | Statement::Update { .. } => {
None => {
format!("Are you sure you want to commit a transaction that will affect {rows_affected} rows?")
},
Some(Statement::Delete(_)) | Some(Statement::Insert(_)) | Some(Statement::Update { .. }) => {
format!(
"Are you sure you want to {} {} rows?",
statement_type_string(&self.statement_type).to_uppercase(),
statement_type_string(self.statement_type.clone()).to_uppercase(),
rows_affected
)
},
Statement::Explain { statement, .. } => {
Some(Statement::Explain { statement, .. }) => {
format!(
"Are you sure you want to run an EXPLAIN ANALYZE that will {} rows?",
statement_type_string(&statement).to_uppercase(),
statement_type_string(Some(*statement.clone())).to_uppercase(),
)
},
_ => {
format!(
"Are you sure you want to use a {} statement?",
statement_type_string(&self.statement_type).to_uppercase()
statement_type_string(self.statement_type.clone()).to_uppercase()
)
},
}

2
src/popups/mod.rs
View File

@@ -4,6 +4,7 @@ use sqlparser::ast::Statement;
use crate::{app::AppState, database::Rows};
pub mod confirm_bypass;
pub mod confirm_export;
pub mod confirm_query;
pub mod confirm_tx;
@@ -21,6 +22,7 @@ pub enum PopUpPayload {
CommitTx,
RollbackTx,
ConfirmQuery(String),
ConfirmBypass(String),
ConfirmExport(bool),
NamedFavorite(String, Vec<String>),
}