Skip to content

Commit e18ba3d

Browse files
LucaCappelletti94LucaCappelletti94
committed
SQLite: Tokenize parameters the way SQLite does
1 parent 0cb62a1 commit e18ba3d

5 files changed

Lines changed: 272 additions & 6 deletions

File tree

‎src/dialect/mod.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1148,6 +1148,14 @@ pub trait Dialect: Debug + Any {
11481148
false
11491149
}
11501150

1151+
/// Returns true if this dialect supports SQLite parameter token forms:
1152+
/// `$AAAA`, `@AAAA`, `:AAAA`, and `#AAAA`, where AAAA may contain
1153+
/// alphanumeric chars and `$`, with optional TCL `::` namespace and
1154+
/// `(subscript)` extensions.
1155+
fn supports_sqlite_variable_syntax(&self) -> bool {
1156+
false
1157+
}
1158+
11511159
/// Returns true if this dialect supports `$` as a prefix for money literals
11521160
/// e.g. `SELECT $123.45` (SQL Server)
11531161
fn supports_dollar_as_money_prefix(&self) -> bool {

‎src/dialect/sqlite.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -129,4 +129,8 @@ impl Dialect for SQLiteDialect {
129129
fn supports_numeric_literal_underscores(&self) -> bool {
130130
true
131131
}
132+
133+
fn supports_sqlite_variable_syntax(&self) -> bool {
134+
true
135+
}
132136
}

‎src/tokenizer.rs‎

Lines changed: 190 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1664,10 +1664,51 @@ impl<'a> Tokenizer<'a> {
16641664
}
16651665
':' => {
16661666
chars.next();
1667-
match chars.peek() {
1668-
Some(':') => self.consume_and_return(chars, Token::DoubleColon),
1669-
Some('=') => self.consume_and_return(chars, Token::Assignment),
1670-
_ => Ok(Some(Token::Colon)),
1667+
if self.dialect.supports_sqlite_variable_syntax() {
1668+
match chars.peek().copied() {
1669+
Some('=') => self.consume_and_return(chars, Token::Assignment),
1670+
Some(c)
1671+
if c.is_ascii_alphanumeric()
1672+
|| c == '_'
1673+
|| c == '$'
1674+
|| c >= '\u{0080}' =>
1675+
{
1676+
let body = self.tokenize_sqlite_variable_body(chars)?;
1677+
Ok(Some(Token::Placeholder(format!(":{body}"))))
1678+
}
1679+
Some(':') => {
1680+
chars.next(); // consume second ':'
1681+
match chars.peek().copied() {
1682+
Some(':') => {
1683+
chars.next(); // consume third ':'
1684+
match chars.peek().copied() {
1685+
Some(c)
1686+
if c.is_ascii_alphanumeric()
1687+
|| c == '_'
1688+
|| c == '$'
1689+
|| c >= '\u{0080}' =>
1690+
{
1691+
let body =
1692+
self.tokenize_sqlite_variable_body(chars)?;
1693+
Ok(Some(Token::Placeholder(format!(":::{body}"))))
1694+
}
1695+
_ => self.tokenizer_error(
1696+
chars.location(),
1697+
"unexpected ':::' without identifier body",
1698+
),
1699+
}
1700+
}
1701+
_ => Ok(Some(Token::DoubleColon)),
1702+
}
1703+
}
1704+
_ => Ok(Some(Token::Colon)),
1705+
}
1706+
} else {
1707+
match chars.peek() {
1708+
Some(':') => self.consume_and_return(chars, Token::DoubleColon),
1709+
Some('=') => self.consume_and_return(chars, Token::Assignment),
1710+
_ => Ok(Some(Token::Colon)),
1711+
}
16711712
}
16721713
}
16731714
';' => self.consume_and_return(chars, Token::SemiColon),
@@ -1758,6 +1799,41 @@ impl<'a> Tokenizer<'a> {
17581799
Some(sch) if self.dialect.is_identifier_start('#') => {
17591800
self.tokenize_identifier_or_keyword([ch, *sch], chars)
17601801
}
1802+
Some(&c)
1803+
if self.dialect.supports_sqlite_variable_syntax()
1804+
&& (c.is_ascii_alphabetic()
1805+
|| c == '_'
1806+
|| c == '$'
1807+
|| c >= '\u{0080}') =>
1808+
{
1809+
let body = self.tokenize_sqlite_variable_body(chars)?;
1810+
Ok(Some(Token::Placeholder(format!("#{body}"))))
1811+
}
1812+
Some(':') if self.dialect.supports_sqlite_variable_syntax() => {
1813+
let mut lookahead = chars.peekable.clone();
1814+
lookahead.next(); // skip first ':'
1815+
if lookahead.peek() == Some(&':') {
1816+
chars.next(); // consume first ':'
1817+
chars.next(); // consume second ':'
1818+
match chars.peek().copied() {
1819+
Some(c)
1820+
if c.is_ascii_alphanumeric()
1821+
|| c == '_'
1822+
|| c == '$'
1823+
|| c >= '\u{0080}' =>
1824+
{
1825+
let body = self.tokenize_sqlite_variable_body(chars)?;
1826+
Ok(Some(Token::Placeholder(format!("#::{body}"))))
1827+
}
1828+
_ => self.tokenizer_error(
1829+
chars.location(),
1830+
"expected identifier after '#::'",
1831+
),
1832+
}
1833+
} else {
1834+
self.start_binop(chars, "#", Token::Sharp)
1835+
}
1836+
}
17611837
_ => self.start_binop(chars, "#", Token::Sharp),
17621838
}
17631839
}
@@ -1802,6 +1878,41 @@ impl<'a> Tokenizer<'a> {
18021878
Some(sch) if self.dialect.is_identifier_start('@') => {
18031879
self.tokenize_identifier_or_keyword([ch, *sch], chars)
18041880
}
1881+
Some(&c)
1882+
if self.dialect.supports_sqlite_variable_syntax()
1883+
&& (c.is_ascii_alphanumeric()
1884+
|| c == '_'
1885+
|| c == '$'
1886+
|| c >= '\u{0080}') =>
1887+
{
1888+
let body = self.tokenize_sqlite_variable_body(chars)?;
1889+
Ok(Some(Token::Placeholder(format!("@{body}"))))
1890+
}
1891+
Some(':') if self.dialect.supports_sqlite_variable_syntax() => {
1892+
let mut lookahead = chars.peekable.clone();
1893+
lookahead.next(); // skip first ':'
1894+
if lookahead.peek() == Some(&':') {
1895+
chars.next(); // consume first ':'
1896+
chars.next(); // consume second ':'
1897+
match chars.peek().copied() {
1898+
Some(c)
1899+
if c.is_ascii_alphanumeric()
1900+
|| c == '_'
1901+
|| c == '$'
1902+
|| c >= '\u{0080}' =>
1903+
{
1904+
let body = self.tokenize_sqlite_variable_body(chars)?;
1905+
Ok(Some(Token::Placeholder(format!("@::{body}"))))
1906+
}
1907+
_ => self.tokenizer_error(
1908+
chars.location(),
1909+
"expected identifier after '@::'",
1910+
),
1911+
}
1912+
} else {
1913+
Ok(Some(Token::AtSign))
1914+
}
1915+
}
18051916
_ => Ok(Some(Token::AtSign)),
18061917
}
18071918
}
@@ -1902,6 +2013,49 @@ impl<'a> Tokenizer<'a> {
19022013
}
19032014
}
19042015

