vibesql_executor/cache/prepared_statement/

plan.rs

//! Query plan caching for prepared statements
//!
//! This module provides cached execution plans that bypass the full query planning
//! pipeline for simple query patterns. When a prepared statement is created, we
//! analyze the AST to detect patterns that can use fast execution paths.
//!
//! ## Supported Patterns
//!
//! ### Primary Key Point Lookup
//!
//! Queries like `SELECT * FROM table WHERE pk_col = ?` are detected and cached with:
//! - Table name
//! - Primary key column indices in schema
//! - Parameter indices for PK values
//! - Projection column indices (or wildcard flag)
//!
//! At execution time, this allows direct index lookup without:
//! - Re-parsing the AST
//! - Checking if it's a simple query
//! - Resolving table/column names
//! - Building schema objects
//!
//! ## Performance Impact
//!
//! For simple point SELECT via prepared statement:
//! - Without plan caching: ~60µs (AST manipulation, pattern detection, table lookup)
//! - With plan caching: ~1-5µs (direct index lookup)
//!
//! ## Cache Invalidation
//!
//! Cached plans are invalidated when:
//! - DDL modifies the referenced table (schema change)
//! - The table is dropped
//! - Indexes are added/removed

use vibesql_ast::{Expression, FromClause, SelectItem, SelectStmt, Statement};

/// Cached execution plan for a prepared statement
#[derive(Debug, Clone)]
pub enum CachedPlan {
    /// Simple primary key point lookup
    /// SELECT [cols] FROM table WHERE pk_col = ? [AND pk_col2 = ? ...]
    PkPointLookup(PkPointLookupPlan),

    /// Query doesn't match any fast-path pattern
    /// Fall back to standard execution
    Standard,
}

impl CachedPlan {
    /// Check if this is a cacheable fast-path plan
    pub fn is_fast_path(&self) -> bool {
        !matches!(self, CachedPlan::Standard)
    }
}

/// Cached plan for primary key point lookup queries
#[derive(Debug, Clone)]
pub struct PkPointLookupPlan {
    /// Table name (normalized to uppercase for case-insensitive matching)
    pub table_name: String,

    /// Primary key column names in order
    pub pk_columns: Vec<String>,

    /// Mapping from parameter index (0-based) to PK column index
    /// e.g., for `WHERE pk1 = ? AND pk2 = ?`, this would be [(0, 0), (1, 1)]
    pub param_to_pk_col: Vec<(usize, usize)>,

    /// Projection type
    pub projection: ProjectionPlan,
}

/// How to project columns from the result
#[derive(Debug, Clone)]
pub enum ProjectionPlan {
    /// SELECT * - return all columns
    Wildcard,

    /// SELECT col1, col2, ... - return specific columns by index
    Columns(Vec<ColumnProjection>),
}

/// A single column in the projection
#[derive(Debug, Clone)]
pub struct ColumnProjection {
    /// Column name in the table
    pub column_name: String,

    /// Output alias (if specified via AS)
    pub alias: Option<String>,
}

/// Analyze a prepared statement and create a cached plan if possible
pub fn analyze_for_plan(stmt: &Statement) -> CachedPlan {
    match stmt {
        Statement::Select(select) => analyze_select(select),
        _ => CachedPlan::Standard,
    }
}

