Expose target table and parameter clause context for plugin queries

internal/cmd/shim.go

@@ -152,6 +152,22 @@ func pluginQueries(r *compiler.Result) []*plugin.Query {
 				Name:    q.InsertIntoTable.Name,
 			}
 		}
+		var ut *plugin.Identifier
+		if q.UpdateTable != nil {
+			ut = &plugin.Identifier{
+				Catalog: q.UpdateTable.Catalog,
+				Schema:  q.UpdateTable.Schema,
+				Name:    q.UpdateTable.Name,
+			}
+		}
+		var dft *plugin.Identifier
+		if q.DeleteFromTable != nil {
+			dft = &plugin.Identifier{
+				Catalog: q.DeleteFromTable.Catalog,
+				Schema:  q.DeleteFromTable.Schema,
+				Name:    q.DeleteFromTable.Name,
+			}
+		}
 		out = append(out, &plugin.Query{
 			Name:            q.Metadata.Name,
 			Cmd:             q.Metadata.Cmd,
@@ -161,6 +177,8 @@ func pluginQueries(r *compiler.Result) []*plugin.Query {
 			Params:          params,
 			Filename:        q.Metadata.Filename,
 			InsertIntoTable: iit,
+			UpdateTable:     ut,
+			DeleteFromTable: dft,
 		})
 	}
 	return out
@@ -218,8 +236,9 @@ func pluginQueryColumn(c *compiler.Column) *plugin.Column {
 
 func pluginQueryParam(p compiler.Parameter) *plugin.Parameter {
 	return &plugin.Parameter{
-		Number: int32(p.Number),
-		Column: pluginQueryColumn(p.Column),
+		Number:  int32(p.Number),
+		Column:  pluginQueryColumn(p.Column),
+		Context: plugin.ParameterContext(p.Context),
 	}
 }
 

internal/compiler/analyze.go

@@ -152,6 +152,18 @@ func (c *Compiler) _analyzeQuery(raw *ast.RawStmt, query string, failfast bool)
 		if err := check(err); err != nil {
 			return nil, err
 		}
+	case *ast.UpdateStmt:
+		if n.Relations != nil && len(n.Relations.Items) > 0 {
+			if rv, ok := n.Relations.Items[0].(*ast.RangeVar); ok {
+				table, _ = ParseTableName(rv)
+			}
+		}
+	case *ast.DeleteStmt:
+		if n.Relations != nil && len(n.Relations.Items) > 0 {
+			if rv, ok := n.Relations.Items[0].(*ast.RangeVar); ok {
+				table, _ = ParseTableName(rv)
+			}
+		}
 	}
 
 	if err := check(validate.FuncCall(c.catalog, c.combo, raw)); err != nil {

internal/compiler/parse.go

@@ -108,10 +108,23 @@ func (c *Compiler) parseQuery(stmt ast.Node, src string, o opts.Parser) (*Query,
 			})
 		}
 
-		// Determine the insert table if applicable
+		// Determine the target table if applicable
 		var table *ast.TableName
