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
21use crate::statements::OptionMap;
22
23/// TRUNCATE TABLE statement.
24#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
25pub struct TruncateTable {
26    table_name: ObjectName,
27    time_ranges: Vec<(sqlparser::ast::Value, sqlparser::ast::Value)>,
28    /// Options for this DDL submission, never persisted as table options.
29    options: OptionMap,
30}
31
32impl Visit for TruncateTable {
33    fn visit<V: Visitor>(&self, visitor: &mut V) -> ::std::ops::ControlFlow<V::Break> {
34        self.table_name.visit(visitor)?;
35        for (start, end) in &self.time_ranges {
36            start.visit(visitor)?;
37            end.visit(visitor)?;
38        }
39        Visit::visit(&self.options, visitor)?;
40        ::std::ops::ControlFlow::Continue(())
41    }
42}
43
44impl VisitMut for TruncateTable {
45    fn visit<V: VisitorMut>(&mut self, visitor: &mut V) -> ::std::ops::ControlFlow<V::Break> {
46        sqlparser::ast::VisitMut::visit(&mut self.table_name, visitor)?;
47        for (start, end) in &mut self.time_ranges {
48            start.visit(visitor)?;
49            end.visit(visitor)?;
50        }
51        VisitMut::visit(&mut self.options, visitor)?;
52        ::std::ops::ControlFlow::Continue(())
53    }
54}
55
56impl TruncateTable {
57    /// Creates a statement for `TRUNCATE TABLE`
58    pub fn new(table_name: ObjectName) -> Self {
59        Self {
60            table_name,
61            time_ranges: Vec::new(),
62            options: OptionMap::default(),
63        }
64    }
65
66    /// Creates a statement for `TRUNCATE TABLE RANGE`
67    pub fn new_with_ranges(
68        table_name: ObjectName,
69        time_ranges: Vec<(sqlparser::ast::Value, sqlparser::ast::Value)>,
70    ) -> Self {
71        Self {
72            table_name,
73            time_ranges,
74            options: OptionMap::default(),
75        }
76    }
77
78    /// Sets the options for this DDL submission.
79    pub fn with_options(mut self, options: OptionMap) -> Self {
80        self.options = options;
81        self
82    }
83
84    /// Returns the options for this DDL submission.
85    pub fn options(&self) -> &OptionMap {
86        &self.options
87    }
88
89    pub fn table_name(&self) -> &ObjectName {
90        &self.table_name
91    }
92
93    pub fn time_ranges(&self) -> &[(sqlparser::ast::Value, sqlparser::ast::Value)] {
94        &self.time_ranges
95    }
96}
97
98impl Display for TruncateTable {
99    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
100        let table_name = self.table_name();
101        write!(f, r#"TRUNCATE TABLE {table_name}"#)?;
102        if !self.time_ranges.is_empty() {
103            write!(f, " FILE RANGE ")?;
104            write!(
105                f,
106                "{}",
107                self.time_ranges
108                    .iter()
109                    .map(|(start, end)| format!("({}, {})", start, end))
110                    .join(", ")
111            )?;
112        }
113        if !self.options.is_empty() {
114            f.write_str(" WITH (")?;
115            self.options.write_sql(f, ", ", "", false)?;
116            f.write_str(")")?;
117        }
118        Ok(())
119    }
120}
121
122#[cfg(test)]
123mod tests {
124    use std::assert_matches;
125
126    use crate::dialect::GreptimeDbDialect;
127    use crate::parser::{ParseOptions, ParserContext};
128    use crate::statements::statement::Statement;
129
130    #[test]
131    fn test_display_for_truncate_table() {
132        let sql = r"truncate table t1;";
133        let stmts: Vec<Statement> =
134            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
135                .unwrap();
136        assert_eq!(1, stmts.len());
137        assert_matches!(&stmts[0], Statement::TruncateTable { .. });
138        match &stmts[0] {
139            Statement::TruncateTable(trunc) => {
140                let new_sql = format!("\n{}", trunc);
141                assert_eq!(
142                    r#"
143TRUNCATE TABLE t1"#,
144                    &new_sql
145                );
146            }
147            _ => {
148                unreachable!();
149            }
150        }
151
152        let sql = r"truncate table t1 file range (1,2);";
153        let stmts: Vec<Statement> =
154            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
155                .unwrap();
156        assert_eq!(1, stmts.len());
157        assert_matches!(&stmts[0], Statement::TruncateTable { .. });
158        match &stmts[0] {
159            Statement::TruncateTable(trunc) => {
160                let new_sql = format!("\n{}", trunc);
161                assert_eq!(
162                    r#"
163TRUNCATE TABLE t1 FILE RANGE (1, 2)"#,
164                    &new_sql
165                );
166            }
167            _ => {
168                unreachable!();
169            }
170        }
171
172        let sql = r"truncate table t1 file range (1,2), (3,4);";
173        let stmts: Vec<Statement> =
174            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
175                .unwrap();
176        assert_eq!(1, stmts.len());
177        assert_matches!(&stmts[0], Statement::TruncateTable { .. });
178        match &stmts[0] {
179            Statement::TruncateTable(trunc) => {
180                let new_sql = format!("\n{}", trunc);
181                assert_eq!(
182                    r#"
183TRUNCATE TABLE t1 FILE RANGE (1, 2), (3, 4)"#,
184                    &new_sql
185                );
186            }
187            _ => {
188                unreachable!();
189            }
190        }
191    }
192}
sql/statements/truncate.rs

sql/statements/
truncate.rs
