featherdb_query/

lib.rs

//! SQL query engine for FeatherDB
//!
//! Provides SQL parsing, planning, optimization, and execution.

pub mod executor;
pub mod expr;
mod metrics;
mod optimizer;
mod parser;
mod planner;
mod prepared;
mod window;

pub use executor::{ExecutionContext, Executor, QueryResult, Row};
pub use expr::{
    AggregateFunction,
    BinaryOp,
    Expr,
    FrameBound,
    FrameUnit,
    SubqueryCompareOp,
    SubqueryType,
    UnaryOp,
    WindowFrame,
    // Window function types
    WindowFunction,
    WindowFunctionType,
    WindowOrderByExpr,
    WindowSortOrder,
};
pub use metrics::{QueryMetrics, QueryMetricsSnapshot};
pub use optimizer::{cost_constants, CostEstimator, Optimizer};
pub use parser::{
    CustomStatement, ParseResult, Parser, ShowGrantsTarget as ParserShowGrantsTarget,
};
pub use planner::{
    CommonTableExpression, IndexBound, IndexRange, JoinType, LogicalPlan, Planner, PlannerContext,
    ShowGrantsTarget, SortOrder,
};
pub use prepared::{
    count_params, count_plan_params, substitute_params, substitute_plan_params, CacheStats,
    PlanCache, PreparedStatement,
};
pub use window::WindowExecutor;

use featherdb_catalog::Catalog;
use featherdb_core::Result;
use featherdb_mvcc::{Snapshot, TransactionManager};
use featherdb_storage::BufferPool;
use std::sync::Arc;

/// Query engine with optional metrics collection
///
/// Provides a high-level interface for executing queries with optional
/// performance tracking. Wraps the parse-plan-optimize-execute pipeline.
pub struct QueryEngine {
    optimizer: Optimizer,
    metrics: Option<Arc<QueryMetrics>>,
}

impl QueryEngine {
    /// Create a new query engine without metrics
    pub fn new() -> Self {
        Self {
            optimizer: Optimizer::new(),
            metrics: None,
        }
    }

    /// Create a new query engine with metrics enabled
    pub fn with_metrics() -> (Self, Arc<QueryMetrics>) {
        let metrics = Arc::new(QueryMetrics::new());
        let engine = Self {
            optimizer: Optimizer::new(),
            metrics: Some(metrics.clone()),
        };
        (engine, metrics)
    }

    /// Enable metrics collection
    pub fn enable_metrics(&mut self) -> Arc<QueryMetrics> {
        let metrics = Arc::new(QueryMetrics::new());
        self.metrics = Some(metrics.clone());
        metrics
    }

    /// Disable metrics collection
    pub fn disable_metrics(&mut self) {
        self.metrics = None;
    }

    /// Get reference to metrics if enabled
    pub fn metrics(&self) -> Option<&Arc<QueryMetrics>> {
        self.metrics.as_ref()
    }

    /// Execute a SQL query with full parse-plan-execute pipeline
    ///
    /// If metrics are enabled, this will track:
    /// - Parse time
    /// - Plan time
    /// - Execution time
    /// - Rows scanned
    /// - Rows returned
    pub fn execute_sql(
        &self,
        sql: &str,
        catalog: &Catalog,
        buffer_pool: &Arc<BufferPool>,
        snapshot: &Snapshot,
        txn_manager: &TransactionManager,
        txn_id: featherdb_core::TransactionId,
    ) -> Result<QueryResult> {
        use std::collections::HashMap;

        if let Some(metrics) = &self.metrics {
            metrics.increment_queries();

            // Parse phase
            let parse_start = std::time::Instant::now();
            let stmt = Parser::parse_one(sql)?;
            metrics.record_parse_time(parse_start.elapsed().as_micros() as u64);

            // Plan phase
            let plan_start = std::time::Instant::now();
            let planner = Planner::new(catalog);
            let plan = planner.plan(&stmt)?;
            let subquery_plans = planner.take_subquery_plans();
            let optimized = self.optimizer.optimize(plan);
            metrics.record_plan_time(plan_start.elapsed().as_micros() as u64);

            // Execute phase (timing happens in execute_with_timing)
            let ctx = ExecutionContext {
                catalog,
                buffer_pool,
                snapshot,
                txn_manager,
                txn_id,
                cte_context: HashMap::new(),
                subquery_plans,
                metrics: Some(metrics),
                api_key_id: None,    // Set by caller (Database)
                auth_enabled: false, // Set by caller (Database)
            };

            Executor::execute_with_timing(&ctx, &optimized)
        } else {
            // No metrics - fast path
            let stmt = Parser::parse_one(sql)?;
            let planner = Planner::new(catalog);
            let plan = planner.plan(&stmt)?;
            let subquery_plans = planner.take_subquery_plans();
            let optimized = self.optimizer.optimize(plan);

            let ctx = ExecutionContext {
                catalog,
                buffer_pool,
                snapshot,
                txn_manager,
                txn_id,
                cte_context: HashMap::new(),
                subquery_plans,
                metrics: None,
                api_key_id: None,    // Set by caller (Database)
                auth_enabled: false, // Set by caller (Database)
            };

            Executor::execute(&ctx, &optimized)
        }
    }

    /// Execute a prepared statement
    pub fn execute_prepared(
        &self,
        plan: &LogicalPlan,
        catalog: &Catalog,
        buffer_pool: &Arc<BufferPool>,
        snapshot: &Snapshot,
        txn_manager: &TransactionManager,
        txn_id: featherdb_core::TransactionId,
    ) -> Result<QueryResult> {
        use std::collections::HashMap;

        let ctx = ExecutionContext {
            catalog,
            buffer_pool,
            snapshot,
            txn_manager,
            txn_id,
            cte_context: HashMap::new(),
            subquery_plans: HashMap::new(),
            metrics: self.metrics.as_ref(),
            api_key_id: None,    // Set by caller (Database)
            auth_enabled: false, // Set by caller (Database)
        };

        if self.metrics.is_some() {
            Executor::execute_with_timing(&ctx, plan)
        } else {
            Executor::execute(&ctx, plan)
        }
    }
}

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

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

    #[test]
    fn test_query_engine_creation() {
        // Create query engine without metrics
        let engine = QueryEngine::new();
        assert!(engine.metrics().is_none());

        // Create query engine with metrics
        let (engine_with_metrics, metrics) = QueryEngine::with_metrics();
        assert!(engine_with_metrics.metrics().is_some());

        // Initial snapshot should be all zeros
        let snapshot = metrics.snapshot();
        assert_eq!(snapshot.queries_executed, 0);
        assert_eq!(snapshot.total_parse_time_us, 0);
        assert_eq!(snapshot.total_plan_time_us, 0);
        assert_eq!(snapshot.total_exec_time_us, 0);
        assert_eq!(snapshot.rows_scanned, 0);
        assert_eq!(snapshot.rows_returned, 0);
    }

    #[test]
    fn test_metrics_enable_disable() {
        let mut engine = QueryEngine::new();
        assert!(engine.metrics().is_none());

        // Enable metrics
        let metrics = engine.enable_metrics();
        assert!(engine.metrics().is_some());

        // Initial snapshot should be all zeros
        let snapshot = metrics.snapshot();
        assert_eq!(snapshot.queries_executed, 0);

        // Disable metrics
        engine.disable_metrics();
        assert!(engine.metrics().is_none());
    }

    #[test]
    fn test_default_query_engine() {
        let engine = QueryEngine::default();
        assert!(engine.metrics().is_none());
    }
}