quill_sql/plan/logical_planner/
logical_planner.rs

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

quill_sql/plan/logical_planner/
logical_planner.rs
