Skip to content

Commit 600ff8b

Browse files
mrmekumrmeku
committed
Expose target table for UPDATE and DELETE queries in plugin proto
Add update_table and delete_from_table fields to the Query message in the plugin codegen proto, matching the existing insert_into_table field. The compiler already resolves the target table for UPDATE and DELETE queries internally; this change wires that information through to plugins. Closes #4317
1 parent ce83d3f commit 600ff8b

File tree

9 files changed

+260
-72
lines changed

9 files changed

+260
-72
lines changed

‎internal/cmd/shim.go‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -152,6 +152,22 @@ func pluginQueries(r *compiler.Result) []*plugin.Query {
152152
Name: q.InsertIntoTable.Name,
153153
}
154154
}
155+
var ut *plugin.Identifier
156+
if q.UpdateTable != nil {
157+
ut = &plugin.Identifier{
158+
Catalog: q.UpdateTable.Catalog,
159+
Schema: q.UpdateTable.Schema,
160+
Name: q.UpdateTable.Name,
161+
}
162+
}
163+
var dft *plugin.Identifier
164+
if q.DeleteFromTable != nil {
165+
dft = &plugin.Identifier{
166+
Catalog: q.DeleteFromTable.Catalog,
167+
Schema: q.DeleteFromTable.Schema,
168+
Name: q.DeleteFromTable.Name,
169+
}
170+
}
155171
out = append(out, &plugin.Query{
156172
Name: q.Metadata.Name,
157173
Cmd: q.Metadata.Cmd,
@@ -161,6 +177,8 @@ func pluginQueries(r *compiler.Result) []*plugin.Query {
161177
Params: params,
162178
Filename: q.Metadata.Filename,
163179
InsertIntoTable: iit,
180+
UpdateTable: ut,
181+
DeleteFromTable: dft,
164182
})
165183
}
166184
return out

‎internal/compiler/analyze.go‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -152,6 +152,18 @@ func (c *Compiler) _analyzeQuery(raw *ast.RawStmt, query string, failfast bool)
152152
if err := check(err); err != nil {
153153
return nil, err
154154
}
155+
case *ast.UpdateStmt:
156+
if n.Relations != nil && len(n.Relations.Items) > 0 {
157+
if rv, ok := n.Relations.Items[0].(*ast.RangeVar); ok {
158+
table, _ = ParseTableName(rv)
159+
}
160+
}
161+
case *ast.DeleteStmt:
162+
if n.Relations != nil && len(n.Relations.Items) > 0 {
163+
if rv, ok := n.Relations.Items[0].(*ast.RangeVar); ok {
164+
table, _ = ParseTableName(rv)
165+
}
166+
}
155167
}
156168

157169
if err := check(validate.FuncCall(c.catalog, c.combo, raw)); err != nil {

‎internal/compiler/parse.go‎

Lines changed: 32 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -108,10 +108,23 @@ func (c *Compiler) parseQuery(stmt ast.Node, src string, o opts.Parser) (*Query,
108108
})
109109
}
110110

