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