quill_sql/plan/logical_planner/
plan_insert.rs

1use crate::error::QuillSQLResult;
2use std::sync::Arc;
3
4use crate::plan::logical_plan::{Insert, LogicalPlan, Values};
5
6use super::LogicalPlanner;
7
8impl<'a> LogicalPlanner<'a> {
9    pub fn plan_insert(
10        &self,
11        table_name: &sqlparser::ast::ObjectName,
12        columns_ident: &Vec<sqlparser::ast::Ident>,
13        source: &sqlparser::ast::Query,
14    ) -> QuillSQLResult<LogicalPlan> {
15        let mut input = self.plan_set_expr(source.body.as_ref())?;
16        let table = self.bind_table_name(table_name)?;
17        let table_schema = self.context.catalog.table_heap(&table)?.schema.clone();
18
19        let projected_schema = if columns_ident.is_empty() {
20            table_schema.clone()
21        } else {
22            let columns: Vec<String> = columns_ident
23                .iter()
24                .map(|ident| ident.value.clone())
25                .collect();
26            let indices = columns
27                .iter()
28                .map(|name| table_schema.index_of(Some(&table), name.as_str()))
29                .collect::<QuillSQLResult<Vec<usize>>>()?;
30
31            Arc::new(table_schema.project(&indices)?)
32        };
33
34        if let LogicalPlan::Values(Values { values, .. }) = input {
35            input = LogicalPlan::Values(Values {
36                values,
37                schema: projected_schema.clone(),
38            })
39        }
40
41        Ok(LogicalPlan::Insert(Insert {
42            table,
43            table_schema,
44            projected_schema,
45            input: Arc::new(input),
46        }))
47    }
48}