111-
// Determine the insert table if applicable
111+
// Determine the target table if applicable
112112
var table *ast.TableName
113-
if insert, ok := expandedRaw.Stmt.(*ast.InsertStmt); ok {
114-
table, _ = ParseTableName(insert.Relation)
113+
switch n := expandedRaw.Stmt.(type) {
114+
case *ast.InsertStmt:
115+
table, _ = ParseTableName(n.Relation)
116+
case *ast.UpdateStmt:
117+
if n.Relations != nil && len(n.Relations.Items) > 0 {
118+
if rv, ok := n.Relations.Items[0].(*ast.RangeVar); ok {
119+
table, _ = ParseTableName(rv)
120+
}
121+
}
122+
case *ast.DeleteStmt:
123+
if n.Relations != nil && len(n.Relations.Items) > 0 {
124+
if rv, ok := n.Relations.Items[0].(*ast.RangeVar); ok {
125+
table, _ = ParseTableName(rv)
126+
}
127+
}
115128
}
116129

117130
anlys = &analysis{
@@ -171,14 +184,22 @@ func (c *Compiler) parseQuery(stmt ast.Node, src string, o opts.Parser) (*Query,
171184

172185
md.Comments = comments
173186

174-
return &Query{
175-
RawStmt: raw,
176-
Metadata: md,
177-
Params: anlys.Parameters,
178-
Columns: anlys.Columns,
179-
SQL: trimmed,
180-
InsertIntoTable: anlys.Table,
181-
}, nil
187+
q := &Query{
188+
RawStmt: raw,
189+
Metadata: md,
190+
Params: anlys.Parameters,
191+
Columns: anlys.Columns,
192+
SQL: trimmed,
193+
}
194+
switch raw.Stmt.(type) {
195+
case *ast.InsertStmt:
196+
q.InsertIntoTable = anlys.Table
197+
case *ast.UpdateStmt:
198+
q.UpdateTable = anlys.Table
199+
case *ast.DeleteStmt:
200+
q.DeleteFromTable = anlys.Table
201+
}
202+
return q, nil
182203
}
183204

184205
func rangeVars(root ast.Node) []*ast.RangeVar {

‎internal/compiler/query.go‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,12 @@ type Query struct {
5151
// Needed for CopyFrom
5252
InsertIntoTable *ast.TableName
5353

54+
// Target table for UPDATE queries
55+
UpdateTable *ast.TableName
56+
57+
// Target table for DELETE queries
58+
DeleteFromTable *ast.TableName
59+
5460
// Needed for vet
5561
RawStmt *ast.RawStmt
5662
}

‎internal/endtoend/testdata/codegen_json/gen/codegen.json‎

Lines changed: 91 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -65079,7 +65079,9 @@
6507965079
],
6508065080
"comments": [],
6508165081
"filename": "query.sql",
65082-
"insert_into_table": null
65082+
"insert_into_table": null,
65083+
"update_table": null,
65084+
"delete_from_table": null
6508365085
},
6508465086
{
6508565087
"text": "SELECT id, name, bio FROM authors\nORDER BY name",
@@ -65168,7 +65170,9 @@
6516865170
"params": [],
6516965171
"comments": [],
6517065172
"filename": "query.sql",
65171-
"insert_into_table": null
65173+
"insert_into_table": null,
65174+
"update_table": null,
65175+
"delete_from_table": null
6517265176
},
6517365177
{
6517465178
"text": "INSERT INTO authors (\n name, bio\n) VALUES (\n $1, $2\n)\nRETURNING id, name, bio",
@@ -65320,7 +65324,84 @@
6532065324
"catalog": "",
6532165325
"schema": "",
6532265326
"name": "authors"
65323-
}
65327+
},
65328+
"update_table": null,
65329+
"delete_from_table": null
65330+
},
65331+
{
65332+
"text": "UPDATE authors SET bio = $1\nWHERE id = $2",
65333+
"name": "UpdateAuthorBio",
65334+
"cmd": ":exec",
65335+
"columns": [],
65336+
"params": [
65337+
{
65338+
"number": 1,
65339+
"column": {
65340+
"name": "bio",
65341+
"not_null": false,
65342+
"is_array": false,
65343+
"comment": "",
65344+
"length": -1,
65345+
"is_named_param": false,
65346+
"is_func_call": false,
65347+
"scope": "",
65348+
"table": {
65349+
"catalog": "",
65350+
"schema": "public",
65351+
"name": "authors"
65352+
},
65353+
"table_alias": "",
65354+
"type": {
65355+
"catalog": "",
65356+
"schema": "",
65357+
"name": "text"
65358+
},
65359+
"is_sqlc_slice": false,
65360+
"embed_table": null,
65361+
"original_name": "bio",
65362+
"unsigned": false,
65363+
"array_dims": 0
65364+
}
65365+
},
65366+
{
65367+
"number": 2,
65368+
"column": {
65369+
"name": "id",
65370+
"not_null": true,
65371+
"is_array": false,
65372+
"comment": "",
65373+
"length": -1,
65374+
"is_named_param": false,
65375+
"is_func_call": false,
65376+
"scope": "",
65377+
"table": {
65378+
"catalog": "",
65379+
"schema": "",
65380+
"name": "authors"
65381+
},
65382+
"table_alias": "",
65383+
"type": {
65384+
"catalog": "",
65385+
"schema": "",
65386+
"name": "bigserial"
65387+
},
65388+
"is_sqlc_slice": false,
65389+
"embed_table": null,
65390+
"original_name": "id",
65391+
"unsigned": false,
65392+
"array_dims": 0
65393+
}
65394+
}
65395+
],
65396+
"comments": [],
65397+
"filename": "query.sql",
65398+
"insert_into_table": null,
65399+
"update_table": {
65400+
"catalog": "",
65401+
"schema": "",
65402+
"name": "authors"
65403+
},
65404+
"delete_from_table": null
6532465405
},
6532565406
{
6532665407
"text": "DELETE FROM authors\nWHERE id = $1",
@@ -65360,7 +65441,13 @@
6536065441
],
6536165442
"comments": [],
6536265443
"filename": "query.sql",
65363-
"insert_into_table": null
65444+
"insert_into_table": null,
65445+
"update_table": null,
65446+
"delete_from_table": {
65447+
"catalog": "",
65448+
"schema": "",
65449+
"name": "authors"
65450+
}
6536465451
}
6536565452
],
6536665453
"sqlc_version": "v1.30.0",