/// Analyze a SELECT statement for caching opportunities
fn analyze_select(stmt: &SelectStmt) -> CachedPlan {
    // Must have exactly one table in FROM (no joins, no subqueries)
    let table_name = match &stmt.from {
        Some(FromClause::Table { name, alias: None, .. }) => name.clone(),
        Some(FromClause::Table { name: _, alias: Some(_), .. }) => {
            // Aliased tables are slightly more complex, skip for now
            return CachedPlan::Standard;
        }
        _ => return CachedPlan::Standard,
    };

    // No complex clauses
    if stmt.with_clause.is_some()
        || stmt.set_operation.is_some()
        || stmt.group_by.is_some()
        || stmt.having.is_some()
        || stmt.distinct
        || stmt.order_by.is_some()
        || stmt.limit.is_some()
        || stmt.offset.is_some()
        || stmt.into_table.is_some()
        || stmt.into_variables.is_some()
    {
        return CachedPlan::Standard;
    }

    // Check SELECT list - must be wildcard or simple column references
    let projection = match analyze_select_list(&stmt.select_list) {
        Some(p) => p,
        None => return CachedPlan::Standard,
    };

    // Must have a WHERE clause with PK equality predicates
    let where_clause = match &stmt.where_clause {
        Some(w) => w,
        None => return CachedPlan::Standard,
    };

    // Extract parameter-to-column mappings from WHERE clause
    let param_mappings = match extract_pk_param_mappings(where_clause) {
        Some(m) if !m.is_empty() => m,
        _ => return CachedPlan::Standard,
    };

    // Build the plan
    // Note: We don't validate PK columns here because we don't have DB access.
    // The Session will validate at prepare time and fall back if not valid.
    let pk_columns: Vec<String> = param_mappings.iter().map(|(_, col)| col.clone()).collect();
    let param_to_pk_col: Vec<(usize, usize)> = param_mappings
        .iter()
        .enumerate()
        .map(|(pk_idx, (param_idx, _))| (*param_idx, pk_idx))
        .collect();

    CachedPlan::PkPointLookup(PkPointLookupPlan {
        table_name: table_name.to_uppercase(),
        pk_columns,
        param_to_pk_col,
        projection,
    })
}

/// Analyze SELECT list and return projection plan if it's simple
fn analyze_select_list(select_list: &[SelectItem]) -> Option<ProjectionPlan> {
    if select_list.len() == 1 {
        if let SelectItem::Wildcard { .. } = &select_list[0] {
            return Some(ProjectionPlan::Wildcard);
        }
    }

    // Check each item is a simple column reference
    let mut columns = Vec::with_capacity(select_list.len());
    for item in select_list {
        match item {
            SelectItem::Wildcard { .. } => {
                // Mixed wildcard and columns - fall back
                return None;
            }
            SelectItem::QualifiedWildcard { .. } => {
                // table.* - fall back for now
                return None;
            }
            SelectItem::Expression { expr, alias } => {
                let column_name = match expr {
                    Expression::ColumnRef { column, table: None } => column.clone(),
                    Expression::ColumnRef { column, table: Some(_) } => {
                        // Qualified column ref - support it
                        column.clone()
                    }
                    _ => {
                        // Complex expression - fall back
                        return None;
                    }
                };
                columns.push(ColumnProjection {
                    column_name,
                    alias: alias.clone(),
                });
            }
        }
    }

    Some(ProjectionPlan::Columns(columns))
}

/// Extract parameter-to-column mappings from WHERE clause
///
/// Returns a list of (parameter_index, column_name) pairs in order.
/// For `WHERE pk1 = ? AND pk2 = ?`, returns [(0, "pk1"), (1, "pk2")]
fn extract_pk_param_mappings(expr: &Expression) -> Option<Vec<(usize, String)>> {
    let mut mappings = Vec::new();
    collect_pk_param_mappings(expr, &mut mappings)?;

    // Sort by parameter index to ensure consistent ordering
    mappings.sort_by_key(|(idx, _)| *idx);

    Some(mappings)
}

/// Recursively collect parameter-to-column mappings
fn collect_pk_param_mappings(
    expr: &Expression,
    mappings: &mut Vec<(usize, String)>,
) -> Option<()> {
    match expr {
        Expression::BinaryOp { left, op, right } => {
            match op {
                vibesql_ast::BinaryOperator::And => {
                    // Recurse into both sides
                    collect_pk_param_mappings(left, mappings)?;
                    collect_pk_param_mappings(right, mappings)?;
                }
                vibesql_ast::BinaryOperator::Equal => {
                    // Check for col = ? pattern
                    if let Some(mapping) = extract_column_placeholder_pair(left, right) {
                        mappings.push(mapping);
                    } else {
                        // Non-placeholder equality - not a fast-path query
                        return None;
                    }
                }
                _ => {
                    // Other operators (OR, <, >, etc.) - not a simple PK lookup
                    return None;
                }
            }
        }
        _ => {
            // Other expression types - not supported
            return None;
        }
    }
    Some(())
}

