diff --git a/core/lib.rs b/core/lib.rs index 592992532..b0c5751b5 100644 --- a/core/lib.rs +++ b/core/lib.rs @@ -33,6 +33,7 @@ mod numeric; static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; use crate::{fast_lock::SpinLock, translate::optimizer::optimize_plan}; +use core::str; pub use error::LimboError; use fallible_iterator::FallibleIterator; pub use io::clock::{Clock, Instant}; @@ -354,6 +355,10 @@ impl Connection { let cmd = parser.next()?; let syms = self.syms.borrow(); let cmd = cmd.expect("Successful parse on nonempty input string should produce a command"); + let byte_offset_end = parser.offset(); + let input = str::from_utf8(&sql.as_bytes()[..byte_offset_end]) + .unwrap() + .trim(); match cmd { Cmd::Stmt(stmt) => { let program = Rc::new(translate::translate( @@ -367,6 +372,7 @@ impl Connection { Rc::downgrade(self), &syms, QueryMode::Normal, + &input, )?); Ok(Statement::new( program, @@ -385,14 +391,22 @@ impl Connection { tracing::trace!("Querying: {}", sql); let mut parser = Parser::new(sql.as_bytes()); let cmd = parser.next()?; + let byte_offset_end = parser.offset(); + let input = str::from_utf8(&sql.as_bytes()[..byte_offset_end]) + .unwrap() + .trim(); match cmd { - Some(cmd) => self.run_cmd(cmd), + Some(cmd) => self.run_cmd(cmd, input), None => Ok(None), } } #[instrument(skip_all, level = Level::TRACE)] - pub(crate) fn run_cmd(self: &Rc, cmd: Cmd) -> Result> { + pub(crate) fn run_cmd( + self: &Rc, + cmd: Cmd, + input: &str, + ) -> Result> { let syms = self.syms.borrow(); match cmd { Cmd::Stmt(ref stmt) | Cmd::Explain(ref stmt) => { @@ -407,6 +421,7 @@ impl Connection { Rc::downgrade(self), &syms, cmd.into(), + input, )?; let stmt = Statement::new( program.into(), @@ -458,6 +473,10 @@ impl Connection { let mut parser = Parser::new(sql.as_bytes()); let cmd = parser.next()?; let syms = self.syms.borrow(); + let byte_offset_end = parser.offset(); + let input = str::from_utf8(&sql.as_bytes()[..byte_offset_end]) + .unwrap() + .trim(); if let Some(cmd) = cmd { match cmd { Cmd::Explain(stmt) => { @@ -472,6 +491,7 @@ impl Connection { Rc::downgrade(self), &syms, QueryMode::Explain, + &input, )?; let _ = std::io::stdout().write_all(program.explain().as_bytes()); } @@ -488,6 +508,7 @@ impl Connection { Rc::downgrade(self), &syms, QueryMode::Normal, + &input, )?; let mut state = @@ -1048,6 +1069,8 @@ impl SymbolTable { pub struct QueryRunner<'a> { parser: Parser<'a>, conn: &'a Rc, + statements: &'a [u8], + last_offset: usize, } impl<'a> QueryRunner<'a> { @@ -1055,6 +1078,8 @@ impl<'a> QueryRunner<'a> { Self { parser: Parser::new(statements), conn, + statements, + last_offset: 0, } } } @@ -1064,7 +1089,14 @@ impl Iterator for QueryRunner<'_> { fn next(&mut self) -> Option { match self.parser.next() { - Ok(Some(cmd)) => Some(self.conn.run_cmd(cmd)), + Ok(Some(cmd)) => { + let byte_offset_end = self.parser.offset(); + let input = str::from_utf8(&self.statements[self.last_offset..byte_offset_end]) + .unwrap() + .trim(); + self.last_offset = byte_offset_end; + Some(self.conn.run_cmd(cmd, &input)) + } Ok(None) => None, Err(err) => { self.parser.finalize(); diff --git a/core/translate/mod.rs b/core/translate/mod.rs index 04b4e52e8..37f222922 100644 --- a/core/translate/mod.rs +++ b/core/translate/mod.rs @@ -62,6 +62,7 @@ pub fn translate( connection: Weak, syms: &SymbolTable, query_mode: QueryMode, + _input: &str, // TODO: going to be used for CREATE VIEW ) -> Result { let change_cnt_on = matches!( stmt,