‎internal/endtoend/testdata/codegen_json/postgresql/query.sql‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,10 @@ INSERT INTO authors (
1414
)
1515
RETURNING *;
1616

17+
-- name: UpdateAuthorBio :exec
18+
UPDATE authors SET bio = $1
19+
WHERE id = $2;
20+
1721
-- name: DeleteAuthor :exec
1822
DELETE FROM authors
1923
WHERE id = $1;

‎internal/endtoend/testdata/process_plugin_sqlc_gen_json/gen/codegen.json‎

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -65081,7 +65081,9 @@
6508165081
],
6508265082
"comments": [],
6508365083
"filename": "query.sql",
65084-
"insert_into_table": null
65084+
"insert_into_table": null,
65085+
"update_table": null,
65086+
"delete_from_table": null
6508565087
},
6508665088
{
6508765089
"text": "SELECT id, name, bio FROM authors\nORDER BY name",
@@ -65170,7 +65172,9 @@
6517065172
"params": [],
6517165173
"comments": [],
6517265174
"filename": "query.sql",
65173-
"insert_into_table": null
65175+
"insert_into_table": null,
65176+
"update_table": null,
65177+
"delete_from_table": null
6517465178
},
6517565179
{
6517665180
"text": "INSERT INTO authors (\n name, bio\n) VALUES (\n $1, $2\n)\nRETURNING id, name, bio",
@@ -65322,7 +65326,9 @@
6532265326
"catalog": "",
6532365327
"schema": "",
6532465328
"name": "authors"
65325-
}
65329+
},
65330+
"update_table": null,
65331+
"delete_from_table": null
6532665332
},
6532765333
{
6532865334
"text": "DELETE FROM authors\nWHERE id = $1",
@@ -65362,7 +65368,13 @@
6536265368
],
6536365369
"comments": [],
6536465370
"filename": "query.sql",
65365-
"insert_into_table": null
65371+
"insert_into_table": null,
65372+
"update_table": null,
65373+
"delete_from_table": {
65374+
"catalog": "",
65375+
"schema": "",
65376+
"name": "authors"
65377+
}
6536665378
}
6536765379
],
6536865380
"sqlc_version": "v1.30.0",

0 commit comments

Comments
 (0)