fix(compiler): Support references to columns in joined tables in UPDATE statements (#1289)

* fix(compiler): Support references to columns in joined tables in UPDATE statements

Author: timstudd

internal/compiler/find_params.go

@@ -117,7 +117,13 @@ func (p paramSearch) Visit(node ast.Node) astutils.Visitor {
 			if !ok {
 				continue
 			}
-			*p.refs = append(*p.refs, paramRef{parent: target, ref: ref, rv: n.Relation})
+			for _, relation := range n.Relations.Items {
+				rv, ok := relation.(*ast.RangeVar)
+				if !ok {
+					continue
+				}
+				*p.refs = append(*p.refs, paramRef{parent: target, ref: ref, rv: rv})
+			}
 			p.seen[ref.Location] = struct{}{}
 		}
 

internal/compiler/output_columns.go

@@ -349,7 +349,7 @@ func sourceTables(qc *QueryCatalog, node ast.Node) ([]*Table, error) {
 		})
 	case *ast.UpdateStmt:
 		list = &ast.List{
-			Items: append(n.FromClause.Items, n.Relation),
+			Items: append(n.FromClause.Items, n.Relations.Items...),
 		}
 	default:
 		return nil, fmt.Errorf("sourceTables: unsupported node type: %T", n)

internal/endtoend/testdata/update_join/mysql/db/db.go

@@ -0,0 +1,37 @@
+CREATE TABLE primary_table (
+  id bigint(20) unsigned NOT NULL AUTO_INCREMENT,
+  user_id bigint(20) unsigned NOT NULL,
+  PRIMARY KEY (id)
+);
+
+CREATE TABLE join_table (
+  id bigint(20) unsigned NOT NULL AUTO_INCREMENT,
+  primary_table_id bigint(20) unsigned NOT NULL,
+  other_table_id bigint(20) unsigned NOT NULL,
+  is_active tinyint(1) NOT NULL DEFAULT '0',
+  PRIMARY KEY (id)
+);
+
+-- name: UpdateJoin :exec
+UPDATE  join_table as jt
+        JOIN primary_table as pt
+            ON jt.primary_table_id = pt.id
+SET     jt.is_active = ?
+WHERE   jt.id = ?
+        AND pt.user_id = ?;
+
+-- name: UpdateLeftJoin :exec
+UPDATE  join_table as jt
+        LEFT JOIN primary_table as pt
+            ON jt.primary_table_id = pt.id
+SET     jt.is_active = ?
+WHERE   jt.id = ?
+        AND pt.user_id = ?;
+
+-- name: UpdateRightJoin :exec
+UPDATE  join_table as jt
+        RIGHT JOIN primary_table as pt
+            ON jt.primary_table_id = pt.id
+SET     jt.is_active = ?
+WHERE   jt.id = ?
+        AND pt.user_id = ?;

internal/endtoend/testdata/update_join/mysql/db/models.go

@@ -0,0 +1,11 @@
+{
+	"version": "1",
+	"packages": [
+		{
+			"path": "db",
+			"engine": "mysql",
+			"schema": "query.sql",
+			"queries": "query.sql"
+		}
+	]
+}