tests_fuzz/translator/mysql/
create_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 datatypes::data_type::ConcreteDataType;
16use sql::statements::concrete_data_type_to_sql_data_type;
17
18use crate::error::{Error, Result};
19use crate::ir::create_expr::ColumnOption;
20use crate::ir::{Column, CreateDatabaseExpr, CreateTableExpr};
21use crate::translator::DslTranslator;
22
23pub struct CreateTableExprTranslator;
24
25impl DslTranslator<CreateTableExpr, String> for CreateTableExprTranslator {
26    type Error = Error;
27
28    fn translate(&self, input: &CreateTableExpr) -> Result<String> {
29        Ok(format!(
30            "CREATE TABLE{}{}(\n{}\n)\n{}{};",
31            Self::create_if_not_exists(input),
32            input.table_name,
33            Self::format_columns(input),
34            Self::format_table_options(input),
35            Self::format_with_clause(input),
36        ))
37    }
38}
39
40impl CreateTableExprTranslator {
41    fn create_if_not_exists(input: &CreateTableExpr) -> &str {
42        if input.if_not_exists {
43            " IF NOT EXISTS "
44        } else {
45            " "
46        }
47    }
48
49    fn format_columns(input: &CreateTableExpr) -> String {
50        let mut output =
51            Vec::with_capacity(input.columns.len() + (!input.primary_keys.is_empty()) as usize);
52        for column in &input.columns {
53            output.push(Self::format_column(column));
54        }
55        if let Some(primary_keys) = Self::format_primary_keys(input) {
56            output.push(primary_keys);
57        }
58        output.join(",\n")
59    }
60
61    fn format_column(column: &Column) -> String {
62        vec![
63            column.name.to_string(),
64            Self::format_column_type(&column.column_type),
65            Self::format_column_options(&column.options),
66        ]
67        .into_iter()
68        .filter(|s| !s.is_empty())
69        .collect::<Vec<_>>()
70        .join(" ")
71    }
72
73    fn format_partition(input: &CreateTableExpr) -> Option<String> {
74        input.partition.as_ref().map(|partition| {
75            format!(
76                "PARTITION ON COLUMNS({}) (\n{}\n)",
77                partition.columns.join(", "),
78                partition
79                    .exprs
80                    .iter()
81                    .map(|expr| expr.to_parser_expr().to_string())
82                    .collect::<Vec<_>>()
83                    .join(",\n")
84            )
85        })
86    }
87
88    fn format_column_type(column_type: &ConcreteDataType) -> String {
89        // Safety: We don't use the `Dictionary` type
90        concrete_data_type_to_sql_data_type(column_type)
91            .unwrap()
92            .to_string()
93    }
94
95    fn format_column_options(options: &[ColumnOption]) -> String {
96        let mut output = Vec::with_capacity(options.len());
97        for option in options {
98            if option != &ColumnOption::PrimaryKey {
99                output.push(option.to_string());
100            }
101        }
102        output.join(" ")
103    }
104
105    fn format_primary_keys(input: &CreateTableExpr) -> Option<String> {
106        if input.primary_keys.is_empty() {
107            None
108        } else {
109            Some(format!(
110                "PRIMARY KEY({})",
111                input
112                    .primary_keys
113                    .iter()
114                    .map(|idx| input.columns[*idx].name.to_string())
115                    .collect::<Vec<_>>()
116                    .join(", ")
117            ))
118        }
119    }
120
121    fn format_table_options(input: &CreateTableExpr) -> String {
122        let mut output = vec![];
123        if let Some(partition) = Self::format_partition(input) {
124            output.push(partition);
125        }
126        if !input.engine.is_empty() {
127            output.push(format!("ENGINE={}", input.engine));
128        }
129
130        output.join("\n")
131    }
132
133    fn format_with_clause(input: &CreateTableExpr) -> String {
134        if input.options.is_empty() {
135            String::new()
136        } else {
137            let mut output = vec![];
138            for (key, value) in &input.options {
139                output.push(format!("\"{key}\" = \"{value}\""));
140            }
141            format!(" with ({})", output.join(",\n"))
142        }
143    }
144}
145
146pub struct CreateDatabaseExprTranslator;
147
148impl DslTranslator<CreateDatabaseExpr, String> for CreateDatabaseExprTranslator {
149    type Error = Error;
150
151    fn translate(&self, input: &CreateDatabaseExpr) -> Result<String> {
152        Ok(format!(
153            "CREATE DATABASE{}{};",
154            Self::create_if_not_exists(input),
155            input.database_name
156        ))
157    }
158}
159
160impl CreateDatabaseExprTranslator {
161    fn create_if_not_exists(input: &CreateDatabaseExpr) -> &str {
162        if input.if_not_exists {
163            " IF NOT EXISTS "
164        } else {
165            " "
166        }
167    }
168}
169
170#[cfg(test)]
171mod tests {
172    use partition::expr::{Operand, PartitionExpr, RestrictedOp};
173
174    use super::CreateTableExprTranslator;
175    use crate::ir::create_expr::{CreateDatabaseExprBuilder, CreateTableExprBuilder, PartitionDef};
176    use crate::test_utils;
177    use crate::translator::DslTranslator;
178
179    #[test]
180    fn test_create_table_expr_translator() {
181        let test_ctx = test_utils::new_test_ctx();
182        let create_table_expr = CreateTableExprBuilder::default()
183            .columns(test_ctx.columns.clone())
184            .table_name("system_metrics")
185            .engine("mito")
186            .primary_keys(vec![0, 1])
187            .partition(PartitionDef {
188                columns: vec!["idc".to_string()],
189                exprs: vec![
190                    PartitionExpr::new(
191                        Operand::Column("idc".to_string()),
192                        RestrictedOp::Lt,
193                        Operand::Value(datatypes::value::Value::Int32(10)),
194                    ),
195                    PartitionExpr::new(
196                        Operand::Expr(PartitionExpr::new(
197                            Operand::Column("idc".to_string()),
198                            RestrictedOp::GtEq,
199                            Operand::Value(datatypes::value::Value::Int32(10)),
200                        )),
201                        RestrictedOp::And,
202                        Operand::Expr(PartitionExpr::new(
203                            Operand::Column("idc".to_string()),
204                            RestrictedOp::Lt,
205                            Operand::Value(datatypes::value::Value::Int32(50)),
206                        )),
207                    ),
208                    PartitionExpr::new(
209                        Operand::Column("idc".to_string()),
210                        RestrictedOp::GtEq,
211                        Operand::Value(datatypes::value::Value::Int32(50)),
212                    ),
213                ],
214            })
215            .build()
216            .unwrap();
217
218        let output = CreateTableExprTranslator
219            .translate(&create_table_expr)
220            .unwrap();
221        assert_eq!(
222            "CREATE TABLE system_metrics(
223host STRING,
224idc STRING,
225cpu_util DOUBLE,
226memory_util DOUBLE,
227disk_util DOUBLE,
228ts TIMESTAMP(3) TIME INDEX,
229PRIMARY KEY(host, idc)
230)
231PARTITION ON COLUMNS(idc) (
232idc < 10,
233idc >= 10 AND idc < 50,
234idc >= 50
235)
236ENGINE=mito;",
237            output
238        );
239    }
240
241    #[test]
242    fn test_create_database_expr_translator() {
243        let create_database_expr = CreateDatabaseExprBuilder::default()
244            .database_name("all_metrics")
245            .if_not_exists(true)
246            .build()
247            .unwrap();
248
249        let output = super::CreateDatabaseExprTranslator
250            .translate(&create_database_expr)
251            .unwrap();
252
253        assert_eq!("CREATE DATABASE IF NOT EXISTS all_metrics;", output);
254    }
255}