From 27e47b9c1f6fd0367412fddd1a0ca1dbc2b0a611 Mon Sep 17 00:00:00 2001
From: Kyle Gray <kyle@sqlc.dev>
Date: Thu, 12 Oct 2023 15:39:07 -0700
Subject: [PATCH] fix(compiler): Search SELECT and UPDATE the same way (#2841)

* fix(compiler): Search SELECT and UPDATE the same way

* Update sqlc.yaml
---
 internal/compiler/output_columns.go           |  7 ++--
 .../endtoend/testdata/join_update/issue.md    |  1 +
 .../join_update/postgresql/pgx/go/db.go       | 32 +++++++++++++++++++
 .../join_update/postgresql/pgx/go/models.go   | 24 ++++++++++++++
 .../postgresql/pgx/go/query.sql.go            | 23 +++++++++++++
 .../join_update/postgresql/pgx/query.sql      |  6 ++++
 .../join_update/postgresql/pgx/schema.sql     | 14 ++++++++
 .../join_update/postgresql/pgx/sqlc.yaml      | 10 ++++++
 8 files changed, 114 insertions(+), 3 deletions(-)
 create mode 100644 internal/endtoend/testdata/join_update/issue.md
 create mode 100644 internal/endtoend/testdata/join_update/postgresql/pgx/go/db.go
 create mode 100644 internal/endtoend/testdata/join_update/postgresql/pgx/go/models.go
 create mode 100644 internal/endtoend/testdata/join_update/postgresql/pgx/go/query.sql.go
 create mode 100644 internal/endtoend/testdata/join_update/postgresql/pgx/query.sql
 create mode 100644 internal/endtoend/testdata/join_update/postgresql/pgx/schema.sql
 create mode 100644 internal/endtoend/testdata/join_update/postgresql/pgx/sqlc.yaml

diff --git a/internal/compiler/output_columns.go b/internal/compiler/output_columns.go
index 104957e18a..8dfca0c5d3 100644
--- a/internal/compiler/output_columns.go
+++ b/internal/compiler/output_columns.go
@@ -507,9 +507,10 @@ func (c *Compiler) sourceTables(qc *QueryCatalog, node ast.Node) ([]*Table, erro
 			return ok
 		})
 	case *ast.UpdateStmt:
-		list = &ast.List{
-			Items: append(n.FromClause.Items, n.Relations.Items...),
-		}
+		var tv tableVisitor
+		astutils.Walk(&tv, n.FromClause)
+		astutils.Walk(&tv, n.Relations)
+		list = &tv.list
 	case *ast.DoStmt:
 		list = &ast.List{}
 	case *ast.CallStmt:
