Merge 'translation: rewrite expressions and properly handle quoted identifiers in UPSERT' from Preston Thorpe

This PR fixes bugs found in the [turso-
go](https://github.com/tursodatabase/turso-go) driver with UPSERT clause
earlier, where `Gorm` will (obviously) use Expr::Variable's as well as
use quotes for `Expr::Qualified` in the tail end of an UPSERT statement.
Example:
```sql
INSERT INTO users (a,b,c) VALUES (?,?,?) ON CONFLICT (`users`.`a`) DO UPDATE SET b = `excluded`.`b`, a = ?;
```
and previously we were not properly calling `rewrite_expr`, which was
not properly setting the anonymous `Expr::Variable` to `__param_N` named
parameter, so it would ignore it completely, then return the wrong # of
parameters.
Also, we didn't handle quoted "`excluded`.`x`", so it would panic in the
optimizer that Qualified should have been rewritten earlier.

Closes #3157

core/translate/insert.rs

@@ -118,45 +118,54 @@ pub fn translate_insert(
     let mut values: Option<Vec<Box<Expr>>> = None;
     let mut upsert_opt: Option<Upsert> = None;
 
-    let inserting_multiple_rows = match &mut body {
-        InsertBody::Select(select, upsert) => {
-            upsert_opt = upsert.as_deref().cloned();
-            match &mut select.body.select {
-                // TODO see how to avoid clone
-                OneSelect::Values(values_expr) if values_expr.len() <= 1 => {
-                    if values_expr.is_empty() {
-                        crate::bail_parse_error!("no values to insert");
-                    }
-                    let mut param_idx = 1;
-                    for expr in values_expr.iter_mut().flat_map(|v| v.iter_mut()) {
-                        match expr.as_mut() {
-                            Expr::Id(name) => {
-                                if name.is_double_quoted() {
-                                    *expr = Expr::Literal(ast::Literal::String(name.to_string()))
-                                        .into();
-                                } else {
-                                    // an INSERT INTO ... VALUES (...) cannot reference columns
-                                    crate::bail_parse_error!("no such column: {name}");
-                                }
-                            }
-                            Expr::Qualified(first_name, second_name) => {
-                                // an INSERT INTO ... VALUES (...) cannot reference columns
-                                crate::bail_parse_error!(
-                                    "no such column: {first_name}.{second_name}"
-                                );
-                            }
-                            _ => {}
-                        }
-                        rewrite_expr(expr, &mut param_idx)?;
-                    }
-                    values = values_expr.pop();
-                    false
+    let mut param_idx = 1;
+    let mut inserting_multiple_rows = false;
+    if let InsertBody::Select(select, upsert) = &mut body {
+        match &mut select.body.select {
+            // TODO see how to avoid clone
+            OneSelect::Values(values_expr) if values_expr.len() <= 1 => {
+                if values_expr.is_empty() {
+                    crate::bail_parse_error!("no values to insert");
+                }
+                for expr in values_expr.iter_mut().flat_map(|v| v.iter_mut()) {
+                    match expr.as_mut() {
+                        Expr::Id(name) => {
+                            if name.is_double_quoted() {
+                                *expr =
+                                    Expr::Literal(ast::Literal::String(name.to_string())).into();
+                            } else {
+                                // an INSERT INTO ... VALUES (...) cannot reference columns
+                                crate::bail_parse_error!("no such column: {name}");
+                            }
+                        }
+                        Expr::Qualified(first_name, second_name) => {
+                            // an INSERT INTO ... VALUES (...) cannot reference columns
+                            crate::bail_parse_error!("no such column: {first_name}.{second_name}");
+                        }
+                        _ => {}
+                    }
+                    rewrite_expr(expr, &mut param_idx)?;
+                }
+                values = values_expr.pop();
+            }
+            _ => inserting_multiple_rows = true,
+        }
+        if let Some(ref mut upsert) = upsert {
+            if let UpsertDo::Set {
+                ref mut sets,
+                ref mut where_clause,
+            } = &mut upsert.do_clause
+            {
+                for set in sets.iter_mut() {
+                    rewrite_expr(set.expr.as_mut(), &mut param_idx)?;
+                }
+                if let Some(ref mut where_expr) = where_clause {
+                    rewrite_expr(where_expr.as_mut(), &mut param_idx)?;
                 }
-                _ => true,
             }
         }
-        InsertBody::DefaultValues => false,
-    };
+        upsert_opt = upsert.as_deref().cloned();
+    }
 
     let halt_label = program.allocate_label();
     let loop_start_label = program.allocate_label();

core/translate/upsert.rs

@@ -39,28 +39,26 @@ pub struct ConflictTarget {
 // Extract `(column, optional_collate)` from an ON CONFLICT target Expr.
 // Accepts: Id, Qualified, DoublyQualified, Parenthesized, Collate
 fn extract_target_key(e: &ast::Expr) -> Option<ConflictTarget> {
+    use ast::Name::{Ident, Quoted};
     match e {
-        // expr COLLATE c: carry c and keep descending into expr
         ast::Expr::Collate(inner, c) => {
             let mut tk = extract_target_key(inner.as_ref())?;
             let cstr = match c {
-                ast::Name::Ident(s) => s.as_str(),
-                _ => return None,
+                Ident(s) | Quoted(s) => s.as_str(),
             };
             tk.collate = Some(cstr.to_ascii_lowercase());
             Some(tk)
         }
         ast::Expr::Parenthesized(v) if v.len() == 1 => extract_target_key(&v[0]),
-        // Bare identifier
-        ast::Expr::Id(ast::Name::Ident(name)) => Some(ConflictTarget {
+
+        ast::Expr::Id(Ident(name)) | ast::Expr::Id(Quoted(name)) => Some(ConflictTarget {
             col_name: normalize_ident(name),
             collate: None,
         }),
-        // t.a or db.t.a
+        // t.a or db.t.a: accept ident or quoted in the column position
         ast::Expr::Qualified(_, col) | ast::Expr::DoublyQualified(_, _, col) => {
             let cname = match col {
-                ast::Name::Ident(s) => s.as_str(),
-                _ => return None,
+                Ident(s) | Quoted(s) => s.as_str(),
             };
             Some(ConflictTarget {
                 col_name: normalize_ident(cname),
@@ -526,37 +524,38 @@ fn rewrite_upsert_expr_in_place(
     conflict_rowid_reg: usize,
     insertion: &Insertion,
 ) -> crate::Result<WalkControl> {
-    use ast::Expr;
+    use ast::{Expr, Name};
 
-    // helper: return the CURRENT-row register for a column (including rowid alias)
     let col_reg = |name: &str| -> Option<usize> {
         if name.eq_ignore_ascii_case("rowid") {
             return Some(conflict_rowid_reg);
         }
-        let (idx, _c) = table.get_column_by_name(&normalize_ident(name))?;
+        let (idx, _) = table.get_column_by_name(name)?;
         Some(current_start + idx)
     };
+
     walk_expr_mut(
         e,
         &mut |expr: &mut ast::Expr| -> crate::Result<WalkControl> {
             match expr {
-                // EXCLUDED.x -> insertion register
-                Expr::Qualified(ns, ast::Name::Ident(c)) => {
-                    if ns.as_str().eq_ignore_ascii_case("excluded") {
-                        let Some(reg) = insertion.get_col_mapping_by_name(c.as_str()) else {
+                // EXCLUDED.x or t.x (t may be quoted)
+                Expr::Qualified(ns, Name::Ident(c) | Name::Quoted(c))
+                | Expr::DoublyQualified(_, ns, Name::Ident(c) | Name::Quoted(c)) => {
+                    let ns = normalize_ident(ns.as_str());
+                    let c = normalize_ident(c.as_str());
+                    if ns.eq_ignore_ascii_case("excluded") {
+                        let Some(reg) = insertion.get_col_mapping_by_name(&c) else {
                             bail_parse_error!("no such column in EXCLUDED: {}", c);
                         };
                         *expr = Expr::Register(reg.register);
-                    } else if ns.as_str().eq_ignore_ascii_case(table_name)
-                    // t.x -> CURRENT, only if t matches the target table name (never "excluded")
-                    {
+                    } else if ns.eq_ignore_ascii_case(table_name) {
                         if let Some(reg) = col_reg(c.as_str()) {
                             *expr = Expr::Register(reg);
                         }
                     }
                 }
                 // Unqualified column id -> CURRENT
-                Expr::Id(ast::Name::Ident(name)) => {
+                Expr::Id(Name::Ident(name)) | Expr::Id(Name::Quoted(name)) => {
                     if let Some(reg) = col_reg(name.as_str()) {
                         *expr = Expr::Register(reg);
                     }

testing/upsert.test

@@ -305,7 +305,30 @@ do_execsql_test_on_specific_db {:memory:} upsert-omitted-target-updates-unique {
     SELECT * FROM ou;
 } {3|y}
 
-# Target qualified with database.table.column should match too (assuming single db).
+do_execsql_test_on_specific_db {:memory:} upsert-current-qualified.1 {
+  CREATE TABLE dq (a INTEGER UNIQUE, b TEXT);
+  INSERT INTO dq VALUES (1,'old');
+  INSERT INTO dq VALUES (1,'new')
+    ON CONFLICT(dq.a) DO UPDATE SET b = dq.b || '-' || excluded.b;
+  SELECT * FROM dq;
+} {1|old-new}
+
+do_execsql_test_on_specific_db {:memory:} upsert-multicol-set.1 {
+  CREATE TABLE dq (a INTEGER UNIQUE, b TEXT);
+  INSERT INTO dq VALUES (1,'old');
+  INSERT INTO dq VALUES (1,'new')
+    ON CONFLICT(a) DO UPDATE SET (`a`,`b`) = (`excluded`.`a`, `excluded`.`b`);
+  SELECT * FROM dq;
+} {1|new}
+
+do_execsql_test_on_specific_db {:memory:} upsert-where-predicate.1 {
+  CREATE TABLE dq (a INTEGER UNIQUE, b TEXT);
+  INSERT INTO dq VALUES (1,'old');
+  INSERT INTO dq VALUES (1,'old')
+    ON CONFLICT(a) DO UPDATE SET b = excluded.b WHERE dq.b != excluded.b;
+  SELECT * FROM dq;
+} {1|old}
+
 do_execsql_test_on_specific_db {:memory:} upsert-doubly-qualified-target {
     CREATE TABLE dq (a UNIQUE, b);
     INSERT INTO dq VALUES (1,'old');