2016+
/// Reads the body of a SQLite parameter variable after its prefix char.
2017+
/// Body chars are alphanumeric, `_`, `$`, and any char above U+007F.
2018+
/// Handles TCL `::` namespace separators and `(subscript)` forms.
2019+
fn tokenize_sqlite_variable_body(&self, chars: &mut State) -> Result<String, TokenizerError> {
2020+
let mut body = peeking_take_while(chars, |c| {
2021+
c.is_ascii_alphanumeric() || c == '_' || c == '$' || c >= '\u{0080}'
2022+
});
2023+
loop {
2024+
if chars.peek() != Some(&':') {
2025+
break;
2026+
}
2027+
chars.next(); // consume first ':'
2028+
if chars.peek() != Some(&':') {
2029+
// Single ':', not a TCL namespace separator
2030+
break;
2031+
}
2032+
chars.next(); // consume second ':'
2033+
body.push_str("::");
2034+
body.push_str(&peeking_take_while(chars, |c| {
2035+
c.is_ascii_alphanumeric() || c == '_' || c == '$' || c >= '\u{0080}'
2036+
}));
2037+
}
2038+
if chars.peek() == Some(&'(') {
2039+
chars.next(); // consume '('
2040+
let sub = peeking_take_while(chars, |c| c != ')' && !c.is_ascii_whitespace());
2041+
match chars.peek() {
2042+
Some(&')') => {
2043+
chars.next(); // consume ')'
2044+
body.push('(');
2045+
body.push_str(&sub);
2046+
body.push(')');
2047+
}
2048+
_ => {
2049+
return self.tokenizer_error(
2050+
chars.location(),
2051+
"unterminated SQLite variable subscript",
2052+
);
2053+
}
2054+
}
2055+
}
2056+
Ok(body)
2057+
}
2058+
19052059
/// Tokenize dollar preceded value (i.e: a string/placeholder)
19062060
fn tokenize_dollar_preceded_value(&self, chars: &mut State) -> Result<Token, TokenizerError> {
19072061
let mut s = String::new();
@@ -1998,6 +2152,38 @@ impl<'a> Tokenizer<'a> {
19982152
}
19992153
}
20002154
} else {
2155+
if self.dialect.supports_sqlite_variable_syntax() {
2156+
loop {
2157+
if chars.peek() != Some(&':') {
2158+
break;
2159+
}
2160+
chars.next(); // consume first ':'
2161+
if chars.peek() != Some(&':') {
2162+
break;
2163+
}
2164+
chars.next(); // consume second ':'
2165+
value.push_str("::");
2166+
value.push_str(&peeking_take_while(chars, |ch| {
2167+
ch.is_ascii_alphanumeric() || ch == '_' || ch == '$' || ch >= '\u{0080}'
2168+
}));
2169+
}
2170+
if !value.is_empty() && chars.peek() == Some(&'(') {
2171+
chars.next(); // consume '('
2172+
let sub =
2173+
peeking_take_while(chars, |c| c != ')' && !c.is_ascii_whitespace());
2174+
if chars.peek() == Some(&')') {
2175+
chars.next(); // consume ')'
2176+
value.push('(');
2177+
value.push_str(&sub);
2178+
value.push(')');
2179+
} else {
2180+
return self.tokenizer_error(
2181+
chars.location(),
2182+
"unterminated SQLite variable subscript",
2183+
);
2184+
}
2185+
}
2186+
}
20012187
return Ok(Token::Placeholder(format!("${value}")));
20022188
}
20032189
}

