Repository: sqlc-dev/sqlc
Pull request #2841 (2023-10-12 21:16:38)

fix(compiler): Search SELECT and UPDATE the same way
Fixes https://github.com/sqlc-dev/sqlc/issues/1100


Patch:
diff --git a/internal/compiler/output_columns.go b/internal/compiler/output_columns.go
index 104957e18..8dfca0c5d 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:


Test patch:
diff --git a/internal/endtoend/testdata/join_update/issue.md b/internal/endtoend/testdata/join_update/issue.md
new file mode 100644
index 000000000..2b957486b
--- /dev/null
+++ b/internal/endtoend/testdata/join_update/issue.md
@@ -0,0 +1,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 000000000..8a010ccc4
--- /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 000000000..6173ad232
--- /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 000000000..700ef45f8
--- /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 000000000..776517029
--- /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 000000000..c91a086a7
--- /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 000000000..5dc63e3f9
--- /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"