vibesql-executor 0.2.0

//! Column position extraction, validation, and wildcard expansion for ORDER BY.
//!
//! This module handles:
//! - Extracting numeric column positions from ORDER BY expressions (`ORDER BY 1`, `ORDER BY +N`)
//! - Validating column positions are within range
//! - Resolving positions to column names for aggregate queries
//! - Wildcard expansion for position resolution

use crate::{errors::ExecutorError, schema::CombinedSchema};

/// Result of extracting a column position from an ORDER BY expression
#[derive(Debug, Clone, Copy)]
pub(crate) enum ColumnPositionResult {
    /// A valid positive column position (1-indexed)
    Position(i64),
    /// A negative column position (always invalid, but we track it for error reporting)
    Negative(i64),
    /// Not a column position expression
    NotAPosition,
}

/// Extracts a numeric column position from an ORDER BY expression.
///
/// Handles:
/// - `ORDER BY N` - Integer literal → Position(N)
/// - `ORDER BY +N` - Unary plus with integer → Position(N) (#4418)
/// - `ORDER BY -N` - Unary minus with integer → Negative(N)
/// - Other expressions → NotAPosition
pub(crate) fn extract_column_position(expr: &vibesql_ast::Expression) -> ColumnPositionResult {
    match expr {
        // Direct integer literal: ORDER BY 1, ORDER BY 2, etc.
        vibesql_ast::Expression::Literal(vibesql_types::SqlValue::Integer(pos)) => {
            ColumnPositionResult::Position(*pos)
        }
        // Unary operator with integer
        vibesql_ast::Expression::UnaryOp { op, expr } => {
            if let vibesql_ast::Expression::Literal(vibesql_types::SqlValue::Integer(pos)) =
                expr.as_ref()
            {
                match op {
                    // ORDER BY +N: treat as ORDER BY N (#4418)
                    vibesql_ast::UnaryOperator::Plus => ColumnPositionResult::Position(*pos),
                    // ORDER BY -N: always invalid
                    vibesql_ast::UnaryOperator::Minus => ColumnPositionResult::Negative(*pos),
                    _ => ColumnPositionResult::NotAPosition,
                }
            } else {
                ColumnPositionResult::NotAPosition
            }
        }
        _ => ColumnPositionResult::NotAPosition,
    }
}

/// Validates a column position and returns an error if out of range.
pub(crate) fn validate_column_position(
    pos: i64,
    column_count: usize,
    term_index: usize,
) -> Result<usize, ExecutorError> {
    if pos <= 0 || (pos as usize) > column_count {
        return Err(ExecutorError::OrderByOutOfRange {
            term_position: term_index + 1, // Convert to 1-indexed for error message
            column_number: pos,
            select_list_len: column_count,
        });
    }
    Ok((pos as usize) - 1) // Convert to 0-indexed
}

/// Resolves a column position to a column name for aggregate queries.
///
/// For aggregate queries, we simply look up the alias/expression at the given position
/// in the SELECT list without wildcard expansion (aggregates don't typically use wildcards).
pub(crate) fn resolve_position_to_column_name(
    idx: usize,
    select_list: &[vibesql_ast::SelectItem],
) -> Option<String> {
    if let vibesql_ast::SelectItem::Expression { expr, alias, .. } = &select_list[idx] {
        Some(if let Some(alias_name) = alias {
            alias_name.clone()
        } else if let vibesql_ast::Expression::ColumnRef(col_id) = expr {
            col_id.column_canonical().to_string()
        } else if let vibesql_ast::Expression::AggregateFunction { name, .. } = expr {
            // For aggregate functions without alias, use the function name (lowercase)
            // This matches the schema column name generated by the aggregate executor
            name.to_lowercase()
        } else {
            format!("col{}", idx + 1)
        })
    } else {
        None
    }
}

/// Result of resolving a column position with wildcard expansion
pub(crate) enum ResolvedPosition<'a> {
    /// The expression at the position (for regular SELECT items)
    Expression(&'a vibesql_ast::Expression),
    /// A column resolved from a wildcard expansion.
    ///
    /// The column is qualified with the canonical name of the table it came
    /// from (when known) so downstream name resolution does not trip the
    /// ambiguity check when the same column name exists in multiple tables
    /// (issue #5231: `SELECT * FROM b, (SELECT x FROM a) ORDER BY 1` must not
    /// fail with "ambiguous column name: x").
    ColumnName {
        /// Canonical table name the column belongs to, if known
        table: Option<String>,
        /// The column name
        column: String,
    },
    /// Position not found (shouldn't happen if validation passed)
    NotFound,
}

