use std::{
collections::{HashMap, HashSet},
sync::Arc,
};
use crate::errors::ExecutorError;
pub type CteResult = (vibesql_catalog::TableSchema, Arc<Vec<vibesql_storage::Row>>);
pub fn execute_ctes<F>(
ctes: &[vibesql_ast::CommonTableExpr],
database: &vibesql_storage::Database,
executor: F,
) -> Result<HashMap<String, CteResult>, ExecutorError>
where
F: Fn(
&vibesql_ast::SelectStmt,
&HashMap<String, CteResult>,
) -> Result<Vec<vibesql_storage::Row>, ExecutorError>,
{
execute_ctes_with_memory_check(ctes, database, executor, |_| Ok(()))
}
pub(super) fn execute_ctes_with_memory_check<F, M>(
ctes: &[vibesql_ast::CommonTableExpr],
database: &vibesql_storage::Database,
executor: F,
memory_check: M,
) -> Result<HashMap<String, CteResult>, ExecutorError>
where
F: Fn(
&vibesql_ast::SelectStmt,
&HashMap<String, CteResult>,
) -> Result<Vec<vibesql_storage::Row>, ExecutorError>,
M: Fn(usize) -> Result<(), ExecutorError>,
{
let mut cte_results = HashMap::new();
for cte in ctes {
let is_recursive = cte.recursive || is_cte_self_referential(cte);
let rows = if is_recursive {
execute_recursive_cte(cte, &cte_results, database, &executor, &memory_check)?
} else {
executor(&cte.query, &cte_results)?
};
let estimated_size = super::helpers::estimate_result_size(&rows);
memory_check(estimated_size)?;
let schema = derive_cte_schema(cte, &rows, database, &cte_results)?;
cte_results.insert(cte.name.clone(), (schema, Arc::new(rows)));
}
Ok(cte_results)
}
pub(super) fn derive_cte_schema(
cte: &vibesql_ast::CommonTableExpr,
rows: &[vibesql_storage::Row],
database: &vibesql_storage::Database,
prior_ctes: &HashMap<String, CteResult>,
) -> Result<vibesql_catalog::TableSchema, ExecutorError> {
if let Some(column_names) = &cte.columns {
if let Some(first_row) = rows.first() {
if first_row.values.len() != column_names.len() {
return Err(ExecutorError::UnsupportedFeature(format!(
"CTE column count mismatch: specified {} columns but query returned {}",
column_names.len(),
first_row.values.len()
)));
}
let columns = column_names
.iter()
.zip(&first_row.values)
.map(|(name, value)| {
let data_type = infer_type_from_value(value);
vibesql_catalog::ColumnSchema::new(name.clone(), data_type, true)
})
.collect();
Ok(cte_pseudo_schema(cte.name.clone(), columns))
} else {
let columns = column_names
.iter()
.map(|name| {
vibesql_catalog::ColumnSchema::new(
name.clone(),
vibesql_types::DataType::Varchar { max_length: Some(255) },
true,
)
})
.collect();
Ok(cte_pseudo_schema(cte.name.clone(), columns))
}
} else {
let mut columns: Vec<vibesql_catalog::ColumnSchema> = Vec::new();
let mut value_idx = 0usize;
for (i, item) in cte.query.select_list.iter().enumerate() {
let names: Vec<String> = match item {
vibesql_ast::SelectItem::Wildcard { .. }
| vibesql_ast::SelectItem::QualifiedWildcard { .. } => {
expand_wildcard_names(item, &cte.query, database, prior_ctes)
.unwrap_or_else(|| vec![format!("col{}", i)])
}
vibesql_ast::SelectItem::Expression { expr, alias, .. } => {
vec![if let Some(a) = alias {
a.clone()
} else {
match expr {
vibesql_ast::Expression::ColumnRef(col_id) => {
col_id.column_canonical().to_string()
}
_ => format!("col{}", i),
}
}]
}
};
for name in names {
let data_type = rows
.first()
.and_then(|first_row| first_row.values.get(value_idx))
.map(infer_type_from_value)
.unwrap_or(vibesql_types::DataType::Varchar { max_length: Some(255) });
columns.push(vibesql_catalog::ColumnSchema::new(name, data_type, true)); value_idx += 1;
}
}
if let Some(first_row) = rows.first() {
if columns.len() != first_row.values.len() {
return Ok(legacy_cte_schema(cte, rows));
}
}
Ok(cte_pseudo_schema(cte.name.clone(), columns))
}
}
fn cte_pseudo_schema(
name: String,
columns: Vec<vibesql_catalog::ColumnSchema>,
) -> vibesql_catalog::TableSchema {
let mut schema = vibesql_catalog::TableSchema::new(name, columns);
schema.set_is_view(true);
schema
}
fn legacy_cte_schema(
cte: &vibesql_ast::CommonTableExpr,
rows: &[vibesql_storage::Row],
) -> vibesql_catalog::TableSchema {
let columns = cte
.query
.select_list
.iter()
.enumerate()
.map(|(i, item)| {
let data_type = rows
.first()
.and_then(|first_row| first_row.values.get(i))
.map(infer_type_from_value)
.unwrap_or(vibesql_types::DataType::Varchar { max_length: Some(255) });
let col_name = match item {
vibesql_ast::SelectItem::Wildcard { .. }
| vibesql_ast::SelectItem::QualifiedWildcard { .. } => format!("col{}", i),
vibesql_ast::SelectItem::Expression { expr, alias, .. } => {
if let Some(a) = alias {
a.clone()
} else {
match expr {
vibesql_ast::Expression::ColumnRef(col_id) => {
col_id.column_canonical().to_string()
}
_ => format!("col{}", i),
}
}
}
};
vibesql_catalog::ColumnSchema::new(col_name, data_type, true) })
.collect();
cte_pseudo_schema(cte.name.clone(), columns)
}
struct WildcardColumn {
name: String,
hidden_for_star: bool,
}
struct WildcardSource {
qualifier: String,
columns: Vec<WildcardColumn>,
}
fn visible_columns(names: Vec<String>) -> Vec<WildcardColumn> {
names.into_iter().map(|name| WildcardColumn { name, hidden_for_star: false }).collect()
}
fn view_column_names(name: &str, database: &vibesql_storage::Database) -> Option<Vec<String>> {
let view = database.catalog.get_view(name)?;
if let Some(cols) = &view.columns {
return Some(cols.clone());
}
collect_select_list_columns(&view.query, database, &HashMap::new())
}
fn expand_wildcard_names(
item: &vibesql_ast::SelectItem,
stmt: &vibesql_ast::SelectStmt,
database: &vibesql_storage::Database,
prior_ctes: &HashMap<String, CteResult>,
) -> Option<Vec<String>> {
match item {
vibesql_ast::SelectItem::Wildcard { alias } => {
if let Some(alias_names) = alias {
return Some(alias_names.clone());
}
let sources = collect_from_sources(stmt.from.as_ref()?, database, prior_ctes)?;
Some(
sources
.into_iter()
.flat_map(|s| s.columns)
.filter(|c| !c.hidden_for_star)
.map(|c| c.name)
.collect(),
)
}
vibesql_ast::SelectItem::QualifiedWildcard { qualifier, alias } => {
if let Some(alias_names) = alias {
return Some(alias_names.clone());
}
let sources = collect_from_sources(stmt.from.as_ref()?, database, prior_ctes)?;
sources
.into_iter()
.find(|s| s.qualifier.eq_ignore_ascii_case(qualifier))
.map(|s| s.columns.into_iter().map(|c| c.name).collect())
}
vibesql_ast::SelectItem::Expression { .. } => None,
}
}
fn collect_from_sources(
from: &vibesql_ast::FromClause,
database: &vibesql_storage::Database,
prior_ctes: &HashMap<String, CteResult>,
) -> Option<Vec<WildcardSource>> {
match from {
vibesql_ast::FromClause::Table { name, alias, column_aliases, .. } => {
let base_columns: Vec<String> = if let Some((schema, _)) =
prior_ctes.get(name).or_else(|| {
prior_ctes.iter().find(|(k, _)| k.eq_ignore_ascii_case(name)).map(|(_, v)| v)
}) {
schema.columns.iter().map(|c| c.name.clone()).collect()
} else if let Some(table) = database.get_table(name) {
table.schema.columns.iter().map(|c| c.name.clone()).collect()
} else if let Some(cols) = view_column_names(name, database) {
cols
} else {
return None;
};
let columns = match column_aliases {
Some(aliases) if aliases.len() == base_columns.len() => aliases.clone(),
Some(_) => return None, None => base_columns,
};
let qualifier = alias.clone().unwrap_or_else(|| name.clone());
Some(vec![WildcardSource { qualifier, columns: visible_columns(columns) }])
}
vibesql_ast::FromClause::Join { left, right, natural, using_columns, alias, .. } => {
if alias.is_some() && (*natural || using_columns.is_some()) {
return None;
}
let mut sources = collect_from_sources(left, database, prior_ctes)?;
let mut right_sources = collect_from_sources(right, database, prior_ctes)?;
if *natural || using_columns.is_some() {
let shared: Vec<String> = if let Some(using) = using_columns {
using.clone()
} else {
let left_visible: Vec<&str> = sources
.iter()
.flat_map(|s| s.columns.iter())
.filter(|c| !c.hidden_for_star)
.map(|c| c.name.as_str())
.collect();
right_sources
.iter()
.flat_map(|s| s.columns.iter())
.filter(|c| !c.hidden_for_star)
.filter(|c| left_visible.iter().any(|l| l.eq_ignore_ascii_case(&c.name)))
.map(|c| c.name.clone())
.collect()
};
for source in &mut right_sources {
for col in &mut source.columns {
if shared.iter().any(|s| s.eq_ignore_ascii_case(&col.name)) {
col.hidden_for_star = true;
}
}
}
}
sources.extend(right_sources);
Some(sources)
}
vibesql_ast::FromClause::Subquery { query, alias, column_aliases } => {
let columns = if let Some(aliases) = column_aliases {
aliases.clone()
} else {
collect_select_list_columns(query, database, prior_ctes)?
};
Some(vec![WildcardSource {
qualifier: alias.clone(),
columns: visible_columns(columns),
}])
}
vibesql_ast::FromClause::Values { rows, alias, column_aliases } => {
let columns = if let Some(aliases) = column_aliases {
aliases.clone()
} else {
let first_row = rows.first()?;
(0..first_row.len()).map(|i| format!("col{}", i)).collect()
};
Some(vec![WildcardSource {
qualifier: alias.clone(),
columns: visible_columns(columns),
}])
}
}
}
pub(crate) fn try_static_select_columns(
stmt: &vibesql_ast::SelectStmt,
database: &vibesql_storage::Database,
) -> Option<Vec<String>> {
collect_select_list_columns(stmt, database, &HashMap::new())
}
fn collect_select_list_columns(
stmt: &vibesql_ast::SelectStmt,
database: &vibesql_storage::Database,
prior_ctes: &HashMap<String, CteResult>,
) -> Option<Vec<String>> {
if let Some(values_rows) = &stmt.values {
let first_row = values_rows.first()?;
return Some((0..first_row.len()).map(|i| format!("col{}", i)).collect());
}
let mut names = Vec::new();
for (i, item) in stmt.select_list.iter().enumerate() {
match item {
vibesql_ast::SelectItem::Wildcard { .. }
| vibesql_ast::SelectItem::QualifiedWildcard { .. } => {
names.extend(expand_wildcard_names(item, stmt, database, prior_ctes)?);
}
vibesql_ast::SelectItem::Expression { expr, alias, .. } => {
names.push(if let Some(a) = alias {
a.clone()
} else {
match expr {
vibesql_ast::Expression::ColumnRef(col_id) => {
col_id.column_canonical().to_string()
}
_ => format!("col{}", i),
}
});
}
}
}
Some(names)
}
fn execute_recursive_cte<F, M>(
cte: &vibesql_ast::CommonTableExpr,
cte_results: &HashMap<String, CteResult>,
database: &vibesql_storage::Database,
executor: &F,
memory_check: &M,
) -> Result<Vec<vibesql_storage::Row>, ExecutorError>
where
F: Fn(
&vibesql_ast::SelectStmt,
&HashMap<String, CteResult>,
) -> Result<Vec<vibesql_storage::Row>, ExecutorError>,
M: Fn(usize) -> Result<(), ExecutorError>,
{
use crate::limits::MAX_RECURSIVE_CTE_ITERATIONS;
let set_op = cte.query.set_operation.as_ref().ok_or_else(|| {
ExecutorError::UnsupportedFeature(format!(
"Recursive CTE '{}' must use UNION ALL",
cte.name
))
})?;
if set_op.op != vibesql_ast::SetOperator::Union {
return Err(ExecutorError::UnsupportedFeature(format!(
"Recursive CTE '{}' must use UNION or UNION ALL (not INTERSECT or EXCEPT)",
cte.name
)));
}
let base_query = vibesql_ast::SelectStmt {
with_clause: cte.query.with_clause.clone(),
distinct: cte.query.distinct,
select_list: cte.query.select_list.clone(),
into_table: cte.query.into_table.clone(),
into_variables: cte.query.into_variables.clone(),
from: cte.query.from.clone(),
where_clause: cte.query.where_clause.clone(),
group_by: cte.query.group_by.clone(),
having: cte.query.having.clone(),
window_definitions: cte.query.window_definitions.clone(),
order_by: cte.query.order_by.clone(),
limit: cte.query.limit.clone(),
offset: cte.query.offset.clone(),
set_operation: None, values: cte.query.values.clone(),
};
let recursive_query = &set_op.right;
if crate::select::window::has_window_functions(&recursive_query.select_list) {
return Err(ExecutorError::SqliteCompatError(
"cannot use window functions in recursive queries".to_string(),
));
}
if let (Some(base_count), Some(recursive_count)) =
(count_stmt_columns(&base_query), count_stmt_columns(recursive_query))
{
if base_count != recursive_count {
return Err(ExecutorError::UnsupportedFeature(
"SELECTs to the left and right of UNION ALL do not have the same number of result columns".to_string()
));
}
}
let mut all_rows = executor(&base_query, cte_results)?;
let mut working_table = all_rows.clone();
let schema = derive_cte_schema(cte, &all_rows, database, cte_results)?;
let mut seen_rows: Option<HashSet<vibesql_storage::RowValues>> = if !set_op.all {
let mut seen = HashSet::with_capacity(all_rows.len());
for row in &all_rows {
seen.insert(row.values.clone());
}
Some(seen)
} else {
None
};
let mut depth = 0;
while !working_table.is_empty() && depth < MAX_RECURSIVE_CTE_ITERATIONS {
depth += 1;
let mut recursive_cte_results = cte_results.clone();
recursive_cte_results
.insert(cte.name.clone(), (schema.clone(), Arc::new(working_table.clone())));
let new_rows = executor(recursive_query, &recursive_cte_results)?;
if new_rows.is_empty() {
break;
}
if depth == 1 && !new_rows.is_empty() && !all_rows.is_empty() {
let base_col_count = all_rows[0].values.len();
let recursive_col_count = new_rows[0].values.len();
if base_col_count != recursive_col_count {
return Err(ExecutorError::UnsupportedFeature(
"SELECTs to the left and right of UNION ALL do not have the same number of result columns".to_string()
));
}
}
let estimated_size = super::helpers::estimate_result_size(&new_rows);
memory_check(estimated_size)?;
let rows_to_add: Vec<vibesql_storage::Row> = if let Some(ref mut seen) = seen_rows {
new_rows.into_iter().filter(|row| seen.insert(row.values.clone())).collect()
} else {
new_rows
};
if rows_to_add.is_empty() {
break;
}
all_rows.extend(rows_to_add.clone());
working_table = rows_to_add;
}
if depth >= MAX_RECURSIVE_CTE_ITERATIONS {
return Err(ExecutorError::UnsupportedFeature(format!(
"Recursive CTE '{}' exceeded maximum iteration limit of {}",
cte.name, MAX_RECURSIVE_CTE_ITERATIONS
)));
}
Ok(all_rows)
}
fn count_explicit_columns(select_list: &[vibesql_ast::SelectItem]) -> Option<usize> {
let mut count = 0;
for item in select_list {
match item {
vibesql_ast::SelectItem::Expression { .. } => count += 1,
vibesql_ast::SelectItem::Wildcard { .. }
| vibesql_ast::SelectItem::QualifiedWildcard { .. } => {
return None;
}
}
}
Some(count)
}
fn count_stmt_columns(stmt: &vibesql_ast::SelectStmt) -> Option<usize> {
if let Some(values_rows) = &stmt.values {
return values_rows.first().map(|row| row.len());
}
count_explicit_columns(&stmt.select_list)
}
fn is_cte_self_referential(cte: &vibesql_ast::CommonTableExpr) -> bool {
let set_op = match &cte.query.set_operation {
Some(op) if op.op == vibesql_ast::SetOperator::Union => op,
_ => return false,
};
stmt_references_table(&set_op.right, &cte.name)
}
fn stmt_references_table(stmt: &vibesql_ast::SelectStmt, table_name: &str) -> bool {
if let Some(from) = &stmt.from {
if from_clause_references_table(from, table_name) {
return true;
}
}
if let Some(where_clause) = &stmt.where_clause {
if expr_references_table(where_clause, table_name) {
return true;
}
}
for item in &stmt.select_list {
if let vibesql_ast::SelectItem::Expression { expr, .. } = item {
if expr_references_table(expr, table_name) {
return true;
}
}
}
false
}
fn from_clause_references_table(from: &vibesql_ast::FromClause, table_name: &str) -> bool {
match from {
vibesql_ast::FromClause::Table { name, .. } => name.eq_ignore_ascii_case(table_name),
vibesql_ast::FromClause::Subquery { query, .. } => stmt_references_table(query, table_name),
vibesql_ast::FromClause::Join { left, right, condition, .. } => {
from_clause_references_table(left, table_name)
|| from_clause_references_table(right, table_name)
|| condition.as_ref().map_or(false, |c| expr_references_table(c, table_name))
}
vibesql_ast::FromClause::Values { .. } => false,
}
}
fn expr_references_table(expr: &vibesql_ast::Expression, table_name: &str) -> bool {
match expr {
vibesql_ast::Expression::ScalarSubquery(subquery) => {
stmt_references_table(subquery, table_name)
}
vibesql_ast::Expression::In { subquery, .. } => stmt_references_table(subquery, table_name),
vibesql_ast::Expression::Exists { subquery, .. } => {
stmt_references_table(subquery, table_name)
}
vibesql_ast::Expression::BinaryOp { left, right, .. } => {
expr_references_table(left, table_name) || expr_references_table(right, table_name)
}
vibesql_ast::Expression::UnaryOp { expr, .. } => expr_references_table(expr, table_name),
vibesql_ast::Expression::Function { args, .. } => {
args.iter().any(|arg| expr_references_table(arg, table_name))
}
vibesql_ast::Expression::AggregateFunction { args, filter, .. } => {
args.iter().any(|arg| expr_references_table(arg, table_name))
|| filter.as_ref().map_or(false, |f| expr_references_table(f, table_name))
}
vibesql_ast::Expression::Case { operand, when_clauses, else_result, .. } => {
operand.as_ref().map_or(false, |o| expr_references_table(o, table_name))
|| when_clauses.iter().any(|when| {
when.conditions.iter().any(|c| expr_references_table(c, table_name))
|| expr_references_table(&when.result, table_name)
})
|| else_result.as_ref().map_or(false, |e| expr_references_table(e, table_name))
}
vibesql_ast::Expression::Between { expr, low, high, .. } => {
expr_references_table(expr, table_name)
|| expr_references_table(low, table_name)
|| expr_references_table(high, table_name)
}
vibesql_ast::Expression::InList { expr, values, .. } => {
expr_references_table(expr, table_name)
|| values.iter().any(|e| expr_references_table(e, table_name))
}
vibesql_ast::Expression::Cast { expr, .. }
| vibesql_ast::Expression::Collate { expr, .. } => expr_references_table(expr, table_name),
vibesql_ast::Expression::Conjunction(exprs)
| vibesql_ast::Expression::Disjunction(exprs) => {
exprs.iter().any(|e| expr_references_table(e, table_name))
}
vibesql_ast::Expression::QuantifiedComparison { expr, subquery, .. } => {
expr_references_table(expr, table_name) || stmt_references_table(subquery, table_name)
}
_ => false,
}
}
pub(super) fn infer_type_from_value(value: &vibesql_types::SqlValue) -> vibesql_types::DataType {
match value {
vibesql_types::SqlValue::Null => vibesql_types::DataType::Varchar { max_length: Some(255) },
vibesql_types::SqlValue::Integer(_) => vibesql_types::DataType::Integer,
vibesql_types::SqlValue::Varchar(_) => {
vibesql_types::DataType::Varchar { max_length: Some(255) }
}
vibesql_types::SqlValue::Character(_) => vibesql_types::DataType::Character { length: 1 },
vibesql_types::SqlValue::Boolean(_) => vibesql_types::DataType::Boolean,
vibesql_types::SqlValue::Float(_) => vibesql_types::DataType::Float { precision: 53 },
vibesql_types::SqlValue::Double(_) => vibesql_types::DataType::DoublePrecision,
vibesql_types::SqlValue::Numeric(_) => {
vibesql_types::DataType::Numeric { precision: 10, scale: 2 }
}
vibesql_types::SqlValue::Real(_) => vibesql_types::DataType::Real,
vibesql_types::SqlValue::Smallint(_) => vibesql_types::DataType::Smallint,
vibesql_types::SqlValue::Bigint(_) => vibesql_types::DataType::Bigint,
vibesql_types::SqlValue::Unsigned(_) => vibesql_types::DataType::Unsigned,
vibesql_types::SqlValue::Date(_) => vibesql_types::DataType::Date,
vibesql_types::SqlValue::Time(_) => vibesql_types::DataType::Time { with_timezone: false },
vibesql_types::SqlValue::Timestamp(_) => {
vibesql_types::DataType::Timestamp { with_timezone: false }
}
vibesql_types::SqlValue::Interval(_) => {
vibesql_types::DataType::Interval {
start_field: vibesql_types::IntervalField::Day,
end_field: None,
}
}
vibesql_types::SqlValue::Vector(v) => {
vibesql_types::DataType::Vector { dimensions: v.len() as u32 }
}
vibesql_types::SqlValue::Blob(_) => vibesql_types::DataType::BinaryLargeObject,
}
}