Merge pull request #1 from mnmandahalf/support-for-specifying-column-…

…names-on-insert support for specifying column names on insert
gtnao · Jul 5, 2024 · e018d96 · e018d96
2 parents b3a044b + 09defe6
commit e018d96
Show file tree

Hide file tree

Showing 4 changed files with 96 additions and 8 deletions.
diff --git a/src/binder.rs b/src/binder.rs
@@ -78,6 +78,7 @@ pub struct BoundSubqueryTableReferenceAST {
 #[derive(Debug, PartialEq, Eq, Clone)]
 pub struct BoundInsertStatementAST {
     pub table_name: String,
+    pub column_names: Option<Vec<String>>,
     pub values: Vec<BoundExpressionAST>,
     pub first_page_id: PageID,
     pub table_schema: Schema,
@@ -440,19 +441,45 @@ impl Binder {
             .lock()
             .map_err(|_| anyhow::anyhow!("lock error"))?
             .get_schema_by_table_name(&statement.table_name, self.txn_id)?;
-        if statement.values.len() != schema.columns.len() {
-            return Err(anyhow::anyhow!(
-                "expected {} values, but got {}",
-                schema.columns.len(),
-                statement.values.len()
-            ));
+        match &statement.column_names {
+            Some(column_names) => {
+                if column_names.len() != statement.values.len() {
+                    return Err(anyhow::anyhow!(
+                        "expected {} values, but got {}",
+                        column_names.len(),
+                        statement.values.len()
+                    ));
+                }
+                if column_names.len() > schema.columns.len() {
+                    return Err(anyhow::anyhow!(
+                        "expected {} values, but got {}",
+                        schema.columns.len(),
+                        column_names.len()
+                    ));
+                }
+                for column_name in column_names {
+                    if !schema.columns.iter().any(|column| column.name == *column_name) {
+                        return Err(anyhow::anyhow!("column {} not found", column_name));
+                    }
+                }
+            }
+            None => {
+                if statement.values.len() != schema.columns.len() {
+                    return Err(anyhow::anyhow!(
+                        "expected {} values, but got {}",
+                        schema.columns.len(),
+                        statement.values.len()
+                    ));
+                }
+            }
         }
         let mut values = Vec::new();
         for value in &statement.values {
             values.push(self.bind_expression(value)?);
         }
         Ok(BoundStatementAST::Insert(BoundInsertStatementAST {
             table_name: statement.table_name.clone(),
+            column_names: statement.column_names.clone(),
             values,
             first_page_id,
             table_schema: schema,
@@ -1283,6 +1310,7 @@ mod tests {
             bound_statement,
             BoundStatementAST::Insert(BoundInsertStatementAST {
                 table_name: "t1".to_string(),
+                column_names: None,
                 values: vec![
                     BoundExpressionAST::Literal(BoundLiteralExpressionAST {
                         value: Value::Integer(IntegerValue(1)),

diff --git a/src/executor/insert_executor.rs b/src/executor/insert_executor.rs
@@ -30,7 +30,20 @@ impl InsertExecutor<'_> {
             .iter()
             .enumerate()
             .map(|(i, c)| {
-                let raw_value = self.plan.values[i].eval(
+                let index;
+                match &self.plan.column_names {
+                    Some(column_names) => {
+                        let position = column_names.iter().position(|x| x == &c.name);
+                        match position {
+                            Some(pos) => index = pos,
+                            None => return Ok(Value::Null)
+                        }
+                    },
+                    None => {
+                        index = i;
+                    }
+                }
+                let raw_value = self.plan.values[index].eval(
                     &vec![&Tuple::new(None, &[])],
                     &vec![&Schema { columns: vec![] }],
                 )?;

diff --git a/src/parser.rs b/src/parser.rs
@@ -95,6 +95,7 @@ pub struct LimitAST {
 #[derive(Debug, PartialEq, Eq, Clone)]
 pub struct InsertStatementAST {
     pub table_name: String,
+    pub column_names: Option<Vec<String>>,
     // TODO: support multiple rows
     pub values: Vec<ExpressionAST>,
 }
@@ -460,6 +461,18 @@ impl Parser {
         self.consume_token_or_error(Token::Keyword(Keyword::Insert))?;
         self.consume_token_or_error(Token::Keyword(Keyword::Into))?;
         let table_name = self.identifier()?;
+        let mut column_names: Option<Vec<String>> = None;
+        if self.consume_token(Token::LeftParen) {
+            let mut names = Vec::new();
+            loop {
+                names.push(self.identifier()?);
+                if !self.consume_token(Token::Comma) {
+                    break;
+                }
+            }
+            self.consume_token_or_error(Token::RightParen)?;
+            column_names = Some(names);
+        }
         self.consume_token_or_error(Token::Keyword(Keyword::Values))?;
         self.consume_token_or_error(Token::LeftParen)?;
         let mut values = Vec::new();
@@ -470,7 +483,7 @@ impl Parser {
             }
         }
         self.consume_token_or_error(Token::RightParen)?;
-        Ok(InsertStatementAST { table_name, values })
+        Ok(InsertStatementAST { table_name, column_names, values })
     }
     fn delete_statement(&mut self) -> Result<DeleteStatementAST> {
         self.consume_token_or_error(Token::Keyword(Keyword::Delete))?;
@@ -956,6 +969,38 @@ mod tests {
             statement,
             StatementAST::Insert(InsertStatementAST {
                 table_name: String::from("users"),
+                column_names: None,
+                values: vec![
+                    ExpressionAST::Literal(LiteralExpressionAST {
+                        value: Value::Integer(IntegerValue(1)),
+                    }),
+                    ExpressionAST::Literal(LiteralExpressionAST {
+                        value: Value::Varchar(VarcharValue(String::from("foo"))),
+                    }),
+                    ExpressionAST::Literal(LiteralExpressionAST {
+                        value: Value::Boolean(BooleanValue(true)),
+                    }),
+                ],
+            })
+        );
+        Ok(())
+    }
+
+    #[test]
+    fn test_parse_insert_with_columns() -> Result<()> {
+        let sql = "INSERT INTO users (id, name, is_deleted) VALUES (1, 'foo', true)";
+        let mut parser = Parser::new(tokenize(&mut sql.chars().peekable())?);
+
+        let statement = parser.parse()?;
+        assert_eq!(
+            statement,
+            StatementAST::Insert(InsertStatementAST {
+                table_name: String::from("users"),
+                column_names: Some(vec![
+                    String::from("id"),
+                    String::from("name"),
+                    String::from("is_deleted"),
+                ]),
                 values: vec![
                     ExpressionAST::Literal(LiteralExpressionAST {
                         value: Value::Integer(IntegerValue(1)),

diff --git a/src/plan.rs b/src/plan.rs
@@ -137,6 +137,7 @@ pub struct EmptyRowPlan {
 pub struct InsertPlan {
     pub first_page_id: PageID,
     pub table_schema: Schema,
+    pub column_names: Option<Vec<String>>,
     pub values: Vec<BoundExpressionAST>,
     pub schema: Schema,
     pub table_name: String,
@@ -382,6 +383,7 @@ impl Planner {
         Plan::Insert(InsertPlan {
             first_page_id: insert_statement.first_page_id,
             table_schema: insert_statement.table_schema.clone(),
+            column_names: insert_statement.column_names.clone(),
             values: insert_statement.values.clone(),
             schema: Schema {
                 columns: vec![Column {