qail_core/

validator.rs

//! Schema validator and fuzzy matching suggestions.
//!
//! Provides compile-time-like validation for Qail against a known schema.
//! Used by CLI, LSP, and the encoder to catch errors before they hit the wire.

use crate::ast::{Expr, Qail};
use std::collections::HashMap;
use strsim::levenshtein;

/// Validation error with structured information.
#[derive(Debug, Clone, PartialEq)]
pub enum ValidationError {
    /// Referenced table does not exist in the schema.
    TableNotFound {
        /// Name of the missing table.
        table: String,
        /// Closest match from known tables.
        suggestion: Option<String>,
    },
    /// Referenced column does not exist in the table.
    ColumnNotFound {
        /// Table the column was looked up in.
        table: String,
        /// Name of the missing column.
        column: String,
        /// Closest match from known columns.
        suggestion: Option<String>,
    },
    /// Type mismatch (future: when schema includes types)
    TypeMismatch {
        /// Table name.
        table: String,
        /// Column name.
        column: String,
        /// Expected type.
        expected: String,
        /// Actual type.
        got: String,
    },
    /// Invalid operator for column type (future)
    InvalidOperator {
        /// Column name.
        column: String,
        /// Operator string.
        operator: String,
        /// Explanation.
        reason: String,
    },
}

impl std::fmt::Display for ValidationError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ValidationError::TableNotFound { table, suggestion } => {
                if let Some(s) = suggestion {
                    write!(f, "Table '{}' not found. Did you mean '{}'?", table, s)
                } else {
                    write!(f, "Table '{}' not found.", table)
                }
            }
            ValidationError::ColumnNotFound {
                table,
                column,
                suggestion,
            } => {
                if let Some(s) = suggestion {
                    write!(
                        f,
                        "Column '{}' not found in table '{}'. Did you mean '{}'?",
                        column, table, s
                    )
                } else {
                    write!(f, "Column '{}' not found in table '{}'.", column, table)
                }
            }
            ValidationError::TypeMismatch {
                table,
                column,
                expected,
                got,
            } => {
                write!(
                    f,
                    "Type mismatch for '{}.{}': expected {}, got {}",
                    table, column, expected, got
                )
            }
            ValidationError::InvalidOperator {
                column,
                operator,
                reason,
            } => {
                write!(
                    f,
                    "Invalid operator '{}' for column '{}': {}",
                    operator, column, reason
                )
            }
        }
    }
}

impl std::error::Error for ValidationError {}

/// Result of validation
pub type ValidationResult = Result<(), Vec<ValidationError>>;

/// Validates query elements against known schema and provides suggestions.
#[derive(Debug, Clone)]
pub struct Validator {
    /// Known table names.
    tables: Vec<String>,
    /// Columns indexed by table name.
    columns: HashMap<String, Vec<String>>,
    /// Column types indexed by table.column.
    column_types: HashMap<String, HashMap<String, String>>,
}

impl Default for Validator {
    fn default() -> Self {
        Self::new()
    }
}

impl Validator {
    /// Create a new Validator with known tables and columns.
    pub fn new() -> Self {
        Self {
            tables: Vec::new(),
            columns: HashMap::new(),
            column_types: HashMap::new(),
        }
    }

    /// Register a table and its columns.
    pub fn add_table(&mut self, table: &str, cols: &[&str]) {
        self.tables.push(table.to_string());
        self.columns.insert(
            table.to_string(),
            cols.iter().map(|s| s.to_string()).collect(),
        );
    }

    /// Register a table name only (no column metadata).
    ///
    /// Useful for views discovered from schema text where build-time parser
    /// does not infer projected columns. Table existence is validated, while
    /// column-level checks are skipped.
    pub fn add_table_name(&mut self, table: &str) {
        self.tables.push(table.to_string());
    }

    /// Register a table with column types (for future type validation).
    ///
    /// # Arguments
    ///
    /// * `table` — Table name to register.
    /// * `cols` — Slice of `(column_name, column_type)` pairs.
    pub fn add_table_with_types(&mut self, table: &str, cols: &[(&str, &str)]) {
        self.tables.push(table.to_string());
        let col_names: Vec<String> = cols.iter().map(|(name, _)| name.to_string()).collect();
        self.columns.insert(table.to_string(), col_names);

        let type_map: HashMap<String, String> = cols
            .iter()
            .map(|(name, typ)| (name.to_string(), typ.to_string()))
            .collect();
        self.column_types.insert(table.to_string(), type_map);
    }

