Skip to content

Commit 18a6169

Browse files
yuhao-suclaude
yuhao-su
and
claude
authored
feat(parser,runner): add let record for binding query results to variables (#279)
* feat(parser,runner): add `let` record for binding query results to variables Add a new `let` record type that executes a SQL query and binds the results to variables for use in subsequent queries. This is useful for capturing dynamic values like auto-generated IDs. Syntax: `let (var1, var2, ...) \n SQL` Features: - Binds query results to named variables stored in RunnerLocals - Requires exactly 1 row with N columns matching N variables - Supports conditions (skipif/onlyif) and connection directives - Variables can be used with substitution syntax ($var, ${var:default}) - Clear error messages for row/column count mismatches Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> Signed-off-by: Yuhao Su <yuhaosu@outlook.com> * fix: remove redundant trim() before split_whitespace() Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> Signed-off-by: Yuhao Su <yuhaosu@outlook.com> * fmt Signed-off-by: Yuhao Su <yuhaosu@outlook.com> * refactor: consolidate let record errors into `LetError` enum Replace duplicated inline error structs (RowCountError, ColumnCountError) and unused TestErrorKind variants (LetRowCountMismatch, LetColumnCountMismatch) with a single `LetError` enum. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Yuhao Su <yuhaosu@outlook.com> * refactor: simplify `let` record syntax by removing parentheses Change `let (var1, var2)` to `let var1, var2` for consistency with other sqllogictest record syntax and reduced visual noise. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Yuhao Su <yuhaosu@outlook.com> * fmt Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Yuhao Su <yuhaosu@outlook.com> --------- Signed-off-by: Yuhao Su <yuhaosu@outlook.com> Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
1 parent 492c9e3 commit 18a6169

6 files changed

Lines changed: 571 additions & 33 deletions

File tree

‎README.md‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -219,6 +219,35 @@ echo "foo" > "$__TEST_DIR__/foo.txt"
219219
> and excaping is also not needed.
220220
> Environment variables are supported by the shell, and special variables are still supported by plain string substitution.
221221
222+
### Extension: Bind query results to variables with `let`
223+
224+
The `let` record allows you to execute a SQL query and bind the results to variables for use in subsequent queries.
225+
This is useful when you need to capture dynamic values (like auto-generated IDs) from the database.
226+
227+
```text
228+
control substitution on
229+
230+
# Execute a query that returns exactly 1 row with 1 column, bind result to 'id'
231+
let id
232+
SELECT id FROM users WHERE name = 'alice'
233+
234+
# Multiple variables can be bound from a single query
235+
let user_id, user_name, user_email
236+
SELECT id, name, email FROM users WHERE id = 1
237+
238+
# Use the bound variable in subsequent queries
239+
query TTT
240+
SELECT $user_id, $user_name, $user_email
241+
----
242+
1 alice alice@example.com
243+
```
244+
245+
**Requirements:**
246+
- The query must return exactly **1 row**.
247+
- The number of columns must match the number of variables.
248+
- `let` requires `control substitution on` to be enabled for variable usage, but the `let` statement itself will execute regardless.
249+
250+
222251
## Used by
223252

224253
- [RisingLight](https://github.com/risinglightdb/risinglight): An OLAP database system for educational purpose

‎sqllogictest/src/parser.rs‎

Lines changed: 176 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -196,6 +196,17 @@ pub enum Record<T: ColumnType> {
196196
Newline,
197197
/// Internally injected record which should not occur in the test file.
198198
Injected(Injected),
199+
/// A let record binds SQL query results to variables.
200+
/// The query must return exactly 1 row with N columns matching the N variable names.
201+
Let {
202+
loc: Location,
203+
conditions: Vec<Condition>,
204+
connection: Connection,
205+
/// Variable names to bind the results to.
206+
variables: Vec<String>,
207+
/// The SQL query to execute.
208+
sql: String,
209+
},
199210
}
200211

201212
impl<T: ColumnType> Record<T> {
@@ -357,6 +368,15 @@ impl<T: ColumnType> std::fmt::Display for Record<T> {
357368
}
358369
Record::Newline => Ok(()), // Display doesn't end with newline
359370
Record::Injected(p) => panic!("unexpected injected record: {p:?}"),
371+
Record::Let {
372+
loc: _,
373+
conditions: _,
374+
connection: _,
375+
variables,
376+
sql,
377+
} => {
378+
write!(f, "let {}\n{sql}\n", variables.join(", "))
379+
}
360380
}
361381
}
362382
}
@@ -700,6 +720,8 @@ pub enum ParseErrorKind {
700720
EmptyIncludeFile(String),
701721
#[error("no such file")]
702722
FileNotFound,
723+
#[error("invalid variable name: {0:?}")]
724+
InvalidVariableName(String),
703725
}
704726

