1use std::fmt;
16
17use serde::Serialize;
18use sqlparser::ast::Query as SpQuery;
19use sqlparser_derive::{Visit, VisitMut};
20
21use crate::error::Error;
22use crate::parsers::with_tql_parser::HybridCteWith;
23
24#[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut, Serialize)]
26pub struct Query {
27 pub inner: SpQuery,
28 pub hybrid_cte: Option<HybridCteWith>,
30}
31
32impl TryFrom<SpQuery> for Query {
33 type Error = Error;
34
35 fn try_from(inner: SpQuery) -> Result<Self, Self::Error> {
36 Ok(Self {
37 inner,
38 hybrid_cte: None,
39 })
40 }
41}
42
43impl TryFrom<Query> for SpQuery {
44 type Error = Error;
45
46 fn try_from(value: Query) -> Result<Self, Self::Error> {
47 Ok(value.inner)
48 }
49}
50
51impl fmt::Display for Query {
52 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
53 if let Some(hybrid_cte) = &self.hybrid_cte {
54 write!(f, "{} ", hybrid_cte)?;
56
57 let mut main_query = self.inner.clone();
59 main_query.with = None;
60 write!(f, "{}", main_query)
61 } else {
62 write!(f, "{}", self.inner)
63 }
64 }
65}
66
67#[cfg(test)]
68mod test {
69
70 use super::Query;
71 use crate::dialect::GreptimeDbDialect;
72 use crate::parser::{ParseOptions, ParserContext};
73 use crate::statements::statement::Statement;
74
75 fn create_query(sql: &str) -> Option<Box<Query>> {
76 match ParserContext::create_with_dialect(
77 sql,
78 &GreptimeDbDialect {},
79 ParseOptions::default(),
80 )
81 .unwrap()
82 .remove(0)
83 {
84 Statement::Query(query) => Some(query),
85 _ => None,
86 }
87 }
88
89 #[test]
90 fn test_query_display() {
91 assert_eq!(
92 create_query("select * from abc where x = 1 and y = 7")
93 .unwrap()
94 .to_string(),
95 "SELECT * FROM abc WHERE x = 1 AND y = 7"
96 );
97 assert_eq!(
98 create_query(
99 "select * from abc left join bcd where abc.a = 1 and bcd.d = 7 and abc.id = bcd.id"
100 )
101 .unwrap()
102 .to_string(),
103 "SELECT * FROM abc LEFT JOIN bcd WHERE abc.a = 1 AND bcd.d = 7 AND abc.id = bcd.id"
104 );
105 assert_eq!(
106 create_query("WITH tql_cte AS (TQL EVAL (0, 100, '5s') up) SELECT * FROM tql_cte")
107 .unwrap()
108 .to_string(),
109 "WITH tql_cte AS (TQL EVAL (0, 100, '5s') up) SELECT * FROM tql_cte"
110 );
111 }
112}