    /// Get list of all table names (for autocomplete).
    pub fn table_names(&self) -> &[String] {
        &self.tables
    }

    /// Get column names for a table (for autocomplete).
    pub fn column_names(&self, table: &str) -> Option<&Vec<String>> {
        self.columns.get(table)
    }

    /// Check if a table exists.
    pub fn table_exists(&self, table: &str) -> bool {
        self.tables.contains(&table.to_string())
    }

    /// Check if a table exists. If not, returns structured error with suggestion.
    pub fn validate_table(&self, table: &str) -> Result<(), ValidationError> {
        if self.tables.contains(&table.to_string()) {
            Ok(())
        } else {
            let suggestion = self.did_you_mean(table, &self.tables);
            Err(ValidationError::TableNotFound {
                table: table.to_string(),
                suggestion,
            })
        }
    }

    /// Check if a column exists in a table. If not, returns a structured error.
    ///
    /// # Arguments
    ///
    /// * `table` — Table to look up.
    /// * `column` — Column name to validate.
    pub fn validate_column(&self, table: &str, column: &str) -> Result<(), ValidationError> {
        // If table doesn't exist, skip column validation (table error takes precedence)
        if !self.tables.contains(&table.to_string()) {
            return Ok(());
        }

        // Always allow *
        if column == "*" {
            return Ok(());
        }

        // Skip SQL expressions — these are not plain column names.
        // Covers: COUNT(*), id::text, leg_ids[1], brand_name AS alias,
        //         count(distinct x), EXCLUDED.col, NOW(), etc.
        if column.contains('(')
            || column.contains('[')
            || column.contains("::")
            || column.contains(" AS ")
            || column.contains(" as ")
            || column.starts_with("distinct ")
            || column.starts_with("DISTINCT ")
        {
            return Ok(());
        }

        // Qualified names like "table.column" — validate against the referenced table
        if column.contains('.') {
            let parts: Vec<&str> = column.split('.').collect();
            if parts.len() == 3 && parts[0] == "public" {
                let public_table = format!("{}.{}", parts[0], parts[1]);
                if self.tables.contains(&public_table) {
                    return self.validate_column(&public_table, parts[2]);
                }
                if self.tables.contains(&parts[1].to_string()) {
                    return self.validate_column(parts[1], parts[2]);
                }
            }
            if parts.len() == 2 {
                // Only validate if the referenced table is known to the validator
                if self.tables.contains(&parts[0].to_string()) {
                    return self.validate_column(parts[0], parts[1]);
                }
                let public_table = format!("public.{}", parts[0]);
                if self.tables.contains(&public_table) {
                    return self.validate_column(&public_table, parts[1]);
                }
            }
            // Unknown table prefix or complex dotted path — allow (might be JSON)
            return Ok(());
        }

        if let Some(cols) = self.columns.get(table) {
            if cols.contains(&column.to_string()) {
                Ok(())
            } else {
                let suggestion = self.did_you_mean(column, cols);
                Err(ValidationError::ColumnNotFound {
                    table: table.to_string(),
                    column: column.to_string(),
                    suggestion,
                })
            }
        } else {
            Ok(())
        }
    }

    /// Extract column name from an Expr for validation.
    fn extract_column_name(expr: &Expr) -> Option<String> {
        match expr {
            Expr::Named(name) => Some(name.clone()),
            Expr::Aliased { name, .. } => Some(name.clone()),
            Expr::Aggregate { col, .. } => Some(col.clone()),
            Expr::Cast { expr, .. } => Self::extract_column_name(expr),
            Expr::JsonAccess { column, .. } => Some(column.clone()),
            _ => None,
        }
    }

    /// Get column type for a table.column
    pub fn get_column_type(&self, table: &str, column: &str) -> Option<&String> {
        self.column_types.get(table)?.get(column)
    }

