sql/statements/
truncate.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::fmt::Display;
16
17use itertools::Itertools;
18use serde::Serialize;
19use sqlparser::ast::{ObjectName, Visit, VisitMut, Visitor, VisitorMut};
20
21/// TRUNCATE TABLE statement.
22#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
23pub struct TruncateTable {
24    table_name: ObjectName,
25    time_ranges: Vec<(sqlparser::ast::Value, sqlparser::ast::Value)>,
26}
27
28impl Visit for TruncateTable {
29    fn visit<V: Visitor>(&self, visitor: &mut V) -> ::std::ops::ControlFlow<V::Break> {
30        self.table_name.visit(visitor)?;
31        for (start, end) in &self.time_ranges {
32            start.visit(visitor)?;
33            end.visit(visitor)?;
34        }
35        ::std::ops::ControlFlow::Continue(())
36    }
37}
38
39impl VisitMut for TruncateTable {
40    fn visit<V: VisitorMut>(&mut self, visitor: &mut V) -> ::std::ops::ControlFlow<V::Break> {
41        sqlparser::ast::VisitMut::visit(&mut self.table_name, visitor)?;
42        for (start, end) in &mut self.time_ranges {
43            start.visit(visitor)?;
44            end.visit(visitor)?;
45        }
46        ::std::ops::ControlFlow::Continue(())
47    }
48}
49
50impl TruncateTable {
51    /// Creates a statement for `TRUNCATE TABLE`
52    pub fn new(table_name: ObjectName) -> Self {
53        Self {
54            table_name,
55            time_ranges: Vec::new(),
56        }
57    }
58
59    /// Creates a statement for `TRUNCATE TABLE RANGE`
60    pub fn new_with_ranges(
61        table_name: ObjectName,
62        time_ranges: Vec<(sqlparser::ast::Value, sqlparser::ast::Value)>,
63    ) -> Self {
64        Self {
65            table_name,
66            time_ranges,
67        }
68    }
69
70    pub fn table_name(&self) -> &ObjectName {
71        &self.table_name
72    }
73
74    pub fn time_ranges(&self) -> &[(sqlparser::ast::Value, sqlparser::ast::Value)] {
75        &self.time_ranges
76    }
77}
78
79impl Display for TruncateTable {
80    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
81        let table_name = self.table_name();
82        write!(f, r#"TRUNCATE TABLE {table_name}"#)?;
83        if self.time_ranges.is_empty() {
84            return Ok(());
85        }
86
87        write!(f, " FILE RANGE ")?;
88        write!(
89            f,
90            "{}",
91            self.time_ranges
92                .iter()
93                .map(|(start, end)| format!("({}, {})", start, end))
94                .join(", ")
95        )
96    }
97}
98
99#[cfg(test)]
100mod tests {
101    use std::assert_matches::assert_matches;
102
103    use crate::dialect::GreptimeDbDialect;
104    use crate::parser::{ParseOptions, ParserContext};
105    use crate::statements::statement::Statement;
106
107    #[test]
108    fn test_display_for_tuncate_table() {
109        let sql = r"truncate table t1;";
110        let stmts: Vec<Statement> =
111            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
112                .unwrap();
113        assert_eq!(1, stmts.len());
114        assert_matches!(&stmts[0], Statement::TruncateTable { .. });
115        match &stmts[0] {
116            Statement::TruncateTable(trunc) => {
117                let new_sql = format!("\n{}", trunc);
118                assert_eq!(
119                    r#"
120TRUNCATE TABLE t1"#,
121                    &new_sql
122                );
123            }
124            _ => {
125                unreachable!();
126            }
127        }
128
129        let sql = r"truncate table t1 file range (1,2);";
130        let stmts: Vec<Statement> =
131            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
132                .unwrap();
133        assert_eq!(1, stmts.len());
134        assert_matches!(&stmts[0], Statement::TruncateTable { .. });
135        match &stmts[0] {
136            Statement::TruncateTable(trunc) => {
137                let new_sql = format!("\n{}", trunc);
138                assert_eq!(
139                    r#"
140TRUNCATE TABLE t1 FILE RANGE (1, 2)"#,
141                    &new_sql
142                );
143            }
144            _ => {
145                unreachable!();
146            }
147        }
148
149        let sql = r"truncate table t1 file range (1,2), (3,4);";
150        let stmts: Vec<Statement> =
151            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
152                .unwrap();
153        assert_eq!(1, stmts.len());
154        assert_matches!(&stmts[0], Statement::TruncateTable { .. });
155        match &stmts[0] {
156            Statement::TruncateTable(trunc) => {
157                let new_sql = format!("\n{}", trunc);
158                assert_eq!(
159                    r#"
160TRUNCATE TABLE t1 FILE RANGE (1, 2), (3, 4)"#,
161                    &new_sql
162                );
163            }
164            _ => {
165                unreachable!();
166            }
167        }
168    }
169}
sql/statements/truncate.rs

sql/statements/
truncate.rs
