feat: add partial truncate (#6602)

* feat: add partial truncate Signed-off-by: discord9 <discord9@163.com> * fix: per review Signed-off-by: discord9 <discord9@163.com> * feat: add proto partial truncate kind Signed-off-by: discord9 <discord9@163.com> * chore: clippy Signed-off-by: discord9 <discord9@163.com> * chore: update branched proto Signed-off-by: discord9 <discord9@163.com> * feat: grpc support truncate WIP sql support Signed-off-by: discord9 <discord9@163.com> * wip: parse truncate range Signed-off-by: discord9 <discord9@163.com> * feat: truncate by range Signed-off-by: discord9 <discord9@163.com> * fix: truncate range display Signed-off-by: discord9 <discord9@163.com> * chore: resolve todo Signed-off-by: discord9 <discord9@163.com> * refactor: per review Signed-off-by: discord9 <discord9@163.com> * test: more invalid parse Signed-off-by: discord9 <discord9@163.com> * chore: per review Signed-off-by: discord9 <discord9@163.com> * refactor: per review Signed-off-by: discord9 <discord9@163.com> * chore: unused Signed-off-by: discord9 <discord9@163.com> * chore: per review Signed-off-by: discord9 <discord9@163.com> * chore: update branch Signed-off-by: discord9 <discord9@163.com> --------- Signed-off-by: discord9 <discord9@163.com>
2026-05-24 17:00:37 +00:00 · 2025-08-04 18:50:27 +08:00
parent 414101fafa
commit 1afa0afc67
25 changed files with 938 additions and 102 deletions
--- a/src/sql/src/parsers/truncate_parser.rs
+++ b/src/sql/src/parsers/truncate_parser.rs
@@ -14,8 +14,9 @@

 use snafu::{ensure, ResultExt};
 use sqlparser::keywords::Keyword;
+use sqlparser::tokenizer::Token;

-use crate::error::{self, InvalidTableNameSnafu, Result};
+use crate::error::{self, InvalidSqlSnafu, InvalidTableNameSnafu, Result, UnexpectedTokenSnafu};
 use crate::parser::ParserContext;
 use crate::statements::statement::Statement;
 use crate::statements::truncate::TruncateTable;
@@ -41,7 +42,92 @@ impl ParserContext<'_> {
            }
        );

-        Ok(Statement::TruncateTable(TruncateTable::new(table_ident)))
+        let have_range = self.parser.parse_keywords(&[Keyword::FILE, Keyword::RANGE]);
+
+        // if no range is specified, we just truncate the table
+        if !have_range {
+            return Ok(Statement::TruncateTable(TruncateTable::new(table_ident)));
+        }
+
+        // parse a list of time ranges consist of (Timestamp, Timestamp),?
+        let mut time_ranges = vec![];
+        loop {
+            let _ = self
+                .parser
+                .expect_token(&sqlparser::tokenizer::Token::LParen)
+                .with_context(|_| error::UnexpectedSnafu {
+                    expected: "a left parenthesis",
+                    actual: self.peek_token_as_string(),
+                })?;
+
+            // parse to values here, no need to valid in parser
+            let start = self
+                .parser
+                .parse_value()
+                .with_context(|e| error::UnexpectedSnafu {
+                    expected: "a timestamp value",
+                    actual: e.to_string(),
+                })?;
+
+            let _ = self
+                .parser
+                .expect_token(&sqlparser::tokenizer::Token::Comma)
+                .with_context(|_| error::UnexpectedSnafu {
+                    expected: "a comma",
+                    actual: self.peek_token_as_string(),
+                })?;
+
+            let end = self
+                .parser
+                .parse_value()
+                .with_context(|_| error::UnexpectedSnafu {
+                    expected: "a timestamp",
+                    actual: self.peek_token_as_string(),
+                })?;
+
+            let _ = self
+                .parser
+                .expect_token(&sqlparser::tokenizer::Token::RParen)
+                .with_context(|_| error::UnexpectedSnafu {
+                    expected: "a right parenthesis",
+                    actual: self.peek_token_as_string(),
+                })?;
+
+            time_ranges.push((start, end));
+
+            let peek = self.parser.peek_token().token;
+
+            match peek {
+                sqlparser::tokenizer::Token::EOF | Token::SemiColon => {
+                    if time_ranges.is_empty() {
+                        return Err(InvalidSqlSnafu {
+                            msg: "TRUNCATE TABLE RANGE must have at least one range".to_string(),
+                        }
+                        .build());
+                    }
+                    break;
+                }
+                Token::Comma => {
+                    self.parser.next_token(); // Consume the comma
+                    let next_peek = self.parser.peek_token().token; // Peek the token after the comma
+                    if matches!(next_peek, Token::EOF | Token::SemiColon) {
+                        break; // Trailing comma, end of statement
+                    }
+                    // Otherwise, continue to parse next range
+                    continue;
+                }
+                _ => UnexpectedTokenSnafu {
+                    expected: "a comma or end of statement",
+                    actual: self.peek_token_as_string(),
+                }
+                .fail()?,
+            }
+        }
+
+        Ok(Statement::TruncateTable(TruncateTable::new_with_ranges(
+            table_ident,
+            time_ranges,
+        )))
    }
 }