-		if insert, ok := expandedRaw.Stmt.(*ast.InsertStmt); ok {
-			table, _ = ParseTableName(insert.Relation)
+		switch n := expandedRaw.Stmt.(type) {
+		case *ast.InsertStmt:
+			table, _ = ParseTableName(n.Relation)
+		case *ast.UpdateStmt:
+			if n.Relations != nil && len(n.Relations.Items) > 0 {
+				if rv, ok := n.Relations.Items[0].(*ast.RangeVar); ok {
+					table, _ = ParseTableName(rv)
+				}
+			}
+		case *ast.DeleteStmt:
+			if n.Relations != nil && len(n.Relations.Items) > 0 {
+				if rv, ok := n.Relations.Items[0].(*ast.RangeVar); ok {
+					table, _ = ParseTableName(rv)
+				}
+			}
 		}
 
 		anlys = &analysis{
@@ -171,14 +184,22 @@ func (c *Compiler) parseQuery(stmt ast.Node, src string, o opts.Parser) (*Query,
 
 	md.Comments = comments
 
-	return &Query{
-		RawStmt:         raw,
-		Metadata:        md,
-		Params:          anlys.Parameters,
-		Columns:         anlys.Columns,
-		SQL:             trimmed,
-		InsertIntoTable: anlys.Table,
-	}, nil
+	q := &Query{
+		RawStmt:  raw,
+		Metadata: md,
+		Params:   anlys.Parameters,
+		Columns:  anlys.Columns,
+		SQL:      trimmed,
+	}
+	switch raw.Stmt.(type) {
+	case *ast.InsertStmt:
+		q.InsertIntoTable = anlys.Table
+	case *ast.UpdateStmt:
+		q.UpdateTable = anlys.Table
+	case *ast.DeleteStmt:
+		q.DeleteFromTable = anlys.Table
+	}
+	return q, nil
 }
 
 func rangeVars(root ast.Node) []*ast.RangeVar {

internal/compiler/query.go

@@ -51,11 +51,31 @@ type Query struct {
 	// Needed for CopyFrom
 	InsertIntoTable *ast.TableName
 
+	// Target table for UPDATE queries
+	UpdateTable *ast.TableName
+
+	// Target table for DELETE queries
+	DeleteFromTable *ast.TableName
+
 	// Needed for vet
 	RawStmt *ast.RawStmt
 }
 
+type ParameterContext int
+
+const (
+	ParameterContextUnspecified ParameterContext = iota
+	ParameterContextSet
+	ParameterContextValues
+	ParameterContextWhere
+	ParameterContextHaving
+	ParameterContextFunctionArg
+	ParameterContextLimit
+	ParameterContextOffset
+)
+
 type Parameter struct {
-	Number int
-	Column *Column
+	Number  int
+	Column  *Column
+	Context ParameterContext
 }

internal/compiler/resolve.go

@@ -126,6 +126,7 @@ func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar,
 					NotNull:      p.NotNull(),
 					IsNamedParam: isNamed,
 				},
+				Context: ParameterContextOffset,
 			})
 
 		case *limitCount:
@@ -139,6 +140,7 @@ func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar,
 					NotNull:      p.NotNull(),
 					IsNamedParam: isNamed,
 				},
+				Context: ParameterContextLimit,
 			})
 
 		case *ast.A_Expr:
@@ -173,6 +175,7 @@ func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar,
 						NotNull:      p.NotNull(),
 						IsSqlcSlice:  p.IsSqlcSlice(),
 					},
+					Context: ParameterContextWhere,
 				})
 				continue
 			}
@@ -246,6 +249,7 @@ func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar,
 								IsNamedParam: isNamed,
 								IsSqlcSlice:  p.IsSqlcSlice(),
 							},
+							Context: ParameterContextWhere,
 						})
 					}
 				}
@@ -311,6 +315,7 @@ func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar,
 							IsNamedParam: isNamed,
 							IsSqlcSlice:  p.IsSqlcSlice(),
 						},
+						Context: ParameterContextWhere,
 					})
 				}
 			}
@@ -383,6 +388,7 @@ func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar,
 							NotNull:      p.NotNull(),
 							IsSqlcSlice:  p.IsSqlcSlice(),
 						},
+						Context: ParameterContextFunctionArg,
 					})
 					continue
 				}
@@ -425,6 +431,7 @@ func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar,
 						IsNamedParam: isNamed,
 						IsSqlcSlice:  p.IsSqlcSlice(),
 					},
+					Context: ParameterContextFunctionArg,
 				})
 			}
 
@@ -498,6 +505,7 @@ func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar,
 						IsNamedParam: isNamed,
 						IsSqlcSlice:  p.IsSqlcSlice(),
 					},
+					Context: ParameterContextSet,
 				})
 			} else {
 				return nil, &sqlerr.Error{
@@ -608,6 +616,7 @@ func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar,
 								IsNamedParam: isNamed,
 								IsSqlcSlice:  p.IsSqlcSlice(),
 							},
+							Context: ParameterContextWhere,
 						})
 					}
 				}

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

@@ -65074,12 +65074,15 @@
             "original_name": "id",
             "unsigned": false,
             "array_dims": 0
-          }
+          },
+          "context": "PARAMETER_CONTEXT_WHERE"
         }
       ],
       "comments": [],
       "filename": "query.sql",
