1use std::str::FromStr;
16
17use chrono::{DateTime, Utc};
18use snafu::{OptionExt, ResultExt};
19use sqlparser::ast::{Ident, Query, Value};
20use sqlparser::dialect::Dialect;
21use sqlparser::keywords::Keyword;
22use sqlparser::parser::{Parser, ParserError, ParserOptions};
23use sqlparser::tokenizer::{Token, TokenWithSpan};
24
25use crate::ast::{Expr, ObjectName};
26use crate::error::{self, InvalidSqlSnafu, Result, SyntaxSnafu};
27use crate::parsers::tql_parser;
28use crate::statements::kill::Kill;
29use crate::statements::statement::Statement;
30use crate::statements::transform_statements;
31
32pub const FLOW: &str = "FLOW";
33
34#[derive(Clone, Debug, Default)]
36pub struct ParseOptions {
37 pub scheduled_time: Option<DateTime<Utc>>,
40}
41
42pub struct ParserContext<'a> {
44 pub(crate) parser: Parser<'a>,
45 pub(crate) sql: &'a str,
46 pub(crate) scheduled_time: Option<DateTime<Utc>>,
48}
49
50impl ParserContext<'_> {
51 pub fn new<'a>(dialect: &'a dyn Dialect, sql: &'a str) -> Result<ParserContext<'a>> {
53 let parser = Parser::new(dialect)
54 .with_options(ParserOptions::new().with_trailing_commas(true))
55 .try_with_sql(sql)
56 .context(SyntaxSnafu)?;
57
58 Ok(ParserContext {
59 parser,
60 sql,
61 scheduled_time: None,
62 })
63 }
64
65 pub fn parser_query(&mut self) -> Result<Box<Query>> {
67 self.parser.parse_query().context(SyntaxSnafu)
68 }
69
70 pub fn create_with_dialect(
72 sql: &str,
73 dialect: &dyn Dialect,
74 opts: ParseOptions,
75 ) -> Result<Vec<Statement>> {
76 let mut stmts: Vec<Statement> = Vec::new();
77
78 let mut parser_ctx = ParserContext::new(dialect, sql)?;
79 parser_ctx.scheduled_time = opts.scheduled_time;
80
81 let mut expecting_statement_delimiter = false;
82 loop {
83 while parser_ctx.parser.consume_token(&Token::SemiColon) {
85 expecting_statement_delimiter = false;
86 }
87
88 if parser_ctx.parser.peek_token() == Token::EOF {
89 break;
90 }
91 if expecting_statement_delimiter {
92 return parser_ctx.unsupported(parser_ctx.peek_token_as_string());
93 }
94
95 let statement = parser_ctx.parse_statement()?;
96 stmts.push(statement);
97 expecting_statement_delimiter = true;
98 }
99
100 transform_statements(&mut stmts)?;
101
102 Ok(stmts)
103 }
104
105 pub fn parse_table_name(sql: &str, dialect: &dyn Dialect) -> Result<ObjectName> {
106 let parser = Parser::new(dialect)
107 .with_options(ParserOptions::new().with_trailing_commas(true))
108 .try_with_sql(sql)
109 .context(SyntaxSnafu)?;
110 ParserContext {
111 parser,
112 sql,
113 scheduled_time: None,
114 }
115 .intern_parse_table_name()
116 }
117
118 pub(crate) fn intern_parse_table_name(&mut self) -> Result<ObjectName> {
119 let raw_table_name =
120 self.parser
121 .parse_object_name(false)
122 .context(error::UnexpectedSnafu {
123 expected: "a table name",
124 actual: self.parser.peek_token().to_string(),
125 })?;
126 Self::canonicalize_object_name(raw_table_name)
127 }
128
129 pub fn parse_function(sql: &str, dialect: &dyn Dialect) -> Result<Expr> {
130 let mut parser = Parser::new(dialect)
131 .with_options(ParserOptions::new().with_trailing_commas(true))
132 .try_with_sql(sql)
133 .context(SyntaxSnafu)?;
134
135 let function_name = parser.parse_identifier().context(SyntaxSnafu)?;
136 parser
137 .parse_function(vec![function_name].into())
138 .context(SyntaxSnafu)
139 }
140
141 pub fn parse_statement(&mut self) -> Result<Statement> {
143 match self.parser.peek_token().token {
144 Token::Word(w) => match w.keyword {
145 Keyword::CREATE => {
146 let _ = self.parser.next_token();
147 self.parse_create()
148 }
149
150 Keyword::EXPLAIN => {
151 let _ = self.parser.next_token();
152 self.parse_explain()
153 }
154
155 Keyword::SHOW => {
156 let _ = self.parser.next_token();
157 self.parse_show()
158 }
159
160 Keyword::DELETE => self.parse_delete(),
161
162 Keyword::DESCRIBE | Keyword::DESC => {
163 let _ = self.parser.next_token();
164 self.parse_describe()
165 }
166
167 Keyword::INSERT => self.parse_insert(),
168
169 Keyword::REPLACE => self.parse_replace(),
170
171 Keyword::SELECT | Keyword::VALUES => self.parse_query(),
172
173 Keyword::WITH => self.parse_with_tql(),
174
175 Keyword::ALTER => self.parse_alter(),
176
177 Keyword::DROP => self.parse_drop(),
178
179 Keyword::COPY => self.parse_copy(),
180
181 Keyword::TRUNCATE => self.parse_truncate(),
182
183 Keyword::COMMENT => self.parse_comment(),
184
185 Keyword::SET => self.parse_set_variables(),
186
187 Keyword::ADMIN => self.parse_admin_command(),
188
189 Keyword::NoKeyword
190 if w.quote_style.is_none() && w.value.to_uppercase() == tql_parser::TQL =>
191 {
192 self.parse_tql(false)
193 }
194
195 Keyword::NoKeyword
196 if w.quote_style.is_none() && w.value.eq_ignore_ascii_case("UNDROP") =>
197 {
198 self.parse_undrop_table()
199 }
200
201 Keyword::DECLARE => self.parse_declare_cursor(),
202
203 Keyword::FETCH => self.parse_fetch_cursor(),
204
205 Keyword::CLOSE => self.parse_close_cursor(),
206
207 Keyword::USE => {
208 let _ = self.parser.next_token();
209
210 let database_name = self.parser.parse_identifier().with_context(|_| {
211 error::UnexpectedSnafu {
212 expected: "a database name",
213 actual: self.peek_token_as_string(),
214 }
215 })?;
216 Ok(Statement::Use(
217 Self::canonicalize_identifier(database_name).value,
218 ))
219 }
220
221 Keyword::KILL => {
222 let _ = self.parser.next_token();
223 let kill = if self.parser.parse_keyword(Keyword::QUERY) {
224 let connection_id_exp =
226 self.parser.parse_number_value().with_context(|_| {
227 error::UnexpectedSnafu {
228 expected: "MySQL numeric connection id",
229 actual: self.peek_token_as_string(),
230 }
231 })?;
232 let Value::Number(s, _) = connection_id_exp.value else {
233 return error::UnexpectedTokenSnafu {
234 expected: "MySQL numeric connection id",
235 actual: connection_id_exp.to_string(),
236 }
237 .fail();
238 };
239
240 let connection_id = u32::from_str(&s).map_err(|_| {
241 error::UnexpectedTokenSnafu {
242 expected: "MySQL numeric connection id",
243 actual: s,
244 }
245 .build()
246 })?;
247 Kill::ConnectionId(connection_id)
248 } else {
249 let process_id_ident =
250 self.parser.parse_literal_string().with_context(|_| {
251 error::UnexpectedSnafu {
252 expected: "process id string literal",
253 actual: self.peek_token_as_string(),
254 }
255 })?;
256 Kill::ProcessId(process_id_ident)
257 };
258
259 Ok(Statement::Kill(kill))
260 }
261
262 _ => self.unsupported(self.peek_token_as_string()),
263 },
264 Token::LParen => self.parse_query(),
265 unexpected => self.unsupported(unexpected.to_string()),
266 }
267 }
268
269 pub fn parse_mysql_prepare_stmt(sql: &str, dialect: &dyn Dialect) -> Result<(String, String)> {
271 ParserContext::new(dialect, sql)?.parse_mysql_prepare()
272 }
273
274 pub fn parse_mysql_execute_stmt(
276 sql: &str,
277 dialect: &dyn Dialect,
278 ) -> Result<(String, Vec<Expr>)> {
279 ParserContext::new(dialect, sql)?.parse_mysql_execute()
280 }
281
282 pub fn parse_mysql_deallocate_stmt(sql: &str, dialect: &dyn Dialect) -> Result<String> {
284 ParserContext::new(dialect, sql)?.parse_deallocate()
285 }
286
287 pub fn unsupported<T>(&self, keyword: String) -> Result<T> {
289 error::UnsupportedSnafu { keyword }.fail()
290 }
291
292 pub(crate) fn expected<T>(&self, expected: &str, found: TokenWithSpan) -> Result<T> {
294 Err(ParserError::ParserError(format!(
295 "Expected {expected}, found: {found}",
296 )))
297 .context(SyntaxSnafu)
298 }
299
300 pub fn matches_keyword(&mut self, expected: Keyword) -> bool {
301 match self.parser.peek_token().token {
302 Token::Word(w) => w.keyword == expected,
303 _ => false,
304 }
305 }
306
307 pub fn consume_token(&mut self, expected: &str) -> bool {
308 if self.peek_token_as_string().to_uppercase() == *expected.to_uppercase() {
309 let _ = self.parser.next_token();
310 true
311 } else {
312 false
313 }
314 }
315
316 #[inline]
317 pub(crate) fn peek_token_as_string(&self) -> String {
318 self.parser.peek_token().to_string()
319 }
320
321 pub fn canonicalize_identifier(ident: Ident) -> Ident {
323 if ident.quote_style.is_some() {
324 ident
325 } else {
326 Ident::new(ident.value.to_lowercase())
327 }
328 }
329
330 pub(crate) fn canonicalize_object_name(object_name: ObjectName) -> Result<ObjectName> {
332 object_name
333 .0
334 .into_iter()
335 .map(|x| {
336 x.as_ident()
337 .cloned()
338 .map(Self::canonicalize_identifier)
339 .with_context(|| InvalidSqlSnafu {
340 msg: format!("not an ident: '{x}'"),
341 })
342 })
343 .collect::<Result<Vec<_>>>()
344 .map(Into::into)
345 }
346
347 pub(crate) fn parse_object_name(&mut self) -> std::result::Result<ObjectName, ParserError> {
352 self.parser.parse_object_name(false)
353 }
354}
355
356#[cfg(test)]
357mod tests {
358
359 use datatypes::prelude::ConcreteDataType;
360 use sqlparser::dialect::MySqlDialect;
361
362 use super::*;
363 use crate::dialect::GreptimeDbDialect;
364 use crate::statements::create::CreateTable;
365 use crate::statements::sql_data_type_to_concrete_data_type;
366
367 fn test_timestamp_precision(sql: &str, expected_type: ConcreteDataType) {
368 match ParserContext::create_with_dialect(
369 sql,
370 &GreptimeDbDialect {},
371 ParseOptions::default(),
372 )
373 .unwrap()
374 .pop()
375 .unwrap()
376 {
377 Statement::CreateTable(CreateTable { columns, .. }) => {
378 let ts_col = columns.first().unwrap();
379 assert_eq!(
380 expected_type,
381 sql_data_type_to_concrete_data_type(ts_col.data_type()).unwrap()
382 );
383 }
384 _ => unreachable!(),
385 }
386 }
387
388 #[test]
389 pub fn test_create_table_with_precision() {
390 test_timestamp_precision(
391 "create table demo (ts timestamp time index, cnt int);",
392 ConcreteDataType::timestamp_millisecond_datatype(),
393 );
394 test_timestamp_precision(
395 "create table demo (ts timestamp(0) time index, cnt int);",
396 ConcreteDataType::timestamp_second_datatype(),
397 );
398 test_timestamp_precision(
399 "create table demo (ts timestamp(3) time index, cnt int);",
400 ConcreteDataType::timestamp_millisecond_datatype(),
401 );
402 test_timestamp_precision(
403 "create table demo (ts timestamp(6) time index, cnt int);",
404 ConcreteDataType::timestamp_microsecond_datatype(),
405 );
406 test_timestamp_precision(
407 "create table demo (ts timestamp(9) time index, cnt int);",
408 ConcreteDataType::timestamp_nanosecond_datatype(),
409 );
410 }
411
412 #[test]
413 #[should_panic]
414 pub fn test_create_table_with_invalid_precision() {
415 test_timestamp_precision(
416 "create table demo (ts timestamp(1) time index, cnt int);",
417 ConcreteDataType::timestamp_millisecond_datatype(),
418 );
419 }
420
421 #[test]
422 pub fn test_parse_table_name() {
423 let table_name = "a.b.c";
424
425 let object_name =
426 ParserContext::parse_table_name(table_name, &GreptimeDbDialect {}).unwrap();
427
428 assert_eq!(object_name.0.len(), 3);
429 assert_eq!(object_name.to_string(), table_name);
430
431 let table_name = "a.b";
432
433 let object_name =
434 ParserContext::parse_table_name(table_name, &GreptimeDbDialect {}).unwrap();
435
436 assert_eq!(object_name.0.len(), 2);
437 assert_eq!(object_name.to_string(), table_name);
438
439 let table_name = "Test.\"public-test\"";
440
441 let object_name =
442 ParserContext::parse_table_name(table_name, &GreptimeDbDialect {}).unwrap();
443
444 assert_eq!(object_name.0.len(), 2);
445 assert_eq!(object_name.to_string(), table_name.to_ascii_lowercase());
446
447 let table_name = "HelloWorld";
448
449 let object_name =
450 ParserContext::parse_table_name(table_name, &GreptimeDbDialect {}).unwrap();
451
452 assert_eq!(object_name.0.len(), 1);
453 assert_eq!(object_name.to_string(), table_name.to_ascii_lowercase());
454 }
455
456 #[test]
457 pub fn test_parse_mysql_prepare_stmt() {
458 let sql = "PREPARE stmt1 FROM 'SELECT * FROM t1 WHERE id = ?';";
459 let (stmt_name, stmt) =
460 ParserContext::parse_mysql_prepare_stmt(sql, &MySqlDialect {}).unwrap();
461 assert_eq!(stmt_name, "stmt1");
462 assert_eq!(stmt, "SELECT * FROM t1 WHERE id = ?");
463
464 let sql = "PREPARE stmt2 FROM \"SELECT * FROM t1 WHERE id = ?\"";
465 let (stmt_name, stmt) =
466 ParserContext::parse_mysql_prepare_stmt(sql, &MySqlDialect {}).unwrap();
467 assert_eq!(stmt_name, "stmt2");
468 assert_eq!(stmt, "SELECT * FROM t1 WHERE id = ?");
469 }
470
471 #[test]
472 pub fn test_parse_mysql_execute_stmt() {
473 let sql = "EXECUTE stmt1 USING 1, 'hello';";
474 let (stmt_name, params) =
475 ParserContext::parse_mysql_execute_stmt(sql, &GreptimeDbDialect {}).unwrap();
476 assert_eq!(stmt_name, "stmt1");
477 assert_eq!(params.len(), 2);
478 assert_eq!(params[0].to_string(), "1");
479 assert_eq!(params[1].to_string(), "'hello'");
480
481 let sql = "EXECUTE stmt2;";
482 let (stmt_name, params) =
483 ParserContext::parse_mysql_execute_stmt(sql, &GreptimeDbDialect {}).unwrap();
484 assert_eq!(stmt_name, "stmt2");
485 assert_eq!(params.len(), 0);
486
487 let sql = "EXECUTE stmt3 USING 231, 'hello', \"2003-03-1\", NULL, ;";
488 let (stmt_name, params) =
489 ParserContext::parse_mysql_execute_stmt(sql, &GreptimeDbDialect {}).unwrap();
490 assert_eq!(stmt_name, "stmt3");
491 assert_eq!(params.len(), 4);
492 assert_eq!(params[0].to_string(), "231");
493 assert_eq!(params[1].to_string(), "'hello'");
494 assert_eq!(params[2].to_string(), "\"2003-03-1\"");
495 assert_eq!(params[3].to_string(), "NULL");
496 }
497
498 #[test]
499 pub fn test_parse_mysql_deallocate_stmt() {
500 let sql = "DEALLOCATE stmt1;";
501 let stmt_name = ParserContext::parse_mysql_deallocate_stmt(sql, &MySqlDialect {}).unwrap();
502 assert_eq!(stmt_name, "stmt1");
503
504 let sql = "DEALLOCATE stmt2";
505 let stmt_name = ParserContext::parse_mysql_deallocate_stmt(sql, &MySqlDialect {}).unwrap();
506 assert_eq!(stmt_name, "stmt2");
507 }
508
509 #[test]
510 pub fn test_parse_kill_query_statement() {
511 use crate::statements::kill::Kill;
512
513 let sql = "KILL QUERY 123";
515 let statements =
516 ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
517 .unwrap();
518
519 assert_eq!(statements.len(), 1);
520 match &statements[0] {
521 Statement::Kill(Kill::ConnectionId(connection_id)) => {
522 assert_eq!(*connection_id, 123);
523 }
524 _ => panic!("Expected Kill::ConnectionId statement"),
525 }
526
527 let sql = "KILL QUERY 999999";
529 let statements =
530 ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
531 .unwrap();
532
533 assert_eq!(statements.len(), 1);
534 match &statements[0] {
535 Statement::Kill(Kill::ConnectionId(connection_id)) => {
536 assert_eq!(*connection_id, 999999);
537 }
538 _ => panic!("Expected Kill::ConnectionId statement"),
539 }
540 }
541
542 #[test]
543 pub fn test_parse_kill_process_statement() {
544 use crate::statements::kill::Kill;
545
546 let sql = "KILL 'process-123'";
548 let statements =
549 ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
550 .unwrap();
551
552 assert_eq!(statements.len(), 1);
553 match &statements[0] {
554 Statement::Kill(Kill::ProcessId(process_id)) => {
555 assert_eq!(process_id, "process-123");
556 }
557 _ => panic!("Expected Kill::ProcessId statement"),
558 }
559
560 let sql = "KILL \"process-456\"";
562 let statements =
563 ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
564 .unwrap();
565
566 assert_eq!(statements.len(), 1);
567 match &statements[0] {
568 Statement::Kill(Kill::ProcessId(process_id)) => {
569 assert_eq!(process_id, "process-456");
570 }
571 _ => panic!("Expected Kill::ProcessId statement"),
572 }
573
574 let sql = "KILL 'f47ac10b-58cc-4372-a567-0e02b2c3d479'";
576 let statements =
577 ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
578 .unwrap();
579
580 assert_eq!(statements.len(), 1);
581 match &statements[0] {
582 Statement::Kill(Kill::ProcessId(process_id)) => {
583 assert_eq!(process_id, "f47ac10b-58cc-4372-a567-0e02b2c3d479");
584 }
585 _ => panic!("Expected Kill::ProcessId statement"),
586 }
587 }
588
589 #[test]
590 pub fn test_parse_kill_statement_errors() {
591 let sql = "KILL QUERY";
593 let result =
594 ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
595 assert!(result.is_err());
596
597 let sql = "KILL QUERY 'not-a-number'";
599 let result =
600 ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
601 assert!(result.is_err());
602
603 let sql = "KILL";
605 let result =
606 ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
607 assert!(result.is_err());
608
609 let sql = "KILL QUERY 4294967296"; let result =
612 ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
613 assert!(result.is_err());
614 }
615
616 #[test]
617 pub fn test_parse_kill_statement_edge_cases() {
618 use crate::statements::kill::Kill;
619
620 let sql = "KILL QUERY 0";
622 let statements =
623 ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
624 .unwrap();
625
626 assert_eq!(statements.len(), 1);
627 match &statements[0] {
628 Statement::Kill(Kill::ConnectionId(connection_id)) => {
629 assert_eq!(*connection_id, 0);
630 }
631 _ => panic!("Expected Kill::ConnectionId statement"),
632 }
633
634 let sql = "KILL QUERY 4294967295"; let statements =
637 ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
638 .unwrap();
639
640 assert_eq!(statements.len(), 1);
641 match &statements[0] {
642 Statement::Kill(Kill::ConnectionId(connection_id)) => {
643 assert_eq!(*connection_id, 4294967295);
644 }
645 _ => panic!("Expected Kill::ConnectionId statement"),
646 }
647
648 let sql = "KILL ''";
650 let statements =
651 ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
652 .unwrap();
653
654 assert_eq!(statements.len(), 1);
655 match &statements[0] {
656 Statement::Kill(Kill::ProcessId(process_id)) => {
657 assert_eq!(process_id, "");
658 }
659 _ => panic!("Expected Kill::ProcessId statement"),
660 }
661 }
662
663 #[test]
664 pub fn test_parse_kill_statement_case_insensitive() {
665 use crate::statements::kill::Kill;
666
667 let sql = "kill query 123";
669 let statements =
670 ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
671 .unwrap();
672
673 assert_eq!(statements.len(), 1);
674 match &statements[0] {
675 Statement::Kill(Kill::ConnectionId(connection_id)) => {
676 assert_eq!(*connection_id, 123);
677 }
678 _ => panic!("Expected Kill::ConnectionId statement"),
679 }
680
681 let sql = "Kill Query 456";
683 let statements =
684 ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
685 .unwrap();
686
687 assert_eq!(statements.len(), 1);
688 match &statements[0] {
689 Statement::Kill(Kill::ConnectionId(connection_id)) => {
690 assert_eq!(*connection_id, 456);
691 }
692 _ => panic!("Expected Kill::ConnectionId statement"),
693 }
694 }
695}