    /// Validate that a Value's type matches the expected column type.
    /// Returns Ok(()) if compatible, Err with TypeMismatch if not.
    pub fn validate_value_type(
        &self,
        table: &str,
        column: &str,
        value: &crate::ast::Value,
    ) -> Result<(), ValidationError> {
        use crate::ast::Value;

        // Get the expected type for this column
        let expected_type = match self.get_column_type(table, column) {
            Some(t) => t.to_uppercase(),
            None => return Ok(()), // Column type unknown, skip validation
        };

        // NULL is always allowed (for nullable columns)
        if matches!(value, Value::Null | Value::NullUuid) {
            return Ok(());
        }

        // Param and NamedParam are runtime values - can't validate at compile time
        if matches!(
            value,
            Value::Param(_)
                | Value::NamedParam(_)
                | Value::Function(_)
                | Value::Subquery(_)
                | Value::Expr(_)
        ) {
            return Ok(());
        }

        // Value::Array is used with IN/NOT IN operators — the array is a container
        // of elements whose type should match the column, not the array itself.
        // Skip type validation for arrays (the element types are checked at runtime).
        if matches!(value, Value::Array(_)) {
            return Ok(());
        }

        // Map Value variant to its type name
        let value_type = match value {
            Value::Bool(_) => "BOOLEAN",
            Value::Int(_) => "INT",
            Value::Float(_) => "FLOAT",
            Value::String(_) => "TEXT",
            Value::Uuid(_) => "UUID",
            Value::Column(_) => return Ok(()), // Column reference, type checked elsewhere
            Value::Interval { .. } => "INTERVAL",
            Value::Timestamp(_) => "TIMESTAMP",
            Value::Bytes(_) => "BYTEA",
            Value::Vector(_) => "VECTOR",
            Value::Json(_) => "JSONB",
            _ => return Ok(()), // Unknown value type, skip
        };

        // Check compatibility
        if !Self::types_compatible(&expected_type, value_type) {
            return Err(ValidationError::TypeMismatch {
                table: table.to_string(),
                column: column.to_string(),
                expected: expected_type,
                got: value_type.to_string(),
            });
        }

        Ok(())
    }

    /// Check if expected column type is compatible with value type.
    /// Allows flexible matching (e.g., INT matches INT4, BIGINT, INTEGER, etc.)
    fn types_compatible(expected: &str, value_type: &str) -> bool {
        let expected = expected.to_uppercase();
        let value_type = value_type.to_uppercase();

        // Exact match
        if expected == value_type {
            return true;
        }

        // Integer family
        let int_types = [
            "INT",
            "INT4",
            "INT8",
            "INTEGER",
            "BIGINT",
            "SMALLINT",
            "SERIAL",
            "BIGSERIAL",
        ];
        if int_types.contains(&expected.as_str()) && value_type == "INT" {
            return true;
        }

        // Float family — includes "DOUBLE PRECISION" which is how PG stores f64
        let float_types = [
            "FLOAT",
            "FLOAT4",
            "FLOAT8",
            "DOUBLE",
            "DOUBLE PRECISION",
            "DECIMAL",
            "NUMERIC",
            "REAL",
        ];
        if float_types.contains(&expected.as_str())
            && (value_type == "FLOAT" || value_type == "INT")
        {
            return true;
        }

        // Text family - TEXT is very flexible in PostgreSQL
        let text_types = ["TEXT", "VARCHAR", "CHAR", "CHARACTER", "CITEXT", "NAME"];
        if text_types.contains(&expected.as_str()) && value_type == "TEXT" {
            return true;
        }

        // Boolean
        if (expected == "BOOLEAN" || expected == "BOOL") && value_type == "BOOLEAN" {
            return true;
        }

        // UUID
        if expected == "UUID" && (value_type == "UUID" || value_type == "TEXT") {
            return true;
        }

        // Timestamp family
        let ts_types = ["TIMESTAMP", "TIMESTAMPTZ", "DATE", "TIME", "TIMETZ"];
        if ts_types.contains(&expected.as_str())
            && (value_type == "TIMESTAMP" || value_type == "TEXT")
        {
            return true;
        }

        // JSONB/JSON accepts almost anything
        if expected == "JSONB" || expected == "JSON" {
            return true;
        }

        // Arrays
        if expected.contains("[]") || expected.starts_with("_") {
            return value_type == "ARRAY";
        }

        false
    }

