featherdb_query/planner/

context.rs

//! Planning context for tracking available columns and tables

use featherdb_catalog::Table;
use featherdb_core::{find_best_match, levenshtein_distance, Error, Result};
use std::sync::Arc;

/// Context for tracking available columns during planning
#[derive(Debug, Clone, Default)]
pub struct PlannerContext {
    /// Available tables and their aliases
    tables: Vec<(String, Arc<Table>)>,
    /// The original SQL query for error messages
    sql: String,
    /// Outer query column names (for detecting correlated subqueries)
    /// This is populated when planning subqueries to detect column references
    /// that belong to the outer query rather than the subquery itself
    outer_columns: Vec<String>,
}

impl PlannerContext {
    /// Create a new context
    pub fn new() -> Self {
        Self::default()
    }

    /// Set the SQL query for error messages
    pub fn with_sql(mut self, sql: impl Into<String>) -> Self {
        self.sql = sql.into();
        self
    }

    /// Add a table to the context
    pub fn add_table(&mut self, alias: Option<String>, table: Arc<Table>) {
        let name = alias.unwrap_or_else(|| table.name.clone());
        self.tables.push((name, table));
    }

    /// Set the outer query columns (for subquery planning)
    pub fn with_outer_columns(mut self, columns: Vec<String>) -> Self {
        self.outer_columns = columns;
        self
    }

    /// Get the current table columns (columns defined in this scope)
    pub fn current_columns(&self) -> Vec<String> {
        self.available_columns()
    }

    /// Check if a column reference is from the outer query
    /// Returns true if the column is in outer_columns but not in current tables
    pub fn is_outer_column(&self, table: Option<&str>, column: &str) -> bool {
        // First check if the column exists in current scope
        if self.find_column(table, column).is_ok() {
            return false; // Column is in current scope, not outer
        }

        // If we have outer columns set and the column matches one, it's an outer reference
        if !self.outer_columns.is_empty() {
            let col_name = if let Some(t) = table {
                format!("{}.{}", t, column)
            } else {
                column.to_string()
            };

            // Check if this column (or unqualified version) is in outer columns
            self.outer_columns
                .iter()
                .any(|c| c == &col_name || c.ends_with(&format!(".{}", column)))
        } else {
            false
        }
    }

    /// Find a column by name, optionally qualified by table (case-insensitive)
    pub fn find_column(&self, table: Option<&str>, column: &str) -> Result<(Arc<Table>, usize)> {
        if let Some(table_name) = table {
            // Qualified reference: find the specific table (case-insensitive)
            for (alias, tbl) in &self.tables {
                if alias.eq_ignore_ascii_case(table_name)
                    || tbl.name.eq_ignore_ascii_case(table_name)
                {
                    return match tbl.get_column_index(column) {
                        Some(idx) => Ok((tbl.clone(), idx)),
                        None => Err(self.column_not_found_error(column, &tbl.name)),
                    };
                }
            }
            Err(self.table_not_found_error(table_name))
        } else {
            // Unqualified reference: search all tables
            let mut found = Vec::new();
            for (alias, tbl) in &self.tables {
                if let Some(idx) = tbl.get_column_index(column) {
                    found.push((alias.clone(), tbl.clone(), idx));
                }
            }

            match found.len() {
                0 => Err(self.column_not_found_in_all_tables_error(column)),
                1 => Ok((found[0].1.clone(), found[0].2)),
                _ => {
                    let tables: Vec<_> = found.iter().map(|(a, _, _)| a.clone()).collect();
                    Err(Error::AmbiguousColumn {
                        column: column.to_string(),
                        tables: tables.join(", "),
                    })
                }
            }
        }
    }

    /// Create a column not found error with suggestions
    fn column_not_found_error(&self, column: &str, table_name: &str) -> Error {
        // Find the table to get its columns
        let suggestion = self
            .tables
            .iter()
            .find(|(_, t)| t.name == table_name)
            .and_then(|(_, t)| {
                let columns: Vec<&str> = t.columns.iter().map(|c| c.name.as_str()).collect();
                find_best_match(column, &columns, 3)
            })
            .map(|s| s.to_string());

        Error::ColumnNotFound {
            column: column.to_string(),
            table: table_name.to_string(),
            suggestion,
        }
    }

    /// Create a column not found error when searching all tables
    fn column_not_found_in_all_tables_error(&self, column: &str) -> Error {
        // Search all tables for a similar column
        let mut best_match: Option<(String, &str)> = None;
        let mut best_distance = usize::MAX;

        for (alias, table) in &self.tables {
            for col in &table.columns {
                let distance = levenshtein_distance(column, &col.name);
                if distance < best_distance && distance <= 3 {
                    best_distance = distance;
                    best_match = Some((alias.clone(), col.name.as_str()));
                }
            }
        }

        let table_name = if self.tables.len() == 1 {
            self.tables[0].0.clone()
        } else {
            "query".to_string()
        };

        let suggestion = best_match.map(|(tbl, col)| {
            if self.tables.len() == 1 {
                col.to_string()
            } else {
                format!("{}.{}", tbl, col)
            }
        });

        Error::ColumnNotFound {
            column: column.to_string(),
            table: table_name,
            suggestion,
        }
    }

    /// Create a table not found error with suggestions
    fn table_not_found_error(&self, table: &str) -> Error {
        let available: Vec<&str> = self.tables.iter().map(|(a, _)| a.as_str()).collect();
        let suggestion = find_best_match(table, &available, 3).map(|s| s.to_string());

        Error::TableNotFound {
            table: table.to_string(),
            suggestion,
        }
    }

    /// Get all available column names (for error messages)
    pub fn available_columns(&self) -> Vec<String> {
        let mut columns = Vec::new();
        for (alias, table) in &self.tables {
            for col in &table.columns {
                if self.tables.len() == 1 {
                    columns.push(col.name.clone());
                } else {
                    columns.push(format!("{}.{}", alias, col.name));
                }
            }
        }
        columns
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use featherdb_catalog::TableBuilder;

    fn create_test_table(name: &str, columns: Vec<&str>) -> Table {
        let mut builder = TableBuilder::new(name);
        for col in columns {
            builder = builder.column(col, featherdb_core::ColumnType::Integer);
        }
        builder.build(0, featherdb_core::PageId::INVALID)
    }

    #[test]
    fn test_planner_context_column_lookup() {
        let users = Arc::new(create_test_table("users", vec!["id", "name", "age"]));
        let orders = Arc::new(create_test_table("orders", vec!["id", "user_id", "total"]));

        let mut ctx = PlannerContext::new();
        ctx.add_table(None, users.clone());
        ctx.add_table(None, orders.clone());

        // Qualified lookup
        let (table, idx) = ctx.find_column(Some("users"), "name").unwrap();
        assert_eq!(table.name, "users");
        assert_eq!(idx, 1);

        // Ambiguous unqualified lookup should error
        let result = ctx.find_column(None, "id");
        assert!(matches!(result, Err(Error::AmbiguousColumn { .. })));

        // Unambiguous unqualified lookup
        let (table, idx) = ctx.find_column(None, "age").unwrap();
        assert_eq!(table.name, "users");
        assert_eq!(idx, 2);
    }

    #[test]
    fn test_planner_context_available_columns() {
        let users = Arc::new(create_test_table("users", vec!["id", "name"]));
        let mut ctx = PlannerContext::new();
        ctx.add_table(None, users);

        let cols = ctx.available_columns();
        assert_eq!(cols, vec!["id", "name"]);
    }
}