mool 0.2.0

A source-first Rust database toolkit for typed SQL queries, migrations, and model metadata
Documentation
//! Parameter planning and SQLx statement binding for typed queries.

use std::any::TypeId;
use std::collections::{HashMap, HashSet};

use crate::argvalue::ArgValue;
use crate::commons::Arguments;
use crate::interfaces::Record;

use super::super::{QueryError, Statement};
use super::GENERATED_PREFIX;
use super::expr::{ColumnRef, ExprNode};
use super::handles::{Table, VarId};
use super::output::SelectAssignment;
use super::plan::{ParamSource, QueryPlan};
use super::source::Source;
use super::validate::{output_column, table_name, validate_identifier};

pub(super) fn statement_from_plan(
    plan: QueryPlan,
    mut args: Arguments<'static>,
) -> Result<Statement, QueryError> {
    validate_unused_binds(&plan)?;
    for name in &plan.bind_order {
        let value = plan
            .values
            .get(name)
            .ok_or_else(|| QueryError::MissingBinding(name.clone()))?;
        value
            .bind_value(&mut args)
            .map_err(|err| QueryError::BindError(err.to_string()))?;
    }
    Ok(Statement::new(&plan.sql, args))
}

pub(super) fn finish_plan(plan: QueryPlan) -> Result<QueryPlan, QueryError> {
    validate_missing_binds(&plan)?;
    validate_unused_binds(&plan)?;
    Ok(plan)
}

fn validate_missing_binds(plan: &QueryPlan) -> Result<(), QueryError> {
    for (name, spec) in &plan.params {
        if spec.source == ParamSource::Var && !plan.values.contains_key(name) {
            return Err(QueryError::MissingBinding(
                spec.display_name
                    .clone()
                    .unwrap_or_else(|| spec.name.clone()),
            ));
        }
    }
    Ok(())
}

pub(super) fn validate_unused_binds(plan: &QueryPlan) -> Result<(), QueryError> {
    for name in plan.values.keys() {
        if name.starts_with(GENERATED_PREFIX) {
            continue;
        }
        match plan.params.get(name) {
            Some(spec) if spec.source == ParamSource::Var => {}
            Some(_) => {}
            None => return Err(QueryError::UnusedBinding(name.clone())),
        }
    }
    Ok(())
}

pub(super) fn collect_expr_binds(
    node: &ExprNode,
    values: &mut HashMap<VarId, ArgValue>,
) -> Result<(), QueryError> {
    match node {
        ExprNode::Column(_) | ExprNode::Value(_) => Ok(()),
        ExprNode::Binary { left, right, .. } | ExprNode::Bool { left, right, .. } => {
            collect_expr_binds(left, values)?;
            collect_expr_binds(right, values)
        }
        ExprNode::Unary { expr, .. } => collect_expr_binds(expr, values),
        ExprNode::Function { args, .. } | ExprNode::Custom { args, .. } => {
            for arg in args {
                collect_expr_binds(arg, values)?;
            }
            Ok(())
        }
        ExprNode::Over { expr, window } => {
            collect_expr_binds(expr, values)?;
            collect_window_binds(window, values)
        }
        ExprNode::InSource { left, source } => {
            collect_expr_binds(left, values)?;
            collect_source_binds(&source.source, values)
        }
        ExprNode::InList { left, values: list } => {
            collect_expr_binds(left, values)?;
            for value in list {
                collect_expr_binds(value, values)?;
            }
            Ok(())
        }
        ExprNode::RelationExists { predicate, .. } => {
            if let Some(predicate) = predicate {
                collect_expr_binds(predicate, values)?;
            }
            Ok(())
        }
        ExprNode::RelationAggregate { expr, .. } => {
            if let Some(expr) = expr {
                collect_expr_binds(expr, values)?;
            }
            Ok(())
        }
        ExprNode::ManyToManyExists { predicate, .. } => {
            if let Some(predicate) = predicate {
                collect_expr_binds(predicate, values)?;
            }
            Ok(())
        }
    }
}

