1use 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 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 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 if need_partible_column {
116 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 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 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 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 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#[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#[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 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 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 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}