fix(compiler): Set RV for TargetList items on updates (#1252)

* fix(compiler): Set RV for TargetList items on updates

* test(compiler): Test CTE updates

Author: andreasgangso

internal/compiler/find_params.go

@@ -107,6 +107,20 @@ func (p paramSearch) Visit(node ast.Node) astutils.Visitor {
 			}
 		}
 
+	case *ast.UpdateStmt:
+		for _, item := range n.TargetList.Items {
+			target, ok := item.(*ast.ResTarget)
+			if !ok {
+				continue
+			}
+			ref, ok := target.Val.(*ast.ParamRef)
+			if !ok {
+				continue
+			}
+			*p.refs = append(*p.refs, paramRef{parent: target, ref: ref, rv: n.Relation})
+			p.seen[ref.Location] = struct{}{}
+		}
+
 	case *ast.RangeVar:
 		p.rangeVar = n
 

internal/endtoend/testdata/update_cte/pgx/go/db.go

@@ -0,0 +1,50 @@
+-- FILE: schema.sql
+
+DROP SCHEMA IF EXISTS td3 CASCADE;
+CREATE SCHEMA td3;
+
+CREATE TABLE td3.codes (
+            id SERIAL PRIMARY KEY,
+            ts_created timestamptz DEFAULT now() NOT NULL,
+            ts_updated timestamptz DEFAULT now() NOT NULL,
+            created_by text NOT NULL,
+            updated_by text NOT NULL,
+            
+            code text,
+            hash text,
+            is_private boolean
+);
+
+
+CREATE TABLE td3.test_codes (
+            id SERIAL PRIMARY KEY,
+            ts_created timestamptz DEFAULT now() NOT NULL,
+            ts_updated timestamptz DEFAULT now() NOT NULL,
+            created_by text NOT NULL,
+            updated_by text NOT NULL,
+
+            test_id integer NOT NULL,
+            code_hash text NOT NULL
+);
+
+-- FILE: query.sql
+
+-- name: UpdateCode :one
+WITH cc AS (
+            UPDATE td3.codes
+            SET
+                created_by = $1,
+                updated_by  = $1,
+                code = $2,
+                hash = $3,
+                is_private = false
+            RETURNING hash
+)
+UPDATE td3.test_codes
+SET
+    created_by = $1,
+    updated_by  = $1,
+    test_id = $4,
+    code_hash = cc.hash
+    FROM cc
+RETURNING *;

internal/endtoend/testdata/update_cte/pgx/go/models.go

@@ -0,0 +1,13 @@
+{
+  "version": "1",
+  "packages": [
+    {
+      "path": "go",
+      "engine": "postgresql",
+      "sql_package": "pgx/v4",
+      "name": "querytest",
+      "schema": "query.sql",
+      "queries": "query.sql"
+    }
+  ]
+}