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
196fn generate_metric_partition(partitions: usize) -> Option<(Column, PartitionDef)> {
197 if partitions <= 1 {
198 return None;
199 }
200
201 let partition_column = Column {
202 name: Ident::new("host"),
203 column_type: ConcreteDataType::string_datatype(),
204 options: vec![ColumnOption::PrimaryKey],
205 };
206 let bounds = generate_partition_bounds(&partition_column.column_type, partitions - 1);
207 let partitions = SimplePartitions::new(partition_column.name.clone(), bounds);
208 let partition_def = PartitionDef {
209 columns: vec![partitions.column_name.clone()],
210 exprs: partitions.generate().unwrap(),
211 };
212
213 Some((partition_column, partition_def))
214}
215
216#[derive(Builder)]
218#[builder(pattern = "owned")]
219pub struct CreatePhysicalTableExprGenerator<R: Rng + 'static> {
220 #[builder(default = "Box::new(WordGenerator)")]
221 name_generator: Box<dyn Random<Ident, R>>,
222 #[builder(default = "false")]
223 if_not_exists: bool,
224 #[builder(default = "0")]
225 partition: usize,
226 #[builder(default, setter(into))]
227 with_clause: HashMap<String, String>,
228}
229
230impl<R: Rng + 'static> Generator<CreateTableExpr, R> for CreatePhysicalTableExprGenerator<R> {
231 type Error = Error;
232
233 fn generate(&self, rng: &mut R) -> Result<CreateTableExpr> {
234 let mut options = HashMap::with_capacity(self.with_clause.len() + 1);
235 options.insert("physical_metric_table".to_string(), Value::from(""));
236 for (key, value) in &self.with_clause {
237 options.insert(key.clone(), Value::from(value.clone()));
238 }
239
240 let mut columns = vec![
241 Column {
242 name: Ident::new("ts"),
243 column_type: ConcreteDataType::timestamp_millisecond_datatype(),
244 options: vec![ColumnOption::TimeIndex],
245 },
246 Column {
247 name: Ident::new("val"),
248 column_type: ConcreteDataType::float64_datatype(),
249 options: vec![],
250 },
251 ];
252
253 let mut partition = None;
254 let mut primary_keys = vec![];
255 if let Some((partition_column, partition_def)) = generate_metric_partition(self.partition) {
256 columns.push(partition_column);
257 partition = Some(partition_def);
258 primary_keys.push(columns.len() - 1);
259 }
260
261 Ok(CreateTableExpr {
262 table_name: self.name_generator.generate(rng),
263 columns,
264 if_not_exists: self.if_not_exists,
265 partition,
266 engine: "metric".to_string(),
267 options,
268 primary_keys,
269 })
270 }
271}
272
273#[derive(Builder)]
275#[builder(pattern = "owned")]
276pub struct CreateLogicalTableExprGenerator<R: Rng + 'static> {
277 physical_table_ctx: TableContextRef,
278 labels: usize,
279 if_not_exists: bool,
280 #[builder(default = "true")]
281 include_partition_column: bool,
282 #[builder(default = "Box::new(WordGenerator)")]
283 name_generator: Box<dyn Random<Ident, R>>,
284}
285
286impl<R: Rng + 'static> Generator<CreateTableExpr, R> for CreateLogicalTableExprGenerator<R> {
287 type Error = Error;
288
289 fn generate(&self, rng: &mut R) -> Result<CreateTableExpr> {
290 ensure!(
292 self.physical_table_ctx.columns.len() >= 2,
293 error::UnexpectedSnafu {
294 violated: "The physical table must have at least two columns"
295 }
296 );
297
298 let logical_table_name = self
300 .physical_table_ctx
301 .generate_unique_table_name(rng, self.name_generator.as_ref());
302 let mut physical_columns = self.physical_table_ctx.columns.clone();
303 if !self.include_partition_column
304 && let Some(partition_def) = &self.physical_table_ctx.partition
305 {
306 physical_columns.retain(|column| !partition_def.columns.contains(&column.name));
307 }
308
309 let mut logical_table = CreateTableExpr {
310 table_name: logical_table_name,
311 columns: physical_columns,
312 if_not_exists: self.if_not_exists,
313 partition: None,
314 engine: "metric".to_string(),
315 options: [(
316 "on_physical_table".to_string(),
317 self.physical_table_ctx.name.value.clone().into(),
318 )]
319 .into(),
320 primary_keys: vec![],
321 };
322
323 let column_names = self.name_generator.choose(rng, self.labels);
324 logical_table.columns.extend(generate_columns(
325 rng,
326 column_names,
327 &StringColumnTypeGenerator,
328 Box::new(primary_key_options_generator),
329 ));
330
331 let mut primary_keys = vec![];
334 for (idx, column) in logical_table.columns.iter().enumerate() {
335 if column.is_primary_key() {
336 primary_keys.push(idx);
337 }
338 }
339 primary_keys.shuffle(rng);
340 logical_table.primary_keys = primary_keys;
341
342 Ok(logical_table)
343 }
344}
345
346#[derive(Builder)]
347#[builder(default, pattern = "owned")]
348pub struct CreateDatabaseExprGenerator<R: Rng + 'static> {
349 #[builder(setter(into))]
350 database_name: String,
351 name_generator: Box<dyn Random<Ident, R>>,
352 if_not_exists: bool,
353}
354
355impl<R: Rng + 'static> Default for CreateDatabaseExprGenerator<R> {
356 fn default() -> Self {
357 Self {
358 database_name: String::new(),
359 name_generator: Box::new(MappedGenerator::new(WordGenerator, random_capitalize_map)),
360 if_not_exists: false,
361 }
362 }
363}
364
365impl<R: Rng + 'static> Generator<CreateDatabaseExpr, R> for CreateDatabaseExprGenerator<R> {
366 type Error = Error;
367
368 fn generate(&self, rng: &mut R) -> Result<CreateDatabaseExpr> {
369 let mut builder = CreateDatabaseExprBuilder::default();
370 builder.if_not_exists(self.if_not_exists);
371 if self.database_name.is_empty() {
372 builder.database_name(self.name_generator.generate(rng));
373 } else {
374 builder.database_name(self.database_name.clone());
375 }
376 builder.build().context(error::BuildCreateDatabaseExprSnafu)
377 }
378}
379
380#[cfg(test)]
381mod tests {
382 use std::sync::Arc;
383
384 use datatypes::data_type::ConcreteDataType;
385 use datatypes::value::Value;
386 use rand::SeedableRng;
387
388 use super::*;
389 use crate::context::TableContext;
390
391 #[test]
392 fn test_float64() {
393 let value = Value::from(0.047318541668048164);
394 assert_eq!("0.047318541668048164", value.to_string());
395 let value: f64 = "0.047318541668048164".parse().unwrap();
396 assert_eq!("0.047318541668048164", value.to_string());
397 }
398
399 #[test]
400 fn test_create_table_expr_generator() {
401 let mut rng = rand::rng();
402
403 let expr = CreateTableExprGeneratorBuilder::default()
404 .columns(10)
405 .partition(3)
406 .if_not_exists(true)
407 .engine("mito2")
408 .build()
409 .unwrap()
410 .generate(&mut rng)
411 .unwrap();
412 assert_eq!(expr.engine, "mito2");
413 assert!(expr.if_not_exists);
414 assert_eq!(expr.columns.len(), 10);
415 assert_eq!(expr.partition.unwrap().exprs.len(), 3);
416
417 let expr = CreateTableExprGeneratorBuilder::default()
418 .columns(10)
419 .partition(1)
420 .build()
421 .unwrap()
422 .generate(&mut rng)
423 .unwrap();
424 assert_eq!(expr.columns.len(), 10);
425 assert!(expr.partition.is_none());
426 }
427
428 #[test]
429 fn test_create_table_expr_generator_deterministic() {
430 let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(0);
431 let expr = CreateTableExprGeneratorBuilder::default()
432 .columns(10)
433 .partition(3)
434 .if_not_exists(true)
435 .engine("mito2")
436 .build()
437 .unwrap()
438 .generate(&mut rng)
439 .unwrap();
440
441 let serialized = serde_json::to_string(&expr).unwrap();
442 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]}"#;
443 assert_eq!(expected, serialized);
444 }
445
446 #[test]
447 fn test_create_logical_table_expr_generator() {
448 let mut rng = rand::rng();
449
450 let physical_table_expr = CreatePhysicalTableExprGeneratorBuilder::default()
451 .if_not_exists(false)
452 .build()
453 .unwrap()
454 .generate(&mut rng)
455 .unwrap();
456 assert_eq!(physical_table_expr.engine, "metric");
457 assert_eq!(physical_table_expr.columns.len(), 2);
458
459 let physical_ts = physical_table_expr.columns.iter().position(|column| {
460 column
461 .options
462 .iter()
463 .any(|option| option == &ColumnOption::TimeIndex)
464 });
465 let physical_ts_name = physical_table_expr.columns[physical_ts.unwrap()]
466 .name
467 .value
468 .clone();
469
470 let physical_table_ctx = Arc::new(TableContext::from(&physical_table_expr));
471
472 let logical_table_expr = CreateLogicalTableExprGeneratorBuilder::default()
473 .physical_table_ctx(physical_table_ctx)
474 .labels(5)
475 .if_not_exists(false)
476 .build()
477 .unwrap()
478 .generate(&mut rng)
479 .unwrap();
480 let logical_ts = logical_table_expr.columns.iter().position(|column| {
481 column
482 .options
483 .iter()
484 .any(|option| option == &ColumnOption::TimeIndex)
485 });
486 let logical_ts_name = logical_table_expr.columns[logical_ts.unwrap()]
487 .name
488 .value
489 .clone();
490
491 assert_eq!(logical_table_expr.engine, "metric");
492 assert_eq!(logical_table_expr.columns.len(), 7);
493 assert_eq!(logical_ts_name, physical_ts_name);
494 assert!(logical_table_expr.columns.iter().all(|column| {
495 column.column_type != ConcreteDataType::string_datatype()
496 || column
497 .options
498 .iter()
499 .any(|option| option == &ColumnOption::PrimaryKey)
500 }));
501 }
502
503 #[test]
504 fn test_create_physical_table_expr_generator_with_partition() {
505 let mut rng = rand::rng();
506 let physical_table_expr = CreatePhysicalTableExprGeneratorBuilder::default()
507 .partition(3)
508 .if_not_exists(false)
509 .build()
510 .unwrap()
511 .generate(&mut rng)
512 .unwrap();
513
514 assert_eq!(physical_table_expr.engine, "metric");
515 assert!(physical_table_expr.partition.is_some());
516 assert_eq!(physical_table_expr.partition.unwrap().exprs.len(), 3);
517 }
518
519 #[test]
520 fn test_create_logical_table_expr_generator_without_partition_column() {
521 let mut rng = rand::rng();
522 let physical_table_expr = CreatePhysicalTableExprGeneratorBuilder::default()
523 .partition(3)
524 .if_not_exists(false)
525 .build()
526 .unwrap()
527 .generate(&mut rng)
528 .unwrap();
529 let partition_columns = physical_table_expr
530 .partition
531 .as_ref()
532 .unwrap()
533 .columns
534 .clone();
535 let physical_table_ctx = Arc::new(TableContext::from(&physical_table_expr));
536
537 let logical_table_expr = CreateLogicalTableExprGeneratorBuilder::default()
538 .physical_table_ctx(physical_table_ctx)
539 .labels(3)
540 .include_partition_column(false)
541 .if_not_exists(false)
542 .build()
543 .unwrap()
544 .generate(&mut rng)
545 .unwrap();
546
547 assert!(
548 logical_table_expr
549 .columns
550 .iter()
551 .all(|column| !partition_columns.contains(&column.name))
552 );
553 }
554
555 #[test]
556 fn test_create_logical_table_expr_generator_deterministic() {
557 let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(0);
558 let physical_table_expr = CreatePhysicalTableExprGeneratorBuilder::default()
559 .if_not_exists(false)
560 .build()
561 .unwrap()
562 .generate(&mut rng)
563 .unwrap();
564 let physical_table_serialized = serde_json::to_string(&physical_table_expr).unwrap();
565 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":[]}"#;
566 assert_eq!(physical_table_expected, physical_table_serialized);
567
568 let physical_table_ctx = Arc::new(TableContext::from(&physical_table_expr));
569
570 let logical_table_expr = CreateLogicalTableExprGeneratorBuilder::default()
571 .physical_table_ctx(physical_table_ctx)
572 .labels(5)
573 .if_not_exists(false)
574 .build()
575 .unwrap()
576 .generate(&mut rng)
577 .unwrap();
578
579 let logical_table_serialized = serde_json::to_string(&logical_table_expr).unwrap();
580 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]}"#;
581 assert_eq!(logical_table_expected, logical_table_serialized);
582 }
583
584 #[test]
585 fn test_create_database_expr_generator() {
586 let mut rng = rand::rng();
587
588 let expr = CreateDatabaseExprGeneratorBuilder::default()
589 .if_not_exists(true)
590 .build()
591 .unwrap()
592 .generate(&mut rng)
593 .unwrap();
594 assert!(expr.if_not_exists);
595 }
596
597 #[test]
598 fn test_create_database_expr_generator_deterministic() {
599 let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(0);
600 let expr = CreateDatabaseExprGeneratorBuilder::default()
601 .if_not_exists(true)
602 .build()
603 .unwrap()
604 .generate(&mut rng)
605 .unwrap();
606
607 let serialized = serde_json::to_string(&expr).unwrap();
608 let expected =
609 r#"{"database_name":{"value":"EXPediTA","quote_style":null},"if_not_exists":true}"#;
610 assert_eq!(expected, serialized);
611 }
612}