alopex_sql/planner/
mod.rs

1//! Query planning module for the Alopex SQL dialect.
2//!
3//! This module provides:
4//! - [`PlannerError`]: Error types for planning phase
5//! - [`ResolvedType`]: Normalized type information for type checking
6//! - [`TypedExpr`]: Type-checked expressions with resolved types
7//! - [`LogicalPlan`]: Logical query plan representation
8//! - [`NameResolver`]: Table and column reference resolution
9//! - [`TypeChecker`]: Expression type inference and validation
10//! - [`Planner`]: Main entry point for converting AST to LogicalPlan
11
12mod error;
13pub mod knn_optimizer;
14pub mod logical_plan;
15pub mod name_resolver;
16pub mod type_checker;
17pub mod typed_expr;
18pub mod types;
19
20#[cfg(test)]
21mod planner_tests;
22
23pub use error::PlannerError;
24pub use knn_optimizer::{KnnPattern, SortDirection, detect_knn_pattern};
25pub use logical_plan::LogicalPlan;
26pub use name_resolver::{NameResolver, ResolvedColumn};
27pub use type_checker::TypeChecker;
28pub use typed_expr::{
29    ProjectedColumn, Projection, SortExpr, TypedAssignment, TypedExpr, TypedExprKind,
30};
31pub use types::ResolvedType;
32
33use crate::ast::ddl::{
34    ColumnConstraint, ColumnDef, CreateIndex, CreateTable, DropIndex, DropTable,
35};
36use crate::ast::dml::{Delete, Insert, OrderByExpr, Select, SelectItem, Update};
37use crate::ast::expr::Literal;
38use crate::ast::{Statement, StatementKind};
39use crate::catalog::{Catalog, ColumnMetadata, IndexMetadata, TableMetadata};
40
41/// The SQL query planner.
42///
43/// The planner converts AST statements into logical plans. It performs:
44/// - Name resolution: Validates table and column references
45/// - Type checking: Infers and validates expression types
46/// - Plan construction: Builds the logical plan tree
47///
48/// # Design Notes
49///
50/// - The planner uses an immutable reference to the catalog (`&C`)
51/// - DDL statements produce plans but don't modify the catalog
52/// - The executor is responsible for applying catalog changes
53///
54/// # Examples
55///
56/// ```
57/// use alopex_sql::catalog::MemoryCatalog;
58/// use alopex_sql::planner::Planner;
59///
60/// let catalog = MemoryCatalog::new();
61/// let planner = Planner::new(&catalog);
62///
63/// // Parse and plan a statement
64/// // let stmt = parser.parse("SELECT * FROM users")?;
65/// // let plan = planner.plan(&stmt)?;
66/// ```
67pub struct Planner<'a, C: Catalog> {
68    catalog: &'a C,
69    name_resolver: NameResolver<'a, C>,
70    type_checker: TypeChecker<'a, C>,
71}
72
73impl<'a, C: Catalog> Planner<'a, C> {
74    /// Create a new planner with the given catalog.
75    pub fn new(catalog: &'a C) -> Self {
76        Self {
77            catalog,
78            name_resolver: NameResolver::new(catalog),
79            type_checker: TypeChecker::new(catalog),
80        }
81    }
82
83    /// Plan a SQL statement.
84    ///
85    /// This is the main entry point for converting an AST statement into a logical plan.
86    ///
87    /// # Errors
88    ///
89    /// Returns a `PlannerError` if:
90    /// - Referenced tables or columns don't exist
91    /// - Type checking fails
92    /// - DDL validation fails (e.g., table already exists for CREATE TABLE)
93    pub fn plan(&self, stmt: &Statement) -> Result<LogicalPlan, PlannerError> {
94        match &stmt.kind {
95            // DDL statements
96            StatementKind::CreateTable(ct) => self.plan_create_table(ct),
97            StatementKind::DropTable(dt) => self.plan_drop_table(dt),
98            StatementKind::CreateIndex(ci) => self.plan_create_index(ci),
99            StatementKind::DropIndex(di) => self.plan_drop_index(di),
100
101            // DML statements
102            StatementKind::Select(sel) => self.plan_select(sel),
103            StatementKind::Insert(ins) => self.plan_insert(ins),
104            StatementKind::Update(upd) => self.plan_update(upd),
105            StatementKind::Delete(del) => self.plan_delete(del),
106        }
107    }
108
109    // ============================================================
110    // DDL Planning Methods (Task 16)
111    // ============================================================
112
113    /// Plan a CREATE TABLE statement.
114    ///
115    /// Validates that the table doesn't already exist (unless IF NOT EXISTS is specified),
116    /// and converts the AST column definitions to catalog metadata.
117    fn plan_create_table(&self, stmt: &CreateTable) -> Result<LogicalPlan, PlannerError> {
118        // Check if table already exists
119        if !stmt.if_not_exists && self.catalog.table_exists(&stmt.name) {
120            return Err(PlannerError::table_already_exists(&stmt.name));
121        }
122
123        // Convert column definitions to metadata
124        let columns: Vec<ColumnMetadata> = stmt
125            .columns
126            .iter()
127            .map(|col| self.convert_column_def(col))
128            .collect();
129
130        // Collect primary key from table constraints
131        let primary_key = Self::extract_primary_key(stmt);
132
133        // Build table metadata
134        // Note: table_id defaults to 0 as placeholder; Executor assigns the actual ID
135        let mut table = TableMetadata::new(stmt.name.clone(), columns);
136        if let Some(pk) = primary_key {
137            table = table.with_primary_key(pk);
138        }
139
140        Ok(LogicalPlan::CreateTable {
141            table,
142            if_not_exists: stmt.if_not_exists,
143            with_options: stmt.with_options.clone(),
144        })
145    }
146
147    /// Convert an AST column definition to catalog column metadata.
148    fn convert_column_def(&self, col: &ColumnDef) -> ColumnMetadata {
149        let data_type = ResolvedType::from_ast(&col.data_type);
150        let mut meta = ColumnMetadata::new(col.name.clone(), data_type);
151
152        // Process constraints
153        for constraint in &col.constraints {
154            meta = Self::apply_column_constraint(meta, constraint);
155        }
156
157        meta
158    }
159
160    /// Apply a column constraint to column metadata.
161    fn apply_column_constraint(
162        mut meta: ColumnMetadata,
163        constraint: &ColumnConstraint,
164    ) -> ColumnMetadata {
165        match constraint {
166            ColumnConstraint::NotNull => {
167                meta.not_null = true;
168            }
169            ColumnConstraint::Null => {
170                meta.not_null = false;
171            }
172            ColumnConstraint::PrimaryKey => {
173                meta.primary_key = true;
174                meta.not_null = true; // PRIMARY KEY implies NOT NULL
175            }
176            ColumnConstraint::Unique => {
177                meta.unique = true;
178            }
179            ColumnConstraint::Default(expr) => {
180                meta.default = Some(expr.clone());
181            }
182            ColumnConstraint::WithSpan { kind, .. } => {
183                meta = Self::apply_column_constraint(meta, kind);
184            }
185        }
186        meta
187    }
188
189    /// Extract primary key columns from table constraints.
190    fn extract_primary_key(stmt: &CreateTable) -> Option<Vec<String>> {
191        use crate::ast::ddl::TableConstraint;
192
193        // First check table-level constraints
194        // Note: Currently only PrimaryKey variant exists; when more variants are added,
195        // this should iterate to find the first PrimaryKey constraint
196        if let Some(TableConstraint::PrimaryKey { columns, .. }) = stmt.constraints.first() {
197            return Some(columns.clone());
198        }
199
200        // Then check column-level PRIMARY KEY constraints
201        let pk_columns: Vec<String> = stmt
202            .columns
203            .iter()
204            .filter(|col| col.constraints.iter().any(Self::is_primary_key_constraint))
205            .map(|col| col.name.clone())
206            .collect();
207
208        if pk_columns.is_empty() {
209            None
210        } else {
211            Some(pk_columns)
212        }
213    }
214
215    /// Check if a column constraint is a PRIMARY KEY constraint.
216    fn is_primary_key_constraint(constraint: &ColumnConstraint) -> bool {
217        match constraint {
218            ColumnConstraint::PrimaryKey => true,
219            ColumnConstraint::WithSpan { kind, .. } => Self::is_primary_key_constraint(kind),
220            _ => false,
221        }
222    }
223
224    /// Plan a DROP TABLE statement.
225    ///
226    /// Validates that the table exists (unless IF EXISTS is specified).
227    fn plan_drop_table(&self, stmt: &DropTable) -> Result<LogicalPlan, PlannerError> {
228        // Check if table exists
229        if !stmt.if_exists && !self.catalog.table_exists(&stmt.name) {
230            return Err(PlannerError::TableNotFound {
231                name: stmt.name.clone(),
232                line: stmt.span.start.line,
233                column: stmt.span.start.column,
234            });
235        }
236
237        Ok(LogicalPlan::DropTable {
238            name: stmt.name.clone(),
239            if_exists: stmt.if_exists,
240        })
241    }
242
243    /// Plan a CREATE INDEX statement.
244    ///
245    /// Validates that:
246    /// - The index doesn't already exist (unless IF NOT EXISTS is specified)
247    /// - The target table exists
248    /// - The target column exists in the table
249    fn plan_create_index(&self, stmt: &CreateIndex) -> Result<LogicalPlan, PlannerError> {
250        // Check if index already exists
251        if !stmt.if_not_exists && self.catalog.index_exists(&stmt.name) {
252            return Err(PlannerError::index_already_exists(&stmt.name));
253        }
254
255        // Validate table exists
256        let table = self.name_resolver.resolve_table(&stmt.table, stmt.span)?;
257
258        // Validate column exists
259        self.name_resolver
260            .resolve_column(table, &stmt.column, stmt.span)?;
261
262        // Build index metadata
263        // Note: index_id is set to 0 as placeholder; Executor assigns the actual ID
264        // Note: column_indices will be resolved by Executor when table schema is available
265        let mut index = IndexMetadata::new(
266            0,
267            stmt.name.clone(),
268            stmt.table.clone(),
269            vec![stmt.column.clone()],
270        );
271
272        if let Some(method) = stmt.method {
273            index = index.with_method(method);
274        }
275
276        let options: Vec<(String, String)> = stmt
277            .options
278            .iter()
279            .map(|opt| (opt.key.clone(), opt.value.clone()))
280            .collect();
281        if !options.is_empty() {
282            index = index.with_options(options);
283        }
284
285        Ok(LogicalPlan::CreateIndex {
286            index,
287            if_not_exists: stmt.if_not_exists,
288        })
289    }
290
291    /// Plan a DROP INDEX statement.
292    ///
293    /// Validates that the index exists (unless IF EXISTS is specified).
294    fn plan_drop_index(&self, stmt: &DropIndex) -> Result<LogicalPlan, PlannerError> {
295        // Check if index exists
296        if !stmt.if_exists && !self.catalog.index_exists(&stmt.name) {
297            return Err(PlannerError::index_not_found(&stmt.name));
298        }
299
300        Ok(LogicalPlan::DropIndex {
301            name: stmt.name.clone(),
302            if_exists: stmt.if_exists,
303        })
304    }
305
306    // ============================================================
307    // DML Planning Methods (Task 17 & 18)
308    // ============================================================
309
310    /// Plan a SELECT statement.
311    ///
312    /// Builds a logical plan tree: Scan -> Filter -> Sort -> Limit
313    /// Each layer is optional and only added if the corresponding clause is present.
314    fn plan_select(&self, stmt: &Select) -> Result<LogicalPlan, PlannerError> {
315        // 1. Resolve the FROM table
316        let table = self
317            .name_resolver
318            .resolve_table(&stmt.from.name, stmt.from.span)?;
319
320        // 2. Build the projection
321        let projection = self.build_projection(&stmt.projection, table)?;
322
323        // 3. Create the base Scan plan
324        let mut plan = LogicalPlan::Scan {
325            table: table.name.clone(),
326            projection,
327        };
328
329        // 4. Add Filter if WHERE clause is present
330        if let Some(ref selection) = stmt.selection {
331            let predicate = self.type_checker.infer_type(selection, table)?;
332
333            // Verify predicate returns Boolean
334            if predicate.resolved_type != ResolvedType::Boolean {
335                return Err(PlannerError::type_mismatch(
336                    "Boolean",
337                    predicate.resolved_type.to_string(),
338                    selection.span,
339                ));
340            }
341
342            plan = LogicalPlan::Filter {
343                input: Box::new(plan),
344                predicate,
345            };
346        }
347
348        // 5. Add Sort if ORDER BY clause is present
349        if !stmt.order_by.is_empty() {
350            let order_by = self.build_sort_exprs(&stmt.order_by, table)?;
351            plan = LogicalPlan::Sort {
352                input: Box::new(plan),
353                order_by,
354            };
355        }
356
357        // 6. Add Limit if LIMIT/OFFSET is present
358        if stmt.limit.is_some() || stmt.offset.is_some() {
359            let limit = self.extract_limit_value(&stmt.limit, stmt.span)?;
360            let offset = self.extract_limit_value(&stmt.offset, stmt.span)?;
361            plan = LogicalPlan::Limit {
362                input: Box::new(plan),
363                limit,
364                offset,
365            };
366        }
367
368        Ok(plan)
369    }
370
371    /// Build the projection for a SELECT statement.
372    ///
373    /// Handles wildcard expansion and expression type checking.
374    fn build_projection(
375        &self,
376        items: &[SelectItem],
377        table: &TableMetadata,
378    ) -> Result<Projection, PlannerError> {
379        // Check for wildcard - if present, expand it
380        if items.len() == 1 && matches!(&items[0], SelectItem::Wildcard { .. }) {
381            let columns = self.name_resolver.expand_wildcard(table);
382            return Ok(Projection::All(columns));
383        }
384
385        // Process each select item
386        let mut projected_columns = Vec::new();
387        for item in items {
388            match item {
389                SelectItem::Wildcard { span } => {
390                    // Wildcard mixed with other items - expand inline
391                    for col in &table.columns {
392                        let column_index = table.get_column_index(&col.name).unwrap();
393                        let typed_expr = TypedExpr::column_ref(
394                            table.name.clone(),
395                            col.name.clone(),
396                            column_index,
397                            col.data_type.clone(),
398                            *span,
399                        );
400                        projected_columns.push(ProjectedColumn::new(typed_expr));
401                    }
402                }
403                SelectItem::Expr { expr, alias, .. } => {
404                    let typed_expr = self.type_checker.infer_type(expr, table)?;
405                    let projected = if let Some(alias) = alias {
406                        ProjectedColumn::with_alias(typed_expr, alias.clone())
407                    } else {
408                        ProjectedColumn::new(typed_expr)
409                    };
410                    projected_columns.push(projected);
411                }
412            }
413        }
414
415        Ok(Projection::Columns(projected_columns))
416    }
417
418    /// Build sort expressions from ORDER BY clause.
419    fn build_sort_exprs(
420        &self,
421        order_by: &[OrderByExpr],
422        table: &TableMetadata,
423    ) -> Result<Vec<SortExpr>, PlannerError> {
424        let mut sort_exprs = Vec::new();
425
426        for order_expr in order_by {
427            let typed_expr = self.type_checker.infer_type(&order_expr.expr, table)?;
428
429            // Determine sort direction (default: ASC)
430            let asc = order_expr.asc.unwrap_or(true);
431
432            // Determine NULLS ordering (default: NULLS LAST for both ASC and DESC)
433            let nulls_first = order_expr.nulls_first.unwrap_or(false);
434
435            sort_exprs.push(SortExpr::new(typed_expr, asc, nulls_first));
436        }
437
438        Ok(sort_exprs)
439    }
440
441    /// Extract a numeric value from a LIMIT or OFFSET expression.
442    ///
443    /// Currently only supports literal integer values.
444    fn extract_limit_value(
445        &self,
446        expr: &Option<crate::ast::expr::Expr>,
447        stmt_span: crate::ast::Span,
448    ) -> Result<Option<u64>, PlannerError> {
449        match expr {
450            None => Ok(None),
451            Some(e) => {
452                // For now, only support literal integers
453                if let crate::ast::expr::ExprKind::Literal(Literal::Number(s)) = &e.kind {
454                    s.parse::<u64>().map(Some).map_err(|_| {
455                        PlannerError::type_mismatch("unsigned integer", s.clone(), e.span)
456                    })
457                } else {
458                    Err(PlannerError::unsupported_feature(
459                        "non-literal LIMIT/OFFSET",
460                        "v0.3.0+",
461                        stmt_span,
462                    ))
463                }
464            }
465        }
466    }
467
468    /// Plan an INSERT statement.
469    ///
470    /// Handles column list specification or implicit column ordering.
471    /// When columns are omitted, uses table definition order from TableMetadata.
472    fn plan_insert(&self, stmt: &Insert) -> Result<LogicalPlan, PlannerError> {
473        // Resolve the target table
474        let table = self.name_resolver.resolve_table(&stmt.table, stmt.span)?;
475
476        // Determine the column list
477        let columns: Vec<String> = if let Some(ref cols) = stmt.columns {
478            // Explicit column list - validate each column exists
479            for col in cols {
480                self.name_resolver.resolve_column(table, col, stmt.span)?;
481            }
482            cols.clone()
483        } else {
484            // Implicit - use all columns in table definition order
485            table.column_names().into_iter().map(String::from).collect()
486        };
487
488        // Validate and type-check each row of values
489        let mut typed_values: Vec<Vec<TypedExpr>> = Vec::new();
490
491        for row in &stmt.values {
492            // Check column count matches
493            if row.len() != columns.len() {
494                return Err(PlannerError::column_value_count_mismatch(
495                    columns.len(),
496                    row.len(),
497                    stmt.span,
498                ));
499            }
500
501            // Type-check each value
502            let typed_row = self.type_check_insert_values(row, &columns, table)?;
503            typed_values.push(typed_row);
504        }
505
506        Ok(LogicalPlan::Insert {
507            table: table.name.clone(),
508            columns,
509            values: typed_values,
510        })
511    }
512
513    /// Type-check INSERT values against column definitions.
514    fn type_check_insert_values(
515        &self,
516        values: &[crate::ast::expr::Expr],
517        columns: &[String],
518        table: &TableMetadata,
519    ) -> Result<Vec<TypedExpr>, PlannerError> {
520        let mut typed_values = Vec::new();
521
522        for (i, value) in values.iter().enumerate() {
523            let column_name = &columns[i];
524            let column_meta = table.get_column(column_name).ok_or_else(|| {
525                PlannerError::column_not_found(column_name, &table.name, value.span)
526            })?;
527
528            // Type-check the value expression
529            let typed_value = self.type_checker.infer_type(value, table)?;
530
531            // Check for NOT NULL constraint violation (except for NULL literal which is allowed if nullable)
532            if column_meta.not_null
533                && matches!(&typed_value.kind, TypedExprKind::Literal(Literal::Null))
534            {
535                return Err(PlannerError::null_constraint_violation(
536                    column_name,
537                    value.span,
538                ));
539            }
540
541            // Validate type compatibility
542            self.validate_type_assignment(&typed_value, &column_meta.data_type, value.span)?;
543
544            typed_values.push(typed_value);
545        }
546
547        Ok(typed_values)
548    }
549
550    /// Validate that a value type can be assigned to a column type.
551    fn validate_type_assignment(
552        &self,
553        value: &TypedExpr,
554        target_type: &ResolvedType,
555        span: crate::ast::Span,
556    ) -> Result<(), PlannerError> {
557        // NULL can be assigned to any nullable column
558        if value.resolved_type == ResolvedType::Null {
559            return Ok(());
560        }
561
562        // Check for exact match or implicit conversion compatibility
563        if self.types_compatible(&value.resolved_type, target_type) {
564            return Ok(());
565        }
566
567        Err(PlannerError::type_mismatch(
568            target_type.to_string(),
569            value.resolved_type.to_string(),
570            span,
571        ))
572    }
573
574    /// Check if two types are compatible for assignment.
575    fn types_compatible(&self, source: &ResolvedType, target: &ResolvedType) -> bool {
576        use ResolvedType::*;
577
578        // Same type is always compatible
579        if source == target {
580            return true;
581        }
582
583        // Numeric promotions
584        match (source, target) {
585            // Integer can be assigned to BigInt, Float, Double
586            (Integer, BigInt) | (Integer, Float) | (Integer, Double) => true,
587            // BigInt can be assigned to Float, Double
588            (BigInt, Float) | (BigInt, Double) => true,
589            // Float can be assigned to Double
590            (Float, Double) => true,
591            // Vector dimensions must match
592            (Vector { dimension: d1, .. }, Vector { dimension: d2, .. }) => d1 == d2,
593            _ => false,
594        }
595    }
596
597    /// Plan an UPDATE statement.
598    ///
599    /// Validates assignments and optional WHERE clause.
600    fn plan_update(&self, stmt: &Update) -> Result<LogicalPlan, PlannerError> {
601        // Resolve the target table
602        let table = self.name_resolver.resolve_table(&stmt.table, stmt.span)?;
603
604        // Process assignments
605        let mut typed_assignments = Vec::new();
606
607        for assignment in &stmt.assignments {
608            // Resolve the column
609            let column_meta =
610                self.name_resolver
611                    .resolve_column(table, &assignment.column, assignment.span)?;
612            let column_index = table.get_column_index(&assignment.column).unwrap();
613
614            // Type-check the value expression
615            let typed_value = self.type_checker.infer_type(&assignment.value, table)?;
616
617            // Check NOT NULL constraint
618            if column_meta.not_null
619                && matches!(&typed_value.kind, TypedExprKind::Literal(Literal::Null))
620            {
621                return Err(PlannerError::null_constraint_violation(
622                    &assignment.column,
623                    assignment.value.span,
624                ));
625            }
626
627            // Validate type compatibility
628            self.validate_type_assignment(
629                &typed_value,
630                &column_meta.data_type,
631                assignment.value.span,
632            )?;
633
634            typed_assignments.push(TypedAssignment::new(
635                assignment.column.clone(),
636                column_index,
637                typed_value,
638            ));
639        }
640
641        // Process optional WHERE clause
642        let filter = if let Some(ref selection) = stmt.selection {
643            let predicate = self.type_checker.infer_type(selection, table)?;
644
645            // Verify predicate returns Boolean
646            if predicate.resolved_type != ResolvedType::Boolean {
647                return Err(PlannerError::type_mismatch(
648                    "Boolean",
649                    predicate.resolved_type.to_string(),
650                    selection.span,
651                ));
652            }
653
654            Some(predicate)
655        } else {
656            None
657        };
658
659        Ok(LogicalPlan::Update {
660            table: table.name.clone(),
661            assignments: typed_assignments,
662            filter,
663        })
664    }
665
666    /// Plan a DELETE statement.
667    ///
668    /// Validates optional WHERE clause.
669    fn plan_delete(&self, stmt: &Delete) -> Result<LogicalPlan, PlannerError> {
670        // Resolve the target table
671        let table = self.name_resolver.resolve_table(&stmt.table, stmt.span)?;
672
673        // Process optional WHERE clause
674        let filter = if let Some(ref selection) = stmt.selection {
675            let predicate = self.type_checker.infer_type(selection, table)?;
676
677            // Verify predicate returns Boolean
678            if predicate.resolved_type != ResolvedType::Boolean {
679                return Err(PlannerError::type_mismatch(
680                    "Boolean",
681                    predicate.resolved_type.to_string(),
682                    selection.span,
683                ));
684            }
685
686            Some(predicate)
687        } else {
688            None
689        };
690
691        Ok(LogicalPlan::Delete {
692            table: table.name.clone(),
693            filter,
694        })
695    }
696}