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

quill_sql/plan/logical_planner/
logical_planner.rs
