From 5efe089d0b775990521757fec5ff7a2d6e073835 Mon Sep 17 00:00:00 2001
From: Clearlove <52417396+Eurekaaw@users.noreply.github.com>
Date: Thu, 16 Mar 2023 09:42:52 -0400
Subject: [PATCH] feat(frontend): ban `update` statements modifying pk columns
 (#8569)

Signed-off-by: Clearlove <yifei.c.wei@gmail.com>
---
 e2e_test/batch/basic/dml.slt.part             | 21 +++++++++++++++++++
 e2e_test/ddl/invalid_operation.slt            |  2 +-
 e2e_test/ddl/table.slt                        |  3 +++
 .../planner_test/tests/testdata/update.yaml   |  4 ++++
 src/frontend/src/binder/update.rs             | 17 +++++++++++++++
 5 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/e2e_test/batch/basic/dml.slt.part b/e2e_test/batch/basic/dml.slt.part
index 8b387b56e570..1b4ad5e459c2 100644
--- a/e2e_test/batch/basic/dml.slt.part
+++ b/e2e_test/batch/basic/dml.slt.part
@@ -70,3 +70,24 @@ select count(*) from t;
 
 statement ok
 drop table t;
+
+statement ok
+create table t (v1 int, v2 int primary key, v3 varchar);
+
+statement ok
+insert into t values (0, 1, 'a'), (1, 2, 'b');
+
+statement ok
+update t set (v1, v3) = (v1+v2, v3||v3);
+
+query IIT
+select * from t order by v1;
+----
+1 1 aa
+3 2 bb
+
+statement error QueryError: Bind error: update modifying the PK column is banned
+update t set (v3, v2) = (v3||v3, v1+v2);
+
+statement ok
+drop table t;
diff --git a/e2e_test/ddl/invalid_operation.slt b/e2e_test/ddl/invalid_operation.slt
index e236acd071a2..e1a7ad90db12 100644
--- a/e2e_test/ddl/invalid_operation.slt
+++ b/e2e_test/ddl/invalid_operation.slt
@@ -264,7 +264,7 @@ SELECT * from v limit 0;
 statement ok
 insert into t values (1);
 
-statement ok
+statement error QueryError: Bind error: update modifying the PK column is banned
 update t set v = 2;
 
 statement ok
diff --git a/e2e_test/ddl/table.slt b/e2e_test/ddl/table.slt
index dacebfe493b3..9a3c83a60e11 100644
--- a/e2e_test/ddl/table.slt
+++ b/e2e_test/ddl/table.slt
@@ -120,6 +120,9 @@ create table "T2" ("V1" int);
 statement ok
 insert into "T2" ("V1") values (1);
 
+statement ok
+drop table "T2"
+
 statement error
 create table C1 (c1 varchar(5));
 
diff --git a/src/frontend/planner_test/tests/testdata/update.yaml b/src/frontend/planner_test/tests/testdata/update.yaml
index c604c1782759..d8ec77dcac35 100644
--- a/src/frontend/planner_test/tests/testdata/update.yaml
+++ b/src/frontend/planner_test/tests/testdata/update.yaml
@@ -60,3 +60,7 @@
         └─BatchExchange { order: [], dist: Single }
           └─BatchFilter { predicate: (t.v1 <> t.v2) }
             └─BatchScan { table: t, columns: [t.v1, t.v2, t._row_id], distribution: UpstreamHashShard(t._row_id) }
+- sql: |
+    create table t (v1 int primary key, v2 int);
+    update t set (v2, v1) = (v1, v2);
+  binder_error: 'Bind error: update modifying the PK column is banned'
diff --git a/src/frontend/src/binder/update.rs b/src/frontend/src/binder/update.rs
index d5a7eb95cede..7a65c617db3b 100644
--- a/src/frontend/src/binder/update.rs
+++ b/src/frontend/src/binder/update.rs
@@ -91,6 +91,11 @@ impl Binder {
             Self::resolve_schema_qualified_name(&self.db_name, name.clone())?;
 
         let table_catalog = self.resolve_dml_table(schema_name.as_deref(), &table_name, false)?;
+        let pk_indices = table_catalog
+            .pk()
+            .iter()
+            .map(|column_order| column_order.column_index)
+            .collect_vec();
         let table_id = table_catalog.id;
         let owner = table_catalog.owner;
         let table_version_id = table_catalog.version_id().expect("table must be versioned");
@@ -131,6 +136,18 @@ impl Binder {
 
             for (id, value) in assignments {
                 let id_expr = self.bind_expr(Expr::Identifier(id.clone()))?;
+                if let Some(id_input_ref) = id_expr.clone().as_input_ref() {
+                    let id_index = id_input_ref.index;
+                    for &pk in &pk_indices {
+                        if id_index == pk {
+                            return Err(ErrorCode::BindError(
+                                "update modifying the PK column is banned".to_owned(),
+                            )
+                            .into());
+                        }
+                    }
+                }
+
                 let value_expr = self.bind_expr(value)?.cast_assign(id_expr.return_type())?;
 
                 match assignment_exprs.entry(id_expr) {