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                wait,
33            }) => {
34                let target_expr = format_partition_expr_sql(target);
35                let into_exprs = into
36                    .iter()
37                    .map(format_partition_expr_sql)
38                    .collect::<Vec<_>>()
39                    .join(",\n  ");
40                let wait_clause = format_wait_clause(*wait);
41                Ok(format!(
42                    "ALTER TABLE {} SPLIT PARTITION (\n  {}\n) INTO (\n  {}\n){};",
43                    table_name, target_expr, into_exprs, wait_clause
44                ))
45            }
46            RepartitionExpr::Merge(MergePartitionExpr {
47                table_name,
48                targets,
49                wait,
50            }) => {
51                let merge_exprs = targets
52                    .iter()
53                    .map(format_partition_expr_sql)
54                    .collect::<Vec<_>>()
55                    .join(",\n  ");
56                let wait_clause = format_wait_clause(*wait);
57                Ok(format!(
58                    "ALTER TABLE {} MERGE PARTITION (\n  {}\n){};",
59                    table_name, merge_exprs, wait_clause
60                ))
61            }
62        }
63    }
64}
65
66fn format_partition_expr_sql(expr: &PartitionExpr) -> String {
67    expr.to_parser_expr().to_string()
68}
69
70fn format_wait_clause(wait: bool) -> String {
71    if wait {
72        String::new()
73    } else {
74        " WITH (\n  WAIT = false\n)".to_string()
75    }
76}
77
78#[cfg(test)]
79mod tests {
80    use datatypes::value::Value;
81    use partition::expr::col;
82
83    use super::RepartitionExprTranslator;
84    use crate::ir::repartition_expr::{MergePartitionExpr, RepartitionExpr, SplitPartitionExpr};
85    use crate::translator::DslTranslator;
86
87    #[test]
88    fn test_translate_split_expr() {
89        let expr = RepartitionExpr::Split(SplitPartitionExpr {
90            table_name: "demo".into(),
91            target: col("id").lt(Value::Int32(10)),
92            into: vec![
93                col("id").lt(Value::Int32(5)),
94                col("id")
95                    .gt_eq(Value::Int32(5))
96                    .and(col("id").lt(Value::Int32(10))),
97            ],
98            wait: true,
99        });
100        let sql = RepartitionExprTranslator.translate(&expr).unwrap();
101        let expected = r#"ALTER TABLE demo SPLIT PARTITION (
102  id < 10
103) INTO (
104  id < 5,
105  id >= 5 AND id < 10
106);"#;
107        assert_eq!(sql, expected);
108    }
109
110    #[test]
111    fn test_translate_merge_expr() {
112        let expr = RepartitionExpr::Merge(MergePartitionExpr {
113            table_name: "demo".into(),
114            targets: vec![
115                col("id").gt_eq(Value::Int32(10)),
116                col("id").gt_eq(Value::Int32(20)),
117            ],
118            wait: true,
119        });
120        let sql = RepartitionExprTranslator.translate(&expr).unwrap();
121        let expected = r#"ALTER TABLE demo MERGE PARTITION (
122  id >= 10,
123  id >= 20
124);"#;
125        assert_eq!(sql, expected);
126    }
127
128    #[test]
129    fn test_translate_split_expr_wait_false() {
130        let expr = RepartitionExpr::Split(SplitPartitionExpr {
131            table_name: "demo".into(),
132            target: col("id").lt(Value::Int32(10)),
133            into: vec![
134                col("id").lt(Value::Int32(5)),
135                col("id")
136                    .gt_eq(Value::Int32(5))
137                    .and(col("id").lt(Value::Int32(10))),
138            ],
139            wait: false,
140        });
141        let sql = RepartitionExprTranslator.translate(&expr).unwrap();
142        let expected = r#"ALTER TABLE demo SPLIT PARTITION (
143  id < 10
144) INTO (
145  id < 5,
146  id >= 5 AND id < 10
147) WITH (
148  WAIT = false
149);"#;
150        assert_eq!(sql, expected);
151    }
152}
tests_fuzz/translator/mysql/repartition_expr.rs

tests_fuzz/translator/mysql/
repartition_expr.rs