    /// Validate an entire Qail against the schema.
    pub fn validate_command(&self, cmd: &Qail) -> ValidationResult {
        let mut errors = Vec::new();

        if let Err(e) = self.validate_table(&cmd.table) {
            errors.push(e);
        }

        // Collect aliases from the SELECT column list so that order_by / having
        // references to computed aliases (e.g. count().alias("route_count"))
        // are not flagged as column-not-found errors.
        let mut aliases: Vec<String> = Vec::new();
        for col in &cmd.columns {
            if let Expr::Aliased { alias, .. } = col {
                aliases.push(alias.clone());
            }
            if let Some(name) = Self::extract_column_name(col)
                && let Err(e) = self.validate_column(&cmd.table, &name)
            {
                errors.push(e);
            }
        }

        for cage in &cmd.cages {
            // Skip column validation for Sort cages — ORDER BY can reference
            // aliases from column_expr (e.g. count().alias("route_count")),
            // which may not exist as physical columns in the primary table.
            if matches!(cage.kind, crate::ast::CageKind::Sort(_)) {
                continue;
            }
            for cond in &cage.conditions {
                if let Some(name) = Self::extract_column_name(&cond.left) {
                    // Skip validation for columns that match a computed alias
                    if aliases.iter().any(|a| a == &name) {
                        continue;
                    }
                    // For join conditions, column might be qualified (table.column)
                    if name.contains('.') {
                        let parts: Vec<&str> = name.split('.').collect();
                        if parts.len() == 2 {
                            if let Err(e) = self.validate_column(parts[0], parts[1]) {
                                errors.push(e);
                            }
                            // Type validation for qualified column
                            if let Err(e) =
                                self.validate_value_type(parts[0], parts[1], &cond.value)
                            {
                                errors.push(e);
                            }
                        }
                    } else {
                        if let Err(e) = self.validate_column(&cmd.table, &name) {
                            errors.push(e);
                        }
                        // Type validation for unqualified column
                        if let Err(e) = self.validate_value_type(&cmd.table, &name, &cond.value) {
                            errors.push(e);
                        }
                    }
                }
            }
        }

        for cond in &cmd.having {
            if let Some(name) = Self::extract_column_name(&cond.left) {
                if name.contains('(') || name == "*" {
                    continue;
                }
                if name.contains('.') {
                    let parts: Vec<&str> = name.split('.').collect();
                    if parts.len() == 2 {
                        if let Err(e) = self.validate_column(parts[0], parts[1]) {
                            errors.push(e);
                        }
                        if let Err(e) = self.validate_value_type(parts[0], parts[1], &cond.value) {
                            errors.push(e);
                        }
                    }
                } else {
                    if let Err(e) = self.validate_column(&cmd.table, &name) {
                        errors.push(e);
                    }
                    if let Err(e) = self.validate_value_type(&cmd.table, &name, &cond.value) {
                        errors.push(e);
                    }
                }
            }
        }

        for join in &cmd.joins {
            if let Err(e) = self.validate_table(&join.table) {
                errors.push(e);
            }

            if let Some(conditions) = &join.on {
                for cond in conditions {
                    if let Some(name) = Self::extract_column_name(&cond.left)
                        && name.contains('.')
                    {
                        let parts: Vec<&str> = name.split('.').collect();
                        if parts.len() == 2
                            && let Err(e) = self.validate_column(parts[0], parts[1])
                        {
                            errors.push(e);
                        }
                    }
                    // Also check right side if it's a column reference
                    if let crate::ast::Value::Column(col_name) = &cond.value
                        && col_name.contains('.')
                    {
                        let parts: Vec<&str> = col_name.split('.').collect();
                        if parts.len() == 2
                            && let Err(e) = self.validate_column(parts[0], parts[1])
                        {
                            errors.push(e);
                        }
                    }
                }
            }
        }

        if let Some(returning) = &cmd.returning {
            for col in returning {
                if let Some(name) = Self::extract_column_name(col)
                    && let Err(e) = self.validate_column(&cmd.table, &name)
                {
                    errors.push(e);
                }
            }
        }

        if errors.is_empty() {
            Ok(())
        } else {
            Err(errors)
        }
    }

