Skip to content

Commit af1e5b0

Browse files
renecannaorenecannao
authored
feat: expose lossless MySQL user variables (#55)
ParserSQL support for lossless typed MySQL user-variable assignments and usage classification.
2 parents a5cd096 + db9d464 commit af1e5b0

19 files changed

Lines changed: 989 additions & 77 deletions

‎Makefile‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ TEST_SRCS = $(TEST_DIR)/test_main.cpp \
4242
$(TEST_DIR)/test_classifier.cpp \
4343
$(TEST_DIR)/test_expression.cpp \
4444
$(TEST_DIR)/test_set.cpp \
45+
$(TEST_DIR)/test_user_variable.cpp \
4546
$(TEST_DIR)/test_select.cpp \
4647
$(TEST_DIR)/test_emitter.cpp \
4748
$(TEST_DIR)/test_stmt_cache.cpp \

‎include/sql_parser/ast.h‎

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33

44
#include "sql_parser/common.h"
55
#include "sql_parser/arena.h"
6+
#include "sql_parser/token.h"
67
#include <cstdint>
78
#include <type_traits>
89

@@ -12,17 +13,25 @@ struct AstNode {
1213
AstNode* first_child;
1314
AstNode* next_sibling;
1415
const char* value_ptr;
16+
const char* source_ptr;
1517
uint32_t value_len;
18+
uint32_t source_len;
1619
NodeType type;
1720
uint16_t flags;
1821

1922
StringRef value() const { return StringRef{value_ptr, value_len}; }
23+
StringRef source() const { return StringRef{source_ptr, source_len}; }
2024

2125
void set_value(StringRef ref) {
2226
value_ptr = ref.ptr;
2327
value_len = ref.len;
2428
}
2529

30+
void set_source(StringRef ref) {
31+
source_ptr = ref.ptr;
32+
source_len = ref.len;
33+
}
34+
2635
void add_child(AstNode* child) {
2736
if (!child) return;
2837
if (!first_child) {
@@ -34,7 +43,7 @@ struct AstNode {
3443
last->next_sibling = child;
3544
}
3645
};
37-
static_assert(sizeof(AstNode) == 32, "AstNode must be 32 bytes");
46+
static_assert(sizeof(AstNode) == 48, "AstNode layout changed unexpectedly");
3847
static_assert(std::is_trivially_copyable_v<AstNode>);
3948

4049
inline AstNode* make_node(Arena& arena, NodeType type, StringRef value = {},
@@ -48,6 +57,13 @@ inline AstNode* make_node(Arena& arena, NodeType type, StringRef value = {},
4857
return node;
4958
}
5059

60+
inline AstNode* make_node_from_token(Arena& arena, NodeType type,
61+
const Token& token, uint16_t flags = 0) {
62+
AstNode* node = make_node(arena, type, token.text, flags);
63+
if (node) node->set_source(token.source);
64+
return node;
65+
}
66+
5167
} // namespace sql_parser
5268

5369
#endif // SQL_PARSER_AST_H

‎include/sql_parser/common.h‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -231,6 +231,12 @@ enum class NodeType : uint16_t {
231231
NODE_SET_ROLE, // SET [LOCAL] ROLE <name>|NONE|DEFAULT
232232
NODE_SET_SESSION_AUTHORIZATION, // SET SESSION AUTHORIZATION <name>|DEFAULT
233233
NODE_SET_CONSTRAINTS, // SET CONSTRAINTS {ALL|<name>[,...]} {DEFERRED|IMMEDIATE}
234+
235+
// MySQL lossless user-variable/literal nodes. Keep appended so existing
236+
// enum values remain stable for consumers that index by NodeType.
237+
NODE_USER_VARIABLE,
238+
NODE_LITERAL_HEX,
239+
NODE_LITERAL_BIT,
234240
};
235241

236242
} // namespace sql_parser

‎include/sql_parser/digest.h‎

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -63,15 +63,17 @@ class Digest {
6363

6464
// Helper: check if a token type is a keyword (not an identifier, literal, or operator)
6565
static bool is_keyword_token(TokenType type) {
66-
// Keywords start at TK_SELECT and go through TK_EXCEPT
67-
return static_cast<uint16_t>(type) >= static_cast<uint16_t>(TokenType::TK_SELECT);
66+
return static_cast<uint16_t>(type) >= static_cast<uint16_t>(TokenType::TK_SELECT) &&
67+
static_cast<uint16_t>(type) <= static_cast<uint16_t>(TokenType::TK_RECURSIVE);
6868
}
6969

7070
// Helper: check if a token type is a literal value that should become ?
7171
static bool is_literal_token(TokenType type) {
7272
return type == TokenType::TK_INTEGER ||
7373
type == TokenType::TK_FLOAT ||
74-
type == TokenType::TK_STRING;
74+
type == TokenType::TK_STRING ||
75+
type == TokenType::TK_HEX_LITERAL ||
76+
type == TokenType::TK_BIT_LITERAL;
7577
}
7678

7779
// Helper: uppercase a character
@@ -104,7 +106,12 @@ class Digest {
104106

105107
// Emit a single token to the string builder, uppercasing keywords, replacing literals with ?
106108
void emit_token(StringBuilder& sb, const Token& t, TokenType prev) {
107-
bool space = (prev != TokenType::TK_EOF) && needs_space_before(prev, t.type);
109+
bool quoted_user_after_account = t.type == TokenType::TK_USER_VARIABLE &&
110+
t.source.len >= 2 &&
111+
(t.source.ptr[1] == '\'' || t.source.ptr[1] == '"' || t.source.ptr[1] == '`') &&
112+
(prev == TokenType::TK_STRING || prev == TokenType::TK_QUESTION);
113+
bool space = (prev != TokenType::TK_EOF) &&
114+
!quoted_user_after_account && needs_space_before(prev, t.type);
108115
if (space) sb.append_char(' ');
109116

110117
if (is_literal_token(t.type)) {
@@ -115,6 +122,13 @@ class Digest {
115122
sb.append(t.text.ptr, t.text.len);
116123
} else if (t.type == TokenType::TK_QUESTION) {
117124
sb.append_char('?');
125+
} else if (t.type == TokenType::TK_USER_VARIABLE) {
126+
if (t.source.len >= 2 &&
127+
(t.source.ptr[1] == '\'' || t.source.ptr[1] == '"' || t.source.ptr[1] == '`')) {
128+
sb.append("@?", 2);
129+
} else {
130+
sb.append(t.source);
131+
}
118132
} else if (t.type == TokenType::TK_COMMA) {
119133
sb.append(",", 1);
120134
} else {

‎include/sql_parser/emitter.h‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -123,6 +123,8 @@ class Emitter {
123123
case NodeType::NODE_ARRAY_SUBSCRIPT: emit_array_subscript(node); break;
124124
case NodeType::NODE_FIELD_ACCESS: emit_field_access(node); break;
125125
case NodeType::NODE_SUBQUERY: emit_subquery(node); break;
126+
case NodeType::NODE_EXPRESSION: emit_parenthesized_expression(node); break;
127+
case NodeType::NODE_USER_VARIABLE: emit_user_variable(node); break;
126128

127129
// ---- Leaf nodes (emit value directly) ----
128130
case NodeType::NODE_PLACEHOLDER:
@@ -131,6 +133,8 @@ class Emitter {
131133
// ---- Leaf nodes (emit value directly) ----
132134
case NodeType::NODE_LITERAL_INT:
133135
case NodeType::NODE_LITERAL_FLOAT:
136+
case NodeType::NODE_LITERAL_HEX:
137+
case NodeType::NODE_LITERAL_BIT:
134138
if (mode_ == EmitMode::DIGEST) { sb_.append_char('?'); break; }
135139
emit_value(node); break;
136140
case NodeType::NODE_LITERAL_NULL:
@@ -152,6 +156,30 @@ class Emitter {
152156
sb_.append(node->value_ptr, node->value_len);
153157
}
154158

159+
void emit_user_variable(const AstNode* node) {
160+
if (mode_ == EmitMode::DIGEST) {
161+
StringRef source = node->source();
162+
if (source.len >= 2 &&
163+
(source.ptr[1] == '\'' || source.ptr[1] == '"' || source.ptr[1] == '`')) {
164+
sb_.append("@?", 2);
165+
return;
166+
}
167+
}
168+
StringRef source = node->source();
169+
if (!source.empty()) {
170+
sb_.append(source.ptr, source.len);
171+
return;
172+
}
173+
sb_.append_char('@');
174+
emit_value(node);
175+
}
176+
177+
void emit_parenthesized_expression(const AstNode* node) {
178+
sb_.append_char('(');
179+
if (node->first_child) emit_node(node->first_child);
180+
sb_.append_char(')');
181+
}
182+
155183
void emit_string_literal(const AstNode* node) {
156184
sb_.append_char('\'');
157185
sb_.append(node->value_ptr, node->value_len);

‎include/sql_parser/expression_parser.h‎

Lines changed: 59 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
#include "sql_parser/tokenizer.h"
77
#include "sql_parser/ast.h"
88
#include "sql_parser/arena.h"
9+
#include "sql_parser/user_variable.h"
910

1011
namespace sql_parser {
1112

@@ -99,19 +100,27 @@ class ExpressionParser {
99100
switch (t.type) {
100101
case TokenType::TK_INTEGER: {
101102
tok_.skip();
102-
return make_node(arena_, NodeType::NODE_LITERAL_INT, t.text);
103+
return make_node_from_token(arena_, NodeType::NODE_LITERAL_INT, t);
103104
}
104105
case TokenType::TK_FLOAT: {
105106
tok_.skip();
106-
return make_node(arena_, NodeType::NODE_LITERAL_FLOAT, t.text);
107+
return make_node_from_token(arena_, NodeType::NODE_LITERAL_FLOAT, t);
108+
}
109+
case TokenType::TK_HEX_LITERAL: {
110+
tok_.skip();
111+
return make_node_from_token(arena_, NodeType::NODE_LITERAL_HEX, t);
112+
}
113+
case TokenType::TK_BIT_LITERAL: {
114+
tok_.skip();
115+
return make_node_from_token(arena_, NodeType::NODE_LITERAL_BIT, t);
107116
}
108117
case TokenType::TK_STRING: {
109118
tok_.skip();
110-
return make_node(arena_, NodeType::NODE_LITERAL_STRING, t.text);
119+
return make_node_from_token(arena_, NodeType::NODE_LITERAL_STRING, t);
111120
}
112121
case TokenType::TK_NULL: {
113122
tok_.skip();
114-
return make_node(arena_, NodeType::NODE_LITERAL_NULL, t.text);
123+
return make_node_from_token(arena_, NodeType::NODE_LITERAL_NULL, t);
115124
}
116125
case TokenType::TK_TRUE:
117126
case TokenType::TK_FALSE: {
@@ -150,6 +159,10 @@ class ExpressionParser {
150159
static_cast<uint32_t>((name.text.ptr + name.text.len) - t.text.ptr)};
151160
return make_node(arena_, NodeType::NODE_COLUMN_REF, full);
152161
}
162+
case TokenType::TK_USER_VARIABLE: {
163+
tok_.skip();
164+
return make_mysql_user_variable_node(arena_, t);
165+
}
153166
case TokenType::TK_DOUBLE_AT: {
154167
// System variable: @@name or @@scope.name
155168
tok_.skip();
@@ -174,19 +187,26 @@ class ExpressionParser {
174187
AstNode* operand = parse(Precedence::UNARY);
175188
if (!operand) return nullptr;
176189
AstNode* node = make_node(arena_, NodeType::NODE_UNARY_OP, t.text);
190+
set_span_through_node_(node, t.source, operand);
177191
node->add_child(operand);
178192
return node;
179193
}
180194
case TokenType::TK_PLUS: {
181195
// Unary plus
182196
tok_.skip();
183-
return parse(Precedence::UNARY);
197+
AstNode* operand = parse(Precedence::UNARY);
198+
if (!operand) return nullptr;
199+
AstNode* node = make_node(arena_, NodeType::NODE_UNARY_OP, t.text);
200+
set_span_through_node_(node, t.source, operand);
201+
node->add_child(operand);
202+
return node;
184203
}
185204
case TokenType::TK_NOT: {
186205
tok_.skip();
187206
AstNode* operand = parse(Precedence::NOT);
188207
if (!operand) return nullptr;
189208
AstNode* node = make_node(arena_, NodeType::NODE_UNARY_OP, t.text);
209+
set_span_through_node_(node, t.source, operand);
190210
node->add_child(operand);
191211
return node;
192212
}
@@ -266,7 +286,12 @@ class ExpressionParser {
266286
return parse_postfix(tuple);
267287
}
268288
if (tok_.peek().type == TokenType::TK_RPAREN) {
269-
tok_.skip();
289+
Token close = tok_.next_token();
290+
AstNode* wrapper = make_node(arena_, NodeType::NODE_EXPRESSION);
291+
wrapper->set_source(StringRef{t.source.ptr,
292+
static_cast<uint32_t>(close.source.ptr + close.source.len - t.source.ptr)});
293+
wrapper->add_child(expr);
294+
return parse_postfix(wrapper);
270295
}
271296
// Check for postfix: (expr).field or (expr)[index]
272297
return parse_postfix(expr);
@@ -287,11 +312,39 @@ class ExpressionParser {
287312
}
288313
}
289314

315+
static void set_span_through_node_(AstNode* node, StringRef start,
316+
const AstNode* end_node) {
317+
if (!node || !start.ptr || !end_node) return;
318+
StringRef end = end_node->source();
319+
if (end.empty()) end = end_node->value();
320+
if (!end.ptr || end.ptr < start.ptr) return;
321+
node->set_source(StringRef{start.ptr,
322+
static_cast<uint32_t>(end.ptr + end.len - start.ptr)});
323+
}
324+
290325
AstNode* parse_identifier_or_function(const Token& name_token) {
291326
// Check for function call: name(
292327
if (tok_.peek().type == TokenType::TK_LPAREN) {
293328
tok_.skip(); // consume (
294329
AstNode* func = make_node(arena_, NodeType::NODE_FUNCTION_CALL, name_token.text);
330+
// CAST uses `CAST(expr AS type)` rather than a comma-separated
331+
// argument list. Model it as a function call so consumers can
332+
// reject or handle the expression without leaving valid input
333+
// unconsumed.
334+
if (name_token.text.equals_ci("CAST", 4)) {
335+
AstNode* arg = parse();
336+
if (!arg || tok_.peek().type != TokenType::TK_AS) return func;
337+
func->add_child(arg);
338+
tok_.skip();
339+
Token type = tok_.next_token();
340+
if (type.type == TokenType::TK_EOF ||
341+
type.type == TokenType::TK_RPAREN) {
342+
return func;
343+
}
344+
func->add_child(make_node(arena_, NodeType::NODE_IDENTIFIER, type.text));
345+
if (tok_.peek().type == TokenType::TK_RPAREN) tok_.skip();
346+
return func;
347+
}
295348
// Parse argument list
296349
if (tok_.peek().type != TokenType::TK_RPAREN) {
297350
while (true) {

‎include/sql_parser/parse_result.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,8 @@ struct ParseResult {
3939
AstNode* ast = nullptr;
4040
ErrorInfo error;
4141
StringRef remaining;
42+
bool full_input = false;
43+
bool has_user_variables = false;
4244

4345
StringRef table_name;
4446
StringRef schema_name;

‎include/sql_parser/parser.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
#include "sql_parser/ast.h"
88
#include "sql_parser/parse_result.h"
99
#include "sql_parser/stmt_cache.h"
10+
#include "sql_parser/user_variable.h"
1011

1112
namespace sql_parser {
1213

‎include/sql_parser/set_parser.h‎

Lines changed: 25 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -352,9 +352,14 @@ class SetParser {
352352
}
353353
} else {
354354
while (tok_.peek().type == TokenType::TK_COMMA) {
355-
tok_.skip();
355+
Token comma = tok_.next_token();
356356
AstNode* next_assign = parse_comma_item();
357-
if (next_assign) root->add_child(next_assign);
357+
if (next_assign) {
358+
root->add_child(next_assign);
359+
} else {
360+
tok_.flag_error_at(comma.source);
361+
break;
362+
}
358363
}
359364
}
360365

@@ -551,7 +556,17 @@ class SetParser {
551556
}
552557

553558
Token var = tok_.peek();
554-
if (var.type == TokenType::TK_AT) {
559+
bool user_variable_target = false;
560+
if (var.type == TokenType::TK_USER_VARIABLE) {
561+
user_variable_target = true;
562+
tok_.skip();
563+
AstNode* variable = make_mysql_user_variable_node(arena_, var);
564+
if (!variable) {
565+
tok_.flag_error_at(var.source);
566+
return nullptr;
567+
}
568+
target->add_child(variable);
569+
} else if (var.type == TokenType::TK_AT) {
555570
// User variable @name. The name may be backtick/double-quoted;
556571
// in that case the source bytes between `@` and the name include
557572
// the opening delimiter (and the closing delimiter sits one past
@@ -624,13 +639,20 @@ class SetParser {
624639

625640
// Expect = or := (MySQL) or TO (PostgreSQL)
626641
Token eq = tok_.peek();
642+
bool has_assignment_operator = false;
627643
if (eq.type == TokenType::TK_EQUAL || eq.type == TokenType::TK_COLON_EQUAL) {
628644
tok_.skip();
645+
has_assignment_operator = true;
629646
} else if constexpr (D == Dialect::PostgreSQL) {
630647
if (eq.type == TokenType::TK_TO) {
631648
tok_.skip();
649+
has_assignment_operator = true;
632650
}
633651
}
652+
if (user_variable_target && !has_assignment_operator) {
653+
tok_.flag_error_at(eq.source);
654+
return nullptr;
655+
}
634656

635657
// Parse RHS expression. If the parser couldn't produce one --
636658
// typically because the input is truncated (`SET x =`), starts

0 commit comments

Comments
 (0)