diff --git a/internal/endtoend/testdata/join_update/issue.md b/internal/endtoend/testdata/join_update/issue.md
new file mode 100644
index 0000000000..2b957486bc
--- /dev/null
+++ b/internal/endtoend/testdata/join_update/issue.md
@@ -0,0 +1 @@
+https://github.com/sqlc-dev/sqlc/issues/1100
diff --git a/internal/endtoend/testdata/join_update/postgresql/pgx/go/db.go b/internal/endtoend/testdata/join_update/postgresql/pgx/go/db.go
new file mode 100644
index 0000000000..8a010ccc48
--- /dev/null
+++ b/internal/endtoend/testdata/join_update/postgresql/pgx/go/db.go
@@ -0,0 +1,32 @@
+// Code generated by sqlc. DO NOT EDIT.
+// versions:
+//   sqlc v1.22.0
+
+package querytest
+
+import (
+	"context"
+
+	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgconn"
+)
+
+type DBTX interface {
+	Exec(context.Context, string, ...interface{}) (pgconn.CommandTag, error)
+	Query(context.Context, string, ...interface{}) (pgx.Rows, error)
+	QueryRow(context.Context, string, ...interface{}) pgx.Row
+}
+
+func New(db DBTX) *Queries {
+	return &Queries{db: db}
+}
+
+type Queries struct {
+	db DBTX
+}
+
+func (q *Queries) WithTx(tx pgx.Tx) *Queries {
+	return &Queries{
+		db: tx,
+	}
+}
diff --git a/internal/endtoend/testdata/join_update/postgresql/pgx/go/models.go b/internal/endtoend/testdata/join_update/postgresql/pgx/go/models.go
new file mode 100644
index 0000000000..6173ad232d
--- /dev/null
+++ b/internal/endtoend/testdata/join_update/postgresql/pgx/go/models.go
@@ -0,0 +1,24 @@
+// Code generated by sqlc. DO NOT EDIT.
+// versions:
+//   sqlc v1.22.0
+
+package querytest
+
+import (
+	"github.com/jackc/pgx/v5/pgtype"
+)
+
+type GroupCalcTotal struct {
+	Npn     pgtype.Text
+	GroupID pgtype.Text
+}
+
+type NpnExternalMap struct {
+	ID  pgtype.Text
+	Npn pgtype.Text
+}
+
+type ProducerGroupAttribute struct {
+	NpnExternalMapID pgtype.Text
+	GroupID          pgtype.Text
+}
diff --git a/internal/endtoend/testdata/join_update/postgresql/pgx/go/query.sql.go b/internal/endtoend/testdata/join_update/postgresql/pgx/go/query.sql.go
new file mode 100644
index 0000000000..700ef45f87
--- /dev/null
+++ b/internal/endtoend/testdata/join_update/postgresql/pgx/go/query.sql.go
@@ -0,0 +1,23 @@
+// Code generated by sqlc. DO NOT EDIT.
+// versions:
+//   sqlc v1.22.0
+// source: query.sql
+
+package querytest
+
+import (
+	"context"
+)
+
+const percentile = `-- name: Percentile :exec
+UPDATE group_calc_totals gct 
+SET npn = nem.npn 
+FROM producer_group_attribute ga 
+JOIN npn_external_map nem ON ga.npn_external_map_id = nem.id 
+WHERE gct.group_id = ga.group_id
+`
+
+func (q *Queries) Percentile(ctx context.Context) error {
+	_, err := q.db.Exec(ctx, percentile)
+	return err
+}
diff --git a/internal/endtoend/testdata/join_update/postgresql/pgx/query.sql b/internal/endtoend/testdata/join_update/postgresql/pgx/query.sql
new file mode 100644
index 0000000000..776517029d
--- /dev/null
+++ b/internal/endtoend/testdata/join_update/postgresql/pgx/query.sql
@@ -0,0 +1,6 @@
+-- name: Percentile :exec
+UPDATE group_calc_totals gct 
+SET npn = nem.npn 
+FROM producer_group_attribute ga 
+JOIN npn_external_map nem ON ga.npn_external_map_id = nem.id 
+WHERE gct.group_id = ga.group_id;
diff --git a/internal/endtoend/testdata/join_update/postgresql/pgx/schema.sql b/internal/endtoend/testdata/join_update/postgresql/pgx/schema.sql
new file mode 100644
index 0000000000..c91a086a71
--- /dev/null
+++ b/internal/endtoend/testdata/join_update/postgresql/pgx/schema.sql
@@ -0,0 +1,14 @@
+CREATE TABLE group_calc_totals (
+  npn  text,
+  group_id text
+);
+
+CREATE TABLE producer_group_attribute (
+  npn_external_map_id text,
+  group_id text
+);
+
+CREATE TABLE npn_external_map (
+  id   text,
+  npn  text
+);
diff --git a/internal/endtoend/testdata/join_update/postgresql/pgx/sqlc.yaml b/internal/endtoend/testdata/join_update/postgresql/pgx/sqlc.yaml
new file mode 100644
index 0000000000..5dc63e3f91
--- /dev/null
+++ b/internal/endtoend/testdata/join_update/postgresql/pgx/sqlc.yaml
@@ -0,0 +1,10 @@
+version: "2"
+sql:
+  - engine: "postgresql"
+    schema: "schema.sql"
+    queries: "query.sql"
+    gen:
+      go:
+        package: "querytest"
+        out: "go"
+        sql_package: "pgx/v5"