Skip to content

Commit 4886e95

Browse files
sundy-lisundy-li
authored
chore(query): declare var in procedure support optional default expr (#18897)
* update * declare support optional default expr * declare support optional default expr * declare support optional default expr * address comments
1 parent b8b06d6 commit 4886e95

File tree

7 files changed

+278
-29
lines changed

7 files changed

+278
-29
lines changed

โ€Žsrc/query/ast/src/ast/statements/script.rsโ€Ž

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ use std::fmt::Formatter;
1818
use crate::ast::Expr;
1919
use crate::ast::Identifier;
2020
use crate::ast::Statement;
21+
use crate::ast::TypeName;
2122
use crate::Span;
2223

2324
const INDENT_DEPTH: usize = 4;
@@ -72,13 +73,26 @@ impl Display for DeclareItem {
7273
pub struct DeclareVar {
7374
pub span: Span,
7475
pub name: Identifier,
75-
pub default: Expr,
76+
pub data_type: Option<TypeName>,
77+
pub default: Option<Expr>,
7678
}
7779

7880
impl Display for DeclareVar {
7981
fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
80-
let DeclareVar { name, default, .. } = self;
81-
write!(f, "{name} := {default}")?;
82+
let DeclareVar {
83+
name,
84+
data_type,
85+
default,
86+
..
87+
} = self;
88+
89+
write!(f, "{name}")?;
90+
if let Some(data_type) = data_type {
91+
write!(f, " {data_type}")?;
92+
}
93+
if let Some(default) = default {
94+
write!(f, " := {default}")?;
95+
}
8296
Ok(())
8397
}
8498
}

โ€Žsrc/query/ast/src/parser/script.rsโ€Ž

Lines changed: 12 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -76,12 +76,13 @@ pub fn declare_item(i: Input) -> IResult<DeclareItem> {
7676
pub fn declare_var(i: Input) -> IResult<DeclareVar> {
7777
map(
7878
consumed(rule! {
79-
#ident ~ ":=" ~ ^#expr
79+
#ident ~ ( #type_name )? ~ ( ( ":=" | DEFAULT ) ~ ^#expr )?
8080
}),
81-
|(span, (name, _, default))| DeclareVar {
81+
|(span, (name, data_type, default))| DeclareVar {
8282
span: transform_span(span.tokens),
8383
name,
84-
default,
84+
data_type,
85+
default: default.map(|(_, default)| default),
8586
},
8687
)(i)
8788
}
@@ -401,15 +402,19 @@ pub fn script_stmt(i: Input) -> IResult<ScriptStatement> {
401402
);
402403

403404
let cursor_stmts = rule!(
404-
#let_cursor_stmt
405-
| #open_cursor_stmt
405+
#open_cursor_stmt
406406
| #fetch_cursor_stmt
407407
| #close_cursor_stmt
408408
);
409409

410-
let assignment_stmts = rule!(
411-
#let_stmt_stmt
410+
let let_stmts = rule!(
411+
#let_cursor_stmt
412+
| #let_stmt_stmt
412413
| #let_var_stmt
414+
);
415+
416+
let assignment_stmts = rule!(
417+
#let_stmts
413418
| #assign_stmt
414419
);
415420

โ€Žsrc/query/ast/tests/it/parser.rsโ€Ž

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1510,8 +1510,12 @@ fn test_script() {
15101510
let file = &mut mint.new_goldenfile("script.txt").unwrap();
15111511

15121512
let cases = &[
1513+
r#"LET cost FLOAT"#,
1514+
r#"LET cost FLOAT default 3.0"#,
1515+
r#"LET cost FLOAT := 100.0"#,
15131516
r#"LET cost := 100.0"#,
15141517
r#"LET t1 RESULTSET := SELECT * FROM numbers(100)"#,
1518+
r#"LET t1 cursor FOR SELECT * FROM numbers(100)"#,
15151519
r#"profit := revenue - cost"#,
15161520
r#"RETURN"#,
15171521
r#"RETURN profit"#,

โ€Žsrc/query/ast/tests/it/testdata/script.txtโ€Ž

Lines changed: 216 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
11
---------- Input ----------
2-
LET cost := 100.0
2+
LET cost FLOAT
33
---------- Output ---------
4-
LET cost := 100.0
4+
LET cost Float32
55
---------- AST ------------
66
LetVar {
77
declare: DeclareVar {
88
span: Some(
9-
4..17,
9+
4..14,
1010
),
1111
name: Identifier {
1212
span: Some(
@@ -16,16 +16,119 @@ LetVar {
1616
quote: None,
1717
ident_type: None,
1818
},
19-
default: Literal {
19+
data_type: Some(
20+
Float32,
21+
),
22+
default: None,
23+
},
24+
}
25+
26+
27+
---------- Input ----------
28+
LET cost FLOAT default 3.0
29+
---------- Output ---------
30+
LET cost Float32 := 3.0
31+
---------- AST ------------
32+
LetVar {
33+
declare: DeclareVar {
34+
span: Some(
35+
4..26,
36+
),
37+
name: Identifier {
2038
span: Some(
21-
12..17,
39+
4..8,
2240
),
23-
value: Decimal256 {
24-
value: 1000,
25-
precision: 76,
26-
scale: 1,
41+
name: "cost",
42+
quote: None,
43+
ident_type: None,
44+
},
45+
data_type: Some(
46+
Float32,
47+
),
48+
default: Some(
49+
Literal {
50+
span: Some(
51+
23..26,
52+
),
53+
value: Decimal256 {
54+
value: 30,
55+
precision: 76,
56+
scale: 1,
57+
},
2758
},
59+
),
60+
},
61+
}
62+
63+
64+
---------- Input ----------
65+
LET cost FLOAT := 100.0
66+
---------- Output ---------
67+
LET cost Float32 := 100.0
68+
---------- AST ------------
69+
LetVar {
70+
declare: DeclareVar {
71+
span: Some(
72+
4..23,
73+
),
74+
name: Identifier {
75+
span: Some(
76+
4..8,
77+
),
78+
name: "cost",
79+
quote: None,
80+
ident_type: None,
81+
},
82+
data_type: Some(
83+
Float32,
84+
),
85+
default: Some(
86+
Literal {
87+
span: Some(
88+
18..23,
89+
),
90+
value: Decimal256 {
91+
value: 1000,
92+
precision: 76,
93+
scale: 1,
94+
},
95+
},
96+
),
97+
},
98+
}
99+
100+
101+
---------- Input ----------
102+
LET cost := 100.0
103+
---------- Output ---------
104+
LET cost := 100.0
105+
---------- AST ------------
106+
LetVar {
107+
declare: DeclareVar {
108+
span: Some(
109+
4..17,
110+
),
111+
name: Identifier {
112+
span: Some(
113+
4..8,
114+
),
115+
name: "cost",
116+
quote: None,
117+
ident_type: None,
28118
},
119+
data_type: None,
120+
default: Some(
121+
Literal {
122+
span: Some(
123+
12..17,
124+
),
125+
value: Decimal256 {
126+
value: 1000,
127+
precision: 76,
128+
scale: 1,
129+
},
130+
},
131+
),
29132
},
30133
}
31134

@@ -120,6 +223,99 @@ LetStatement {
120223
}
121224

122225

226+
---------- Input ----------
227+
LET t1 cursor FOR SELECT * FROM numbers(100)
228+
---------- Output ---------
229+
LET t1 CURSOR FOR SELECT * FROM numbers(100)
230+
---------- AST ------------
231+
LetCursor {
232+
declare: DeclareCursor {
233+
span: Some(
234+
4..44,
235+
),
236+
name: Identifier {
237+
span: Some(
238+
4..6,
239+
),
240+
name: "t1",
241+
quote: None,
242+
ident_type: None,
243+
},
244+
stmt: Some(
245+
Query(
246+
Query {
247+
span: Some(
248+
18..44,
249+
),
250+
with: None,
251+
body: Select(
252+
SelectStmt {
253+
span: Some(
254+
18..44,
255+
),
256+
hints: None,
257+
distinct: false,
258+
top_n: None,
259+
select_list: [
260+
StarColumns {
261+
qualified: [
262+
Star(
263+
Some(
264+
25..26,
265+
),
266+
),
267+
],
268+
column_filter: None,
269+
},
270+
],
271+
from: [
272+
TableFunction {
273+
span: Some(
274+
32..44,
275+
),
276+
lateral: false,
277+
name: Identifier {
278+
span: Some(
279+
32..39,
280+
),
281+
name: "numbers",
282+
quote: None,
283+
ident_type: None,
284+
},
285+
params: [
286+
Literal {
287+
span: Some(
288+
40..43,
289+
),
290+
value: UInt64(
291+
100,
292+
),
293+
},
294+
],
295+
named_params: [],
296+
alias: None,
297+
sample: None,
298+
},
299+
],
300+
selection: None,
301+
group_by: None,
302+
having: None,
303+
window_list: None,
304+
qualify: None,
305+
},
306+
),
307+
order_by: [],
308+
limit: [],
309+
offset: None,
310+
ignore_result: false,
311+
},
312+
),
313+
),
314+
resultset: None,
315+
},
316+
}
317+
318+
123319
---------- Input ----------
124320
profit := revenue - cost
125321
---------- Output ---------
@@ -2094,14 +2290,17 @@ ScriptBlock {
20942290
quote: None,
20952291
ident_type: None,
20962292
},
2097-
default: Literal {
2098-
span: Some(
2099-
17..18,
2100-
),
2101-
value: UInt64(
2102-
1,
2103-
),
2104-
},
2293+
data_type: None,
2294+
default: Some(
2295+
Literal {
2296+
span: Some(
2297+
17..18,
2298+
),
2299+
value: UInt64(
2300+
1,
2301+
),
2302+
},
2303+
),
21052304
},
21062305
),
21072306
],

โ€Žsrc/query/script/src/compiler.rsโ€Ž

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -89,7 +89,9 @@ impl Compiler {
8989
&declare.name.name,
9090
&mut self.ref_allocator,
9191
);
92-
output.append(&mut self.compile_expr(&declare.default, to_var.clone())?);
92+
if let Some(default_expr) = &declare.default {
93+
output.append(&mut self.compile_expr(default_expr, to_var.clone())?);
94+
}
9395
self.declare_ref(&declare.name, RefItem::Var(to_var))?;
9496
}
9597
ScriptStatement::LetStatement { declare } => {

โ€Žsrc/query/service/src/interpreters/interpreter_procedure_call.rsโ€Ž

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,8 @@ impl Interpreter for CallProcedureInterpreter {
7777
declare: DeclareVar {
7878
span: None,
7979
name: Identifier::from_name(None, arg_name),
80-
default: arg.clone(),
80+
data_type: None,
81+
default: Some(arg.clone()),
8182
},
8283
});
8384
}

0 commit comments

Comments
ย (0)