/// Resolves a column position to an expression or column name, handling wildcard expansion.
///
/// This handles SELECT lists that may contain wildcards (`*`) or qualified wildcards (`table.*`)
/// which expand to multiple columns. The position is 0-indexed.
pub(crate) fn resolve_position_with_wildcards<'a>(
    target_col: usize,
    select_list: &'a [vibesql_ast::SelectItem],
    schema: Option<&CombinedSchema>,
) -> ResolvedPosition<'a> {
    let mut current_col = 0;

    for item in select_list {
        match item {
            vibesql_ast::SelectItem::Wildcard { .. } => {
                // Count how many columns this wildcard expands to
                let wildcard_cols = if let Some(s) = schema {
                    s.table_schemas.values().map(|(_, ts)| ts.columns.len()).sum()
                } else {
                    1
                };

                if target_col < current_col + wildcard_cols {
                    // The position is within this wildcard's expanded columns
                    if let Some(s) = schema {
                        let offset_in_wildcard = target_col - current_col;
                        // Collect all columns from all tables in schema order,
                        // tracking which table each column belongs to so the
                        // result can be table-qualified (issue #5231)
                        let mut all_columns: Vec<(usize, String, String)> = Vec::new();
                        for (table_id, (start_idx, table_schema)) in &s.table_schemas {
                            for (i, col) in table_schema.columns.iter().enumerate() {
                                all_columns.push((
                                    start_idx + i,
                                    table_id.table_canonical().to_string(),
                                    col.name.clone(),
                                ));
                            }
                        }
                        all_columns.sort_by_key(|(idx, _, _)| *idx);

                        if offset_in_wildcard < all_columns.len() {
                            let (_, table, column) = all_columns[offset_in_wildcard].clone();
                            return ResolvedPosition::ColumnName { table: Some(table), column };
                        }
                    }
                    return ResolvedPosition::NotFound;
                }
                current_col += wildcard_cols;
            }
            vibesql_ast::SelectItem::QualifiedWildcard { qualifier, .. } => {
                // Count how many columns this qualified wildcard expands to
                let qualified_cols = if let Some(s) = schema {
                    s.get_table(qualifier).map(|(_, ts)| ts.columns.len()).unwrap_or(1)
                } else {
                    1
                };

                if target_col < current_col + qualified_cols {
                    // The position is within this qualified wildcard's columns
                    if let Some(s) = schema {
                        if let Some((_, table_schema)) = s.get_table(qualifier) {
                            let offset = target_col - current_col;
                            if offset < table_schema.columns.len() {
                                // Qualify with the (canonicalized) qualifier so
                                // downstream resolution is unambiguous (#5231).
                                // get_table resolves via TableIdentifier::unquoted,
                                // so lowercase folding matches its lookup.
                                return ResolvedPosition::ColumnName {
                                    table: Some(qualifier.to_ascii_lowercase()),
                                    column: table_schema.columns[offset].name.clone(),
                                };
                            }
                        }
                    }
                    return ResolvedPosition::NotFound;
                }
                current_col += qualified_cols;
            }
            vibesql_ast::SelectItem::Expression { expr, .. } => {
                if target_col == current_col {
                    return ResolvedPosition::Expression(expr);
                }
                current_col += 1;
            }
        }
    }

    ResolvedPosition::NotFound
}

/// Count the number of columns in a SELECT list after wildcard expansion
///
/// This correctly handles:
/// - Regular expressions: count as 1 column each
/// - Wildcard (`SELECT *`): expands to all columns from the schema
/// - Qualified wildcards (`SELECT table.*`): expands to all columns from that table
///
/// # Arguments
/// * `select_list` - The SELECT list items
/// * `schema` - Optional schema for wildcard expansion. If None, wildcards count as 1.
pub(crate) fn count_select_columns(
    select_list: &[vibesql_ast::SelectItem],
    schema: Option<&CombinedSchema>,
) -> usize {
    let mut count = 0;
    for item in select_list {
        match item {
            vibesql_ast::SelectItem::Wildcard { .. } => {
                // Expand wildcard to all columns from schema
                if let Some(s) = schema {
                    // Count all columns from all tables in the schema
                    for (_, table_schema) in s.table_schemas.values() {
                        count += table_schema.columns.len();
                    }
                } else {
                    // No schema available, count as 1
                    count += 1;
                }
            }
            vibesql_ast::SelectItem::QualifiedWildcard { qualifier, .. } => {
                // Expand qualified wildcard to all columns from the specified table
                if let Some(s) = schema {
                    if let Some((_, table_schema)) = s.get_table(qualifier) {
                        count += table_schema.columns.len();
                    } else {
                        // Table not found, count as 1 (error will be caught elsewhere)
                        count += 1;
                    }
                } else {
                    // No schema available, count as 1
                    count += 1;
                }
            }
            vibesql_ast::SelectItem::Expression { .. } => {
                count += 1;
            }
        }
    }
    count
}