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::DECLARE => self.parse_declare_cursor(),
196
197 Keyword::FETCH => self.parse_fetch_cursor(),
198
199 Keyword::CLOSE => self.parse_close_cursor(),
200
201 Keyword::USE => {
202 let _ = self.parser.next_token();
203
204 let database_name = self.parser.parse_identifier().with_context(|_| {
205 error::UnexpectedSnafu {
206 expected: "a database name",
207 actual: self.peek_token_as_string(),
208 }
209 })?;
210 Ok(Statement::Use(
211 Self::canonicalize_identifier(database_name).value,
212 ))
213 }
214
215 Keyword::KILL => {
216 let _ = self.parser.next_token();
217 let kill = if self.parser.parse_keyword(Keyword::QUERY) {
218 let connection_id_exp =
220 self.parser.parse_number_value().with_context(|_| {
221 error::UnexpectedSnafu {
222 expected: "MySQL numeric connection id",
223 actual: self.peek_token_as_string(),
224 }
225 })?;
226 let Value::Number(s, _) = connection_id_exp.value else {
227 return error::UnexpectedTokenSnafu {
228 expected: "MySQL numeric connection id",
229 actual: connection_id_exp.to_string(),
230 }
231 .fail();
232 };
233
234 let connection_id = u32::from_str(&s).map_err(|_| {
235 error::UnexpectedTokenSnafu {
236 expected: "MySQL numeric connection id",
237 actual: s,
238 }
239 .build()
240 })?;
241 Kill::ConnectionId(connection_id)
242 } else {
243 let process_id_ident =
244 self.parser.parse_literal_string().with_context(|_| {
245 error::UnexpectedSnafu {
246 expected: "process id string literal",
247 actual: self.peek_token_as_string(),
248 }
249 })?;
250 Kill::ProcessId(process_id_ident)
251 };
252
253 Ok(Statement::Kill(kill))
254 }
255
256 _ => self.unsupported(self.peek_token_as_string()),
257 },
258 Token::LParen => self.parse_query(),
259 unexpected => self.unsupported(unexpected.to_string()),
260 }
261 }
262
263 pub fn parse_mysql_prepare_stmt(sql: &str, dialect: &dyn Dialect) -> Result<(String, String)> {
265 ParserContext::new(dialect, sql)?.parse_mysql_prepare()
266 }
267
268 pub fn parse_mysql_execute_stmt(
270 sql: &str,
271 dialect: &dyn Dialect,
272 ) -> Result<(String, Vec<Expr>)> {
273 ParserContext::new(dialect, sql)?.parse_mysql_execute()
274 }
275
276 pub fn parse_mysql_deallocate_stmt(sql: &str, dialect: &dyn Dialect) -> Result<String> {
278 ParserContext::new(dialect, sql)?.parse_deallocate()
279 }
280
281 pub fn unsupported<T>(&self, keyword: String) -> Result<T> {
283 error::UnsupportedSnafu { keyword }.fail()
284 }
285
286 pub(crate) fn expected<T>(&self, expected: &str, found: TokenWithSpan) -> Result<T> {
288 Err(ParserError::ParserError(format!(
289 "Expected {expected}, found: {found}",
290 )))
291 .context(SyntaxSnafu)
292 }
293
294 pub fn matches_keyword(&mut self, expected: Keyword) -> bool {
295 match self.parser.peek_token().token {
296 Token::Word(w) => w.keyword == expected,
297 _ => false,
298 }
299 }
300
301 pub fn consume_token(&mut self, expected: &str) -> bool {
302 if self.peek_token_as_string().to_uppercase() == *expected.to_uppercase() {
303 let _ = self.parser.next_token();
304 true
305 } else {
306 false
307 }
308 }
309
310 #[inline]
311 pub(crate) fn peek_token_as_string(&self) -> String {
312 self.parser.peek_token().to_string()
313 }
314
315 pub fn canonicalize_identifier(ident: Ident) -> Ident {
317 if ident.quote_style.is_some() {
318 ident
319 } else {
320 Ident::new(ident.value.to_lowercase())
321 }
322 }
323
324 pub(crate) fn canonicalize_object_name(object_name: ObjectName) -> Result<ObjectName> {
326 object_name
327 .0
328 .into_iter()
329 .map(|x| {
330 x.as_ident()
331 .cloned()
332 .map(Self::canonicalize_identifier)
333 .with_context(|| InvalidSqlSnafu {
334 msg: format!("not an ident: '{x}'"),
335 })
336 })
337 .collect::<Result<Vec<_>>>()
338 .map(Into::into)
339 }
340
341 pub(crate) fn parse_object_name(&mut self) -> std::result::Result<ObjectName, ParserError> {
346 self.parser.parse_object_name(false)
347 }
348}
349
350#[cfg(test)]
351mod tests {
352
353 use datatypes::prelude::ConcreteDataType;
354 use sqlparser::dialect::MySqlDialect;
355
356 use super::*;
357 use crate::dialect::GreptimeDbDialect;
358 use crate::statements::create::CreateTable;
359 use crate::statements::sql_data_type_to_concrete_data_type;
360
361 fn test_timestamp_precision(sql: &str, expected_type: ConcreteDataType) {
362 match ParserContext::create_with_dialect(
363 sql,
364 &GreptimeDbDialect {},
365 ParseOptions::default(),
366 )
367 .unwrap()
368 .pop()
369 .unwrap()
370 {
371 Statement::CreateTable(CreateTable { columns, .. }) => {
372 let ts_col = columns.first().unwrap();
373 assert_eq!(
374 expected_type,
375 sql_data_type_to_concrete_data_type(ts_col.data_type(), &Default::default())
376 .unwrap()
377 );
378 }
379 _ => unreachable!(),
380 }
381 }
382
383 #[test]
384 pub fn test_create_table_with_precision() {
385 test_timestamp_precision(
386 "create table demo (ts timestamp time index, cnt int);",
387 ConcreteDataType::timestamp_millisecond_datatype(),
388 );
389 test_timestamp_precision(
390 "create table demo (ts timestamp(0) time index, cnt int);",
391 ConcreteDataType::timestamp_second_datatype(),
392 );
393 test_timestamp_precision(
394 "create table demo (ts timestamp(3) time index, cnt int);",
395 ConcreteDataType::timestamp_millisecond_datatype(),
396 );
397 test_timestamp_precision(
398 "create table demo (ts timestamp(6) time index, cnt int);",
399 ConcreteDataType::timestamp_microsecond_datatype(),
400 );
401 test_timestamp_precision(
402 "create table demo (ts timestamp(9) time index, cnt int);",
403 ConcreteDataType::timestamp_nanosecond_datatype(),
404 );
405 }
406
407 #[test]
408 #[should_panic]
409 pub fn test_create_table_with_invalid_precision() {
410 test_timestamp_precision(
411 "create table demo (ts timestamp(1) time index, cnt int);",
412 ConcreteDataType::timestamp_millisecond_datatype(),
413 );
414 }
415
416 #[test]
417 pub fn test_parse_table_name() {
418 let table_name = "a.b.c";
419
420 let object_name =
421 ParserContext::parse_table_name(table_name, &GreptimeDbDialect {}).unwrap();
422
423 assert_eq!(object_name.0.len(), 3);
424 assert_eq!(object_name.to_string(), table_name);
425
426 let table_name = "a.b";
427
428 let object_name =
429 ParserContext::parse_table_name(table_name, &GreptimeDbDialect {}).unwrap();
430
431 assert_eq!(object_name.0.len(), 2);
432 assert_eq!(object_name.to_string(), table_name);
433
434 let table_name = "Test.\"public-test\"";
435
436 let object_name =
437 ParserContext::parse_table_name(table_name, &GreptimeDbDialect {}).unwrap();
438
439 assert_eq!(object_name.0.len(), 2);
440 assert_eq!(object_name.to_string(), table_name.to_ascii_lowercase());
441
442 let table_name = "HelloWorld";
443
444 let object_name =
445 ParserContext::parse_table_name(table_name, &GreptimeDbDialect {}).unwrap();
446
447 assert_eq!(object_name.0.len(), 1);
448 assert_eq!(object_name.to_string(), table_name.to_ascii_lowercase());
449 }
450
451 #[test]
452 pub fn test_parse_mysql_prepare_stmt() {
453 let sql = "PREPARE stmt1 FROM 'SELECT * FROM t1 WHERE id = ?';";
454 let (stmt_name, stmt) =
455 ParserContext::parse_mysql_prepare_stmt(sql, &MySqlDialect {}).unwrap();
456 assert_eq!(stmt_name, "stmt1");
457 assert_eq!(stmt, "SELECT * FROM t1 WHERE id = ?");
458
459 let sql = "PREPARE stmt2 FROM \"SELECT * FROM t1 WHERE id = ?\"";
460 let (stmt_name, stmt) =
461 ParserContext::parse_mysql_prepare_stmt(sql, &MySqlDialect {}).unwrap();
462 assert_eq!(stmt_name, "stmt2");
463 assert_eq!(stmt, "SELECT * FROM t1 WHERE id = ?");
464 }
465
466 #[test]
467 pub fn test_parse_mysql_execute_stmt() {
468 let sql = "EXECUTE stmt1 USING 1, 'hello';";
469 let (stmt_name, params) =
470 ParserContext::parse_mysql_execute_stmt(sql, &GreptimeDbDialect {}).unwrap();
471 assert_eq!(stmt_name, "stmt1");
472 assert_eq!(params.len(), 2);
473 assert_eq!(params[0].to_string(), "1");
474 assert_eq!(params[1].to_string(), "'hello'");
475
476 let sql = "EXECUTE stmt2;";
477 let (stmt_name, params) =
478 ParserContext::parse_mysql_execute_stmt(sql, &GreptimeDbDialect {}).unwrap();
479 assert_eq!(stmt_name, "stmt2");
480 assert_eq!(params.len(), 0);
481
482 let sql = "EXECUTE stmt3 USING 231, 'hello', \"2003-03-1\", NULL, ;";
483 let (stmt_name, params) =
484 ParserContext::parse_mysql_execute_stmt(sql, &GreptimeDbDialect {}).unwrap();
485 assert_eq!(stmt_name, "stmt3");
486 assert_eq!(params.len(), 4);
487 assert_eq!(params[0].to_string(), "231");
488 assert_eq!(params[1].to_string(), "'hello'");
489 assert_eq!(params[2].to_string(), "\"2003-03-1\"");
490 assert_eq!(params[3].to_string(), "NULL");
491 }
492
493 #[test]
494 pub fn test_parse_mysql_deallocate_stmt() {
495 let sql = "DEALLOCATE stmt1;";
496 let stmt_name = ParserContext::parse_mysql_deallocate_stmt(sql, &MySqlDialect {}).unwrap();
497 assert_eq!(stmt_name, "stmt1");
498
499 let sql = "DEALLOCATE stmt2";
500 let stmt_name = ParserContext::parse_mysql_deallocate_stmt(sql, &MySqlDialect {}).unwrap();
501 assert_eq!(stmt_name, "stmt2");
502 }
503
504 #[test]
505 pub fn test_parse_kill_query_statement() {
506 use crate::statements::kill::Kill;
507
508 let sql = "KILL QUERY 123";
510 let statements =
511 ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
512 .unwrap();
513
514 assert_eq!(statements.len(), 1);
515 match &statements[0] {
516 Statement::Kill(Kill::ConnectionId(connection_id)) => {
517 assert_eq!(*connection_id, 123);
518 }
519 _ => panic!("Expected Kill::ConnectionId statement"),
520 }
521
522 let sql = "KILL QUERY 999999";
524 let statements =
525 ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
526 .unwrap();
527
528 assert_eq!(statements.len(), 1);
529 match &statements[0] {
530 Statement::Kill(Kill::ConnectionId(connection_id)) => {
531 assert_eq!(*connection_id, 999999);
532 }
533 _ => panic!("Expected Kill::ConnectionId statement"),
534 }
535 }
536
537 #[test]
538 pub fn test_parse_kill_process_statement() {
539 use crate::statements::kill::Kill;
540
541 let sql = "KILL 'process-123'";
543 let statements =
544 ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
545 .unwrap();
546
547 assert_eq!(statements.len(), 1);
548 match &statements[0] {
549 Statement::Kill(Kill::ProcessId(process_id)) => {
550 assert_eq!(process_id, "process-123");
551 }
552 _ => panic!("Expected Kill::ProcessId statement"),
553 }
554
555 let sql = "KILL \"process-456\"";
557 let statements =
558 ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
559 .unwrap();
560
561 assert_eq!(statements.len(), 1);
562 match &statements[0] {
563 Statement::Kill(Kill::ProcessId(process_id)) => {
564 assert_eq!(process_id, "process-456");
565 }
566 _ => panic!("Expected Kill::ProcessId statement"),
567 }
568
569 let sql = "KILL 'f47ac10b-58cc-4372-a567-0e02b2c3d479'";
571 let statements =
572 ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
573 .unwrap();
574
575 assert_eq!(statements.len(), 1);
576 match &statements[0] {
577 Statement::Kill(Kill::ProcessId(process_id)) => {
578 assert_eq!(process_id, "f47ac10b-58cc-4372-a567-0e02b2c3d479");
579 }
580 _ => panic!("Expected Kill::ProcessId statement"),
581 }
582 }
583
584 #[test]
585 pub fn test_parse_kill_statement_errors() {
586 let sql = "KILL QUERY";
588 let result =
589 ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
590 assert!(result.is_err());
591
592 let sql = "KILL QUERY 'not-a-number'";
594 let result =
595 ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
596 assert!(result.is_err());
597
598 let sql = "KILL";
600 let result =
601 ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
602 assert!(result.is_err());
603
604 let sql = "KILL QUERY 4294967296"; let result =
607 ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
608 assert!(result.is_err());
609 }
610
611 #[test]
612 pub fn test_parse_kill_statement_edge_cases() {
613 use crate::statements::kill::Kill;
614
615 let sql = "KILL QUERY 0";
617 let statements =
618 ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
619 .unwrap();
620
621 assert_eq!(statements.len(), 1);
622 match &statements[0] {
623 Statement::Kill(Kill::ConnectionId(connection_id)) => {
624 assert_eq!(*connection_id, 0);
625 }
626 _ => panic!("Expected Kill::ConnectionId statement"),
627 }
628
629 let sql = "KILL QUERY 4294967295"; let statements =
632 ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
633 .unwrap();
634
635 assert_eq!(statements.len(), 1);
636 match &statements[0] {
637 Statement::Kill(Kill::ConnectionId(connection_id)) => {
638 assert_eq!(*connection_id, 4294967295);
639 }
640 _ => panic!("Expected Kill::ConnectionId statement"),
641 }
642
643 let sql = "KILL ''";
645 let statements =
646 ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
647 .unwrap();
648
649 assert_eq!(statements.len(), 1);
650 match &statements[0] {
651 Statement::Kill(Kill::ProcessId(process_id)) => {
652 assert_eq!(process_id, "");
653 }
654 _ => panic!("Expected Kill::ProcessId statement"),
655 }
656 }
657
658 #[test]
659 pub fn test_parse_kill_statement_case_insensitive() {
660 use crate::statements::kill::Kill;
661
662 let sql = "kill query 123";
664 let statements =
665 ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
666 .unwrap();
667
668 assert_eq!(statements.len(), 1);
669 match &statements[0] {
670 Statement::Kill(Kill::ConnectionId(connection_id)) => {
671 assert_eq!(*connection_id, 123);
672 }
673 _ => panic!("Expected Kill::ConnectionId statement"),
674 }
675
676 let sql = "Kill Query 456";
678 let statements =
679 ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
680 .unwrap();
681
682 assert_eq!(statements.len(), 1);
683 match &statements[0] {
684 Statement::Kill(Kill::ConnectionId(connection_id)) => {
685 assert_eq!(*connection_id, 456);
686 }
687 _ => panic!("Expected Kill::ConnectionId statement"),
688 }
689 }
690}