feat(frontend): ban update statements modifying pk columns (risingw…

…avelabs#8569)

Signed-off-by: Clearlove <yifei.c.wei@gmail.com>

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;

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

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));
 

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'

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) {