[arrow-datafusion] branch master updated: Validate assignment target column existence for UPDATE statements (#5069)

[al...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

alamb pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git


The following commit(s) were added to refs/heads/master by this push:
     new 14e153e75 Validate assignment target column existence for UPDATE statements (#5069)
14e153e75 is described below

commit 14e153e755b845295a78ec3833ebf792d3365f3e
Author: Marko Grujic <ma...@gmail.com>
AuthorDate: Thu Jan 26 14:33:45 2023 +0100

    Validate assignment target column existence for UPDATE statements (#5069)
---
 datafusion/sql/Cargo.toml                |  3 +++
 datafusion/sql/src/statement.rs          |  9 +++++----
 datafusion/sql/tests/integration_test.rs | 14 ++++++++++++++
 3 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/datafusion/sql/Cargo.toml b/datafusion/sql/Cargo.toml
index 5676a23c6..bec70291a 100644
--- a/datafusion/sql/Cargo.toml
+++ b/datafusion/sql/Cargo.toml
@@ -42,3 +42,6 @@ datafusion-common = { path = "../common", version = "16.0.0" }
 datafusion-expr = { path = "../expr", version = "16.0.0" }
 log = "^0.4"
 sqlparser = "0.30"
+
+[dev-dependencies]
+rstest = "*"
diff --git a/datafusion/sql/src/statement.rs b/datafusion/sql/src/statement.rs
index 1d31e316a..3368b59b3 100644
--- a/datafusion/sql/src/statement.rs
+++ b/datafusion/sql/src/statement.rs
@@ -670,7 +670,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
 
         // Overwrite with assignment expressions
         let mut planner_context = PlannerContext::new();
-        let assign_map = assignments
+        let mut assign_map = assignments
             .iter()
             .map(|assign| {
                 let col_name: &Ident = assign
@@ -678,11 +678,12 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
                     .iter()
                     .last()
                     .ok_or(DataFusionError::Plan("Empty column id".to_string()))?;
+                // Validate that the assignment target column exists
+                table_schema.field_with_unqualified_name(&col_name.value)?;
                 Ok((col_name.value.clone(), assign.value.clone()))
             })
-            .collect::<Result<Vec<_>>>()?;
-        let mut assign_map: HashMap<String, ast::Expr> =
-            HashMap::from_iter(assign_map.into_iter());
+            .collect::<Result<HashMap<String, Expr>>>()?;
+
         let values = values
             .into_iter()
             .map(|(k, v)| {
diff --git a/datafusion/sql/tests/integration_test.rs b/datafusion/sql/tests/integration_test.rs
index 9a1d5a883..e48095526 100644
--- a/datafusion/sql/tests/integration_test.rs
+++ b/datafusion/sql/tests/integration_test.rs
@@ -34,6 +34,8 @@ use datafusion_expr::{AggregateUDF, ScalarUDF};
 use datafusion_sql::parser::DFParser;
 use datafusion_sql::planner::{ContextProvider, ParserOptions, SqlToRel};
 
+use rstest::rstest;
+
 #[test]
 fn parse_decimals() {
     let test_data = [
@@ -183,6 +185,18 @@ Dml: op=[Update] table=[person]
     quick_test(sql, plan);
 }
 
+#[rstest]
+#[case::missing_assignement_target("UPDATE person SET doesnotexist = true")]
+#[case::missing_assignement_expression("UPDATE person SET age = doesnotexist + 42")]
+#[case::missing_selection_expression(
+    "UPDATE person SET age = 42 WHERE doesnotexist = true"
+)]
+#[test]
+fn update_column_does_not_exist(#[case] sql: &str) {
+    let err = logical_plan(sql).expect_err("query should have failed");
+    assert_field_not_found(err, "doesnotexist");
+}
+
 #[test]
 fn plan_delete() {
     let sql = "delete from person where id=1";