tests_fuzz/generator/
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 std::collections::HashMap;
16
17use datatypes::data_type::ConcreteDataType;
18use datatypes::value::Value;
19use derive_builder::Builder;
20use rand::Rng;
21use rand::seq::SliceRandom;
22use snafu::{ResultExt, ensure};
23
24use super::Generator;
25use crate::context::TableContextRef;
26use crate::error::{self, Error, Result};
27use crate::fake::{MappedGenerator, WordGenerator, random_capitalize_map};
28use crate::generator::{ColumnOptionGenerator, ConcreteDataTypeGenerator, Random};
29use crate::ir::create_expr::{
30    ColumnOption, CreateDatabaseExprBuilder, CreateTableExprBuilder, PartitionDef,
31};
32use crate::ir::partition_expr::SimplePartitions;
33use crate::ir::{
34    Column, ColumnTypeGenerator, CreateDatabaseExpr, CreateTableExpr, Ident,
35    PartibleColumnTypeGenerator, StringColumnTypeGenerator, TsColumnTypeGenerator,
36    column_options_generator, generate_columns, generate_partition_bounds,
37    partible_column_options_generator, primary_key_options_generator, ts_column_options_generator,
38};
39
40#[derive(Builder)]
41#[builder(default, pattern = "owned")]
42pub struct CreateTableExprGenerator<R: Rng + 'static> {
43    columns: usize,
44    #[builder(setter(into))]
45    engine: String,
46    partition: usize,
47    if_not_exists: bool,
48    #[builder(setter(into))]
49    name: Ident,
50    #[builder(setter(into))]
51    with_clause: HashMap<String, String>,
52    name_generator: Box<dyn Random<Ident, R>>,
53    ts_column_type_generator: ConcreteDataTypeGenerator<R>,
54    column_type_generator: ConcreteDataTypeGenerator<R>,
55    partible_column_type_generator: ConcreteDataTypeGenerator<R>,
56    partible_column_options_generator: ColumnOptionGenerator<R>,
57    column_options_generator: ColumnOptionGenerator<R>,
58    ts_column_options_generator: ColumnOptionGenerator<R>,
59}
60
61const DEFAULT_ENGINE: &str = "mito";
62
63impl<R: Rng + 'static> Default for CreateTableExprGenerator<R> {
64    fn default() -> Self {
65        Self {
66            columns: 0,
67            engine: DEFAULT_ENGINE.to_string(),
68            if_not_exists: false,
69            partition: 0,
70            name: Ident::new(""),
71            with_clause: HashMap::default(),
72            name_generator: Box::new(MappedGenerator::new(WordGenerator, random_capitalize_map)),
73            ts_column_type_generator: Box::new(TsColumnTypeGenerator),
74            column_type_generator: Box::new(ColumnTypeGenerator),
75            partible_column_type_generator: Box::new(PartibleColumnTypeGenerator),
76            partible_column_options_generator: Box::new(partible_column_options_generator),
77            column_options_generator: Box::new(column_options_generator),
78            ts_column_options_generator: Box::new(ts_column_options_generator),
79        }
80    }
81}
82
83impl<R: Rng + 'static> Generator<CreateTableExpr, R> for CreateTableExprGenerator<R> {
84    type Error = Error;
85
86    /// Generates the [CreateTableExpr].
87    fn generate(&self, rng: &mut R) -> Result<CreateTableExpr> {
88        ensure!(
89            self.columns != 0,
90            error::UnexpectedSnafu {
91                violated: "The columns must larger than zero"
92            }
93        );
94
95        let mut builder = CreateTableExprBuilder::default();
96        let mut columns = Vec::with_capacity(self.columns);
97        let mut primary_keys = vec![];
98        let need_partible_column = self.partition > 1;
99        let mut column_names = self.name_generator.choose(rng, self.columns);
100
101        if self.columns == 1 {
102            // Generates the ts column.
103            // Safety: columns must large than 0.
104            let name = column_names.pop().unwrap();
105            let column = generate_columns(
106                rng,
107                vec![name.clone()],
108                self.ts_column_type_generator.as_ref(),
109                self.ts_column_options_generator.as_ref(),
110            )
111            .remove(0);
112            columns.push(column);
113        } else {
114            // Generates the partible column.
115            if need_partible_column {
116                // Safety: columns must large than 0.
117                let name = column_names.pop().unwrap();
118                let column = generate_columns(
119                    rng,
120                    vec![name.clone()],
121                    self.partible_column_type_generator.as_ref(),
122                    self.partible_column_options_generator.as_ref(),
123                )
124                .remove(0);
125
126                // Generates partition bounds.
127                let partition_def = generate_partition_def(
128                    self.partition,
129                    column.column_type.clone(),
130                    name.clone(),
131                );
132                builder.partition(partition_def);
133                columns.push(column);
134            }
135            // Generates the ts column.
136            // Safety: columns must large than 1.
137            let name = column_names.pop().unwrap();
138            columns.extend(generate_columns(
139                rng,
140                vec![name],
141                self.ts_column_type_generator.as_ref(),
142                self.ts_column_options_generator.as_ref(),
143            ));
144            // Generates rest columns
145            columns.extend(generate_columns(
146                rng,
147                column_names,
148                self.column_type_generator.as_ref(),
149                self.column_options_generator.as_ref(),
150            ));
151        }
152
153        for (idx, column) in columns.iter().enumerate() {
154            if column.is_primary_key() {
155                primary_keys.push(idx);
156            }
157        }
158        // Shuffles the primary keys.
159        primary_keys.shuffle(rng);
160
161        builder.columns(columns);
162        builder.primary_keys(primary_keys);
163        builder.engine(self.engine.clone());
164        builder.if_not_exists(self.if_not_exists);
165        if self.name.is_empty() {
166            builder.table_name(self.name_generator.generate(rng));
167        } else {
168            builder.table_name(self.name.clone());
169        }
170        if !self.with_clause.is_empty() {
171            let mut options = HashMap::new();
172            for (key, value) in &self.with_clause {
173                options.insert(key.clone(), Value::from(value.clone()));
174            }
175            builder.options(options);
176        }
177        builder.build().context(error::BuildCreateTableExprSnafu)
178    }
179}
180
181fn generate_partition_def(
182    partitions: usize,
183    column_type: ConcreteDataType,
184    column_name: Ident,
185) -> PartitionDef {
186    let bounds = generate_partition_bounds(&column_type, partitions - 1);
187    let partitions = SimplePartitions::new(column_name.clone(), bounds);
188    let partition_exprs = partitions.generate().unwrap();
189
190    PartitionDef {
191        columns: vec![column_name.clone()],
192        exprs: partition_exprs,
193    }
194}
195
196/// Generate a physical table with 2 columns: ts of TimestampType::Millisecond as time index and val of Float64Type.
197#[derive(Builder)]
198#[builder(pattern = "owned")]
199pub struct CreatePhysicalTableExprGenerator<R: Rng + 'static> {
200    #[builder(default = "Box::new(WordGenerator)")]
201    name_generator: Box<dyn Random<Ident, R>>,
202    #[builder(default = "false")]
203    if_not_exists: bool,
204    #[builder(default, setter(into))]
205    with_clause: HashMap<String, String>,
206}
207
208impl<R: Rng + 'static> Generator<CreateTableExpr, R> for CreatePhysicalTableExprGenerator<R> {
209    type Error = Error;
210
211    fn generate(&self, rng: &mut R) -> Result<CreateTableExpr> {
212        let mut options = HashMap::with_capacity(self.with_clause.len() + 1);
213        options.insert("physical_metric_table".to_string(), Value::from(""));
214        for (key, value) in &self.with_clause {
215            options.insert(key.clone(), Value::from(value.clone()));
216        }
217
218        Ok(CreateTableExpr {
219            table_name: self.name_generator.generate(rng),
220            columns: vec![
221                Column {
222                    name: Ident::new("ts"),
223                    column_type: ConcreteDataType::timestamp_millisecond_datatype(),
224                    options: vec![ColumnOption::TimeIndex],
225                },
226                Column {
227                    name: Ident::new("val"),
228                    column_type: ConcreteDataType::float64_datatype(),
229                    options: vec![],
230                },
231            ],
232            if_not_exists: self.if_not_exists,
233            partition: None,
234            engine: "metric".to_string(),
235            options,
236            primary_keys: vec![],
237        })
238    }
239}
240
241/// Generate a logical table based on an existing physical table.
242#[derive(Builder)]
243#[builder(pattern = "owned")]
244pub struct CreateLogicalTableExprGenerator<R: Rng + 'static> {
245    physical_table_ctx: TableContextRef,
246    labels: usize,
247    if_not_exists: bool,
248    #[builder(default = "Box::new(WordGenerator)")]
249    name_generator: Box<dyn Random<Ident, R>>,
250}
251
252impl<R: Rng + 'static> Generator<CreateTableExpr, R> for CreateLogicalTableExprGenerator<R> {
253    type Error = Error;
254
255    fn generate(&self, rng: &mut R) -> Result<CreateTableExpr> {
256        // Currently we mock the usage of GreptimeDB as Prometheus' backend, the physical table must have two columns.
257        ensure!(
258            self.physical_table_ctx.columns.len() == 2,
259            error::UnexpectedSnafu {
260                violated: "The physical table must have two columns"
261            }
262        );
263
264        // Generates the logical table columns based on the physical table.
265        let logical_table_name = self
266            .physical_table_ctx
267            .generate_unique_table_name(rng, self.name_generator.as_ref());
268        let mut logical_table = CreateTableExpr {
269            table_name: logical_table_name,
270            columns: self.physical_table_ctx.columns.clone(),
271            if_not_exists: self.if_not_exists,
272            partition: None,
273            engine: "metric".to_string(),
274            options: [(
275                "on_physical_table".to_string(),
276                self.physical_table_ctx.name.value.clone().into(),
277            )]
278            .into(),
279            primary_keys: vec![],
280        };
281
282        let column_names = self.name_generator.choose(rng, self.labels);
283        logical_table.columns.extend(generate_columns(
284            rng,
285            column_names,
286            &StringColumnTypeGenerator,
287            Box::new(primary_key_options_generator),
288        ));
289
290        // Currently only the `primary key` option is kept in physical table,
291        // so we only keep the `primary key` option in the logical table for fuzz test.
292        let mut primary_keys = vec![];
293        for (idx, column) in logical_table.columns.iter().enumerate() {
294            if column.is_primary_key() {
295                primary_keys.push(idx);
296            }
297        }
298        primary_keys.shuffle(rng);
299        logical_table.primary_keys = primary_keys;
300
301        Ok(logical_table)
302    }
303}
304
305#[derive(Builder)]
306#[builder(default, pattern = "owned")]
307pub struct CreateDatabaseExprGenerator<R: Rng + 'static> {
308    #[builder(setter(into))]
309    database_name: String,
310    name_generator: Box<dyn Random<Ident, R>>,
311    if_not_exists: bool,
312}
313
314impl<R: Rng + 'static> Default for CreateDatabaseExprGenerator<R> {
315    fn default() -> Self {
316        Self {
317            database_name: String::new(),
318            name_generator: Box::new(MappedGenerator::new(WordGenerator, random_capitalize_map)),
319            if_not_exists: false,
320        }
321    }
322}
323
324impl<R: Rng + 'static> Generator<CreateDatabaseExpr, R> for CreateDatabaseExprGenerator<R> {
325    type Error = Error;
326
327    fn generate(&self, rng: &mut R) -> Result<CreateDatabaseExpr> {
328        let mut builder = CreateDatabaseExprBuilder::default();
329        builder.if_not_exists(self.if_not_exists);
330        if self.database_name.is_empty() {
331            builder.database_name(self.name_generator.generate(rng));
332        } else {
333            builder.database_name(self.database_name.clone());
334        }
335        builder.build().context(error::BuildCreateDatabaseExprSnafu)
336    }
337}
338
339#[cfg(test)]
340mod tests {
341    use std::sync::Arc;
342
343    use datatypes::data_type::ConcreteDataType;
344    use datatypes::value::Value;
345    use rand::SeedableRng;
346
347    use super::*;
348    use crate::context::TableContext;
349
350    #[test]
351    fn test_float64() {
352        let value = Value::from(0.047318541668048164);
353        assert_eq!("0.047318541668048164", value.to_string());
354        let value: f64 = "0.047318541668048164".parse().unwrap();
355        assert_eq!("0.047318541668048164", value.to_string());
356    }
357
358    #[test]
359    fn test_create_table_expr_generator() {
360        let mut rng = rand::rng();
361
362        let expr = CreateTableExprGeneratorBuilder::default()
363            .columns(10)
364            .partition(3)
365            .if_not_exists(true)
366            .engine("mito2")
367            .build()
368            .unwrap()
369            .generate(&mut rng)
370            .unwrap();
371        assert_eq!(expr.engine, "mito2");
372        assert!(expr.if_not_exists);
373        assert_eq!(expr.columns.len(), 10);
374        assert_eq!(expr.partition.unwrap().exprs.len(), 3);
375
376        let expr = CreateTableExprGeneratorBuilder::default()
377            .columns(10)
378            .partition(1)
379            .build()
380            .unwrap()
381            .generate(&mut rng)
382            .unwrap();
383        assert_eq!(expr.columns.len(), 10);
384        assert!(expr.partition.is_none());
385    }
386
387    #[test]
388    fn test_create_table_expr_generator_deterministic() {
389        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(0);
390        let expr = CreateTableExprGeneratorBuilder::default()
391            .columns(10)
392            .partition(3)
393            .if_not_exists(true)
394            .engine("mito2")
395            .build()
396            .unwrap()
397            .generate(&mut rng)
398            .unwrap();
399
400        let serialized = serde_json::to_string(&expr).unwrap();
401        let expected = r#"{"table_name":{"value":"quasi","quote_style":null},"columns":[{"name":{"value":"mOLEsTIAs","quote_style":null},"column_type":{"Float64":{}},"options":["PrimaryKey","Null"]},{"name":{"value":"CUMQUe","quote_style":null},"column_type":{"Timestamp":{"Second":null}},"options":["TimeIndex"]},{"name":{"value":"NaTus","quote_style":null},"column_type":{"Int64":{}},"options":[]},{"name":{"value":"EXPeDITA","quote_style":null},"column_type":{"Float64":{}},"options":[]},{"name":{"value":"ImPEDiT","quote_style":null},"column_type":{"Float32":{}},"options":[{"DefaultValue":{"Float32":0.56425774}}]},{"name":{"value":"ADIpisci","quote_style":null},"column_type":{"Float32":{}},"options":["PrimaryKey"]},{"name":{"value":"deBITIs","quote_style":null},"column_type":{"Float32":{}},"options":[{"DefaultValue":{"Float32":0.31315368}}]},{"name":{"value":"toTaM","quote_style":null},"column_type":{"Int32":{}},"options":["NotNull"]},{"name":{"value":"QuI","quote_style":null},"column_type":{"Float32":{}},"options":[{"DefaultValue":{"Float32":0.39941502}}]},{"name":{"value":"INVeNtOre","quote_style":null},"column_type":{"Boolean":null},"options":["PrimaryKey"]}],"if_not_exists":true,"partition":{"columns":[{"value":"mOLEsTIAs","quote_style":null}],"exprs":[{"lhs":{"Column":"mOLEsTIAs"},"op":"Lt","rhs":{"Value":{"Float64":5.992310449541053e+307}}},{"lhs":{"Expr":{"lhs":{"Column":"mOLEsTIAs"},"op":"GtEq","rhs":{"Value":{"Float64":5.992310449541053e+307}}}},"op":"And","rhs":{"Expr":{"lhs":{"Column":"mOLEsTIAs"},"op":"Lt","rhs":{"Value":{"Float64":1.1984620899082105e+308}}}}},{"lhs":{"Column":"mOLEsTIAs"},"op":"GtEq","rhs":{"Value":{"Float64":1.1984620899082105e+308}}}]},"engine":"mito2","options":{},"primary_keys":[0,5,9]}"#;
402        assert_eq!(expected, serialized);
403    }
404
405    #[test]
406    fn test_create_logical_table_expr_generator() {
407        let mut rng = rand::rng();
408
409        let physical_table_expr = CreatePhysicalTableExprGeneratorBuilder::default()
410            .if_not_exists(false)
411            .build()
412            .unwrap()
413            .generate(&mut rng)
414            .unwrap();
415        assert_eq!(physical_table_expr.engine, "metric");
416        assert_eq!(physical_table_expr.columns.len(), 2);
417
418        let physical_ts = physical_table_expr.columns.iter().position(|column| {
419            column
420                .options
421                .iter()
422                .any(|option| option == &ColumnOption::TimeIndex)
423        });
424        let physical_ts_name = physical_table_expr.columns[physical_ts.unwrap()]
425            .name
426            .value
427            .clone();
428
429        let physical_table_ctx = Arc::new(TableContext::from(&physical_table_expr));
430
431        let logical_table_expr = CreateLogicalTableExprGeneratorBuilder::default()
432            .physical_table_ctx(physical_table_ctx)
433            .labels(5)
434            .if_not_exists(false)
435            .build()
436            .unwrap()
437            .generate(&mut rng)
438            .unwrap();
439        let logical_ts = logical_table_expr.columns.iter().position(|column| {
440            column
441                .options
442                .iter()
443                .any(|option| option == &ColumnOption::TimeIndex)
444        });
445        let logical_ts_name = logical_table_expr.columns[logical_ts.unwrap()]
446            .name
447            .value
448            .clone();
449
450        assert_eq!(logical_table_expr.engine, "metric");
451        assert_eq!(logical_table_expr.columns.len(), 7);
452        assert_eq!(logical_ts_name, physical_ts_name);
453        assert!(logical_table_expr.columns.iter().all(|column| {
454            column.column_type != ConcreteDataType::string_datatype()
455                || column
456                    .options
457                    .iter()
458                    .any(|option| option == &ColumnOption::PrimaryKey)
459        }));
460    }
461
462    #[test]
463    fn test_create_logical_table_expr_generator_deterministic() {
464        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(0);
465        let physical_table_expr = CreatePhysicalTableExprGeneratorBuilder::default()
466            .if_not_exists(false)
467            .build()
468            .unwrap()
469            .generate(&mut rng)
470            .unwrap();
471        let physical_table_serialized = serde_json::to_string(&physical_table_expr).unwrap();
472        let physical_table_expected = r#"{"table_name":{"value":"expedita","quote_style":null},"columns":[{"name":{"value":"ts","quote_style":null},"column_type":{"Timestamp":{"Millisecond":null}},"options":["TimeIndex"]},{"name":{"value":"val","quote_style":null},"column_type":{"Float64":{}},"options":[]}],"if_not_exists":false,"partition":null,"engine":"metric","options":{"physical_metric_table":{"String":""}},"primary_keys":[]}"#;
473        assert_eq!(physical_table_expected, physical_table_serialized);
474
475        let physical_table_ctx = Arc::new(TableContext::from(&physical_table_expr));
476
477        let logical_table_expr = CreateLogicalTableExprGeneratorBuilder::default()
478            .physical_table_ctx(physical_table_ctx)
479            .labels(5)
480            .if_not_exists(false)
481            .build()
482            .unwrap()
483            .generate(&mut rng)
484            .unwrap();
485
486        let logical_table_serialized = serde_json::to_string(&logical_table_expr).unwrap();
487        let logical_table_expected = r#"{"table_name":{"value":"impedit","quote_style":null},"columns":[{"name":{"value":"ts","quote_style":null},"column_type":{"Timestamp":{"Millisecond":null}},"options":["TimeIndex"]},{"name":{"value":"val","quote_style":null},"column_type":{"Float64":{}},"options":[]},{"name":{"value":"totam","quote_style":null},"column_type":{"String":{"size_type":"Utf8"}},"options":["PrimaryKey"]},{"name":{"value":"cumque","quote_style":null},"column_type":{"String":{"size_type":"Utf8"}},"options":["PrimaryKey"]},{"name":{"value":"natus","quote_style":null},"column_type":{"String":{"size_type":"Utf8"}},"options":["PrimaryKey"]},{"name":{"value":"molestias","quote_style":null},"column_type":{"String":{"size_type":"Utf8"}},"options":["PrimaryKey"]},{"name":{"value":"qui","quote_style":null},"column_type":{"String":{"size_type":"Utf8"}},"options":["PrimaryKey"]}],"if_not_exists":false,"partition":null,"engine":"metric","options":{"on_physical_table":{"String":"expedita"}},"primary_keys":[4,2,3,6,5]}"#;
488        assert_eq!(logical_table_expected, logical_table_serialized);
489    }
490
491    #[test]
492    fn test_create_database_expr_generator() {
493        let mut rng = rand::rng();
494
495        let expr = CreateDatabaseExprGeneratorBuilder::default()
496            .if_not_exists(true)
497            .build()
498            .unwrap()
499            .generate(&mut rng)
500            .unwrap();
501        assert!(expr.if_not_exists);
502    }
503
504    #[test]
505    fn test_create_database_expr_generator_deterministic() {
506        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(0);
507        let expr = CreateDatabaseExprGeneratorBuilder::default()
508            .if_not_exists(true)
509            .build()
510            .unwrap()
511            .generate(&mut rng)
512            .unwrap();
513
514        let serialized = serde_json::to_string(&expr).unwrap();
515        let expected =
516            r#"{"database_name":{"value":"EXPediTA","quote_style":null},"if_not_exists":true}"#;
517        assert_eq!(expected, serialized);
518    }
519}