rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2025-12-22 22:20:02 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: add create alter table expr translator (#3203)

Browse Source 
* feat: add create table expr translator

* feat: add alter table expr translator

* refactor: expose mod

* refactor: expr generator

* chore: ignore typos check for lorem_words

* feat: add string map helper functions

* chore: remove unit tests
This commit is contained in:
Weny Xu
2024-01-23 11:53:42 +09:00
committed by GitHub
parent 31787f4bfd
commit 364754afa2
17 changed files with 1231 additions and 262 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
Cargo.lock generated
View File

@@ -2875,12 +2875,6 @@ dependencies = [
"syn 1.0.109",
]
[[package]]
name = "deunicode"
version = "1.4.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3ae2a35373c5c74340b79ae6780b498b2b183915ec5dacf263aac5a099bf485a"
[[package]]
name = "diff"
version = "0.1.13"
@@ -3154,17 +3148,6 @@ dependencies = [
"rand",
]
[[package]]
name = "faker_rand"
version = "0.1.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "300d2ddbf2245b5b5e723995e0961033121b4fc2be9045fb661af82bd739ffb6"
dependencies = [
"deunicode",
"lazy_static",
"rand",
]
[[package]]
name = "fallible-iterator"
version = "0.2.0"
@@ -9551,15 +9534,21 @@ dependencies = [
"common-error",
"common-macro",
"common-query",
"common-telemetry",
"datatypes",
"derive_builder 0.12.0",
"faker_rand",
"dotenv",
"lazy_static",
"partition",
"rand",
"rand_chacha",
"serde",
"serde_json",
"snafu",
"sql",
"sqlparser 0.38.0 (git+https://github.com/GreptimeTeam/sqlparser-rs.git?rev=6a93567ae38d42be5c8d08b13c8ff4dde26502ef)",
"sqlx",
"tokio",
]
[[package]]

15
tests-fuzz/Cargo.toml
View File

@@ -9,12 +9,25 @@ async-trait = { workspace = true }
common-error = { workspace = true }
common-macro = { workspace = true }
common-query = { workspace = true }
common-telemetry = { workspace = true }
datatypes = { workspace = true }
derive_builder = { workspace = true }
faker_rand = "0.1"
lazy_static = { workspace = true }
partition = { workspace = true }
rand = { workspace = true }
rand_chacha = "0.3.1"
serde = { workspace = true }
serde_json = { workspace = true }
snafu = { workspace = true }
sql = { workspace = true }
sqlparser.workspace = true
[dev-dependencies]
dotenv = "0.15"
sqlx = { version = "0.6", features = [
"runtime-tokio-rustls",
"mysql",
"postgres",
"chrono",
] }
tokio = { workspace = true }

6
tests-fuzz/src/context.rs
View File

@@ -26,7 +26,7 @@ pub struct TableContext {
pub columns: Vec<Column>,
// GreptimeDB specific options
pub partitions: Vec<PartitionDef>,
pub partition: Option<PartitionDef>,
pub primary_keys: Vec<usize>,
}
@@ -35,7 +35,7 @@ impl From<&CreateTableExpr> for TableContext {
CreateTableExpr {
name,
columns,
partitions,
partition,
primary_keys,
..
}: &CreateTableExpr,
@@ -43,7 +43,7 @@ impl From<&CreateTableExpr> for TableContext {
Self {
name: name.to_string(),
columns: columns.clone(),
partitions: partitions.clone(),
partition: partition.clone(),
primary_keys: primary_keys.clone(),
}
}

249
tests-fuzz/src/data/lorem_words Normal file
View File

@@ -0,0 +1,249 @@
alias
consequatur
aut
perferendis
sit
voluptatem
accusantium
doloremque
aperiam
eaque
ipsa
quae
ab
illo
inventore
veritatis
et
quasi
architecto
beatae
vitae
dicta
sunt
explicabo
aspernatur
aut
odit
aut
fugit
sed
quia
consequuntur
magni
dolores
eos
qui
ratione
voluptatem
sequi
nesciunt
neque
dolorem
ipsum
quia
dolor
sit
amet
consectetur
adipisci
velit
sed
quia
non
numquam
eius
modi
tempora
incidunt
ut
labore
et
dolore
magnam
aliquam
quaerat
voluptatem
ut
enim
ad
minima
veniam
quis
nostrum
exercitationem
ullam
corporis
nemo
enim
ipsam
voluptatem
quia
voluptas
sit
suscipit
laboriosam
nisi
ut
aliquid
ex
ea
commodi
consequatur
quis
autem
vel
eum
iure
reprehenderit
qui
in
ea
voluptate
velit
esse
quam
nihil
molestiae
et
iusto
odio
dignissimos
ducimus
qui
blanditiis
praesentium
laudantium
totam
rem
voluptatum
deleniti
atque
corrupti
quos
dolores
et
quas
molestias
excepturi
sint
occaecati
cupiditate
non
provident
sed
ut
perspiciatis
unde
omnis
iste
natus
error
similique
sunt
in
culpa
qui
officia
deserunt
mollitia
animi
id
est
laborum
et
dolorum
fuga
et
harum
quidem
rerum
facilis
est
et
expedita
distinctio
nam
libero
tempore
cum
soluta
nobis
est
eligendi
optio
cumque
nihil
impedit
quo
porro
quisquam
est
qui
minus
id
quod
maxime
placeat
facere
possimus
omnis
voluptas
assumenda
est
omnis
dolor
repellendus
temporibus
autem
quibusdam
et
aut
consequatur
vel
illum
qui
dolorem
eum
fugiat
quo
voluptas
nulla
pariatur
at
vero
eos
et
accusamus
officiis
debitis
aut
rerum
necessitatibus
saepe
eveniet
ut
et
voluptates
repudiandae
sint
et
molestiae
non
recusandae
itaque
earum
rerum
hic
tenetur
a
sapiente
delectus
ut
aut
reiciendis
voluptatibus
maiores
doloribus
asperiores
repellat

174
tests-fuzz/src/fake.rs Normal file
View File

@@ -0,0 +1,174 @@
// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use std::collections::HashSet;
use lazy_static::lazy_static;
use rand::seq::{IteratorRandom, SliceRandom};
use rand::Rng;
use crate::generator::Random;
use crate::impl_random;
lazy_static! {
pub static ref LOREM_WORDS: Vec<String> = include_str!("data/lorem_words")
.lines()
.map(String::from)
.collect();
}
/// Modified from https://github.com/ucarion/faker_rand/blob/ea70c660e1ecd7320156eddb31d2830a511f8842/src/lib.rs
macro_rules! faker_impl_from_values {
($name: ident, $values: expr) => {
impl rand::distributions::Distribution<$name> for rand::distributions::Standard {
fn sample<R: rand::Rng + ?Sized>(&self, rng: &mut R) -> $name {
$name($values[rng.gen_range(0..$values.len())].clone())
}
}
impl std::fmt::Display for $name {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.0)
}
}
};
}
pub struct Word(String);
faker_impl_from_values!(Word, LOREM_WORDS);
pub struct WordGenerator;
impl_random!(String, WordGenerator, LOREM_WORDS);
pub type WordMapFn<R> = Box<dyn Fn(&mut R, String) -> String>;
pub struct MapWordGenerator<R: Rng> {
base: WordGenerator,
map: WordMapFn<R>,
}
pub fn random_capitalize_map<R: Rng + 'static>(rng: &mut R, s: String) -> String {
let mut v = s.chars().collect::<Vec<_>>();
let select = rng.gen_range(0..s.len());
for idx in (0..s.len()).choose_multiple(rng, select) {
v[idx] = v[idx].to_uppercase().next().unwrap();
}
v.into_iter().collect::<String>()
}
lazy_static! {
static ref KEYWORDS_SET: HashSet<&'static str> = sqlparser::keywords::ALL_KEYWORDS
.iter()
.cloned()
.collect::<HashSet<_>>();
}
/// Returns true if it's a keyword.
pub fn is_keyword(word: impl AsRef<str>) -> bool {
KEYWORDS_SET.contains(word.as_ref())
}
/// Returns true if it contains uppercase char.
pub fn contain_uppercase_char(s: &str) -> bool {
s.chars().any(|c| c.is_uppercase())
}
/// Returns true if it's a keyword or contains uppercase char.
pub fn is_keyword_or_contain_uppercase(s: &str) -> bool {
is_keyword(s.to_uppercase()) || contain_uppercase_char(s)
}
pub fn make_backtick_map<R: Rng + 'static, F: Fn(&str) -> bool>(
f: F,
) -> impl Fn(&mut R, String) -> String {
move |_rng, s| -> String {
let need = f(&s);
if need {
format!("`{s}`")
} else {
s
}
}
}
pub fn make_quote_map<R: Rng + 'static, F: Fn(&str) -> bool>(
f: F,
) -> impl Fn(&mut R, String) -> String {
move |_rng, s| -> String {
let need = f(&s);
if need {
format!("\"{s}\"")
} else {
s
}
}
}
/// Adds backticks if it contains uppercase chars.
pub fn auto_backtick_map<R: Rng + 'static>(_rng: &mut R, s: String) -> String {
let need = s.chars().any(|c| c.is_uppercase());
if need {
format!("`{s}`")
} else {
s
}
}
/// Adds backticks if it contains uppercase chars.
pub fn uppercase_and_keyword_backtick_map<R: Rng + 'static>(rng: &mut R, s: String) -> String {
make_backtick_map(is_keyword_or_contain_uppercase)(rng, s)
}
/// Adds quotes if it contains uppercase chars.
pub fn auto_quote_map<R: Rng + 'static>(rng: &mut R, s: String) -> String {
make_quote_map(contain_uppercase_char)(rng, s)
}
/// Adds quotes if it contains uppercase chars.
pub fn uppercase_and_keyword_quote_map<R: Rng + 'static>(rng: &mut R, s: String) -> String {
make_quote_map(is_keyword_or_contain_uppercase)(rng, s)
}
pub fn merge_two_word_map_fn<R: Rng>(
f1: impl Fn(&mut R, String) -> String,
f2: impl Fn(&mut R, String) -> String,
) -> impl Fn(&mut R, String) -> String {
move |rng, s| -> String {
let s = f1(rng, s);
f2(rng, s)
}
}
impl<R: Rng> MapWordGenerator<R> {
pub fn new(map: WordMapFn<R>) -> Self {
Self {
base: WordGenerator,
map,
}
}
}
impl<R: Rng> Random<String, R> for MapWordGenerator<R> {
fn choose(&self, rng: &mut R, amount: usize) -> Vec<String> {
self.base
.choose(rng, amount)
.into_iter()
.map(|s| (self.map)(rng, s))
.collect()
}
}