pub(super) fn collect_source_binds(
    source: &Source,
    values: &mut HashMap<VarId, ArgValue>,
) -> Result<(), QueryError> {
    match source {
        Source::Subquery(subquery) => subquery.data.scope.collect_binds_into(values),
        Source::Cte(_) | Source::Table(_) => Ok(()),
    }
}

pub(super) fn collect_source_ctes(source: &Source, used: &mut HashSet<String>) {
    if let Source::Cte(cte) = source {
        used.insert(cte.data.name.to_string());
    }
}

pub(super) fn collect_expr_ctes(node: &ExprNode, used: &mut HashSet<String>) {
    match node {
        ExprNode::Column(_) | ExprNode::Value(_) => {}
        ExprNode::Binary { left, right, .. } | ExprNode::Bool { left, right, .. } => {
            collect_expr_ctes(left, used);
            collect_expr_ctes(right, used);
        }
        ExprNode::Unary { expr, .. } => collect_expr_ctes(expr, used),
        ExprNode::Function { args, .. } | ExprNode::Custom { args, .. } => {
            for arg in args {
                collect_expr_ctes(arg, used);
            }
        }
        ExprNode::Over { expr, window } => {
            collect_expr_ctes(expr, used);
            collect_window_ctes(window, used);
        }
        ExprNode::InSource { left, source } => {
            collect_expr_ctes(left, used);
            collect_source_ctes(&source.source, used);
        }
        ExprNode::InList { left, values } => {
            collect_expr_ctes(left, used);
            for value in values {
                collect_expr_ctes(value, used);
            }
        }
        ExprNode::RelationExists { predicate, .. } => {
            if let Some(predicate) = predicate {
                collect_expr_ctes(predicate, used);
            }
        }
        ExprNode::RelationAggregate { expr, .. } => {
            if let Some(expr) = expr {
                collect_expr_ctes(expr, used);
            }
        }
        ExprNode::ManyToManyExists { predicate, .. } => {
            if let Some(predicate) = predicate {
                collect_expr_ctes(predicate, used);
            }
        }
    }
}

fn collect_window_binds(
    window: &super::window::WindowSpec,
    values: &mut HashMap<VarId, ArgValue>,
) -> Result<(), QueryError> {
    for expr in &window.partitions {
        collect_expr_binds(expr, values)?;
    }
    for order in &window.orders {
        collect_expr_binds(&order.expr, values)?;
    }
    if let Some(frame) = &window.frame {
        collect_bound_binds(&frame.start, values)?;
        collect_bound_binds(&frame.end, values)?;
    }
    Ok(())
}

fn collect_bound_binds(
    bound: &super::window::FrameBound,
    values: &mut HashMap<VarId, ArgValue>,
) -> Result<(), QueryError> {
    if let Some(expr) = &bound.expr {
        collect_expr_binds(expr, values)?;
    }
    Ok(())
}

fn collect_window_ctes(window: &super::window::WindowSpec, used: &mut HashSet<String>) {
    for expr in &window.partitions {
        collect_expr_ctes(expr, used);
    }
    for order in &window.orders {
        collect_expr_ctes(&order.expr, used);
    }
    if let Some(frame) = &window.frame {
        collect_bound_ctes(&frame.start, used);
        collect_bound_ctes(&frame.end, used);
    }
}

fn collect_bound_ctes(bound: &super::window::FrameBound, used: &mut HashSet<String>) {
    if let Some(expr) = &bound.expr {
        collect_expr_ctes(expr, used);
    }
}

pub(super) fn validate_cte_usage(
    defined: &HashSet<String>,
    used: &HashSet<String>,
) -> Result<(), QueryError> {
    for name in defined {
        if !used.contains(name) {
            return Err(QueryError::BindError(format!("unused CTE '{}'", name)));
        }
    }
    for name in used {
        if !defined.contains(name) {
            return Err(QueryError::BindError(format!(
                "CTE '{}' is not registered",
                name
            )));
        }
    }
    Ok(())
}

