Skip to main content

featherdb_query/planner/
mod.rs

1//! Query planner - converts SQL AST to logical plans
2
3// Submodules
4mod context;
5mod cte;
6mod ddl;
7mod dml;
8mod expr_conv;
9mod plan;
10mod select;
11mod subquery;
12mod types;
13mod util;
14mod window;
15
16// Re-exports
17pub use context::PlannerContext;
18pub use cte::CommonTableExpression;
19pub use plan::{LogicalPlan, ShowGrantsTarget};
20pub use types::{IndexBound, IndexRange, JoinType, SortOrder};
21
22use featherdb_catalog::Catalog;
23use featherdb_core::{Error, Permissions, QueryError, Result};
24use sqlparser::ast::{Ident, Statement};
25use std::cell::{Cell, RefCell};
26use std::collections::HashMap;
27
28/// Query planner with enhanced error messages
29pub struct Planner<'a> {
30    catalog: &'a Catalog,
31    /// Current SQL being planned (for error messages)
32    current_sql: Option<String>,
33    /// Tracked subquery plans (subquery_id -> LogicalPlan)
34    subquery_plans: RefCell<HashMap<usize, LogicalPlan>>,
35    /// CTE definitions visible to the current planning scope.
36    /// When a WITH clause is encountered, these are propagated into
37    /// subquery planning so that scalar/IN/EXISTS subqueries can
38    /// reference CTEs defined in the outer query.
39    pub(crate) cte_scope: RefCell<Vec<CommonTableExpression>>,
40    /// Per-planner-instance subquery counter. Starts at 0 for each
41    /// new Planner so subquery IDs are deterministic for the same SQL.
42    subquery_counter: Cell<usize>,
43}
44
45impl<'a> Planner<'a> {
46    /// Create a new planner
47    pub fn new(catalog: &'a Catalog) -> Self {
48        Planner {
49            catalog,
50            current_sql: None,
51            subquery_plans: RefCell::new(HashMap::new()),
52            cte_scope: RefCell::new(Vec::new()),
53            subquery_counter: Cell::new(0),
54        }
55    }
56
57    /// Allocate the next subquery ID (deterministic per planner instance)
58    pub(super) fn next_subquery_id(&self) -> usize {
59        let id = self.subquery_counter.get();
60        self.subquery_counter.set(id + 1);
61        id
62    }
63
64    /// Take the collected subquery plans out of the planner
65    pub fn take_subquery_plans(&self) -> HashMap<usize, LogicalPlan> {
66        self.subquery_plans.borrow_mut().drain().collect()
67    }
68
69    /// Plan a statement with SQL context for better error messages
70    pub fn plan_with_sql(&mut self, stmt: &Statement, sql: &str) -> Result<LogicalPlan> {
71        self.current_sql = Some(sql.to_string());
72        let result = self.plan(stmt);
73        self.current_sql = None;
74        result
75    }
76
77    /// Create a rich query error
78    #[allow(dead_code)]
79    fn query_error(&self, message: &str, suggestion: Option<&str>) -> Error {
80        let sql = self.current_sql.as_deref().unwrap_or("");
81        let mut err = QueryError::new(message, sql, 1, 1);
82        if let Some(sugg) = suggestion {
83            err = err.with_suggestion(sugg);
84        }
85        Error::Query(err)
86    }
87
88    /// Plan a SQL statement
89    pub fn plan(&self, stmt: &Statement) -> Result<LogicalPlan> {
90        match stmt {
91            Statement::Query(query) => self.plan_query(query),
92            Statement::Insert {
93                table_name,
94                columns,
95                source,
96                ..
97            } => self.plan_insert_fields(table_name, columns, source.as_deref()),
98            Statement::Update {
99                table,
100                assignments,
101                selection,
102                ..
103            } => self.plan_update(table, assignments, selection.as_ref()),
104            Statement::Delete {
105                from, selection, ..
106            } => self.plan_delete_fields(from, selection.as_ref()),
107            Statement::CreateTable {
108                name,
109                columns,
110                constraints,
111                ..
112            } => self.plan_create_table_fields(name, columns, constraints),
113            Statement::CreateView {
114                name,
115                columns,
116                query,
117                or_replace: _,
118                ..
119            } => {
120                let view_name = name.0.first().map(|i| i.value.clone()).unwrap_or_default();
121                let col_names: Vec<String> = columns.iter().map(|c| c.value.clone()).collect();
122                let query_sql = format!("{}", query);
123                let query_plan = self.plan_query(query)?;
124                Ok(LogicalPlan::CreateView {
125                    name: view_name,
126                    query: Box::new(query_plan),
127                    query_sql,
128                    columns: col_names,
129                })
130            }
131            Statement::Drop {
132                object_type,
133                names,
134                if_exists,
135                ..
136            } => self.plan_drop(object_type, names, *if_exists),
137            Statement::Explain {
138                statement,
139                verbose,
140                analyze,
141                ..
142            } => self.plan_explain(statement, *verbose, *analyze),
143            Statement::AlterTable {
144                name, operations, ..
145            } => self.plan_alter_table(name, operations),
146            Statement::Grant {
147                privileges,
148                objects,
149                grantees,
150                ..
151            } => self.plan_grant(privileges, objects, grantees),
152            Statement::Revoke {
153                privileges,
154                objects,
155                grantees,
156                ..
157            } => self.plan_revoke(privileges, objects, grantees),
158            Statement::CreateIndex {
159                name,
160                table_name,
161                columns,
162                unique,
163                if_not_exists,
164                ..
165            } => {
166                let idx_name = name
167                    .as_ref()
168                    .map(|n| {
169                        n.0.iter()
170                            .map(|i| i.value.clone())
171                            .collect::<Vec<_>>()
172                            .join(".")
173                    })
174                    .unwrap_or_default();
175                let tbl_name = table_name
176                    .0
177                    .iter()
178                    .map(|i| i.value.clone())
179                    .collect::<Vec<_>>()
180                    .join(".");
181                // Verify table exists
182                let table = self.catalog.get_table(&tbl_name)?;
183                // Check for duplicate index name (unless IF NOT EXISTS)
184                if table.indexes.iter().any(|i| i.name == idx_name) {
185                    if *if_not_exists {
186                        return Ok(LogicalPlan::EmptyRelation);
187                    }
188                    return Err(Error::InvalidQuery {
189                        message: format!(
190                            "Index '{}' already exists on table '{}'",
191                            idx_name, tbl_name
192                        ),
193                    });
194                }
195                // Resolve column names
196                let col_names: Vec<String> = columns.iter().map(|c| c.expr.to_string()).collect();
197                // Verify columns exist
198                for col_name in &col_names {
199                    if table.get_column_index(col_name).is_none() {
200                        return Err(Error::ColumnNotFound {
201                            column: col_name.clone(),
202                            table: tbl_name.clone(),
203                            suggestion: None,
204                        });
205                    }
206                }
207                Ok(LogicalPlan::CreateIndex {
208                    index_name: idx_name,
209                    table_name: tbl_name,
210                    columns: col_names,
211                    unique: *unique,
212                })
213            }
214            _ => Err(Error::Unsupported {
215                feature: format!("Statement: {:?}", stmt),
216            }),
217        }
218    }
219
220    /// Plan a GRANT statement
221    fn plan_grant(
222        &self,
223        privileges: &sqlparser::ast::Privileges,
224        objects: &sqlparser::ast::GrantObjects,
225        grantees: &[Ident],
226    ) -> Result<LogicalPlan> {
227        let perms = self.convert_privileges(privileges)?;
228        let table = self.extract_table_from_grant_objects(objects)?;
229        let grantee = self.extract_grantee_from_idents(grantees)?;
230
231        Ok(LogicalPlan::Grant {
232            privileges: perms,
233            table,
234            grantee,
235        })
236    }
237
238    /// Plan a REVOKE statement
239    fn plan_revoke(
240        &self,
241        privileges: &sqlparser::ast::Privileges,
242        objects: &sqlparser::ast::GrantObjects,
243        grantees: &[Ident],
244    ) -> Result<LogicalPlan> {
245        let perms = self.convert_privileges(privileges)?;
246        let table = self.extract_table_from_grant_objects(objects)?;
247        let grantee = self.extract_grantee_from_idents(grantees)?;
248
249        Ok(LogicalPlan::Revoke {
250            privileges: perms,
251            table,
252            grantee,
253        })
254    }
255
256    /// Convert sqlparser Privileges to our Permissions bitflags
257    fn convert_privileges(&self, privileges: &sqlparser::ast::Privileges) -> Result<Permissions> {
258        match privileges {
259            sqlparser::ast::Privileges::All { .. } => Ok(Permissions::ALL),
260            sqlparser::ast::Privileges::Actions(actions) => {
261                let mut perms = Permissions::empty();
262                for action in actions {
263                    match action {
264                        sqlparser::ast::Action::Select { .. } => perms |= Permissions::SELECT,
265                        sqlparser::ast::Action::Insert { .. } => perms |= Permissions::INSERT,
266                        sqlparser::ast::Action::Update { .. } => perms |= Permissions::UPDATE,
267                        sqlparser::ast::Action::Delete => perms |= Permissions::DELETE,
268                        other => {
269                            return Err(Error::Unsupported {
270                                feature: format!("Permission: {:?}", other),
271                            });
272                        }
273                    }
274                }
275                Ok(perms)
276            }
277        }
278    }
279
280    /// Extract table name from GrantObjects (or None for wildcard)
281    fn extract_table_from_grant_objects(
282        &self,
283        objects: &sqlparser::ast::GrantObjects,
284    ) -> Result<Option<String>> {
285        match objects {
286            sqlparser::ast::GrantObjects::Tables(tables) => {
287                if tables.is_empty() {
288                    return Ok(None);
289                }
290                // Get the first table name
291                let table_name = tables
292                    .first()
293                    .and_then(|t| t.0.first())
294                    .map(|ident| ident.value.clone());
295                Ok(table_name)
296            }
297            sqlparser::ast::GrantObjects::AllTablesInSchema { .. } => {
298                // Treat as wildcard for now
299                Ok(None)
300            }
301            other => Err(Error::Unsupported {
302                feature: format!("Grant object type: {:?}", other),
303            }),
304        }
305    }
306
307    /// Extract grantee (API key ID) from the idents list
308    fn extract_grantee_from_idents(&self, grantees: &[Ident]) -> Result<String> {
309        if grantees.is_empty() {
310            return Err(Error::InvalidQuery {
311                message: "GRANT/REVOKE requires a grantee (API key ID)".to_string(),
312            });
313        }
314
315        // Get the first grantee as the API key ID
316        Ok(grantees[0].value.clone())
317    }
318
319    /// Plan a SHOW GRANTS statement (custom parsing required)
320    pub fn plan_show_grants(&self, target: ShowGrantsTarget) -> Result<LogicalPlan> {
321        Ok(LogicalPlan::ShowGrants { target })
322    }
323
324    /// Plan an EXPLAIN statement
325    fn plan_explain(
326        &self,
327        statement: &Statement,
328        verbose: bool,
329        analyze: bool,
330    ) -> Result<LogicalPlan> {
331        // Plan the inner statement
332        let inner_plan = self.plan(statement)?;
333
334        Ok(LogicalPlan::Explain {
335            input: Box::new(inner_plan),
336            verbose,
337            analyze,
338        })
339    }
340}
341
342#[cfg(test)]
343mod tests {
344    use super::*;
345    use crate::Parser;
346    use featherdb_catalog::TableBuilder;
347    use featherdb_core::ColumnType;
348
349    fn create_test_catalog() -> Catalog {
350        let catalog = Catalog::new();
351
352        let users = TableBuilder::new("users")
353            .column_with(
354                "id",
355                ColumnType::Integer,
356                vec![featherdb_catalog::ColumnConstraint::PrimaryKey],
357            )
358            .column_with("name", ColumnType::Text { max_len: None }, vec![])
359            .column("age", ColumnType::Integer)
360            .build(0, featherdb_core::PageId::INVALID);
361
362        catalog.create_table(users).unwrap();
363
364        catalog
365    }
366
367    #[test]
368    fn test_plan_select() {
369        let catalog = create_test_catalog();
370        let planner = Planner::new(&catalog);
371        let stmt = Parser::parse_one("SELECT * FROM users").unwrap();
372        let plan = planner.plan(&stmt).unwrap();
373
374        matches!(plan, LogicalPlan::Project { .. });
375    }
376
377    #[test]
378    fn test_plan_select_where() {
379        let catalog = create_test_catalog();
380        let planner = Planner::new(&catalog);
381        let stmt = Parser::parse_one("SELECT * FROM users WHERE age > 18").unwrap();
382        let plan = planner.plan(&stmt).unwrap();
383
384        matches!(plan, LogicalPlan::Project { .. });
385    }
386
387    #[test]
388    fn test_plan_aggregate() {
389        let catalog = create_test_catalog();
390        let planner = Planner::new(&catalog);
391        let stmt = Parser::parse_one("SELECT COUNT(*) FROM users").unwrap();
392        let plan = planner.plan(&stmt).unwrap();
393
394        matches!(plan, LogicalPlan::Project { .. });
395    }
396
397    #[test]
398    fn test_table_not_found_suggestion() {
399        let catalog = Catalog::new();
400        let users = TableBuilder::new("users")
401            .column("id", ColumnType::Integer)
402            .build(0, featherdb_core::PageId::INVALID);
403        catalog.create_table(users).unwrap();
404
405        let planner = Planner::new(&catalog);
406        let stmt = Parser::parse_one("SELECT * FROM user").unwrap();
407        let result = planner.plan(&stmt);
408
409        assert!(result.is_err());
410        if let Err(Error::TableNotFound { table, suggestion }) = result {
411            assert_eq!(table, "user");
412            assert_eq!(suggestion, Some("users".to_string()));
413        } else {
414            panic!("Expected TableNotFound error");
415        }
416    }
417}