diff --git a/src/tests/sqlsmith/src/lib.rs b/src/tests/sqlsmith/src/lib.rs index 6362ed7243ba0..b5f6adf7fa2db 100644 --- a/src/tests/sqlsmith/src/lib.rs +++ b/src/tests/sqlsmith/src/lib.rs @@ -128,10 +128,12 @@ impl<'a, R: Rng> SqlGenerator<'a, R> { } } - fn add_relation_to_context(&mut self, table: Table) { - let mut bound_columns = table.get_qualified_columns(); - self.bound_columns.append(&mut bound_columns); - self.bound_relations.push(table); + fn add_relations_to_context(&mut self, mut tables: Vec) { + for rel in &tables { + let mut bound_columns = rel.get_qualified_columns(); + self.bound_columns.append(&mut bound_columns); + } + self.bound_relations.append(&mut tables); } fn gen_stmt(&mut self) -> Statement { @@ -204,7 +206,7 @@ impl<'a, R: Rng> SqlGenerator<'a, R> { } /// Generates a query with local context. - /// Used by `WITH`, (and perhaps subquery should use this too) + /// Used by `WITH`, subquery fn gen_local_query(&mut self) -> (Query, Vec) { let old_ctxt = self.new_local_context(); let t = self.gen_query(); @@ -350,7 +352,9 @@ impl<'a, R: Rng> SqlGenerator<'a, R> { .expect("with tables should not be empty"); vec![create_table_with_joins_from_table(with_table)] } else { - vec![self.gen_from_relation()] + let (rel, tables) = self.gen_from_relation(); + self.add_relations_to_context(tables); + vec![rel] }; if self.is_mview { @@ -359,12 +363,15 @@ impl<'a, R: Rng> SqlGenerator<'a, R> { assert!(!self.tables.is_empty()); return from; } - + let mut lateral_contexts = vec![]; for _ in 0..self.tables.len() { if self.flip_coin() { - from.push(self.gen_from_relation()); + let (table_with_join, mut table) = self.gen_from_relation(); + from.push(table_with_join); + lateral_contexts.append(&mut table); } } + self.add_relations_to_context(lateral_contexts); from } diff --git a/src/tests/sqlsmith/src/relation.rs b/src/tests/sqlsmith/src/relation.rs index ed855ccce3e88..c535d9777dc24 100644 --- a/src/tests/sqlsmith/src/relation.rs +++ b/src/tests/sqlsmith/src/relation.rs @@ -32,28 +32,31 @@ fn create_join_on_clause(left: String, right: String) -> Expr { impl<'a, R: Rng> SqlGenerator<'a, R> { /// A relation specified in the FROM clause. - pub(crate) fn gen_from_relation(&mut self) -> TableWithJoins { - match self.rng.gen_range(0..=9) { + pub(crate) fn gen_from_relation(&mut self) -> (TableWithJoins, Vec
) { + let range = if self.can_recurse() { 10 } else { 9 }; + match self.rng.gen_range(0..=range) { 0..=7 => self.gen_simple_table(), 8..=8 => self.gen_time_window_func(), // TODO: Enable after resolving: . 9..=9 => self.gen_equijoin_clause(), - // TODO: Currently `gen_subquery` will cause panic due to some wrong assertions. - 10..=10 => self.gen_subquery(), + 10..=10 => self.gen_table_subquery(), _ => unreachable!(), } } - fn gen_simple_table(&mut self) -> TableWithJoins { - let (relation, _) = self.gen_simple_table_factor(); + fn gen_simple_table(&mut self) -> (TableWithJoins, Vec
) { + let (relation, _, table) = self.gen_simple_table_factor(); - TableWithJoins { - relation, - joins: vec![], - } + ( + TableWithJoins { + relation, + joins: vec![], + }, + table, + ) } - fn gen_simple_table_factor(&mut self) -> (TableFactor, Vec) { + fn gen_simple_table_factor(&mut self) -> (TableFactor, Vec, Vec
) { let alias = self.gen_table_name_with_prefix("t"); let mut table = self.tables.choose(&mut self.rng).unwrap().clone(); let table_factor = TableFactor::Table { @@ -66,11 +69,10 @@ impl<'a, R: Rng> SqlGenerator<'a, R> { }; table.name = alias; // Rename the table. let columns = table.get_qualified_columns(); - self.add_relation_to_context(table); - (table_factor, columns) + (table_factor, columns, vec![table]) } - fn gen_table_factor(&mut self) -> (TableFactor, Vec) { + fn gen_table_factor(&mut self) -> (TableFactor, Vec, Vec
) { let current_context = self.new_local_context(); let factor = self.gen_table_factor_inner(); self.restore_context(current_context); @@ -79,14 +81,14 @@ impl<'a, R: Rng> SqlGenerator<'a, R> { /// Generates a table factor, and provides bound columns. /// Generated column names should be qualified by table name. - fn gen_table_factor_inner(&mut self) -> (TableFactor, Vec) { + fn gen_table_factor_inner(&mut self) -> (TableFactor, Vec, Vec
) { // TODO: TableFactor::Derived, TableFactor::TableFunction, TableFactor::NestedJoin self.gen_simple_table_factor() } - fn gen_equijoin_clause(&mut self) -> TableWithJoins { - let (left_factor, left_columns) = self.gen_table_factor(); - let (right_factor, right_columns) = self.gen_table_factor(); + fn gen_equijoin_clause(&mut self) -> (TableWithJoins, Vec
) { + let (left_factor, left_columns, mut left_table) = self.gen_table_factor(); + let (right_factor, right_columns, mut right_table) = self.gen_table_factor(); let mut available_join_on_columns = vec![]; for left_column in &left_columns { @@ -113,15 +115,18 @@ impl<'a, R: Rng> SqlGenerator<'a, R> { relation: right_factor, join_operator: JoinOperator::Inner(JoinConstraint::On(join_on_expr)), }; - - TableWithJoins { - relation: left_factor, - joins: vec![right_factor_with_join], - } + left_table.append(&mut right_table); + ( + TableWithJoins { + relation: left_factor, + joins: vec![right_factor_with_join], + }, + left_table, + ) } - fn gen_subquery(&mut self) -> TableWithJoins { - let (subquery, columns) = self.gen_query(); + fn gen_table_subquery(&mut self) -> (TableWithJoins, Vec
) { + let (subquery, columns) = self.gen_local_query(); let alias = self.gen_table_name_with_prefix("sq"); let table = Table { name: alias.clone(), @@ -138,7 +143,7 @@ impl<'a, R: Rng> SqlGenerator<'a, R> { }, joins: vec![], }; - self.add_relation_to_context(table); - relation + + (relation, vec![table]) } } diff --git a/src/tests/sqlsmith/src/time_window.rs b/src/tests/sqlsmith/src/time_window.rs index 7c0296ed419eb..6adf60679298b 100644 --- a/src/tests/sqlsmith/src/time_window.rs +++ b/src/tests/sqlsmith/src/time_window.rs @@ -25,7 +25,7 @@ use crate::{Column, Expr, SqlGenerator, Table}; impl<'a, R: Rng> SqlGenerator<'a, R> { /// Generates time window functions. - pub(crate) fn gen_time_window_func(&mut self) -> TableWithJoins { + pub(crate) fn gen_time_window_func(&mut self) -> (TableWithJoins, Vec
) { match self.flip_coin() { true => self.gen_hop(), false => self.gen_tumble(), @@ -34,7 +34,7 @@ impl<'a, R: Rng> SqlGenerator<'a, R> { /// Generates `TUMBLE`. /// TUMBLE(data: TABLE, timecol: COLUMN, size: INTERVAL, offset?: INTERVAL) - fn gen_tumble(&mut self) -> TableWithJoins { + fn gen_tumble(&mut self) -> (TableWithJoins, Vec
) { let tables = find_tables_with_timestamp_cols(self.tables.clone()); let (source_table_name, time_cols, schema) = tables .choose(&mut self.rng) @@ -52,14 +52,13 @@ impl<'a, R: Rng> SqlGenerator<'a, R> { let relation = create_tvf("tumble", alias, args); let table = Table::new(table_name, schema.clone()); - self.add_relation_to_context(table); - relation + (relation, vec![table]) } /// Generates `HOP`. /// HOP(data: TABLE, timecol: COLUMN, slide: INTERVAL, size: INTERVAL, offset?: INTERVAL) - fn gen_hop(&mut self) -> TableWithJoins { + fn gen_hop(&mut self) -> (TableWithJoins, Vec
) { let tables = find_tables_with_timestamp_cols(self.tables.clone()); let (source_table_name, time_cols, schema) = tables .choose(&mut self.rng) @@ -84,9 +83,8 @@ impl<'a, R: Rng> SqlGenerator<'a, R> { let relation = create_tvf("hop", alias, args); let table = Table::new(table_name, schema.clone()); - self.add_relation_to_context(table); - relation + (relation, vec![table]) } }