pub(super) fn insert_bind(
    values: &mut HashMap<VarId, ArgValue>,
    id: VarId,
    value: ArgValue,
) -> Result<(), QueryError> {
    if values.contains_key(&id) {
        return Err(QueryError::BindError(format!(
            "duplicate binding for var {}",
            id.value()
        )));
    }
    values.insert(id, value);
    Ok(())
}

pub(super) fn validate_output_assignments(
    exprs: &[SelectAssignment],
    columns: &[String],
    record: TypeId,
    record_name: &'static str,
) -> Result<(), QueryError> {
    let mut seen = HashSet::new();
    for assignment in exprs {
        let name = assignment.target.name.as_ref();
        validate_identifier(name).map_err(|_| QueryError::InvalidProjection(name.to_string()))?;
        validate_output_record(assignment, record, record_name)?;
        let target_exists = columns.iter().any(|column| {
            column == name || matches!(output_column(column), Ok(output) if output == name)
        });
        if !target_exists || !seen.insert(name.to_string()) {
            return Err(QueryError::InvalidProjection(name.to_string()));
        }
    }
    Ok(())
}

fn validate_output_record(
    assignment: &SelectAssignment,
    record: TypeId,
    record_name: &'static str,
) -> Result<(), QueryError> {
    if assignment.target.record == record {
        return Ok(());
    }
    Err(QueryError::InvalidProjection(format!(
        "{} belongs to {}, expected {}",
        assignment.target.name, assignment.target.record_name, record_name
    )))
}

pub(super) fn bind_columns<T>() -> Result<Vec<String>, QueryError>
where
    T: Record,
{
    let columns = T::record_bind_column_names();
    if columns.is_empty() {
        return Err(QueryError::BindError("no bindable columns".to_string()));
    }
    for column in &columns {
        validate_identifier(column)?;
    }
    Ok(columns)
}

pub(super) fn validate_bind_columns(table: &Table, columns: &[String]) -> Result<(), QueryError> {
    let Some(known) = table.data.columns.as_ref() else {
        return Ok(());
    };
    for column in columns {
        if !known.iter().any(|known| known == column) {
            return Err(QueryError::BindError(format!(
                "column '{}' is not writable for {}",
                column,
                table_name(table.data.schema.as_deref(), table.data.name.as_ref())
            )));
        }
    }
    Ok(())
}

pub(super) fn upsert_update_columns<'a>(
    columns: &'a [String],
    conflict: &[ColumnRef],
) -> Result<Vec<&'a str>, QueryError> {
    let mut update_columns = Vec::with_capacity(columns.len());
    for column in columns {
        validate_identifier(column)?;
        if !conflict
            .iter()
            .any(|conflict| conflict.name.as_ref() == column)
        {
            update_columns.push(column.as_str());
        }
    }
    Ok(update_columns)
}

pub(super) fn bind_rows<T>(rows: &[T], col_count: usize) -> Result<Arguments<'static>, QueryError>
where
    T: Record,
{
    if rows.is_empty() {
        return Err(QueryError::BindError(
            "cannot insert empty list".to_string(),
        ));
    }
    let mut args = Arguments::default();
    for row in rows {
        bind_row(row, col_count, &mut args)?;
    }
    Ok(args)
}

pub(super) fn bind_row<T>(
    row: &T,
    col_count: usize,
    args: &mut Arguments<'static>,
) -> Result<(), QueryError>
where
    T: Record,
{
    use sqlx::Arguments as _;

    let before = args.len();
    row.record_bind_values(args)
        .map_err(|err| QueryError::BindError(err.to_string()))?;
    let bound = args.len().saturating_sub(before);
    if bound != col_count {
        return Err(QueryError::BindCountMismatch {
            expected: col_count,
            got: bound,
        });
    }
    Ok(())
}