40
tests-fuzz/src/generator.rs
View File

@@ -17,16 +17,46 @@ pub mod create_expr;
use std::fmt;
use datatypes::data_type::ConcreteDataType;
use rand::Rng;
use crate::error::Error;
use crate::ir::create_expr::ColumnOption;
use crate::ir::{AlterTableExpr, CreateTableExpr};
pub type CreateTableExprGenerator =
Box<dyn Generator<CreateTableExpr, Error = Error> + Sync + Send>;
pub type CreateTableExprGenerator<R> =
Box<dyn Generator<CreateTableExpr, R, Error = Error> + Sync + Send>;
pub type AlterTableExprGenerator = Box<dyn Generator<AlterTableExpr, Error = Error> + Sync + Send>;
pub type AlterTableExprGenerator<R> =
Box<dyn Generator<AlterTableExpr, R, Error = Error> + Sync + Send>;
pub(crate) trait Generator<T> {
pub type ColumnOptionGenerator<R> = Box<dyn Fn(&mut R, &ConcreteDataType) -> Vec<ColumnOption>>;
pub type ConcreteDataTypeGenerator<R> = Box<dyn Random<ConcreteDataType, R>>;
pub trait Generator<T, R: Rng> {
type Error: Sync + Send + fmt::Debug;
fn generate(&self) -> Result<T, Self::Error>;
fn generate(&self, rng: &mut R) -> Result<T, Self::Error>;
}
pub trait Random<T, R: Rng> {
/// Generates a random element.
fn gen(&self, rng: &mut R) -> T {
self.choose(rng, 1).remove(0)
}
/// Uniformly sample `amount` distinct elements.
fn choose(&self, rng: &mut R, amount: usize) -> Vec<T>;
}
#[macro_export]
macro_rules! impl_random {
($type: ident, $value:ident, $values: ident) => {
impl<R: Rng> Random<$type, R> for $value {
fn choose(&self, rng: &mut R, amount: usize) -> Vec<$type> {
$values.choose_multiple(rng, amount).cloned().collect()
}
}
};
}

142
tests-fuzz/src/generator/alter_expr.rs
View File

@@ -12,47 +12,41 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::f64::consts::E;
use std::sync::Arc;
use std::marker::PhantomData;
use common_query::AddColumnLocation;
use faker_rand::lorem::Word;
use rand::{random, Rng};
use derive_builder::Builder;
use rand::Rng;
use snafu::ensure;
use crate::context::TableContextRef;
use crate::error::{self, Error, Result};
use crate::generator::Generator;
use crate::fake::WordGenerator;
use crate::generator::{ColumnOptionGenerator, ConcreteDataTypeGenerator, Generator, Random};
use crate::ir::alter_expr::{AlterTableExpr, AlterTableOperation};
use crate::ir::{droppable_columns, Column};
use crate::ir::{
column_options_generator, droppable_columns, generate_columns, ColumnTypeGenerator,
};
/// Generates the [AlterTableOperation::AddColumn] of [AlterTableExpr].
pub struct AlterExprAddColumnGenerator {
#[derive(Builder)]
#[builder(pattern = "owned")]
pub struct AlterExprAddColumnGenerator<R: Rng + 'static> {
table_ctx: TableContextRef,
#[builder(default)]
location: bool,
#[builder(default = "Box::new(WordGenerator)")]
name_generator: Box<dyn Random<String, R>>,
#[builder(default = "Box::new(column_options_generator)")]
column_options_generator: ColumnOptionGenerator<R>,
#[builder(default = "Box::new(ColumnTypeGenerator)")]
column_type_generator: ConcreteDataTypeGenerator<R>,
}
impl AlterExprAddColumnGenerator {
/// Returns an [AlterExprAddColumnGenerator].
pub fn new(table_ctx: &TableContextRef) -> Self {
Self {
table_ctx: table_ctx.clone(),
location: false,
}
}
/// Sets true to generate alter expr with a specific location.
pub fn with_location(mut self, v: bool) -> Self {
self.location = v;
self
}
}
impl Generator<AlterTableExpr> for AlterExprAddColumnGenerator {
impl<R: Rng + 'static> Generator<AlterTableExpr, R> for AlterExprAddColumnGenerator<R> {
type Error = Error;
fn generate(&self) -> Result<AlterTableExpr> {
let mut rng = rand::thread_rng();
fn generate(&self, rng: &mut R) -> Result<AlterTableExpr> {
let with_location = self.location && rng.gen::<bool>();
let location = if with_location {
let use_first = rng.gen::<bool>();
@@ -71,7 +65,14 @@ impl Generator<AlterTableExpr> for AlterExprAddColumnGenerator {
None
};
let column = rng.gen::<Column>();
let name = self.name_generator.gen(rng);
let column = generate_columns(
rng,
vec![name],
self.column_type_generator.as_ref(),
self.column_options_generator.as_ref(),
)
.remove(0);
Ok(AlterTableExpr {
name: self.table_ctx.name.to_string(),
alter_options: AlterTableOperation::AddColumn { column, location },
@@ -80,24 +81,18 @@ impl Generator<AlterTableExpr> for AlterExprAddColumnGenerator {
}
/// Generates the [AlterTableOperation::DropColumn] of [AlterTableExpr].
pub struct AlterExprDropColumnGenerator {
#[derive(Builder)]
#[builder(pattern = "owned")]
pub struct AlterExprDropColumnGenerator<R> {
table_ctx: TableContextRef,
#[builder(default)]
_phantom: PhantomData<R>,
}
impl AlterExprDropColumnGenerator {
/// Returns an [AlterExprDropColumnGenerator].
pub fn new(table_ctx: &TableContextRef) -> Self {
Self {
table_ctx: table_ctx.clone(),
}
}
}
impl Generator<AlterTableExpr> for AlterExprDropColumnGenerator {
impl<R: Rng> Generator<AlterTableExpr, R> for AlterExprDropColumnGenerator<R> {
type Error = Error;
fn generate(&self) -> Result<AlterTableExpr> {
let mut rng = rand::thread_rng();
fn generate(&self, rng: &mut R) -> Result<AlterTableExpr> {
let droppable = droppable_columns(&self.table_ctx.columns);
ensure!(!droppable.is_empty(), error::DroppableColumnsSnafu);
let name = droppable[rng.gen_range(0..droppable.len())]
@@ -110,25 +105,20 @@ impl Generator<AlterTableExpr> for AlterExprDropColumnGenerator {
}
}
pub struct AlterExprRenameGenerator {
/// Generates the [AlterTableOperation::RenameTable] of [AlterTableExpr].
#[derive(Builder)]
#[builder(pattern = "owned")]
pub struct AlterExprRenameGenerator<R: Rng> {
table_ctx: TableContextRef,
#[builder(default = "Box::new(WordGenerator)")]
name_generator: Box<dyn Random<String, R>>,
}
impl AlterExprRenameGenerator {
/// Returns an [AlterExprRenameGenerator].
pub fn new(table_ctx: &TableContextRef) -> Self {
Self {
table_ctx: table_ctx.clone(),
}
}
}
impl Generator<AlterTableExpr> for AlterExprRenameGenerator {
impl<R: Rng> Generator<AlterTableExpr, R> for AlterExprRenameGenerator<R> {
type Error = Error;
fn generate(&self) -> Result<AlterTableExpr> {
let mut rng = rand::thread_rng();
let mut new_table_name = rng.gen::<Word>().to_string();
fn generate(&self, rng: &mut R) -> Result<AlterTableExpr> {
let new_table_name = self.name_generator.gen(rng);
Ok(AlterTableExpr {
name: self.table_ctx.name.to_string(),
alter_options: AlterTableOperation::RenameTable { new_table_name },
@@ -140,29 +130,53 @@ impl Generator<AlterTableExpr> for AlterExprRenameGenerator {
mod tests {
use std::sync::Arc;
use rand::SeedableRng;
use super::*;
use crate::context::TableContext;
use crate::generator::create_expr::CreateTableExprGenerator;
use crate::generator::create_expr::CreateTableExprGeneratorBuilder;
use crate::generator::Generator;
#[test]
fn test_alter_table_expr_generator() {
let create_expr = CreateTableExprGenerator::default()
fn test_alter_table_expr_generator_deterministic() {
let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(0);
let create_expr = CreateTableExprGeneratorBuilder::default()
.columns(10)
.generate()
.build()
.unwrap()
.generate(&mut rng)
.unwrap();
let table_ctx = Arc::new(TableContext::from(&create_expr));
let alter_expr = AlterExprAddColumnGenerator::new(&table_ctx)
.generate()
let expr = AlterExprAddColumnGeneratorBuilder::default()
.table_ctx(table_ctx.clone())
.build()
.unwrap()
.generate(&mut rng)
.unwrap();
let serialized = serde_json::to_string(&expr).unwrap();
let expected = r#"{"name":"DigNissIMOS","alter_options":{"AddColumn":{"column":{"name":"sit","column_type":{"Boolean":null},"options":["PrimaryKey"]},"location":null}}}"#;
assert_eq!(expected, serialized);
let alter_expr = AlterExprRenameGenerator::new(&table_ctx)
.generate()
let expr = AlterExprRenameGeneratorBuilder::default()
.table_ctx(table_ctx.clone())
.build()
.unwrap()
.generate(&mut rng)
.unwrap();
let serialized = serde_json::to_string(&expr).unwrap();
let expected = r#"{"name":"DigNissIMOS","alter_options":{"RenameTable":{"new_table_name":"excepturi"}}}"#;
assert_eq!(expected, serialized);
let alter_expr = AlterExprDropColumnGenerator::new(&table_ctx)
.generate()
let expr = AlterExprDropColumnGeneratorBuilder::default()
.table_ctx(table_ctx)
.build()
.unwrap()
.generate(&mut rng)
.unwrap();
let serialized = serde_json::to_string(&expr).unwrap();
let expected =
r#"{"name":"DigNissIMOS","alter_options":{"DropColumn":{"name":"INVentORE"}}}"#;
assert_eq!(expected, serialized);
}
}

248
tests-fuzz/src/generator/create_expr.rs
View File

@@ -12,7 +12,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use faker_rand::lorem::Word;
use derive_builder::Builder;
use partition::partition::{PartitionBound, PartitionDef};
use rand::seq::SliceRandom;
use rand::Rng;
@@ -20,61 +20,60 @@ use snafu::{ensure, ResultExt};
use super::Generator;
use crate::error::{self, Error, Result};
use crate::ir::create_expr::{ColumnOption, CreateTableExprBuilder};
use crate::fake::{random_capitalize_map, MapWordGenerator};
use crate::generator::{ColumnOptionGenerator, ConcreteDataTypeGenerator, Random};
use crate::ir::create_expr::CreateTableExprBuilder;
use crate::ir::{
self, generate_random_value, Column, CreateTableExpr, PartibleColumn, TsColumn,
PARTIBLE_DATA_TYPES,
column_options_generator, generate_columns, generate_random_value,
partible_column_options_generator, ts_column_options_generator, ColumnTypeGenerator,
CreateTableExpr, PartibleColumnTypeGenerator, TsColumnTypeGenerator,
};
pub struct CreateTableExprGenerator {
#[derive(Builder)]
#[builder(default, pattern = "owned")]
pub struct CreateTableExprGenerator<R: Rng + 'static> {
columns: usize,
#[builder(setter(into))]
engine: String,
partition: usize,
if_not_exists: bool,
#[builder(setter(into))]
name: String,
name_generator: Box<dyn Random<String, R>>,
ts_column_type_generator: ConcreteDataTypeGenerator<R>,
column_type_generator: ConcreteDataTypeGenerator<R>,
partible_column_type_generator: ConcreteDataTypeGenerator<R>,
partible_column_options_generator: ColumnOptionGenerator<R>,
column_options_generator: ColumnOptionGenerator<R>,
ts_column_options_generator: ColumnOptionGenerator<R>,
}
const DEFAULT_ENGINE: &str = "mito";
impl Default for CreateTableExprGenerator {
impl<R: Rng + 'static> Default for CreateTableExprGenerator<R> {
fn default() -> Self {
Self {
columns: 0,
engine: DEFAULT_ENGINE.to_string(),
if_not_exists: false,
partition: 0,
name: String::new(),
name_generator: Box::new(MapWordGenerator::new(Box::new(random_capitalize_map))),
ts_column_type_generator: Box::new(TsColumnTypeGenerator),
column_type_generator: Box::new(ColumnTypeGenerator),
partible_column_type_generator: Box::new(PartibleColumnTypeGenerator),
partible_column_options_generator: Box::new(partible_column_options_generator),
column_options_generator: Box::new(column_options_generator),
ts_column_options_generator: Box::new(ts_column_options_generator),
}
}
}
impl CreateTableExprGenerator {
/// Sets column num.
pub fn columns(mut self, columns: usize) -> Self {
self.columns = columns;
self
}
/// Sets the `if_not_exists`.
pub fn create_is_not_exists(mut self, v: bool) -> Self {
self.if_not_exists = v;
self
}
/// Sets the `engine`.
pub fn engine(mut self, engine: &str) -> Self {
self.engine = engine.to_string();
self
}
/// Sets the partition num.
/// If there is no primary key column,
/// it appends a primary key atomically.
pub fn partitions(mut self, partition: usize) -> Self {
self.partition = partition;
self
}
impl<R: Rng + 'static> Generator<CreateTableExpr, R> for CreateTableExprGenerator<R> {
type Error = Error;
/// Generates the [CreateTableExpr].
pub fn generate(&self) -> Result<CreateTableExpr> {
fn generate(&self, rng: &mut R) -> Result<CreateTableExpr> {
ensure!(
self.columns != 0,
error::UnexpectedSnafu {
@@ -83,100 +82,153 @@ impl CreateTableExprGenerator {
);
let mut builder = CreateTableExprBuilder::default();
let mut columns = Vec::with_capacity(self.columns + 1);
let mut columns = Vec::with_capacity(self.columns);
let mut primary_keys = vec![];
let mut rng = rand::thread_rng();
let need_partible_column = self.partition > 1;
let mut column_names = self.name_generator.choose(rng, self.columns);
// Generates columns.
for i in 0..self.columns {
let mut column = rng.gen::<Column>();
let is_primary_key = column
.options
.iter()
.any(|option| option == &ColumnOption::PrimaryKey);
if is_primary_key {
primary_keys.push(i);
if self.columns == 1 {
// Generates the ts column.
// Safety: columns must large than 0.
let name = column_names.pop().unwrap();
let column = generate_columns(
rng,
vec![name.to_string()],
self.ts_column_type_generator.as_ref(),
self.ts_column_options_generator.as_ref(),
)
.remove(0);
if need_partible_column {
// Generates partition bounds.
let mut partition_bounds = Vec::with_capacity(self.partition);
for _ in 0..self.partition - 1 {
partition_bounds.push(PartitionBound::Value(generate_random_value(
rng,
&column.column_type,
)));
partition_bounds.sort();
}
partition_bounds.push(PartitionBound::MaxValue);
builder.partition(PartitionDef::new(vec![name], partition_bounds));
}
columns.push(column);
} else {
// Generates the partible column.
if need_partible_column {
// Safety: columns must large than 0.
let name = column_names.pop().unwrap();
let column = generate_columns(
rng,
vec![name.to_string()],
self.partible_column_type_generator.as_ref(),
self.partible_column_options_generator.as_ref(),
)
.remove(0);
// Generates partition bounds.
let mut partition_bounds = Vec::with_capacity(self.partition);
for _ in 0..self.partition - 1 {
partition_bounds.push(PartitionBound::Value(generate_random_value(
rng,
&column.column_type,
)));
partition_bounds.sort();
}
partition_bounds.push(PartitionBound::MaxValue);
builder.partition(PartitionDef::new(vec![name], partition_bounds));
columns.push(column);
}
// Generates the ts column.
// Safety: columns must large than 1.
let name = column_names.pop().unwrap();
columns.extend(generate_columns(
rng,
vec![name],
self.ts_column_type_generator.as_ref(),
self.ts_column_options_generator.as_ref(),
));
// Generates rest columns
columns.extend(generate_columns(
rng,
column_names,
self.column_type_generator.as_ref(),
self.column_options_generator.as_ref(),
));
}
for (idx, column) in columns.iter().enumerate() {
if column.is_primary_key() {
primary_keys.push(idx);
}
}
// Shuffles the primary keys.
primary_keys.shuffle(&mut rng);
let partitions = if self.partition > 1 {
// Finds a partible primary keys.
let partible_primary_keys = primary_keys
.iter()
.flat_map(|i| {
if PARTIBLE_DATA_TYPES.contains(&columns[*i].column_type) {
Some(*i)
} else {
None
}
})
.collect::<Vec<_>>();
// Generates the partitions.
if partible_primary_keys.is_empty() {
columns.push(rng.gen::<PartibleColumn>().0);
primary_keys.push(columns.len() - 1);
}
let primary_key_idx = primary_keys[rng.gen_range(0..primary_keys.len())];
let primary_column = &columns[primary_key_idx];
// Generates partition bounds.
let mut partition_bounds = Vec::with_capacity(self.partition);
for _ in 0..self.partition - 1 {
partition_bounds.push(PartitionBound::Value(generate_random_value(
&primary_column.column_type,
)));
partition_bounds.sort();
}
partition_bounds.push(PartitionBound::MaxValue);
vec![PartitionDef::new(
vec![primary_column.name.to_string()],
partition_bounds,
)]
} else {
vec![]
};
// Generates ts column.
columns.push(rng.gen::<TsColumn>().0);
primary_keys.shuffle(rng);
builder.columns(columns);
builder.primary_keys(primary_keys);
builder.engine(self.engine.to_string());
builder.if_not_exists(self.if_not_exists);
builder.name(rng.gen::<Word>().to_string());
builder.partitions(partitions);
if self.name.is_empty() {
builder.name(self.name_generator.gen(rng));
} else {
builder.name(self.name.to_string());
}
builder.build().context(error::BuildCreateTableExprSnafu)
}
}
#[cfg(test)]
mod tests {
use rand::SeedableRng;
use super::*;
#[test]
fn test_create_table_expr_generator() {
let expr = CreateTableExprGenerator::default()
let mut rng = rand::thread_rng();
let expr = CreateTableExprGeneratorBuilder::default()
.columns(10)
.partitions(3)
.create_is_not_exists(true)
.partition(3)
.if_not_exists(true)
.engine("mito2")
.generate()
.build()
.unwrap()
.generate(&mut rng)
.unwrap();
assert_eq!(expr.engine, "mito2");
assert!(expr.if_not_exists);
assert!(expr.columns.len() >= 11);
assert_eq!(expr.partitions[0].partition_bounds().len(), 3);
assert_eq!(expr.columns.len(), 10);
assert_eq!(expr.partition.unwrap().partition_bounds().len(), 3);
let expr = CreateTableExprGenerator::default()
let expr = CreateTableExprGeneratorBuilder::default()
.columns(10)
.partitions(1)
.generate()
.partition(1)
.build()
.unwrap()
.generate(&mut rng)
.unwrap();
assert_eq!(expr.columns.len(), 11);
assert_eq!(expr.partitions.len(), 0);
assert_eq!(expr.columns.len(), 10);
assert!(expr.partition.is_none());
}
#[test]
fn test_create_table_expr_generator_deterministic() {
let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(0);
let expr = CreateTableExprGeneratorBuilder::default()
.columns(10)
.partition(3)
.if_not_exists(true)
.engine("mito2")
.build()
.unwrap()
.generate(&mut rng)
.unwrap();
let serialized = serde_json::to_string(&expr).unwrap();
let expected = r#"{"name":"iN","columns":[{"name":"CUlpa","column_type":{"Int16":{}},"options":["PrimaryKey","NotNull"]},{"name":"dEBiTiS","column_type":{"Timestamp":{"Second":null}},"options":["TimeIndex"]},{"name":"HArum","column_type":{"Int16":{}},"options":["NotNull"]},{"name":"NObIS","column_type":{"Int32":{}},"options":["PrimaryKey"]},{"name":"IMPEDiT","column_type":{"Int16":{}},"options":[{"DefaultValue":{"Int16":-25151}}]},{"name":"bLanDITIis","column_type":{"Boolean":null},"options":[{"DefaultValue":{"Boolean":true}}]},{"name":"Dolores","column_type":{"Float32":{}},"options":["PrimaryKey"]},{"name":"eSt","column_type":{"Float32":{}},"options":[{"DefaultValue":{"Float32":0.9152612}}]},{"name":"INVentORE","column_type":{"Int64":{}},"options":["PrimaryKey"]},{"name":"aDIpiSci","column_type":{"Float64":{}},"options":["Null"]}],"if_not_exists":true,"partition":{"partition_columns":["CUlpa"],"partition_bounds":[{"Value":{"Int16":15966}},{"Value":{"Int16":31925}},"MaxValue"]},"engine":"mito2","options":{},"primary_keys":[6,0,8,3]}"#;
assert_eq!(expected, serialized);
}
}

167
tests-fuzz/src/ir.rs
View File

@@ -22,15 +22,17 @@ pub use create_expr::CreateTableExpr;
use datatypes::data_type::ConcreteDataType;
use datatypes::value::Value;
use derive_builder::Builder;
use faker_rand::lorem::Word;
use lazy_static::lazy_static;
use rand::distributions::{Distribution, Standard};
use rand::seq::SliceRandom;
use rand::Rng;
use serde::{Deserialize, Serialize};
use crate::generator::Random;
use crate::impl_random;
use crate::ir::create_expr::ColumnOption;
lazy_static! {
static ref DATA_TYPES: Vec<ConcreteDataType> = vec![
pub static ref DATA_TYPES: Vec<ConcreteDataType> = vec![
ConcreteDataType::boolean_datatype(),
ConcreteDataType::int16_datatype(),
ConcreteDataType::int32_datatype(),
@@ -38,7 +40,7 @@ lazy_static! {
ConcreteDataType::float32_datatype(),
ConcreteDataType::float64_datatype(),
];
static ref TS_DATA_TYPES: Vec<ConcreteDataType> = vec![
pub static ref TS_DATA_TYPES: Vec<ConcreteDataType> = vec![
ConcreteDataType::timestamp_nanosecond_datatype(),
ConcreteDataType::timestamp_microsecond_datatype(),
ConcreteDataType::timestamp_millisecond_datatype(),
@@ -56,25 +58,37 @@ lazy_static! {
];
}
impl_random!(ConcreteDataType, ColumnTypeGenerator, DATA_TYPES);
impl_random!(ConcreteDataType, TsColumnTypeGenerator, TS_DATA_TYPES);
impl_random!(
ConcreteDataType,
PartibleColumnTypeGenerator,
PARTIBLE_DATA_TYPES
);
pub struct ColumnTypeGenerator;
pub struct TsColumnTypeGenerator;
pub struct PartibleColumnTypeGenerator;
/// Generates a random [Value].
pub fn generate_random_value(datatype: &ConcreteDataType) -> Value {
pub fn generate_random_value<R: Rng>(rng: &mut R, datatype: &ConcreteDataType) -> Value {
match datatype {
&ConcreteDataType::Boolean(_) => Value::from(rand::random::<bool>()),
ConcreteDataType::Int16(_) => Value::from(rand::random::<i16>()),
ConcreteDataType::Int32(_) => Value::from(rand::random::<i32>()),
ConcreteDataType::Int64(_) => Value::from(rand::random::<i64>()),
ConcreteDataType::Float32(_) => Value::from(rand::random::<f32>()),
ConcreteDataType::Float64(_) => Value::from(rand::random::<f64>()),
ConcreteDataType::String(_) => Value::from(rand::random::<char>().to_string()),
ConcreteDataType::Date(_) => Value::from(rand::random::<i32>()),
ConcreteDataType::DateTime(_) => Value::from(rand::random::<i64>()),
&ConcreteDataType::Boolean(_) => Value::from(rng.gen::<bool>()),
ConcreteDataType::Int16(_) => Value::from(rng.gen::<i16>()),
ConcreteDataType::Int32(_) => Value::from(rng.gen::<i32>()),
ConcreteDataType::Int64(_) => Value::from(rng.gen::<i64>()),
ConcreteDataType::Float32(_) => Value::from(rng.gen::<f32>()),
ConcreteDataType::Float64(_) => Value::from(rng.gen::<f64>()),
ConcreteDataType::String(_) => Value::from(rng.gen::<char>().to_string()),
ConcreteDataType::Date(_) => Value::from(rng.gen::<i32>()),
ConcreteDataType::DateTime(_) => Value::from(rng.gen::<i64>()),
_ => unimplemented!("unsupported type: {datatype}"),
}
}
/// The IR column.
#[derive(Debug, Builder, Clone, Serialize, Deserialize)]
#[derive(Debug, Builder, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
pub struct Column {
#[builder(setter(into))]
pub name: String,
@@ -83,6 +97,22 @@ pub struct Column {
pub options: Vec<ColumnOption>,
}
impl Column {
/// Returns true if it's [ColumnOption::TimeIndex] [Column].
pub fn is_time_index(&self) -> bool {
self.options
.iter()
.any(|opt| opt == &ColumnOption::TimeIndex)
}
/// Returns true if it's the [ColumnOption::PrimaryKey] [Column].
pub fn is_primary_key(&self) -> bool {
self.options
.iter()
.any(|opt| opt == &ColumnOption::PrimaryKey)
}
}
/// Returns droppable columns. i.e., non-primary key columns, non-ts columns.
pub fn droppable_columns(columns: &[Column]) -> Vec<&Column> {
columns
@@ -95,59 +125,78 @@ pub fn droppable_columns(columns: &[Column]) -> Vec<&Column> {
.collect::<Vec<_>>()
}
impl Distribution<Column> for Standard {
fn sample<R: rand::prelude::Rng + ?Sized>(&self, rng: &mut R) -> Column {
let column_type = DATA_TYPES[rng.gen_range(0..DATA_TYPES.len())].clone();
// 0 -> NULL
// 1 -> NOT NULL
// 2 -> DEFAULT VALUE
// 3 -> PRIMARY KEY
// 4 -> EMPTY
let option_idx = rng.gen_range(0..5);
let options = match option_idx {
0 => vec![ColumnOption::Null],
1 => vec![ColumnOption::NotNull],
2 => vec![ColumnOption::DefaultValue(generate_random_value(
&column_type,
))],
3 => vec![ColumnOption::PrimaryKey],
_ => vec![],
};
Column {
name: rng.gen::<Word>().to_string(),
/// Generates [ColumnOption] for [Column].
pub fn column_options_generator<R: Rng>(
rng: &mut R,
column_type: &ConcreteDataType,
) -> Vec<ColumnOption> {
// 0 -> NULL
// 1 -> NOT NULL
// 2 -> DEFAULT VALUE
// 3 -> PRIMARY KEY
// 4 -> EMPTY
let option_idx = rng.gen_range(0..5);
match option_idx {
0 => vec![ColumnOption::Null],
1 => vec![ColumnOption::NotNull],
2 => vec![ColumnOption::DefaultValue(generate_random_value(
rng,
column_type,
options,
}
))],
3 => vec![ColumnOption::PrimaryKey],
_ => vec![],
}
}
/// The IR ts column.
pub struct TsColumn(pub Column);
impl Distribution<TsColumn> for Standard {
fn sample<R: rand::prelude::Rng + ?Sized>(&self, rng: &mut R) -> TsColumn {
let column_type = TS_DATA_TYPES[rng.gen_range(0..TS_DATA_TYPES.len())].clone();
TsColumn(Column {
name: rng.gen::<Word>().to_string(),
column_type,
options: vec![],
})
/// Generates [ColumnOption] for Partible [Column].
pub fn partible_column_options_generator<R: Rng + 'static>(
rng: &mut R,
column_type: &ConcreteDataType,
) -> Vec<ColumnOption> {
// 0 -> NULL
// 1 -> NOT NULL
// 2 -> DEFAULT VALUE
// 3 -> PRIMARY KEY
let option_idx = rng.gen_range(0..4);
match option_idx {
0 => vec![ColumnOption::PrimaryKey, ColumnOption::Null],
1 => vec![ColumnOption::PrimaryKey, ColumnOption::NotNull],
2 => vec![
ColumnOption::PrimaryKey,
ColumnOption::DefaultValue(generate_random_value(rng, column_type)),
],
3 => vec![ColumnOption::PrimaryKey],
_ => unreachable!(),
}
}
/// The IR partible column.
pub struct PartibleColumn(pub Column);
/// Generates [ColumnOption] for ts [Column].
pub fn ts_column_options_generator<R: Rng + 'static>(
_: &mut R,
_: &ConcreteDataType,
) -> Vec<ColumnOption> {
vec![ColumnOption::TimeIndex]
}
impl Distribution<PartibleColumn> for Standard {
fn sample<R: rand::prelude::Rng + ?Sized>(&self, rng: &mut R) -> PartibleColumn {
let column_type = PARTIBLE_DATA_TYPES[rng.gen_range(0..PARTIBLE_DATA_TYPES.len())].clone();
PartibleColumn(Column {
name: rng.gen::<Word>().to_string(),
column_type,
options: vec![],
/// Generates columns with given `names`.
pub fn generate_columns<R: Rng + 'static>(
rng: &mut R,
names: impl IntoIterator<Item = String>,
types: &(impl Random<ConcreteDataType, R> + ?Sized),
options: impl Fn(&mut R, &ConcreteDataType) -> Vec<ColumnOption>,
) -> Vec<Column> {
names
.into_iter()
.map(|name| {
let column_type = types.gen(rng);
let options = options(rng, &column_type);
Column {
name,
options,
column_type,
}
})
}
.collect()
}
#[cfg(test)]

2
tests-fuzz/src/ir/alter_expr.rs
View File

@@ -18,7 +18,7 @@ use serde::{Deserialize, Serialize};
use crate::ir::Column;
#[derive(Debug, Builder, Clone)]
#[derive(Debug, Builder, Clone, Serialize, Deserialize)]
pub struct AlterTableExpr {
pub name: String,
pub alter_options: AlterTableOperation,

6
tests-fuzz/src/ir/create_expr.rs
View File

@@ -23,7 +23,7 @@ use serde::{Deserialize, Serialize};
use crate::ir::Column;
// The column options
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Hash)]
pub enum ColumnOption {
Null,
NotNull,
@@ -49,6 +49,7 @@ impl Display for ColumnOption {
/// A naive create table expr builder.
#[derive(Debug, Builder, Clone, Serialize, Deserialize)]
pub struct CreateTableExpr {
#[builder(setter(into))]
pub name: String,
pub columns: Vec<Column>,
#[builder(default)]
@@ -56,7 +57,8 @@ pub struct CreateTableExpr {
// GreptimeDB specific options
#[builder(default, setter(into))]
pub partitions: Vec<PartitionDef>,
pub partition: Option<PartitionDef>,
#[builder(default, setter(into))]
pub engine: String,
#[builder(default, setter(into))]
pub options: HashMap<String, Value>,

18
tests-fuzz/src/lib.rs
View File

@@ -12,12 +12,12 @@
// See the License for the specific language governing permissions and
// limitations under the License.
pub(crate) mod context;
pub(crate) mod error;
pub(crate) mod executor;
// TODO(weny): removes it.
#[allow(unused)]
pub(crate) mod generator;
pub(crate) mod ir;
pub(crate) mod table_creator;
pub(crate) mod translator;
#![feature(associated_type_bounds)]
pub mod context;
pub mod error;
pub mod executor;
pub mod fake;
pub mod generator;
pub mod ir;
pub mod translator;

4
tests-fuzz/src/translator.rs
View File

@@ -12,9 +12,11 @@
// See the License for the specific language governing permissions and
// limitations under the License.
pub mod greptime;
use std::fmt;
pub(crate) trait DslTranslator<T, U> {
pub trait DslTranslator<T, U> {
type Error: Sync + Send + fmt::Debug;
fn translate(&self, input: &T) -> Result<U, Self::Error>;

3
tests-fuzz/src/table_creator.rs → tests-fuzz/src/translator/greptime.rs
View File

@@ -11,3 +11,6 @@
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
pub mod alter_expr;
pub mod create_expr;

159
tests-fuzz/src/translator/greptime/alter_expr.rs Normal file
View File

@@ -0,0 +1,159 @@
// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use common_query::AddColumnLocation;
use datatypes::data_type::ConcreteDataType;
use sql::statements::concrete_data_type_to_sql_data_type;
use crate::error::{Error, Result};
use crate::ir::alter_expr::AlterTableOperation;
use crate::ir::create_expr::ColumnOption;
use crate::ir::{AlterTableExpr, Column};
use crate::translator::DslTranslator;
pub struct AlterTableExprTranslator;
impl DslTranslator<AlterTableExpr, String> for AlterTableExprTranslator {
type Error = Error;
fn translate(&self, input: &AlterTableExpr) -> Result<String> {
Ok(match &input.alter_options {
AlterTableOperation::AddColumn { column, location } => {
Self::format_add_column(&input.name, column, location)
}
AlterTableOperation::DropColumn { name } => Self::format_drop(&input.name, name),
AlterTableOperation::RenameTable { new_table_name } => {
Self::format_rename(&input.name, new_table_name)
}
})
}
}
impl AlterTableExprTranslator {
fn format_drop(name: &str, column: &str) -> String {
format!("ALTER TABLE {name} DROP COLUMN {column};")
}
fn format_rename(name: &str, new_name: &str) -> String {
format!("ALTER TABLE {name} RENAME {new_name};")
}
fn format_add_column(
name: &str,
column: &Column,
location: &Option<AddColumnLocation>,
) -> String {
format!(
"{};",
vec![
format!(
"ALTER TABLE {name} ADD COLUMN {}",
Self::format_column(column)
),
Self::format_location(location).unwrap_or_default(),
]
.into_iter()
.filter(|s| !s.is_empty())
.collect::<Vec<_>>()
.join(" ")
)
}
fn format_location(location: &Option<AddColumnLocation>) -> Option<String> {
location.as_ref().map(|location| match location {
AddColumnLocation::First => "FIRST".to_string(),
AddColumnLocation::After { column_name } => format!("AFTER {column_name}"),
})
}
fn format_column(column: &Column) -> String {
vec![
column.name.to_string(),
Self::format_column_type(&column.column_type),
Self::format_column_options(&column.options),
]
.into_iter()
.filter(|s| !s.is_empty())
.collect::<Vec<_>>()
.join(" ")
}
fn format_column_type(column_type: &ConcreteDataType) -> String {
// Safety: We don't use the `Dictionary` type
concrete_data_type_to_sql_data_type(column_type)
.unwrap()
.to_string()
}
fn format_column_options(options: &[ColumnOption]) -> String {
options
.iter()
.map(|option| option.to_string())
.collect::<Vec<_>>()
.join(" ")
}
}
#[cfg(test)]
mod tests {
use common_query::AddColumnLocation;
use datatypes::data_type::ConcreteDataType;
use super::AlterTableExprTranslator;
use crate::ir::alter_expr::AlterTableOperation;
use crate::ir::create_expr::ColumnOption;
use crate::ir::{AlterTableExpr, Column};
use crate::translator::DslTranslator;
#[test]
fn test_alter_table_expr() {
let alter_expr = AlterTableExpr {
name: "test".to_string(),
alter_options: AlterTableOperation::AddColumn {
column: Column {
name: "host".to_string(),
column_type: ConcreteDataType::string_datatype(),
options: vec![ColumnOption::PrimaryKey],
},
location: Some(AddColumnLocation::First),
},
};
let output = AlterTableExprTranslator.translate(&alter_expr).unwrap();
assert_eq!(
"ALTER TABLE test ADD COLUMN host STRING PRIMARY KEY FIRST;",
output
);
let alter_expr = AlterTableExpr {
name: "test".to_string(),
alter_options: AlterTableOperation::RenameTable {
new_table_name: "foo".to_string(),
},
};
let output = AlterTableExprTranslator.translate(&alter_expr).unwrap();
assert_eq!("ALTER TABLE test RENAME foo;", output);
let alter_expr = AlterTableExpr {
name: "test".to_string(),
alter_options: AlterTableOperation::DropColumn {
name: "foo".to_string(),
},
};
let output = AlterTableExprTranslator.translate(&alter_expr).unwrap();
assert_eq!("ALTER TABLE test DROP COLUMN foo;", output);
}
}

233
tests-fuzz/src/translator/greptime/create_expr.rs Normal file
View File

@@ -0,0 +1,233 @@
// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use datatypes::data_type::ConcreteDataType;
use datatypes::value::Value;
use partition::partition::PartitionBound;
use sql::statements::concrete_data_type_to_sql_data_type;
use crate::error::{Error, Result};
use crate::ir::create_expr::ColumnOption;
use crate::ir::{Column, CreateTableExpr};
use crate::translator::DslTranslator;
pub struct CreateTableExprTranslator;
impl DslTranslator<CreateTableExpr, String> for CreateTableExprTranslator {
type Error = Error;
fn translate(&self, input: &CreateTableExpr) -> Result<String> {
Ok(format!(
"CREATE TABLE{}{}(\n{}\n)\n{};",
Self::create_if_not_exists(input),
input.name,
Self::format_columns(input),
Self::format_table_options(input)
))
}
}
impl CreateTableExprTranslator {
fn create_if_not_exists(input: &CreateTableExpr) -> &str {
if input.if_not_exists {
" IF NOT EXISTS "
} else {
" "
}
}
fn format_columns(input: &CreateTableExpr) -> String {
let mut output =
Vec::with_capacity(input.columns.len() + (!input.primary_keys.is_empty()) as usize);
for column in &input.columns {
output.push(Self::format_column(column));
}
if let Some(primary_keys) = Self::format_primary_keys(input) {
output.push(primary_keys);
}
output.join(",\n")
}
fn format_column(column: &Column) -> String {
vec![
column.name.to_string(),
Self::format_column_type(&column.column_type),
Self::format_column_options(&column.options),
]
.into_iter()
.filter(|s| !s.is_empty())
.collect::<Vec<_>>()
.join(" ")
}
fn format_partition(input: &CreateTableExpr) -> Option<String> {
input.partition.as_ref().map(|partition| {
format!(
"PARTITION BY RANGE COLUMNS({}) (\n{}\n)",
partition.partition_columns().join(", "),
partition
.partition_bounds()
.iter()
.enumerate()
.map(|(i, bound)| format!(
"PARTITION r{} VALUES LESS THAN ({})",
i,
Self::format_partition_bound(bound)
))
.collect::<Vec<_>>()
.join(",\n")
)
})
}
fn format_partition_bound(bound: &PartitionBound) -> String {
match bound {
PartitionBound::Value(v) => match v {
Value::String(v) => format!("'{}'", v.as_utf8()),
_ => format!("{v}"),
},
PartitionBound::MaxValue => "MAXVALUE".to_string(),
}
}
fn format_column_type(column_type: &ConcreteDataType) -> String {
// Safety: We don't use the `Dictionary` type
concrete_data_type_to_sql_data_type(column_type)
.unwrap()
.to_string()
}
fn format_column_options(options: &[ColumnOption]) -> String {
let mut output = Vec::with_capacity(options.len());
for option in options {
if option != &ColumnOption::PrimaryKey {
output.push(option.to_string());
}
}
output.join(" ")
}
fn format_primary_keys(input: &CreateTableExpr) -> Option<String> {
if input.primary_keys.is_empty() {
None
} else {
Some(format!(
"PRIMARY KEY({})",
input
.primary_keys
.iter()
.map(|idx| input.columns[*idx].name.to_string())
.collect::<Vec<_>>()
.join(", ")
))
}
}
fn format_table_options(input: &CreateTableExpr) -> String {
let mut output = vec![];
if !input.engine.is_empty() {
output.push(format!("ENGINE={}", input.engine));
}
if let Some(partition) = Self::format_partition(input) {
output.push(partition);
}
output.join("\n")
}
}
#[cfg(test)]
mod tests {
use datatypes::data_type::ConcreteDataType;
use datatypes::value::Value;
use partition::partition::{PartitionBound, PartitionDef};
use super::CreateTableExprTranslator;
use crate::ir::create_expr::{ColumnOption, CreateTableExprBuilder};
use crate::ir::Column;
use crate::translator::DslTranslator;
#[test]
fn test_create_table_expr_translator() {
let create_table_expr = CreateTableExprBuilder::default()
.columns(vec![
Column {
name: "host".to_string(),
column_type: ConcreteDataType::string_datatype(),
options: vec![ColumnOption::PrimaryKey],
},
Column {
name: "idc".to_string(),
column_type: ConcreteDataType::string_datatype(),
options: vec![ColumnOption::PrimaryKey],
},
Column {
name: "cpu_util".to_string(),
column_type: ConcreteDataType::float64_datatype(),
options: vec![],
},
Column {
name: "memory_util".to_string(),
column_type: ConcreteDataType::float64_datatype(),
options: vec![],
},
Column {
name: "disk_util".to_string(),
column_type: ConcreteDataType::float64_datatype(),
options: vec![],
},
Column {
name: "ts".to_string(),
column_type: ConcreteDataType::timestamp_millisecond_datatype(),
options: vec![ColumnOption::TimeIndex],
},
])
.name("system_metrics")
.engine("mito")
.primary_keys(vec![0, 1])
.partition(PartitionDef::new(
vec!["idc".to_string()],
vec![
PartitionBound::Value(Value::String("a".into())),
PartitionBound::Value(Value::String("f".into())),
PartitionBound::MaxValue,
],
))
.build()
.unwrap();
let output = CreateTableExprTranslator
.translate(&create_table_expr)
.unwrap();
assert_eq!(
"CREATE TABLE system_metrics(
host STRING,
idc STRING,
cpu_util DOUBLE,
memory_util DOUBLE,
disk_util DOUBLE,
ts TIMESTAMP(3) TIME INDEX,
PRIMARY KEY(host, idc)
)
ENGINE=mito
PARTITION BY RANGE COLUMNS(idc) (
PARTITION r0 VALUES LESS THAN ('a'),
PARTITION r1 VALUES LESS THAN ('f'),
PARTITION r2 VALUES LESS THAN (MAXVALUE)
);",
output
);
}
}

2
typos.toml
View File

@@ -2,4 +2,4 @@
ue = "ue"
datas = "datas"
[files]
extend-exclude = ["corrupted"]
extend-exclude = ["corrupted","tests-fuzz/src/data/lorem_words"]