tests_fuzz/translator/mysql/
repartition_expr.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 partition::expr::PartitionExpr;
16
17use crate::error::Result;
18use crate::ir::repartition_expr::{MergePartitionExpr, RepartitionExpr, SplitPartitionExpr};
19use crate::translator::DslTranslator;
20
21pub struct RepartitionExprTranslator;
22
23impl DslTranslator<RepartitionExpr, String> for RepartitionExprTranslator {
24    type Error = crate::error::Error;
25
26    fn translate(&self, input: &RepartitionExpr) -> Result<String> {
27        match input {
28            RepartitionExpr::Split(SplitPartitionExpr {
29                table_name,
30                target,
31                into,
32            }) => {
33                let target_expr = format_partition_expr_sql(target);
34                let into_exprs = into
35                    .iter()
36                    .map(format_partition_expr_sql)
37                    .collect::<Vec<_>>()
38                    .join(",\n  ");
39                Ok(format!(
40                    "ALTER TABLE {} SPLIT PARTITION (\n  {}\n) INTO (\n  {}\n);",
41                    table_name, target_expr, into_exprs
42                ))
43            }
44            RepartitionExpr::Merge(MergePartitionExpr {
45                table_name,
46                targets,
47            }) => {
48                let merge_exprs = targets
49                    .iter()
50                    .map(format_partition_expr_sql)
51                    .collect::<Vec<_>>()
52                    .join(",\n  ");
53                Ok(format!(
54                    "ALTER TABLE {} MERGE PARTITION (\n  {}\n);",
55                    table_name, merge_exprs
56                ))
57            }
58        }
59    }
60}
61
62fn format_partition_expr_sql(expr: &PartitionExpr) -> String {
63    expr.to_parser_expr().to_string()
64}
65
66#[cfg(test)]
67mod tests {
68    use datatypes::value::Value;
69    use partition::expr::col;
70
71    use super::RepartitionExprTranslator;
72    use crate::ir::repartition_expr::{MergePartitionExpr, RepartitionExpr, SplitPartitionExpr};
73    use crate::translator::DslTranslator;
74
75    #[test]
76    fn test_translate_split_expr() {
77        let expr = RepartitionExpr::Split(SplitPartitionExpr {
78            table_name: "demo".into(),
79            target: col("id").lt(Value::Int32(10)),
80            into: vec![
81                col("id").lt(Value::Int32(5)),
82                col("id")
83                    .gt_eq(Value::Int32(5))
84                    .and(col("id").lt(Value::Int32(10))),
85            ],
86        });
87        let sql = RepartitionExprTranslator.translate(&expr).unwrap();
88        let expected = r#"ALTER TABLE demo SPLIT PARTITION (
89  id < 10
90) INTO (
91  id < 5,
92  id >= 5 AND id < 10
93);"#;
94        assert_eq!(sql, expected);
95    }
96
97    #[test]
98    fn test_translate_merge_expr() {
99        let expr = RepartitionExpr::Merge(MergePartitionExpr {
100            table_name: "demo".into(),
101            targets: vec![
102                col("id").gt_eq(Value::Int32(10)),
103                col("id").gt_eq(Value::Int32(20)),
104            ],
105        });
106        let sql = RepartitionExprTranslator.translate(&expr).unwrap();
107        let expected = r#"ALTER TABLE demo MERGE PARTITION (
108  id >= 10,
109  id >= 20
110);"#;
111        assert_eq!(sql, expected);
112    }
113}