mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-05-24 17:00:37 +00:00
f99b08796d
* feat: add insert into expr generator & translator * feat: add select expr generator & translator * chore: apply suggestions from CR * fix: fix unit tests
73 lines
2.1 KiB
Rust
73 lines
2.1 KiB
Rust
// 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::marker::PhantomData;
|
|
|
|
use derive_builder::Builder;
|
|
use rand::seq::SliceRandom;
|
|
use rand::Rng;
|
|
|
|
use crate::context::TableContextRef;
|
|
use crate::error::{Error, Result};
|
|
use crate::generator::Generator;
|
|
use crate::ir::select_expr::{Direction, SelectExpr};
|
|
|
|
#[derive(Builder)]
|
|
#[builder(pattern = "owned")]
|
|
pub struct SelectExprGenerator<R: Rng + 'static> {
|
|
table_ctx: TableContextRef,
|
|
#[builder(default = "8192")]
|
|
max_limit: usize,
|
|
#[builder(default)]
|
|
_phantom: PhantomData<R>,
|
|
}
|
|
|
|
impl<R: Rng + 'static> Generator<SelectExpr, R> for SelectExprGenerator<R> {
|
|
type Error = Error;
|
|
|
|
fn generate(&self, rng: &mut R) -> Result<SelectExpr> {
|
|
let selection = rng.gen_range(1..self.table_ctx.columns.len());
|
|
let mut selected_columns = self
|
|
.table_ctx
|
|
.columns
|
|
.choose_multiple(rng, selection)
|
|
.cloned()
|
|
.collect::<Vec<_>>();
|
|
selected_columns.shuffle(rng);
|
|
|
|
let order_by_selection = rng.gen_range(1..selection);
|
|
|
|
let order_by = selected_columns
|
|
.choose_multiple(rng, order_by_selection)
|
|
.map(|c| c.name.to_string())
|
|
.collect::<Vec<_>>();
|
|
|
|
let limit = rng.gen_range(1..self.max_limit);
|
|
|
|
let direction = if rng.gen_bool(1.0 / 2.0) {
|
|
Direction::Asc
|
|
} else {
|
|
Direction::Desc
|
|
};
|
|
|
|
Ok(SelectExpr {
|
|
table_name: self.table_ctx.name.to_string(),
|
|
columns: selected_columns,
|
|
order_by,
|
|
direction,
|
|
limit,
|
|
})
|
|
}
|
|
}
|