risingwave_sqlsmith/sql_gen/
utils.rs

1// Copyright 2025 RisingWave Labs
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
15//! Internal utilities for sql gen.
16use std::mem;
17
18use rand::Rng;
19use risingwave_sqlparser::ast::{
20    FunctionArg, FunctionArgExpr, TableAlias, TableFactor, TableWithJoins,
21};
22
23use crate::sql_gen::{Column, Expr, Ident, ObjectName, SqlGenerator, Table};
24
25type Context = (Vec<Column>, Vec<Table>);
26
27/// Context utils
28impl<R: Rng> SqlGenerator<'_, R> {
29    pub(crate) fn add_relations_to_context(&mut self, mut tables: Vec<Table>) {
30        for rel in &tables {
31            let mut bound_columns = rel.get_qualified_columns();
32            self.bound_columns.append(&mut bound_columns);
33        }
34        self.bound_relations.append(&mut tables);
35    }
36
37    pub(crate) fn new_local_context(&mut self) -> Context {
38        let current_bound_relations = mem::take(&mut self.bound_relations);
39        let current_bound_columns = mem::take(&mut self.bound_columns);
40        (current_bound_columns, current_bound_relations)
41    }
42
43    pub(crate) fn restore_context(&mut self, (old_cols, old_rels): Context) {
44        self.bound_relations = old_rels;
45        self.bound_columns = old_cols;
46    }
47
48    pub(crate) fn clone_local_context(&mut self) -> Context {
49        let current_bound_relations = self.bound_relations.clone();
50        let current_bound_columns = self.bound_columns.clone();
51        (current_bound_columns, current_bound_relations)
52    }
53}
54
55/// Gen utils
56impl<R: Rng> SqlGenerator<'_, R> {
57    pub(crate) fn gen_table_name_with_prefix(&mut self, prefix: &str) -> String {
58        format!("{}_{}", prefix, &self.gen_relation_id())
59    }
60
61    fn gen_relation_id(&mut self) -> u32 {
62        let id = self.relation_id;
63        self.relation_id += 1;
64        id
65    }
66
67    pub(crate) fn gen_table_alias_with_prefix(&mut self, prefix: &str) -> TableAlias {
68        let name = &self.gen_table_name_with_prefix(prefix);
69        create_table_alias(name)
70    }
71}
72
73pub(crate) fn create_table_factor_from_table(table: &Table) -> TableFactor {
74    TableFactor::Table {
75        name: ObjectName(vec![Ident::new_unchecked(&table.name)]),
76        alias: None,
77        as_of: None,
78    }
79}
80
81pub(crate) fn create_table_with_joins_from_table(table: &Table) -> TableWithJoins {
82    TableWithJoins {
83        relation: create_table_factor_from_table(table),
84        joins: vec![],
85    }
86}
87
88pub(crate) fn create_table_alias(table_name: &str) -> TableAlias {
89    TableAlias {
90        name: table_name.into(),
91        columns: vec![],
92    }
93}
94
95pub(crate) fn create_args(arg_exprs: Vec<Expr>) -> Vec<FunctionArg> {
96    arg_exprs
97        .into_iter()
98        .map(create_function_arg_from_expr)
99        .collect()
100}
101
102/// Create `FunctionArg` from an `Expr`.
103fn create_function_arg_from_expr(expr: Expr) -> FunctionArg {
104    FunctionArg::Unnamed(FunctionArgExpr::Expr(expr))
105}
risingwave_sqlsmith/sql_gen/utils.rs

risingwave_sqlsmith/sql_gen/
utils.rs