/// Extract (param_index, column_name) from an equality expression
fn extract_column_placeholder_pair(left: &Expression, right: &Expression) -> Option<(usize, String)> {
    // Try col = ?
    if let Expression::ColumnRef { column, .. } = left {
        if let Expression::Placeholder(idx) = right {
            return Some((*idx, column.clone()));
        }
    }

    // Try ? = col
    if let Expression::Placeholder(idx) = left {
        if let Expression::ColumnRef { column, .. } = right {
            return Some((*idx, column.clone()));
        }
    }

    None
}

#[cfg(test)]
mod tests {
    use super::*;
    use vibesql_parser::Parser;

    fn parse_to_plan(sql: &str) -> CachedPlan {
        let stmt = Parser::parse_sql(sql).unwrap();
        analyze_for_plan(&stmt)
    }

    #[test]
    fn test_simple_pk_lookup() {
        let plan = parse_to_plan("SELECT * FROM users WHERE id = ?");
        match plan {
            CachedPlan::PkPointLookup(p) => {
                assert_eq!(p.table_name, "USERS");
                // Parser normalizes identifiers to uppercase
                assert_eq!(p.pk_columns, vec!["ID"]);
                assert_eq!(p.param_to_pk_col, vec![(0, 0)]);
                assert!(matches!(p.projection, ProjectionPlan::Wildcard));
            }
            _ => panic!("Expected PkPointLookup"),
        }
    }

    #[test]
    fn test_composite_pk_lookup() {
        let plan = parse_to_plan("SELECT * FROM orders WHERE customer_id = ? AND order_id = ?");
        match plan {
            CachedPlan::PkPointLookup(p) => {
                assert_eq!(p.table_name, "ORDERS");
                // Parser normalizes identifiers to uppercase
                assert_eq!(p.pk_columns, vec!["CUSTOMER_ID", "ORDER_ID"]);
                assert_eq!(p.param_to_pk_col.len(), 2);
            }
            _ => panic!("Expected PkPointLookup"),
        }
    }

    #[test]
    fn test_projected_columns() {
        let plan = parse_to_plan("SELECT name, email FROM users WHERE id = ?");
        match plan {
            CachedPlan::PkPointLookup(p) => {
                match p.projection {
                    ProjectionPlan::Columns(cols) => {
                        assert_eq!(cols.len(), 2);
                        // Parser normalizes identifiers to uppercase
                        assert_eq!(cols[0].column_name, "NAME");
                        assert_eq!(cols[1].column_name, "EMAIL");
                    }
                    _ => panic!("Expected Columns projection"),
                }
            }
            _ => panic!("Expected PkPointLookup"),
        }
    }

    #[test]
    fn test_not_cacheable_join() {
        let plan = parse_to_plan("SELECT * FROM users u JOIN orders o ON u.id = o.user_id WHERE u.id = ?");
        assert!(matches!(plan, CachedPlan::Standard));
    }

    #[test]
    fn test_not_cacheable_aggregate() {
        let plan = parse_to_plan("SELECT COUNT(*) FROM users WHERE id = ?");
        assert!(matches!(plan, CachedPlan::Standard));
    }

    #[test]
    fn test_not_cacheable_order_by() {
        let plan = parse_to_plan("SELECT * FROM users WHERE id = ? ORDER BY name");
        assert!(matches!(plan, CachedPlan::Standard));
    }

    #[test]
    fn test_not_cacheable_or() {
        let plan = parse_to_plan("SELECT * FROM users WHERE id = ? OR name = ?");
        assert!(matches!(plan, CachedPlan::Standard));
    }

    #[test]
    fn test_not_cacheable_literal() {
        let plan = parse_to_plan("SELECT * FROM users WHERE id = 1");
        assert!(matches!(plan, CachedPlan::Standard));
    }
}