diff --git a/bindings/nodejs/index.d.ts b/bindings/nodejs/index.d.ts index 05e4f3bd..904e52b2 100644 --- a/bindings/nodejs/index.d.ts +++ b/bindings/nodejs/index.d.ts @@ -34,6 +34,8 @@ export class Connection { exec(sql: string): Promise /** Execute a SQL query, and only return the first row. */ queryRow(sql: string): Promise + /** Execute a SQL query and fetch all data into the result */ + queryAll(sql: string): Promise> /** Execute a SQL query, and return all rows. */ queryIter(sql: string): Promise /** Execute a SQL query, and return all rows with schema and stats. */ diff --git a/bindings/nodejs/src/lib.rs b/bindings/nodejs/src/lib.rs index 51741204..8e470fff 100644 --- a/bindings/nodejs/src/lib.rs +++ b/bindings/nodejs/src/lib.rs @@ -332,6 +332,19 @@ impl Connection { .map_err(format_napi_error) } + /// Execute a SQL query and fetch all data into the result + #[napi] + pub async fn query_all(&self, sql: String) -> Result> { + Ok(self + .0 + .query_all(&sql) + .await + .map_err(format_napi_error)? + .into_iter() + .map(|row| Row(row)) + .collect()) + } + /// Execute a SQL query, and return all rows. #[napi] pub async fn query_iter(&self, sql: String) -> Result { diff --git a/bindings/python/src/asyncio.rs b/bindings/python/src/asyncio.rs index 3b8eb3d2..28725d21 100644 --- a/bindings/python/src/asyncio.rs +++ b/bindings/python/src/asyncio.rs @@ -76,6 +76,20 @@ impl AsyncDatabendConnection { }) } + pub fn query_all<'p>(&self, py: Python<'p>, sql: String) -> PyResult<&'p PyAny> { + let this = self.0.clone(); + future_into_py(py, async move { + let rows: Vec = this + .query_all(&sql) + .await + .map_err(DriverError::new)? + .into_iter() + .map(Row::new) + .collect(); + Ok(rows) + }) + } + pub fn query_iter<'p>(&'p self, py: Python<'p>, sql: String) -> PyResult<&'p PyAny> { let this = self.0.clone(); future_into_py(py, async move { diff --git a/bindings/python/src/blocking.rs b/bindings/python/src/blocking.rs index 9075adf2..ca2cbee6 100644 --- a/bindings/python/src/blocking.rs +++ b/bindings/python/src/blocking.rs @@ -78,6 +78,14 @@ impl BlockingDatabendConnection { Ok(ret.map(Row::new)) } + pub fn query_all(&self, py: Python, sql: String) -> PyResult> { + let this = self.0.clone(); + let rows = wait_for_future(py, async move { + this.query_all(&sql).await.map_err(DriverError::new) + })?; + Ok(rows.into_iter().map(Row::new).collect()) + } + pub fn query_iter(&self, py: Python, sql: String) -> PyResult { let this = self.0.clone(); let it = wait_for_future(py, async { diff --git a/cli/src/ast/tokenizer.rs b/cli/src/ast/tokenizer.rs index f690b063..6ac92384 100644 --- a/cli/src/ast/tokenizer.rs +++ b/cli/src/ast/tokenizer.rs @@ -122,7 +122,10 @@ pub enum TokenKind { #[regex(r#"'([^'\\]|\\.|'')*'"#)] QuotedString, - #[regex(r#"@([^\s`;'"])+"#)] + #[regex(r#"\$\$([^\$]|(\$[^\$]))*\$\$"#)] + CodeString, + + #[regex(r#"@([^\s`;'"()]|\\\s|\\'|\\"|\\\\)+"#)] AtString, #[regex(r"[xX]'[a-fA-F0-9]*'")] @@ -907,7 +910,12 @@ impl TokenKind { pub fn is_literal(&self) -> bool { matches!( self, - LiteralInteger | LiteralFloat | QuotedString | PGLiteralHex | MySQLLiteralHex + LiteralInteger + | CodeString + | LiteralFloat + | QuotedString + | PGLiteralHex + | MySQLLiteralHex ) } @@ -916,6 +924,7 @@ impl TokenKind { self, Ident | QuotedString + | CodeString | PGLiteralHex | MySQLLiteralHex | LiteralInteger diff --git a/cli/src/session.rs b/cli/src/session.rs index f240ffd0..e127f2df 100644 --- a/cli/src/session.rs +++ b/cli/src/session.rs @@ -57,7 +57,6 @@ pub struct Session { settings: Settings, query: String, - in_comment_block: bool, keywords: Arc>, } @@ -108,7 +107,6 @@ impl Session { is_repl, settings, query: String::new(), - in_comment_block: false, keywords: Arc::new(keywords), }) } @@ -316,7 +314,6 @@ impl Session { } pub fn append_query(&mut self, line: &str) -> Vec { - let line = line.trim(); if line.is_empty() { return vec![]; } @@ -338,63 +335,56 @@ impl Session { } } - self.query.push(' '); - + // consume self.query and get the result let mut queries = Vec::new(); - let mut tokenizer = Tokenizer::new(line); - let mut in_comment = false; - let mut start = 0; - let mut comment_block_start = 0; - - while let Some(Ok(token)) = tokenizer.next() { - match token.kind { - TokenKind::SemiColon => { - if in_comment || self.in_comment_block { - continue; - } else { - let mut sql = self.query.trim().to_owned(); - if sql.is_empty() { + + if !self.query.is_empty() { + self.query.push('\n'); + } + self.query.push_str(line); + + 'Parser: loop { + let mut tokenizer = Tokenizer::new(&self.query); + + let mut in_comment = false; + let mut in_comment_block = false; + + while let Some(Ok(token)) = tokenizer.next() { + match token.kind { + TokenKind::SemiColon => { + if in_comment_block || in_comment { continue; } - sql.push(';'); - queries.push(sql); - self.query.clear(); + // push to current and continue the tokenizer + let (sql, remain) = self.query.split_at(token.span.end); + if !sql.is_empty() { + queries.push(sql.to_string()); + } + self.query = remain.to_string(); + continue 'Parser; } - } - TokenKind::Comment => { - in_comment = true; - } - TokenKind::EOI => { - in_comment = false; - } - TokenKind::Newline => { - in_comment = false; - self.query.push('\n'); - } - TokenKind::CommentBlockStart => { - if !self.in_comment_block { - comment_block_start = token.span.start; + TokenKind::Comment => { + if in_comment_block { + continue; + } + in_comment = true; } - self.in_comment_block = true; - } - TokenKind::CommentBlockEnd => { - self.in_comment_block = false; - self.query - .push_str(&line[comment_block_start..token.span.end]); - } - _ => { - if !in_comment && !self.in_comment_block { - self.query.push_str(&line[start..token.span.end]); + TokenKind::Newline => { + in_comment = false; + } + TokenKind::CommentBlockStart => { + in_comment_block = true; + } + TokenKind::CommentBlockEnd => { + in_comment_block = false; } + _ => {} } } - start = token.span.end; + break; } - if self.in_comment_block { - self.query.push_str(&line[comment_block_start..]); - } queries } diff --git a/cli/tests/00-base.result b/cli/tests/00-base.result index e61bb947..c6b8b920 100644 --- a/cli/tests/00-base.result +++ b/cli/tests/00-base.result @@ -6,6 +6,11 @@ a 1 true [1,2] 3 [] {} with comment +1 +2 +" +a" +3 3.00 3.00 0.0000000170141183460469231731687303715884105727000 -0.0000000170141183460469231731687303715884105727000 Asia/Shanghai 0 0.00 diff --git a/cli/tests/00-base.sql b/cli/tests/00-base.sql index fc43f478..14c9a2aa 100644 --- a/cli/tests/00-base.sql +++ b/cli/tests/00-base.sql @@ -16,6 +16,16 @@ select [], {}; select /* ignore this block */ 'with comment'; +select 1; select 2; select ' +a'; select 3; + +-- enable it after we support code string in databend +-- select $$aa$$; +-- select $$ +-- def add(a, b): +-- a + b +-- $$; + /* ignore this block /* /* select 'in comment block'; */ diff --git a/driver/src/conn.rs b/driver/src/conn.rs index e49e47f9..f7a58e54 100644 --- a/driver/src/conn.rs +++ b/driver/src/conn.rs @@ -106,6 +106,10 @@ pub trait Connection: DynClone + Send + Sync { async fn exec(&self, sql: &str) -> Result; async fn query_row(&self, sql: &str) -> Result>; + async fn query_all(&self, sql: &str) -> Result> { + let rows = self.query_iter(sql).await?; + rows.collect().await + } async fn query_iter(&self, sql: &str) -> Result; async fn query_iter_ext(&self, sql: &str) -> Result; diff --git a/driver/src/flight_sql.rs b/driver/src/flight_sql.rs index 66429282..db62361c 100644 --- a/driver/src/flight_sql.rs +++ b/driver/src/flight_sql.rs @@ -121,9 +121,9 @@ impl Connection for FlightSQLConnection { _file_format_options: Option>, _copy_options: Option>, ) -> Result { - return Err(Error::Protocol( + Err(Error::Protocol( "LOAD DATA unavailable for FlightSQL".to_string(), - )); + )) } async fn load_file( @@ -133,15 +133,15 @@ impl Connection for FlightSQLConnection { _format_options: BTreeMap<&str, &str>, _copy_options: Option>, ) -> Result { - return Err(Error::Protocol( + Err(Error::Protocol( "LOAD FILE unavailable for FlightSQL".to_string(), - )); + )) } async fn stream_load(&self, _sql: &str, _data: Vec>) -> Result { - return Err(Error::Protocol( + Err(Error::Protocol( "STREAM LOAD unavailable for FlightSQL".to_string(), - )); + )) } }