-      "insert_into_table": null
+      "insert_into_table": null,
+      "update_table": null,
+      "delete_from_table": null
     },
     {
       "text": "SELECT id, name, bio FROM authors\nORDER BY name",
@@ -65168,7 +65171,9 @@
       "params": [],
       "comments": [],
       "filename": "query.sql",
-      "insert_into_table": null
+      "insert_into_table": null,
+      "update_table": null,
+      "delete_from_table": null
     },
     {
       "text": "INSERT INTO authors (\n          name, bio\n) VALUES (\n  $1, $2\n)\nRETURNING id, name, bio",
@@ -65282,7 +65287,8 @@
             "original_name": "name",
             "unsigned": false,
             "array_dims": 0
-          }
+          },
+          "context": "PARAMETER_CONTEXT_SET"
         },
         {
           "number": 2,
@@ -65311,7 +65317,8 @@
             "original_name": "bio",
             "unsigned": false,
             "array_dims": 0
-          }
+          },
+          "context": "PARAMETER_CONTEXT_SET"
         }
       ],
       "comments": [],
@@ -65320,7 +65327,86 @@
         "catalog": "",
         "schema": "",
         "name": "authors"
-      }
+      },
+      "update_table": null,
+      "delete_from_table": null
+    },
+    {
+      "text": "UPDATE authors SET bio = $1\nWHERE id = $2",
+      "name": "UpdateAuthorBio",
+      "cmd": ":exec",
+      "columns": [],
+      "params": [
+        {
+          "number": 1,
+          "column": {
+            "name": "bio",
+            "not_null": false,
+            "is_array": false,
+            "comment": "",
+            "length": -1,
+            "is_named_param": false,
+            "is_func_call": false,
+            "scope": "",
+            "table": {
+              "catalog": "",
+              "schema": "public",
+              "name": "authors"
+            },
+            "table_alias": "",
+            "type": {
+              "catalog": "",
+              "schema": "",
+              "name": "text"
+            },
+            "is_sqlc_slice": false,
+            "embed_table": null,
+            "original_name": "bio",
+            "unsigned": false,
+            "array_dims": 0
+          },
+          "context": "PARAMETER_CONTEXT_SET"
+        },
+        {
+          "number": 2,
+          "column": {
+            "name": "id",
+            "not_null": true,
+            "is_array": false,
+            "comment": "",
+            "length": -1,
+            "is_named_param": false,
+            "is_func_call": false,
+            "scope": "",
+            "table": {
+              "catalog": "",
+              "schema": "",
+              "name": "authors"
+            },
+            "table_alias": "",
+            "type": {
+              "catalog": "",
+              "schema": "",
+              "name": "bigserial"
+            },
+            "is_sqlc_slice": false,
+            "embed_table": null,
+            "original_name": "id",
+            "unsigned": false,
+            "array_dims": 0
+          },
+          "context": "PARAMETER_CONTEXT_WHERE"
+        }
+      ],
+      "comments": [],
+      "filename": "query.sql",
+      "insert_into_table": null,
+      "update_table": {
+        "catalog": "",
+        "schema": "",
+        "name": "authors"
+      },
+      "delete_from_table": null
     },
     {
       "text": "DELETE FROM authors\nWHERE id = $1",
@@ -65355,12 +65441,19 @@
             "original_name": "id",
             "unsigned": false,
             "array_dims": 0
-          }
+          },
+          "context": "PARAMETER_CONTEXT_WHERE"
         }
       ],
       "comments": [],
       "filename": "query.sql",
-      "insert_into_table": null
+      "insert_into_table": null,
+      "update_table": null,
+      "delete_from_table": {
+        "catalog": "",
+        "schema": "",
+        "name": "authors"
+      }
     }
   ],
   "sqlc_version": "v1.30.0",

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

@@ -14,6 +14,10 @@ INSERT INTO authors (
 )
 RETURNING *;
 
+-- name: UpdateAuthorBio :exec
+UPDATE authors SET bio = $1
+WHERE id = $2;
+
 -- name: DeleteAuthor :exec
 DELETE FROM authors
 WHERE id = $1;