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 partition::expr::{Operand, PartitionExpr, RestrictedOp};
21use rand::Rng;
22use rand::seq::SliceRandom;
23use snafu::{ResultExt, ensure};
24
25use super::Generator;
26use crate::context::TableContextRef;
27use crate::error::{self, Error, Result};
28use crate::fake::{MappedGenerator, WordGenerator, random_capitalize_map};
29use crate::generator::{ColumnOptionGenerator, ConcreteDataTypeGenerator, Random};
30use crate::ir::create_expr::{
31    ColumnOption, CreateDatabaseExprBuilder, CreateTableExprBuilder, PartitionDef,
32};
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.to_string());
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.to_string(), Value::from(value.to_string()));
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 mut partition_exprs = Vec::with_capacity(partitions);
188
189    let first_bound = bounds[0].clone();
190    partition_exprs.push(PartitionExpr::new(
191        Operand::Column(column_name.to_string()),
192        RestrictedOp::Lt,
193        Operand::Value(first_bound),
194    ));
195    for bound_idx in 1..bounds.len() {
196        partition_exprs.push(PartitionExpr::new(
197            Operand::Expr(PartitionExpr::new(
198                Operand::Column(column_name.to_string()),
199                RestrictedOp::GtEq,
200                Operand::Value(bounds[bound_idx - 1].clone()),
201            )),
202            RestrictedOp::And,
203            Operand::Expr(PartitionExpr::new(
204                Operand::Column(column_name.to_string()),
205                RestrictedOp::Lt,
206                Operand::Value(bounds[bound_idx].clone()),
207            )),
208        ));
209    }
210    let last_bound = bounds.last().unwrap().clone();
211    partition_exprs.push(PartitionExpr::new(
212        Operand::Column(column_name.to_string()),
213        RestrictedOp::GtEq,
214        Operand::Value(last_bound),
215    ));
216
217    PartitionDef {
218        columns: vec![column_name.to_string()],
219        exprs: partition_exprs,
220    }
221}
222
223/// Generate a physical table with 2 columns: ts of TimestampType::Millisecond as time index and val of Float64Type.
224#[derive(Builder)]
225#[builder(pattern = "owned")]
226pub struct CreatePhysicalTableExprGenerator<R: Rng + 'static> {
227    #[builder(default = "Box::new(WordGenerator)")]
228    name_generator: Box<dyn Random<Ident, R>>,
229    #[builder(default = "false")]
230    if_not_exists: bool,
231    #[builder(default, setter(into))]
232    with_clause: HashMap<String, String>,
233}
234
235impl<R: Rng + 'static> Generator<CreateTableExpr, R> for CreatePhysicalTableExprGenerator<R> {
236    type Error = Error;
237
238    fn generate(&self, rng: &mut R) -> Result<CreateTableExpr> {
239        let mut options = HashMap::with_capacity(self.with_clause.len() + 1);
240        options.insert("physical_metric_table".to_string(), Value::from(""));
241        for (key, value) in &self.with_clause {
242            options.insert(key.to_string(), Value::from(value.to_string()));
243        }
244
245        Ok(CreateTableExpr {
246            table_name: self.name_generator.generate(rng),
247            columns: vec![
248                Column {
249                    name: Ident::new("ts"),
250                    column_type: ConcreteDataType::timestamp_millisecond_datatype(),
251                    options: vec![ColumnOption::TimeIndex],
252                },
253                Column {
254                    name: Ident::new("val"),
255                    column_type: ConcreteDataType::float64_datatype(),
256                    options: vec![],
257                },
258            ],
259            if_not_exists: self.if_not_exists,
260            partition: None,
261            engine: "metric".to_string(),
262            options,
263            primary_keys: vec![],
264        })
265    }
266}
267
268/// Generate a logical table based on an existing physical table.
269#[derive(Builder)]
270#[builder(pattern = "owned")]
271pub struct CreateLogicalTableExprGenerator<R: Rng + 'static> {
272    physical_table_ctx: TableContextRef,
273    labels: usize,
274    if_not_exists: bool,
275    #[builder(default = "Box::new(WordGenerator)")]
276    name_generator: Box<dyn Random<Ident, R>>,
277}
278
279impl<R: Rng + 'static> Generator<CreateTableExpr, R> for CreateLogicalTableExprGenerator<R> {
280    type Error = Error;
281
282    fn generate(&self, rng: &mut R) -> Result<CreateTableExpr> {
283        // Currently we mock the usage of GreptimeDB as Prometheus' backend, the physical table must have two columns.
284        ensure!(
285            self.physical_table_ctx.columns.len() == 2,
286            error::UnexpectedSnafu {
287                violated: "The physical table must have two columns"
288            }
289        );
290
291        // Generates the logical table columns based on the physical table.
292        let logical_table_name = self
293            .physical_table_ctx
294            .generate_unique_table_name(rng, self.name_generator.as_ref());
295        let mut logical_table = CreateTableExpr {
296            table_name: logical_table_name,
297            columns: self.physical_table_ctx.columns.clone(),
298            if_not_exists: self.if_not_exists,
299            partition: None,
300            engine: "metric".to_string(),
301            options: [(
302                "on_physical_table".to_string(),
303                self.physical_table_ctx.name.value.clone().into(),
304            )]
305            .into(),
306            primary_keys: vec![],
307        };
308
309        let column_names = self.name_generator.choose(rng, self.labels);
310        logical_table.columns.extend(generate_columns(
311            rng,
312            column_names,
313            &StringColumnTypeGenerator,
314            Box::new(primary_key_options_generator),
315        ));
316
317        // Currently only the `primary key` option is kept in physical table,
318        // so we only keep the `primary key` option in the logical table for fuzz test.
319        let mut primary_keys = vec![];
320        for (idx, column) in logical_table.columns.iter().enumerate() {
321            if column.is_primary_key() {
322                primary_keys.push(idx);
323            }
324        }
325        primary_keys.shuffle(rng);
326        logical_table.primary_keys = primary_keys;
327
328        Ok(logical_table)
329    }
330}
331
332#[derive(Builder)]
333#[builder(default, pattern = "owned")]
334pub struct CreateDatabaseExprGenerator<R: Rng + 'static> {
335    #[builder(setter(into))]
336    database_name: String,
337    name_generator: Box<dyn Random<Ident, R>>,
338    if_not_exists: bool,
339}
340
341impl<R: Rng + 'static> Default for CreateDatabaseExprGenerator<R> {
342    fn default() -> Self {
343        Self {
344            database_name: String::new(),
345            name_generator: Box::new(MappedGenerator::new(WordGenerator, random_capitalize_map)),
346            if_not_exists: false,
347        }
348    }
349}
350
351impl<R: Rng + 'static> Generator<CreateDatabaseExpr, R> for CreateDatabaseExprGenerator<R> {
352    type Error = Error;
353
354    fn generate(&self, rng: &mut R) -> Result<CreateDatabaseExpr> {
355        let mut builder = CreateDatabaseExprBuilder::default();
356        builder.if_not_exists(self.if_not_exists);
357        if self.database_name.is_empty() {
358            builder.database_name(self.name_generator.generate(rng));
359        } else {
360            builder.database_name(self.database_name.to_string());
361        }
362        builder.build().context(error::BuildCreateDatabaseExprSnafu)
363    }
364}
365
366#[cfg(test)]
367mod tests {
368    use std::sync::Arc;
369
370    use datatypes::data_type::ConcreteDataType;
371    use datatypes::value::Value;
372    use rand::SeedableRng;
373
374    use super::*;
375    use crate::context::TableContext;
376
377    #[test]
378    fn test_float64() {
379        let value = Value::from(0.047318541668048164);
380        assert_eq!("0.047318541668048164", value.to_string());
381        let value: f64 = "0.047318541668048164".parse().unwrap();
382        assert_eq!("0.047318541668048164", value.to_string());
383    }
384
385    #[test]
386    fn test_create_table_expr_generator() {
387        let mut rng = rand::rng();
388
389        let expr = CreateTableExprGeneratorBuilder::default()
390            .columns(10)
391            .partition(3)
392            .if_not_exists(true)
393            .engine("mito2")
394            .build()
395            .unwrap()
396            .generate(&mut rng)
397            .unwrap();
398        assert_eq!(expr.engine, "mito2");
399        assert!(expr.if_not_exists);
400        assert_eq!(expr.columns.len(), 10);
401        assert_eq!(expr.partition.unwrap().exprs.len(), 3);
402
403        let expr = CreateTableExprGeneratorBuilder::default()
404            .columns(10)
405            .partition(1)
406            .build()
407            .unwrap()
408            .generate(&mut rng)
409            .unwrap();
410        assert_eq!(expr.columns.len(), 10);
411        assert!(expr.partition.is_none());
412    }
413
414    #[test]
415    fn test_create_table_expr_generator_deterministic() {
416        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(0);
417        let expr = CreateTableExprGeneratorBuilder::default()
418            .columns(10)
419            .partition(3)
420            .if_not_exists(true)
421            .engine("mito2")
422            .build()
423            .unwrap()
424            .generate(&mut rng)
425            .unwrap();
426
427        let serialized = serde_json::to_string(&expr).unwrap();
428        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":["mOLEsTIAs"],"exprs":[{"lhs":{"Column":"mOLEsTIAs"},"op":"Lt","rhs":{"Value":{"Float64":5.992310449541053e307}}},{"lhs":{"Expr":{"lhs":{"Column":"mOLEsTIAs"},"op":"GtEq","rhs":{"Value":{"Float64":5.992310449541053e307}}}},"op":"And","rhs":{"Expr":{"lhs":{"Column":"mOLEsTIAs"},"op":"Lt","rhs":{"Value":{"Float64":1.1984620899082105e308}}}}},{"lhs":{"Column":"mOLEsTIAs"},"op":"GtEq","rhs":{"Value":{"Float64":1.1984620899082105e308}}}]},"engine":"mito2","options":{},"primary_keys":[0,5,9]}"#;
429        assert_eq!(expected, serialized);
430    }
431
432    #[test]
433    fn test_create_logical_table_expr_generator() {
434        let mut rng = rand::rng();
435
436        let physical_table_expr = CreatePhysicalTableExprGeneratorBuilder::default()
437            .if_not_exists(false)
438            .build()
439            .unwrap()
440            .generate(&mut rng)
441            .unwrap();
442        assert_eq!(physical_table_expr.engine, "metric");
443        assert_eq!(physical_table_expr.columns.len(), 2);
444
445        let physical_ts = physical_table_expr.columns.iter().position(|column| {
446            column
447                .options
448                .iter()
449                .any(|option| option == &ColumnOption::TimeIndex)
450        });
451        let physical_ts_name = physical_table_expr.columns[physical_ts.unwrap()]
452            .name
453            .value
454            .to_string();
455
456        let physical_table_ctx = Arc::new(TableContext::from(&physical_table_expr));
457
458        let logical_table_expr = CreateLogicalTableExprGeneratorBuilder::default()
459            .physical_table_ctx(physical_table_ctx)
460            .labels(5)
461            .if_not_exists(false)
462            .build()
463            .unwrap()
464            .generate(&mut rng)
465            .unwrap();
466        let logical_ts = logical_table_expr.columns.iter().position(|column| {
467            column
468                .options
469                .iter()
470                .any(|option| option == &ColumnOption::TimeIndex)
471        });
472        let logical_ts_name = logical_table_expr.columns[logical_ts.unwrap()]
473            .name
474            .value
475            .to_string();
476
477        assert_eq!(logical_table_expr.engine, "metric");
478        assert_eq!(logical_table_expr.columns.len(), 7);
479        assert_eq!(logical_ts_name, physical_ts_name);
480        assert!(logical_table_expr.columns.iter().all(|column| {
481            column.column_type != ConcreteDataType::string_datatype()
482                || column
483                    .options
484                    .iter()
485                    .any(|option| option == &ColumnOption::PrimaryKey)
486        }));
487    }
488
489    #[test]
490    fn test_create_logical_table_expr_generator_deterministic() {
491        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(0);
492        let physical_table_expr = CreatePhysicalTableExprGeneratorBuilder::default()
493            .if_not_exists(false)
494            .build()
495            .unwrap()
496            .generate(&mut rng)
497            .unwrap();
498        let physical_table_serialized = serde_json::to_string(&physical_table_expr).unwrap();
499        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":[]}"#;
500        assert_eq!(physical_table_expected, physical_table_serialized);
501
502        let physical_table_ctx = Arc::new(TableContext::from(&physical_table_expr));
503
504        let logical_table_expr = CreateLogicalTableExprGeneratorBuilder::default()
505            .physical_table_ctx(physical_table_ctx)
506            .labels(5)
507            .if_not_exists(false)
508            .build()
509            .unwrap()
510            .generate(&mut rng)
511            .unwrap();
512
513        let logical_table_serialized = serde_json::to_string(&logical_table_expr).unwrap();
514        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":null},"options":["PrimaryKey"]},{"name":{"value":"cumque","quote_style":null},"column_type":{"String":null},"options":["PrimaryKey"]},{"name":{"value":"natus","quote_style":null},"column_type":{"String":null},"options":["PrimaryKey"]},{"name":{"value":"molestias","quote_style":null},"column_type":{"String":null},"options":["PrimaryKey"]},{"name":{"value":"qui","quote_style":null},"column_type":{"String":null},"options":["PrimaryKey"]}],"if_not_exists":false,"partition":null,"engine":"metric","options":{"on_physical_table":{"String":"expedita"}},"primary_keys":[4,2,3,6,5]}"#;
515        assert_eq!(logical_table_expected, logical_table_serialized);
516    }
517
518    #[test]
519    fn test_create_database_expr_generator() {
520        let mut rng = rand::rng();
521
522        let expr = CreateDatabaseExprGeneratorBuilder::default()
523            .if_not_exists(true)
524            .build()
525            .unwrap()
526            .generate(&mut rng)
527            .unwrap();
528        assert!(expr.if_not_exists);
529    }
530
531    #[test]
532    fn test_create_database_expr_generator_deterministic() {
533        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(0);
534        let expr = CreateDatabaseExprGeneratorBuilder::default()
535            .if_not_exists(true)
536            .build()
537            .unwrap()
538            .generate(&mut rng)
539            .unwrap();
540
541        let serialized = serde_json::to_string(&expr).unwrap();
542        let expected =
543            r#"{"database_name":{"value":"EXPediTA","quote_style":null},"if_not_exists":true}"#;
544        assert_eq!(expected, serialized);
545    }
546}
tests_fuzz/generator/create_expr.rs

tests_fuzz/generator/
create_expr.rs
