fix(compiler): Support references to columns in joined tables in UPDA…

…TE statements (#1289)

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

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/query.sql

@@ -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/sqlc.json

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

internal/endtoend/testdata/update_join/postgresql/query.sql

@@ -0,0 +1,19 @@
+CREATE TABLE primary_table (
+  id         INT PRIMARY KEY,
+  user_id    INT NOT NULL
+);
+
+CREATE TABLE join_table (
+  id                INT PRIMARY KEY,
+  primary_table_id  INT NOT NULL,
+  other_table_id    INT NOT NULL,
+  is_active         BOOLEAN NOT NULL
+);
+
+-- name: UpdateJoin :exec
+UPDATE  join_table
+SET     is_active = $1
+FROM    primary_table
+WHERE   join_table.id = $2
+        AND primary_table.user_id = $3
+        AND join_table.primary_table_id = primary_table.id;

internal/endtoend/testdata/update_join/postgresql/sqlc.json

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

internal/engine/dolphin/convert.go

@@ -540,7 +540,7 @@ func (c *cc) convertUpdateStmt(n *pcast.UpdateStmt) *ast.UpdateStmt {
 		panic("expected one range var")
 	}
 
-	var rangeVar *ast.RangeVar
+	relations := &ast.List{}
 	switch rel := rels.Items[0].(type) {
 
 	// Special case for joins in updates
@@ -549,10 +549,16 @@ func (c *cc) convertUpdateStmt(n *pcast.UpdateStmt) *ast.UpdateStmt {
 		if !ok {
 			panic("expected range var")
 		}
-		rangeVar = left
+		relations.Items = append(relations.Items, left)
+
+		right, ok := rel.Rarg.(*ast.RangeVar)
+		if !ok {
+			panic("expected range var")
+		}
+		relations.Items = append(relations.Items, right)
 
 	case *ast.RangeVar:
-		rangeVar = rel
+		relations.Items = append(relations.Items, rel)
 
 	default:
 		panic("expected range var")
@@ -564,7 +570,7 @@ func (c *cc) convertUpdateStmt(n *pcast.UpdateStmt) *ast.UpdateStmt {
 		list.Items = append(list.Items, c.convertAssignment(a))
 	}
 	return &ast.UpdateStmt{
-		Relation:      rangeVar,
+		Relations:     relations,
 		TargetList:    list,
 		WhereClause:   c.convert(n.Where),
 		FromClause:    &ast.List{},

internal/engine/postgresql/convert.go

@@ -1,3 +1,4 @@
+//go:build !windows
 // +build !windows
 
 package postgresql
@@ -2765,8 +2766,11 @@ func convertUpdateStmt(n *pg.UpdateStmt) *ast.UpdateStmt {
 	if n == nil {
 		return nil
 	}
+
 	return &ast.UpdateStmt{
-		Relation:      convertRangeVar(n.Relation),
+		Relations: &ast.List{
+			Items: []ast.Node{convertRangeVar(n.Relation)},
+		},
 		TargetList:    convertSlice(n.TargetList),
 		WhereClause:   convertNode(n.WhereClause),
 		FromClause:    convertSlice(n.FromClause),
@@ -2780,10 +2784,10 @@ func convertVacuumStmt(n *pg.VacuumStmt) *ast.VacuumStmt {
 		return nil
 	}
 	return &ast.VacuumStmt{
-	// FIXME: The VacuumStmt node has changed quite a bit
-	// Options:  n.Options
-	// Relation: convertRangeVar(n.Relation),
-	// VaCols:   convertSlice(n.VaCols),
+		// FIXME: The VacuumStmt node has changed quite a bit
+		// Options:  n.Options
+		// Relation: convertRangeVar(n.Relation),
+		// VaCols:   convertSlice(n.VaCols),
 	}
 }
 

internal/sql/ast/update_stmt.go

@@ -1,7 +1,7 @@
 package ast
 
 type UpdateStmt struct {
-	Relation      *RangeVar
+	Relations     *List
 	TargetList    *List
 	WhereClause   Node
 	FromClause    *List

internal/sql/astutils/rewrite.go

@@ -1141,7 +1141,7 @@ func (a *application) apply(parent ast.Node, name string, iter *iterator, n ast.
 		// pass
 
 	case *ast.UpdateStmt:
-		a.apply(n, "Relation", nil, n.Relation)
+		a.apply(n, "Relations", nil, n.Relations)
 		a.apply(n, "TargetList", nil, n.TargetList)
 		a.apply(n, "WhereClause", nil, n.WhereClause)
 		a.apply(n, "FromClause", nil, n.FromClause)

internal/sql/astutils/walk.go

@@ -2008,8 +2008,8 @@ func Walk(f Visitor, node ast.Node) {
 		// pass
 
 	case *ast.UpdateStmt:
-		if n.Relation != nil {
-			Walk(f, n.Relation)
+		if n.Relations != nil {
+			Walk(f, n.Relations)
 		}
 		if n.TargetList != nil {
 			Walk(f, n.TargetList)