    /// Find the best match with Levenshtein distance within threshold.
    fn did_you_mean(&self, input: &str, candidates: &[impl AsRef<str>]) -> Option<String> {
        let mut best_match = None;
        let mut min_dist = usize::MAX;

        for cand in candidates {
            let cand_str = cand.as_ref();
            let dist = levenshtein(input, cand_str);

            // Dynamic threshold based on length
            let threshold = match input.len() {
                0..=2 => 0, // Precise match only for very short strings
                3..=5 => 2, // Allow 2 char diff (e.g. usr -> users)
                _ => 3,     // Allow 3 char diff for longer strings
            };

            if dist <= threshold && dist < min_dist {
                min_dist = dist;
                best_match = Some(cand_str.to_string());
            }
        }

        best_match
    }
}

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

    #[test]
    fn test_did_you_mean_table() {
        let mut v = Validator::new();
        v.add_table("users", &["id", "name"]);
        v.add_table("orders", &["id", "total"]);

        assert!(v.validate_table("users").is_ok());

        let err = v.validate_table("usr").unwrap_err();
        assert!(
            matches!(err, ValidationError::TableNotFound { suggestion: Some(ref s), .. } if s == "users")
        );

        let err = v.validate_table("usrs").unwrap_err();
        assert!(
            matches!(err, ValidationError::TableNotFound { suggestion: Some(ref s), .. } if s == "users")
        );
    }

    #[test]
    fn test_did_you_mean_column() {
        let mut v = Validator::new();
        v.add_table("users", &["email", "password"]);

        assert!(v.validate_column("users", "email").is_ok());
        assert!(v.validate_column("users", "*").is_ok());

        let err = v.validate_column("users", "emial").unwrap_err();
        assert!(
            matches!(err, ValidationError::ColumnNotFound { suggestion: Some(ref s), .. } if s == "email")
        );
    }

    #[test]
    fn test_qualified_column_name() {
        let mut v = Validator::new();
        v.add_table("users", &["id", "name"]);
        v.add_table("profiles", &["user_id", "avatar"]);

        // Qualified names should pass through
        assert!(v.validate_column("users", "users.id").is_ok());
        assert!(v.validate_column("users", "profiles.user_id").is_ok());
    }

    #[test]
    fn test_validate_command() {
        let mut v = Validator::new();
        v.add_table("users", &["id", "email", "name"]);

        let cmd = Qail::get("users").columns(["id", "email"]);
        assert!(v.validate_command(&cmd).is_ok());

        let cmd = Qail::get("users").columns(["id", "emial"]); // typo
        let errors = v.validate_command(&cmd).unwrap_err();
        assert_eq!(errors.len(), 1);
        assert!(
            matches!(&errors[0], ValidationError::ColumnNotFound { column, .. } if column == "emial")
        );
    }

    #[test]
    fn test_validate_having_columns() {
        let mut v = Validator::new();
        v.add_table("orders", &["id", "status", "total"]);

        let mut cmd = Qail::get("orders");
        cmd.having.push(crate::ast::Condition {
            left: Expr::Named("totl".to_string()),
            op: crate::ast::Operator::Eq,
            value: crate::ast::Value::Int(1),
            is_array_unnest: false,
        });

        let errors = v.validate_command(&cmd).unwrap_err();
        assert!(errors.iter().any(
            |e| matches!(e, ValidationError::ColumnNotFound { column, .. } if column == "totl")
        ));
    }

    #[test]
    fn test_error_display() {
        let err = ValidationError::TableNotFound {
            table: "usrs".to_string(),
            suggestion: Some("users".to_string()),
        };
        assert_eq!(
            err.to_string(),
            "Table 'usrs' not found. Did you mean 'users'?"
        );

        let err = ValidationError::ColumnNotFound {
            table: "users".to_string(),
            column: "emial".to_string(),
            suggestion: Some("email".to_string()),
        };
        assert_eq!(
            err.to_string(),
            "Column 'emial' not found in table 'users'. Did you mean 'email'?"
        );
    }
}