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::create_expr::PartitionDef;
19use crate::ir::repartition_expr::{
20    AlterTablePartitionsExpr, MergePartitionExpr, RepartitionExpr, SplitPartitionExpr,
21};
22use crate::translator::DslTranslator;
23
24pub struct RepartitionExprTranslator;
25
26impl DslTranslator<RepartitionExpr, String> for RepartitionExprTranslator {
27    type Error = crate::error::Error;
28
29    fn translate(&self, input: &RepartitionExpr) -> Result<String> {
30        match input {
31            RepartitionExpr::Split(SplitPartitionExpr {
32                table_name,
33                target,
34                into,
35                wait,
36            }) => {
37                let target_expr = format_partition_expr_sql(target);
38                let into_exprs = into
39                    .iter()
40                    .map(format_partition_expr_sql)
41                    .collect::<Vec<_>>()
42                    .join(",\n  ");
43                let wait_clause = format_wait_clause(*wait);
44                Ok(format!(
45                    "ALTER TABLE {} SPLIT PARTITION (\n  {}\n) INTO (\n  {}\n){};",
46                    table_name, target_expr, into_exprs, wait_clause
47                ))
48            }
49            RepartitionExpr::Merge(MergePartitionExpr {
50                table_name,
51                targets,
52                wait,
53            }) => {
54                let merge_exprs = targets
55                    .iter()
56                    .map(format_partition_expr_sql)
57                    .collect::<Vec<_>>()
58                    .join(",\n  ");
59                let wait_clause = format_wait_clause(*wait);
60                Ok(format!(
61                    "ALTER TABLE {} MERGE PARTITION (\n  {}\n){};",
62                    table_name, merge_exprs, wait_clause
63                ))
64            }
65            RepartitionExpr::AlterPartitions(AlterTablePartitionsExpr {
66                table_name,
67                partition,
68                wait,
69            }) => {
70                let partition_clause = format_partition_clause(partition);
71                let wait_clause = format_wait_clause(*wait);
72                Ok(format!(
73                    "ALTER TABLE {} {}{};",
74                    table_name, partition_clause, wait_clause
75                ))
76            }
77        }
78    }
79}
80
81fn format_partition_clause(partition: &PartitionDef) -> String {
82    let columns = partition
83        .columns
84        .iter()
85        .map(|column| column.to_string())
86        .collect::<Vec<_>>()
87        .join(", ");
88    let exprs = partition
89        .exprs
90        .iter()
91        .map(format_partition_expr_sql)
92        .collect::<Vec<_>>()
93        .join(",\n  ");
94    format!("PARTITION ON COLUMNS ({columns}) (\n  {exprs}\n)")
95}
96
97fn format_partition_expr_sql(expr: &PartitionExpr) -> String {
98    expr.to_parser_expr().to_string()
99}
100
101fn format_wait_clause(wait: bool) -> String {
102    if wait {
103        String::new()
104    } else {
105        " WITH (\n  WAIT = false\n)".to_string()
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use datatypes::value::Value;
112    use partition::expr::col;
113    use sql::dialect::GreptimeDbDialect;
114    use sql::parser::{ParseOptions, ParserContext};
115
116    use super::RepartitionExprTranslator;
117    use crate::ir::Ident;
118    use crate::ir::create_expr::PartitionDef;
119    use crate::ir::repartition_expr::{
120        AlterTablePartitionsExpr, MergePartitionExpr, RepartitionExpr, SplitPartitionExpr,
121    };
122    use crate::translator::DslTranslator;
123
124    #[test]
125    fn test_translate_split_expr() {
126        let expr = RepartitionExpr::Split(SplitPartitionExpr {
127            table_name: "demo".into(),
128            target: col("id").lt(Value::Int32(10)),
129            into: vec![
130                col("id").lt(Value::Int32(5)),
131                col("id")
132                    .gt_eq(Value::Int32(5))
133                    .and(col("id").lt(Value::Int32(10))),
134            ],
135            wait: true,
136        });
137        let sql = RepartitionExprTranslator.translate(&expr).unwrap();
138        let expected = r#"ALTER TABLE demo SPLIT PARTITION (
139  id < 10
140) INTO (
141  id < 5,
142  id >= 5 AND id < 10
143);"#;
144        assert_eq!(sql, expected);
145    }
146
147    #[test]
148    fn test_translate_merge_expr() {
149        let expr = RepartitionExpr::Merge(MergePartitionExpr {
150            table_name: "demo".into(),
151            targets: vec![
152                col("id").gt_eq(Value::Int32(10)),
153                col("id").gt_eq(Value::Int32(20)),
154            ],
155            wait: true,
156        });
157        let sql = RepartitionExprTranslator.translate(&expr).unwrap();
158        let expected = r#"ALTER TABLE demo MERGE PARTITION (
159  id >= 10,
160  id >= 20
161);"#;
162        assert_eq!(sql, expected);
163    }
164
165    #[test]
166    fn test_translate_split_expr_wait_false() {
167        let expr = RepartitionExpr::Split(SplitPartitionExpr {
168            table_name: "demo".into(),
169            target: col("id").lt(Value::Int32(10)),
170            into: vec![
171                col("id").lt(Value::Int32(5)),
172                col("id")
173                    .gt_eq(Value::Int32(5))
174                    .and(col("id").lt(Value::Int32(10))),
175            ],
176            wait: false,
177        });
178        let sql = RepartitionExprTranslator.translate(&expr).unwrap();
179        let expected = r#"ALTER TABLE demo SPLIT PARTITION (
180  id < 10
181) INTO (
182  id < 5,
183  id >= 5 AND id < 10
184) WITH (
185  WAIT = false
186);"#;
187        assert_eq!(sql, expected);
188    }
189
190    #[test]
191    fn test_translate_alter_table_partitions_expr() {
192        let expr = RepartitionExpr::AlterPartitions(AlterTablePartitionsExpr {
193            table_name: "demo".into(),
194            partition: PartitionDef {
195                columns: vec![Ident::new("id")],
196                exprs: vec![
197                    col("id").lt(Value::Int32(10)),
198                    col("id")
199                        .gt_eq(Value::Int32(10))
200                        .and(col("id").lt(Value::Int32(20))),
201                    col("id").gt_eq(Value::Int32(20)),
202                ],
203            },
204            wait: true,
205        });
206        let sql = RepartitionExprTranslator.translate(&expr).unwrap();
207        let expected = r#"ALTER TABLE demo PARTITION ON COLUMNS (id) (
208  id < 10,
209  id >= 10 AND id < 20,
210  id >= 20
211);"#;
212        assert_eq!(sql, expected);
213        assert_repartition_sql_parseable(&sql);
214    }
215
216    #[test]
217    fn test_translate_alter_table_partitions_expr_wait_false() {
218        let expr = RepartitionExpr::AlterPartitions(AlterTablePartitionsExpr {
219            table_name: "demo".into(),
220            partition: PartitionDef {
221                columns: vec![Ident::new("host")],
222                exprs: vec![
223                    col("host").lt(Value::from("m")),
224                    col("host").gt_eq(Value::from("m")),
225                ],
226            },
227            wait: false,
228        });
229        let sql = RepartitionExprTranslator.translate(&expr).unwrap();
230        let expected = r#"ALTER TABLE demo PARTITION ON COLUMNS (host) (
231  host < 'm',
232  host >= 'm'
233) WITH (
234  WAIT = false
235);"#;
236        assert_eq!(sql, expected);
237        assert_repartition_sql_parseable(&sql);
238    }
239
240    fn assert_repartition_sql_parseable(sql: &str) {
241        let statements =
242            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
243                .unwrap();
244        assert_eq!(statements.len(), 1);
245    }
246}
tests_fuzz/translator/mysql/repartition_expr.rs

tests_fuzz/translator/mysql/
repartition_expr.rs
