Skip to content

Commit 4c3041b

Browse files
sundy-lisundy-li
authored
feat: add query_all function (#355)
* feat: add query_all function * feat: update * feat: update * feat: update * update * update * update
1 parent 25b1195 commit 4c3041b

File tree

10 files changed

+112
-57
lines changed

10 files changed

+112
-57
lines changed

‎bindings/nodejs/index.d.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,8 @@ export class Connection {
3434
exec(sql: string): Promise<number>
3535
/** Execute a SQL query, and only return the first row. */
3636
queryRow(sql: string): Promise<Row | null>
37+
/** Execute a SQL query and fetch all data into the result */
38+
queryAll(sql: string): Promise<Array<Row>>
3739
/** Execute a SQL query, and return all rows. */
3840
queryIter(sql: string): Promise<RowIterator>
3941
/** Execute a SQL query, and return all rows with schema and stats. */

‎bindings/nodejs/src/lib.rs‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -332,6 +332,19 @@ impl Connection {
332332
.map_err(format_napi_error)
333333
}
334334

335+
/// Execute a SQL query and fetch all data into the result
336+
#[napi]
337+
pub async fn query_all(&self, sql: String) -> Result<Vec<Row>> {
338+
Ok(self
339+
.0
340+
.query_all(&sql)
341+
.await
342+
.map_err(format_napi_error)?
343+
.into_iter()
344+
.map(|row| Row(row))
345+
.collect())
346+
}
347+
335348
/// Execute a SQL query, and return all rows.
336349
#[napi]
337350
pub async fn query_iter(&self, sql: String) -> Result<RowIterator> {

‎bindings/python/src/asyncio.rs‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,20 @@ impl AsyncDatabendConnection {
7676
})
7777
}
7878

79+
pub fn query_all<'p>(&self, py: Python<'p>, sql: String) -> PyResult<&'p PyAny> {
80+
let this = self.0.clone();
81+
future_into_py(py, async move {
82+
let rows: Vec<Row> = this
83+
.query_all(&sql)
84+
.await
85+
.map_err(DriverError::new)?
86+
.into_iter()
87+
.map(Row::new)
88+
.collect();
89+
Ok(rows)
90+
})
91+
}
92+
7993
pub fn query_iter<'p>(&'p self, py: Python<'p>, sql: String) -> PyResult<&'p PyAny> {
8094
let this = self.0.clone();
8195
future_into_py(py, async move {

‎bindings/python/src/blocking.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,14 @@ impl BlockingDatabendConnection {
7878
Ok(ret.map(Row::new))
7979
}
8080

81+
pub fn query_all(&self, py: Python, sql: String) -> PyResult<Vec<Row>> {
82+
let this = self.0.clone();
83+
let rows = wait_for_future(py, async move {
84+
this.query_all(&sql).await.map_err(DriverError::new)
85+
})?;
86+
Ok(rows.into_iter().map(Row::new).collect())
87+
}
88+
8189
pub fn query_iter(&self, py: Python, sql: String) -> PyResult<RowIterator> {
8290
let this = self.0.clone();
8391
let it = wait_for_future(py, async {

‎cli/src/ast/tokenizer.rs‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -122,7 +122,10 @@ pub enum TokenKind {
122122
#[regex(r#"'([^'\\]|\\.|'')*'"#)]
123123
QuotedString,
124124

125-
#[regex(r#"@([^\s`;'"])+"#)]
125+
#[regex(r#"\$\$([^\$]|(\$[^\$]))*\$\$"#)]
126+
CodeString,
127+
128+
#[regex(r#"@([^\s`;'"()]|\\\s|\\'|\\"|\\\\)+"#)]
126129
AtString,
127130

128131
#[regex(r"[xX]'[a-fA-F0-9]*'")]
@@ -907,7 +910,12 @@ impl TokenKind {
907910
pub fn is_literal(&self) -> bool {
908911
matches!(
909912
self,
910-
LiteralInteger | LiteralFloat | QuotedString | PGLiteralHex | MySQLLiteralHex
913+
LiteralInteger
914+
| CodeString
915+
| LiteralFloat
916+
| QuotedString
917+
| PGLiteralHex
918+
| MySQLLiteralHex
911919
)
912920
}
913921

@@ -916,6 +924,7 @@ impl TokenKind {
916924
self,
917925
Ident
918926
| QuotedString
927+
| CodeString
919928
| PGLiteralHex
920929
| MySQLLiteralHex
921930
| LiteralInteger

‎cli/src/session.rs‎

Lines changed: 39 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,6 @@ pub struct Session {
5757

5858
settings: Settings,
5959
query: String,
60-
in_comment_block: bool,
6160

6261
keywords: Arc<Vec<String>>,
6362
}
@@ -108,7 +107,6 @@ impl Session {
108107
is_repl,
109108
settings,
110109
query: String::new(),
111-
in_comment_block: false,
112110
keywords: Arc::new(keywords),
113111
})
114112
}
@@ -316,7 +314,6 @@ impl Session {
316314
}
317315

318316
pub fn append_query(&mut self, line: &str) -> Vec<String> {
319-
let line = line.trim();
320317
if line.is_empty() {
321318
return vec![];
322319
}
@@ -338,63 +335,56 @@ impl Session {
338335
}
339336
}
340337

341-
self.query.push(' ');
342-
338+
// consume self.query and get the result
343339
let mut queries = Vec::new();
344-
let mut tokenizer = Tokenizer::new(line);
345-
let mut in_comment = false;
346-
let mut start = 0;
347-
let mut comment_block_start = 0;
348-
349-
while let Some(Ok(token)) = tokenizer.next() {
350-
match token.kind {
351-
TokenKind::SemiColon => {
352-
if in_comment || self.in_comment_block {
353-
continue;
354-
} else {
355-
let mut sql = self.query.trim().to_owned();
356-
if sql.is_empty() {
340+
341+
if !self.query.is_empty() {
342+
self.query.push('\n');
343+
}
344+
self.query.push_str(line);
345+
346+
'Parser: loop {
347+
let mut tokenizer = Tokenizer::new(&self.query);
348+
349+
let mut in_comment = false;
350+
let mut in_comment_block = false;
351+
352+
while let Some(Ok(token)) = tokenizer.next() {
353+
match token.kind {
354+
TokenKind::SemiColon => {
355+
if in_comment_block || in_comment {
357356
continue;
358357
}
359-
sql.push(';');
360358

361-
queries.push(sql);
362-
self.query.clear();
359+
// push to current and continue the tokenizer
360+
let (sql, remain) = self.query.split_at(token.span.end);
361+
if !sql.is_empty() {
362+
queries.push(sql.to_string());
363+
}
364+
self.query = remain.to_string();
365+
continue 'Parser;
363366
}
364-
}
365-
TokenKind::Comment => {
366-
in_comment = true;
367-
}
368-
TokenKind::EOI => {
369-
in_comment = false;
370-
}
371-
TokenKind::Newline => {
372-
in_comment = false;
373-
self.query.push('\n');
374-
}
375-
TokenKind::CommentBlockStart => {
376-
if !self.in_comment_block {
377-
comment_block_start = token.span.start;
367+
TokenKind::Comment => {
368+
if in_comment_block {
369+
continue;
370+
}
371+
in_comment = true;
378372
}
379-
self.in_comment_block = true;
380-
}
381-
TokenKind::CommentBlockEnd => {
382-
self.in_comment_block = false;
383-
self.query
384-
.push_str(&line[comment_block_start..token.span.end]);
385-
}
386-
_ => {
387-
if !in_comment && !self.in_comment_block {
388-
self.query.push_str(&line[start..token.span.end]);
373+
TokenKind::Newline => {
374+
in_comment = false;
375+
}
376+
TokenKind::CommentBlockStart => {
377+
in_comment_block = true;
378+
}
379+
TokenKind::CommentBlockEnd => {
380+
in_comment_block = false;
389381
}
382+
_ => {}
390383
}
391384
}
392-
start = token.span.end;
385+
break;
393386
}
394387

395-
if self.in_comment_block {
396-
self.query.push_str(&line[comment_block_start..]);
397-
}
398388
queries
399389
}
400390

‎cli/tests/00-base.result‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,11 @@ a 1 true [1,2]
66
3
77
[] {}
88
with comment
9+
1
10+
2
11+
"
12+
a"
13+
3
914
3.00 3.00 0.0000000170141183460469231731687303715884105727000 -0.0000000170141183460469231731687303715884105727000
1015
Asia/Shanghai
1116
0 0.00

‎cli/tests/00-base.sql‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,16 @@ select [], {};
1616

1717
select /* ignore this block */ 'with comment';
1818

19+
select 1; select 2; select '
20+
a'; select 3;
21+
22+
-- enable it after we support code string in databend
23+
-- select $$aa$$;
24+
-- select $$
25+
-- def add(a, b):
26+
-- a + b
27+
-- $$;
28+
1929
/* ignore this block /* /*
2030
select 'in comment block';
2131
*/

‎driver/src/conn.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,10 @@ pub trait Connection: DynClone + Send + Sync {
106106

107107
async fn exec(&self, sql: &str) -> Result<i64>;
108108
async fn query_row(&self, sql: &str) -> Result<Option<Row>>;
109+
async fn query_all(&self, sql: &str) -> Result<Vec<Row>> {
110+
let rows = self.query_iter(sql).await?;
111+
rows.collect().await
112+
}
109113
async fn query_iter(&self, sql: &str) -> Result<RowIterator>;
110114
async fn query_iter_ext(&self, sql: &str) -> Result<RowStatsIterator>;
111115

‎driver/src/flight_sql.rs‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -121,9 +121,9 @@ impl Connection for FlightSQLConnection {
121121
_file_format_options: Option<BTreeMap<&str, &str>>,
122122
_copy_options: Option<BTreeMap<&str, &str>>,
123123
) -> Result<ServerStats> {
124-
return Err(Error::Protocol(
124+
Err(Error::Protocol(
125125
"LOAD DATA unavailable for FlightSQL".to_string(),
126-
));
126+
))
127127
}
128128

129129
async fn load_file(
@@ -133,15 +133,15 @@ impl Connection for FlightSQLConnection {
133133
_format_options: BTreeMap<&str, &str>,
134134
_copy_options: Option<BTreeMap<&str, &str>>,
135135
) -> Result<ServerStats> {
136-
return Err(Error::Protocol(
136+
Err(Error::Protocol(
137137
"LOAD FILE unavailable for FlightSQL".to_string(),
138-
));
138+
))
139139
}
140140

141141
async fn stream_load(&self, _sql: &str, _data: Vec<Vec<&str>>) -> Result<ServerStats> {
142-
return Err(Error::Protocol(
142+
Err(Error::Protocol(
143143
"STREAM LOAD unavailable for FlightSQL".to_string(),
144-
));
144+
))
145145
}
146146
}
147147

0 commit comments

Comments
 (0)