705727
impl ParseErrorKind {
@@ -976,6 +998,45 @@ fn parse_inner<T: ColumnType>(loc: &Location, script: &str) -> Result<Vec<Record
976998
})?,
977999
});
9781000
}
1001+
["let", rest @ ..] => {
1002+
// Parse: let var1, var2, ... followed by SQL
1003+
// Join the rest of the tokens and parse the variable list
1004+
let rest_str = rest.join(" ");
1005+
let rest_str = rest_str.trim();
1006+
1007+
if rest_str.is_empty() {
1008+
return Err(ParseErrorKind::InvalidLine(line.into()).at(loc));
1009+
}
1010+
1011+
// Extract variable names separated by commas
1012+
let variables: Vec<String> = rest_str
1013+
.split(',')
1014+
.map(|s| s.trim().to_string())
1015+
.filter(|s| !s.is_empty())
1016+
.collect();
1017+
1018+
if variables.is_empty() {
1019+
return Err(ParseErrorKind::InvalidLine(line.into()).at(loc));
1020+
}
1021+
1022+
// Validate variable names (must be valid identifiers)
1023+
for var in &variables {
1024+
if !is_valid_variable_name(var) {
1025+
return Err(ParseErrorKind::InvalidVariableName(var.clone()).at(loc));
1026+
}
1027+
}
1028+
1029+
// Parse the SQL body (following lines until empty line)
1030+
let (sql, _has_results) = parse_lines(&mut lines, &loc, None)?;
1031+
1032+
records.push(Record::Let {
1033+
loc,
1034+
conditions: std::mem::take(&mut conditions),
1035+
connection: std::mem::take(&mut connection),
1036+
variables,
1037+
sql,
1038+
});
1039+
}
9791040
_ => return Err(ParseErrorKind::InvalidLine(line.into()).at(loc)),
9801041
}
9811042
}
@@ -1029,6 +1090,21 @@ fn parse_file_inner<T: ColumnType>(loc: Location) -> Result<Vec<Record<T>>, Pars
10291090
Ok(records)
10301091
}
10311092