@@ -53,6 +139,145 @@ mod tests {
    use crate::dialect::GreptimeDbDialect;
    use crate::parser::ParseOptions;

+    #[test]
+    pub fn test_parse_truncate_with_ranges() {
+        let sql = r#"TRUNCATE foo FILE RANGE (0, 20)"#;
+        let mut stmts =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+        assert_eq!(
+            stmts.pop().unwrap(),
+            Statement::TruncateTable(TruncateTable::new_with_ranges(
+                ObjectName(vec![Ident::new("foo")]),
+                vec![(
+                    sqlparser::ast::Value::Number("0".to_string(), false),
+                    sqlparser::ast::Value::Number("20".to_string(), false)
+                )]
+            ))
+        );
+
+        let sql = r#"TRUNCATE TABLE foo FILE RANGE ("2000-01-01 00:00:00+00:00", "2000-01-01 00:00:00+00:00"), (2,33)"#;
+        let mut stmts =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+        assert_eq!(
+            stmts.pop().unwrap(),
+            Statement::TruncateTable(TruncateTable::new_with_ranges(
+                ObjectName(vec![Ident::new("foo")]),
+                vec![
+                    (
+                        sqlparser::ast::Value::DoubleQuotedString(
+                            "2000-01-01 00:00:00+00:00".to_string()
+                        ),
+                        sqlparser::ast::Value::DoubleQuotedString(
+                            "2000-01-01 00:00:00+00:00".to_string()
+                        )
+                    ),
+                    (
+                        sqlparser::ast::Value::Number("2".to_string(), false),
+                        sqlparser::ast::Value::Number("33".to_string(), false)
+                    )
+                ]
+            ))
+        );
+
+        let sql = "TRUNCATE TABLE my_schema.foo FILE RANGE (1, 2), (3, 4),";
+        let mut stmts =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+        assert_eq!(
+            stmts.pop().unwrap(),
+            Statement::TruncateTable(TruncateTable::new_with_ranges(
+                ObjectName(vec![Ident::new("my_schema"), Ident::new("foo")]),
+                vec![
+                    (
+                        sqlparser::ast::Value::Number("1".to_string(), false),
+                        sqlparser::ast::Value::Number("2".to_string(), false)
+                    ),
+                    (
+                        sqlparser::ast::Value::Number("3".to_string(), false),
+                        sqlparser::ast::Value::Number("4".to_string(), false)
+                    )
+                ]
+            ))
+        );
+
+        let sql = "TRUNCATE my_schema.foo FILE RANGE (1,2),";
+        let mut stmts =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+        assert_eq!(
+            stmts.pop().unwrap(),
+            Statement::TruncateTable(TruncateTable::new_with_ranges(
+                ObjectName(vec![Ident::new("my_schema"), Ident::new("foo")]),
+                vec![(
+                    sqlparser::ast::Value::Number("1".to_string(), false),
+                    sqlparser::ast::Value::Number("2".to_string(), false)
+                )]
+            ))
+        );
+
+        let sql = "TRUNCATE TABLE my_catalog.my_schema.foo FILE RANGE (1,2);";
+        let mut stmts =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+        assert_eq!(
+            stmts.pop().unwrap(),
+            Statement::TruncateTable(TruncateTable::new_with_ranges(
+                ObjectName(vec![
+                    Ident::new("my_catalog"),
+                    Ident::new("my_schema"),
+                    Ident::new("foo")
+                ]),
+                vec![(
+                    sqlparser::ast::Value::Number("1".to_string(), false),
+                    sqlparser::ast::Value::Number("2".to_string(), false)
+                )]
+            ))
+        );
+
+        let sql = "TRUNCATE drop FILE RANGE (1,2)";
+        let mut stmts =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+        assert_eq!(
+            stmts.pop().unwrap(),
+            Statement::TruncateTable(TruncateTable::new_with_ranges(
+                ObjectName(vec![Ident::new("drop")]),
+                vec![(
+                    sqlparser::ast::Value::Number("1".to_string(), false),
+                    sqlparser::ast::Value::Number("2".to_string(), false)
+                )]
+            ))
+        );
+
+        let sql = "TRUNCATE `drop`";
+        let mut stmts =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+        assert_eq!(
+            stmts.pop().unwrap(),
+            Statement::TruncateTable(TruncateTable::new(ObjectName(vec![Ident::with_quote(
+                '`', "drop"
+            ),])))
+        );
+
+        let sql = "TRUNCATE \"drop\" FILE RANGE (\"1\", \"2\")";
+        let mut stmts =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+        assert_eq!(
+            stmts.pop().unwrap(),
+            Statement::TruncateTable(TruncateTable::new_with_ranges(
+                ObjectName(vec![Ident::with_quote('"', "drop")]),
+                vec![(
+                    sqlparser::ast::Value::DoubleQuotedString("1".to_string()),
+                    sqlparser::ast::Value::DoubleQuotedString("2".to_string())
+                )]
+            ))
+        );
+    }
+
    #[test]
    pub fn test_parse_truncate() {
        let sql = "TRUNCATE foo";
@@ -153,5 +378,82 @@ mod tests {
        let result =
            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
        assert!(result.is_err(), "result is: {result:?}");
+
+        let sql = "TRUNCATE TABLE foo RANGE";
+        let result =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
+        assert!(
+            result.is_err()
+                && format!("{result:?}").contains("SQL statement is not supported, keyword: RANGE"),
+            "result is: {result:?}"
+        );
+
+        let sql = "TRUNCATE TABLE foo FILE";
+        let result =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
+        assert!(
+            result.is_err()
+                && format!("{result:?}").contains("SQL statement is not supported, keyword: FILE"),
+            "result is: {result:?}"
+        );
+
+        let sql = "TRUNCATE TABLE foo FILE RANGE";
+        let result =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
+        assert!(
+            result.is_err() && format!("{result:?}").contains("expected: 'a left parenthesis'"),
+            "result is: {result:?}"
+        );
+
+        let sql = "TRUNCATE TABLE foo FILE RANGE (";
+        let result =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
+        assert!(
+            result.is_err() && format!("{result:?}").contains("expected: 'a timestamp value'"),
+            "result is: {result:?}"
+        );
+
+        let sql = "TRUNCATE TABLE foo FILE RANGE ()";
+        let result =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
+        assert!(
+            result.is_err() && format!("{result:?}").contains("expected: 'a timestamp value'"),
+            "result is: {result:?}"
+        );
+
+        let sql = "TRUNCATE TABLE foo FILE RANGE (1 2) (3 4)";
+        let result =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
+        assert!(
+            result.is_err() && format!("{result:?}").contains("expected: 'a comma'"),
+            "result is: {result:?}"
+        );
+
+        let sql = "TRUNCATE TABLE foo FILE RANGE (,),(3,4)";
+        let result =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
+        assert!(
+            result.is_err() && format!("{result:?}").contains("Expected: a value, found: ,"),
+            "result is: {result:?}"
+        );
+
+        let sql = "TRUNCATE TABLE foo FILE RANGE (1,2) (3,4)";
+        let result =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
+        assert!(
+            result.is_err()
+                && format!("{result:?}")
+                    .contains("expected: 'a comma or end of statement', found: ("),
+            "result is: {result:?}"
+        );
+
+        let sql = "TRUNCATE TABLE foo FILE RANGE (1,2),,,,,,,,,(3,4)";
+        let result =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
+        assert!(
+            result.is_err()
+                && format!("{result:?}").contains("expected: 'a left parenthesis', found: ,"),
+            "result is: {result:?}"
+        );
    }
 }
--- a/src/sql/src/statements/truncate.rs
+++ b/src/sql/src/statements/truncate.rs
@@ -14,31 +14,85 @@

 use std::fmt::Display;

+use itertools::Itertools;
 use serde::Serialize;
-use sqlparser::ast::ObjectName;
-use sqlparser_derive::{Visit, VisitMut};
+use sqlparser::ast::{ObjectName, Visit, VisitMut, Visitor, VisitorMut};

 /// TRUNCATE TABLE statement.
-#[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut, Serialize)]
+#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
 pub struct TruncateTable {
    table_name: ObjectName,
+    time_ranges: Vec<(sqlparser::ast::Value, sqlparser::ast::Value)>,
+}
+
+impl Visit for TruncateTable {
+    fn visit<V: Visitor>(&self, visitor: &mut V) -> ::std::ops::ControlFlow<V::Break> {
+        self.table_name.visit(visitor)?;
+        for (start, end) in &self.time_ranges {
+            start.visit(visitor)?;
+            end.visit(visitor)?;
+        }
+        ::std::ops::ControlFlow::Continue(())
+    }
+}
+
+impl VisitMut for TruncateTable {
+    fn visit<V: VisitorMut>(&mut self, visitor: &mut V) -> ::std::ops::ControlFlow<V::Break> {
+        sqlparser::ast::VisitMut::visit(&mut self.table_name, visitor)?;
+        for (start, end) in &mut self.time_ranges {
+            start.visit(visitor)?;
+            end.visit(visitor)?;
+        }
+        ::std::ops::ControlFlow::Continue(())
+    }
 }

 impl TruncateTable {
    /// Creates a statement for `TRUNCATE TABLE`
    pub fn new(table_name: ObjectName) -> Self {
-        Self { table_name }
+        Self {
+            table_name,
+            time_ranges: Vec::new(),
+        }
+    }
+
+    /// Creates a statement for `TRUNCATE TABLE RANGE`
+    pub fn new_with_ranges(
+        table_name: ObjectName,
+        time_ranges: Vec<(sqlparser::ast::Value, sqlparser::ast::Value)>,
+    ) -> Self {
+        Self {
+            table_name,
+            time_ranges,
+        }
    }

    pub fn table_name(&self) -> &ObjectName {
        &self.table_name
    }
+
+    pub fn time_ranges(&self) -> &[(sqlparser::ast::Value, sqlparser::ast::Value)] {
+        &self.time_ranges
+    }
 }

 impl Display for TruncateTable {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let table_name = self.table_name();
-        write!(f, r#"TRUNCATE TABLE {table_name}"#)
+        write!(f, r#"TRUNCATE TABLE {table_name}"#)?;
+        if self.time_ranges.is_empty() {
+            return Ok(());
+        }
+
+        write!(f, " FILE RANGE ")?;
+        write!(
+            f,
+            "{}",
+            self.time_ranges
+                .iter()
+                .map(|(start, end)| format!("({}, {})", start, end))
+                .join(", ")
+        )
    }
 }

@@ -71,5 +125,45 @@ TRUNCATE TABLE t1"#,
                unreachable!();
            }
        }
+
+        let sql = r"truncate table t1 file range (1,2);";
+        let stmts: Vec<Statement> =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+        assert_eq!(1, stmts.len());
+        assert_matches!(&stmts[0], Statement::TruncateTable { .. });
+        match &stmts[0] {
+            Statement::TruncateTable(trunc) => {
+                let new_sql = format!("\n{}", trunc);
+                assert_eq!(
+                    r#"
+TRUNCATE TABLE t1 FILE RANGE (1, 2)"#,
+                    &new_sql
+                );
+            }
+            _ => {
+                unreachable!();
+            }
+        }
+
+        let sql = r"truncate table t1 file range (1,2), (3,4);";
+        let stmts: Vec<Statement> =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+        assert_eq!(1, stmts.len());
+        assert_matches!(&stmts[0], Statement::TruncateTable { .. });
+        match &stmts[0] {
+            Statement::TruncateTable(trunc) => {
+                let new_sql = format!("\n{}", trunc);
+                assert_eq!(
+                    r#"
+TRUNCATE TABLE t1 FILE RANGE (1, 2), (3, 4)"#,
+                    &new_sql
+                );
+            }
+            _ => {
+                unreachable!();
+            }
+        }
    }
 }