quill_sql/plan/logical_planner/
planner.rs

1use crate::error::{QuillSQLError, QuillSQLResult};
2
3use crate::catalog::Catalog;
4use crate::plan::logical_plan::{
5    Analyze, LogicalPlan, OrderByExpr, TransactionModes, TransactionScope,
6};
7use crate::sql::ast;
8use crate::sql::ast::Value;
9use crate::utils::table_ref::TableReference;
10use sqlparser::ast::{Ident, ObjectType};
11
12pub struct PlannerContext<'a> {
13    pub catalog: &'a Catalog,
14}
15
16pub struct LogicalPlanner<'a> {
17    pub context: PlannerContext<'a>,
18}
19impl<'a> LogicalPlanner<'a> {
20    pub fn plan(&mut self, stmt: &ast::Statement) -> QuillSQLResult<LogicalPlan> {
21        match stmt {
22            ast::Statement::CreateTable {
23                name,
24                columns,
25                if_not_exists,
26                engine,
27                ..
28            } => self.plan_create_table(name, columns, *if_not_exists, engine.as_deref()),
29            ast::Statement::CreateIndex {
30                name,
31                table_name,
32                using,
33                columns,
34                ..
35            } => self.plan_create_index(name, table_name, using.as_ref(), columns),
36            ast::Statement::Drop {
37                object_type,
38                if_exists,
39                names,
40                cascade,
41                restrict: _,
42                purge,
43            } => match object_type {
44                ObjectType::Table => self.plan_drop_table(names, *if_exists, *cascade, *purge),
45                ObjectType::Index => self.plan_drop_index(names, *if_exists, *cascade, *purge),
46                other => Err(QuillSQLError::NotSupport(format!(
47                    "DROP {} is not supported",
48                    other
49                ))),
50            },
51            ast::Statement::Query(query) => self.plan_query(query),
52            ast::Statement::Insert {
53                table_name,
54                columns,
55                source,
56                ..
57            } => self.plan_insert(table_name, columns, source),
58            ast::Statement::Update {
59                table,
60                assignments,
61                selection,
62                ..
63            } => self.plan_update(table, assignments, selection),
64            ast::Statement::Delete {
65                tables,
66                from,
67                using,
68                selection,
69                returning,
70            } => {
71                if !tables.is_empty() {
72                    return Err(QuillSQLError::Plan(
73                        "DELETE with table aliases is not supported".to_string(),
74                    ));
75                }
76                if using.is_some() {
77                    return Err(QuillSQLError::Plan(
78                        "DELETE USING clause is not supported".to_string(),
79                    ));
80                }
81                if returning.is_some() {
82                    return Err(QuillSQLError::Plan(
83                        "DELETE RETURNING is not supported".to_string(),
84                    ));
85                }
86                if from.len() != 1 {
87                    return Err(QuillSQLError::Plan(
88                        "DELETE must target exactly one table".to_string(),
89                    ));
90                }
91                self.plan_delete(&from[0], selection)
92            }
93            ast::Statement::Explain { statement, .. } => self.plan_explain(statement),
94            ast::Statement::StartTransaction { modes, .. } => self.plan_begin_transaction(modes),
95            ast::Statement::Commit { .. } => Ok(LogicalPlan::CommitTransaction),
96            ast::Statement::Rollback { .. } => Ok(LogicalPlan::RollbackTransaction),
97            ast::Statement::SetTransaction {
98                modes,
99                snapshot,
100                session,
101            } => self.plan_set_transaction(*session, snapshot, modes),
102            ast::Statement::Analyze {
103                table_name,
104                partitions,
105                for_columns,
106                columns,
107                cache_metadata,
108                noscan,
109                compute_statistics: _,
110            } => self.plan_analyze(
111                table_name,
112                partitions,
113                *for_columns,
114                columns,
115                *cache_metadata,
116                *noscan,
117            ),
118            _ => unimplemented!(),
119        }
120    }
121
122    fn plan_begin_transaction(
123        &self,
124        modes: &[ast::TransactionMode],
125    ) -> QuillSQLResult<LogicalPlan> {
126        Ok(LogicalPlan::BeginTransaction(TransactionModes::from_modes(
127            modes,
128        )))
129    }
130
131    fn plan_set_transaction(
132        &self,
133        session_scope: bool,
134        snapshot: &Option<Value>,
135        modes: &[ast::TransactionMode],
136    ) -> QuillSQLResult<LogicalPlan> {
137        if snapshot.is_some() {
138            return Err(QuillSQLError::Plan(
139                "SET TRANSACTION SNAPSHOT is not supported".to_string(),
140            ));
141        };
142        let logical_scope = if session_scope {
143            TransactionScope::Session
144        } else {
145            TransactionScope::Transaction
146        };
147        Ok(LogicalPlan::SetTransaction {
148            scope: logical_scope,
149            modes: TransactionModes::from_modes(modes),
150        })
151    }
152
153    pub fn bind_order_by_expr(&self, order_by: &ast::OrderByExpr) -> QuillSQLResult<OrderByExpr> {
154        let expr = self.bind_expr(&order_by.expr)?;
155        Ok(OrderByExpr {
156            expr: Box::new(expr),
157            asc: order_by.asc.unwrap_or(true),
158            nulls_first: order_by.nulls_first.unwrap_or(false),
159        })
160    }
161
162    pub fn bind_table_name(&self, table_name: &ast::ObjectName) -> QuillSQLResult<TableReference> {
163        match table_name.0.as_slice() {
164            [table] => Ok(TableReference::Bare {
165                table: table.value.clone(),
166            }),
167            [schema, table] => Ok(TableReference::Partial {
168                schema: schema.value.clone(),
169                table: table.value.clone(),
170            }),
171            [catalog, schema, table] => Ok(TableReference::Full {
172                catalog: catalog.value.clone(),
173                schema: schema.value.clone(),
174                table: table.value.clone(),
175            }),
176            _ => Err(QuillSQLError::Plan(format!(
177                "Fail to plan table name: {}",
178                table_name
179            ))),
180        }
181    }
182
183    fn plan_analyze(
184        &self,
185        table_name: &ast::ObjectName,
186        partitions: &Option<Vec<ast::Expr>>,
187        for_columns: bool,
188        columns: &[Ident],
189        cache_metadata: bool,
190        noscan: bool,
191    ) -> QuillSQLResult<LogicalPlan> {
192        if partitions.as_ref().is_some_and(|p| !p.is_empty()) {
193            return Err(QuillSQLError::Plan(
194                "ANALYZE PARTITION is not supported".to_string(),
195            ));
196        }
197        if for_columns || !columns.is_empty() {
198            return Err(QuillSQLError::Plan(
199                "ANALYZE FOR COLUMNS is not supported".to_string(),
200            ));
201        }
202        if cache_metadata || noscan {
203            return Err(QuillSQLError::Plan(
204                "ANALYZE options NOSCAN/CACHE METADATA are not supported".to_string(),
205            ));
206        }
207        let table_ref = self.bind_table_name(table_name)?;
208        Ok(LogicalPlan::Analyze(Analyze { table: table_ref }))
209    }
210}
quill_sql/plan/logical_planner/planner.rs

quill_sql/plan/logical_planner/
planner.rs