1093+
/// Check if a variable name is valid.
1094+
/// A valid variable name starts with a letter or underscore, and contains only alphanumeric
1095+
/// characters and underscores.
1096+
fn is_valid_variable_name(name: &str) -> bool {
1097+
if name.is_empty() {
1098+
return false;
1099+
}
1100+
let mut chars = name.chars();
1101+
let first = chars.next().unwrap();
1102+
if !first.is_ascii_alphabetic() && first != '_' {
1103+
return false;
1104+
}
1105+
chars.all(|c| c.is_ascii_alphanumeric() || c == '_')
1106+
}
1107+
10321108
/// Parse one or more lines until empty line or a delimiter.
10331109
fn parse_lines<'a>(
10341110
lines: &mut impl Iterator<Item = (usize, &'a str)>,
@@ -1336,6 +1412,7 @@ select * from foo;
13361412
Record::Subtest { loc, .. } => normalize_loc(loc),
13371413
Record::Halt { loc, .. } => normalize_loc(loc),
13381414
Record::HashThreshold { loc, .. } => normalize_loc(loc),
1415+
Record::Let { loc, .. } => normalize_loc(loc),
13391416
// even though these variants don't include a
13401417
// location include them in this match statement
13411418
// so if new variants are added, this match
@@ -1389,4 +1466,103 @@ select * from foo;
13891466
fn test_query_retry() {
13901467
parse_roundtrip::<DefaultColumnType>("../tests/no_run/query_retry.slt")
13911468
}
1469+
1470+
#[test]
1471+
fn test_let_parsing() {
1472+
let script = "\
1473+
let id
1474+
SELECT 1
1475+
";
1476+
let records = parse::<DefaultColumnType>(script).unwrap();
1477+
assert_eq!(records.len(), 1);
1478+
match &records[0] {
1479+
Record::Let { variables, sql, .. } => {
1480+
assert_eq!(variables, &["id".to_string()]);
1481+
assert_eq!(sql, "SELECT 1");
1482+
}
1483+
_ => panic!("expected Let record"),
1484+
}
1485+
}
1486+
1487+
#[test]
1488+
fn test_let_parsing_multiple_vars() {
1489+
let script = "\
1490+
let id, name, value
1491+
SELECT 1, 'hello', 42
1492+
";
1493+
let records = parse::<DefaultColumnType>(script).unwrap();
1494+
assert_eq!(records.len(), 1);
1495+
match &records[0] {
1496+
Record::Let { variables, sql, .. } => {
1497+
assert_eq!(
1498+
variables,
1499+
&["id".to_string(), "name".to_string(), "value".to_string()]
1500+
);
1501+
assert_eq!(sql, "SELECT 1, 'hello', 42");
1502+
}
1503+
_ => panic!("expected Let record"),
1504+
}
1505+
}
1506+
1507+
#[test]
1508+
fn test_let_parsing_roundtrip() {
1509+
let script = "\
1510+
let id
1511+
SELECT 1
1512+
";
1513+
let records = parse::<DefaultColumnType>(script).unwrap();
1514+
let unparsed = records[0].to_string();
1515+
let reparsed = parse::<DefaultColumnType>(&unparsed).unwrap();
1516+
assert_eq!(records.len(), reparsed.len());
1517+
match (&records[0], &reparsed[0]) {
1518+
(
1519+
Record::Let {
1520+
variables: v1,
1521+
sql: s1,
1522+
..
1523+
},
1524+
Record::Let {
1525+
variables: v2,
1526+
sql: s2,
1527+
..
1528+
},
1529+
) => {
1530+
assert_eq!(v1, v2);
1531+
assert_eq!(s1, s2);
1532+
}
1533+
_ => panic!("expected Let records"),
1534+
}
1535+
}
1536+
1537+
#[test]
1538+
fn test_let_parsing_error_empty_vars() {
1539+
let script = "\
1540+
let
1541+
SELECT 1
1542+
";
1543+
let err = parse::<DefaultColumnType>(script).unwrap_err();
1544+
assert!(matches!(err.kind(), ParseErrorKind::InvalidLine(_)));
1545+
}
1546+
1547+
#[test]
1548+
fn test_let_parsing_error_invalid_var_name() {
1549+
let script = "\
1550+
let 123invalid
1551+
SELECT 1
1552+
";
1553+
let err = parse::<DefaultColumnType>(script).unwrap_err();
1554+
assert!(matches!(err.kind(), ParseErrorKind::InvalidVariableName(_)));
1555+
}
1556+
1557+
#[test]
1558+
fn test_is_valid_variable_name() {
1559+
assert!(is_valid_variable_name("foo"));
1560+
assert!(is_valid_variable_name("_bar"));
1561+
assert!(is_valid_variable_name("foo123"));
1562+
assert!(is_valid_variable_name("__TEST_DIR__"));
1563+
assert!(!is_valid_variable_name(""));
1564+
assert!(!is_valid_variable_name("123"));
1565+
assert!(!is_valid_variable_name("foo-bar"));
1566+
assert!(!is_valid_variable_name("foo bar"));
1567+
}
13921568
}

0 commit comments

Comments
 (0)