‎tests/sqlparser_common.rs‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10257,7 +10257,8 @@ fn parse_grant() {
1025710257
verified_stmt("GRANT SELECT ON ALL MATERIALIZED VIEWS IN SCHEMA db1.sc1 TO ROLE role1");
1025810258
verified_stmt("GRANT SELECT ON ALL EXTERNAL TABLES IN SCHEMA db1.sc1 TO ROLE role1");
1025910259
verified_stmt("GRANT USAGE ON ALL FUNCTIONS IN SCHEMA db1.sc1 TO ROLE role1");
10260-
verified_stmt("GRANT USAGE ON SCHEMA sc1 TO a:b");
10260+
all_dialects_except(|d| d.supports_sqlite_variable_syntax())
10261+
.verified_stmt("GRANT USAGE ON SCHEMA sc1 TO a:b");
1026110262
verified_stmt("GRANT USAGE ON SCHEMA sc1 TO GROUP group1");
1026210263
verified_stmt("GRANT OWNERSHIP ON ALL TABLES IN SCHEMA DEV_STAS_ROGOZHIN TO ROLE ANALYST");
1026310264
verified_stmt("GRANT OWNERSHIP ON ALL TABLES IN SCHEMA DEV_STAS_ROGOZHIN TO ROLE ANALYST COPY CURRENT GRANTS");
@@ -20135,7 +20136,9 @@ fn parse_stage_table_factor() {
2013520136
supported.verified_stmt("SELECT * FROM @stage");
2013620137
supported.verified_stmt("SELECT * FROM @stage, my_table");
2013720138

20138-
let unsupported = all_dialects_where(|d| !d.supports_stages() && !d.is_identifier_start('@'));
20139+
let unsupported = all_dialects_where(|d| {
20140+
!d.supports_stages() && !d.is_identifier_start('@') && !d.supports_sqlite_variable_syntax()
20141+
});
2013920142
assert_eq!(
2014020143
unsupported
2014120144
.parse_sql_statements("SELECT * FROM @stage")

‎tests/sqlparser_sqlite.rs‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -971,3 +971,68 @@ fn sqlite_and_generic() -> TestedDialects {
971971
Box::new(GenericDialect {}),
972972
])
973973
}
974+
975+
#[test]
976+
fn parse_sqlite_variable_tokens() {
977+
// $, @, :, # parameter forms with SQLite tokenizer rules
978+
sqlite().verified_stmt("SELECT #a");
979+
sqlite().verified_stmt("SELECT @a");
980+
sqlite().verified_stmt("SELECT :a");
981+
sqlite().verified_stmt("SELECT $a");
982+
983+
// digit-first body (allowed for @, :, not #)
984+
sqlite().verified_stmt("SELECT @2t");
985+
sqlite().verified_stmt("SELECT :9c");
986+
987+
// $ in variable body
988+
sqlite().verified_stmt("SELECT @a$");
989+
sqlite().verified_stmt("SELECT :a$");
990+
991+
// subscript suffix
992+
sqlite().verified_stmt("SELECT :a()");
993+
sqlite().verified_stmt("SELECT :a(x)");
994+
sqlite().verified_stmt("SELECT @a()");
995+
996+
// TCL :: namespace separator
997+
sqlite().verified_stmt("SELECT :::6");
998+
sqlite().verified_stmt("SELECT :a::b");
999+
sqlite().verified_stmt("SELECT $a::b");
1000+
1001+
// TCL :: namespace at start of body
1002+
sqlite().verified_stmt("SELECT @::b");
1003+
sqlite().verified_stmt("SELECT #::b");
1004+
1005+
// dollar variable with :: (repro from fuzzer: SELECT ($$::))
1006+
sqlite().verified_stmt("SELECT ($$::)");
1007+
1008+
// non-ASCII chars in variable body (any char >= U+0080)
1009+
sqlite().verified_stmt("SELECT @\u{00e9}");
1010+
sqlite().verified_stmt("SELECT :\u{00e9}");
1011+
sqlite().verified_stmt("SELECT #\u{00e9}");
1012+
sqlite().verified_stmt("SELECT :::\u{00e9}");
1013+
1014+
// implicit alias forms from repros
1015+
sqlite().one_statement_parses_to("SELECT @2t i", "SELECT @2t AS i");
1016+
sqlite().one_statement_parses_to("SELECT :9c x", "SELECT :9c AS x");
1017+
1018+
// non-SQLite dialects reject the SQLite-specific forms
1019+
// :::6 tokenizes as DoubleColon + Colon + Number in GenericDialect (no valid expr)
1020+
assert!(
1021+
TestedDialects::new(vec![Box::new(GenericDialect {})])
1022+
.parse_sql_statements("SELECT :::6")
1023+
.is_err(),
1024+
"GenericDialect should reject :::6"
1025+
);
1026+
// $$ begins a dollar-quoted string in GenericDialect, making ($$::) unterminated
1027+
assert!(
1028+
TestedDialects::new(vec![Box::new(GenericDialect {})])
1029+
.parse_sql_statements("SELECT ($$::)")
1030+
.is_err(),
1031+
"GenericDialect should reject ($$::)"
1032+
);
1033+
// SQLite rejects subscripts containing whitespace
1034+
assert!(
1035+
sqlite().parse_sql_statements("SELECT $a(x y)").is_err(),
1036+
"SQLiteDialect should reject $a(x y)"
1037+
);
1038+
}

0 commit comments

Comments
 (0)