Add a config flag to add a context key with the current query name

Multiple people have asked for more involved features like stats, but
with this simple context injection people can implement their own DBTX
and do whatever they want from there.

My use case is also measuring query durations per query

internal/cmd/shim.go

@@ -96,6 +96,7 @@ func pluginGoCode(s config.SQLGo) *plugin.GoCode {
 		EmitPointersForNullTypes:    s.EmitPointersForNullTypes,
 		EmitEnumValidMethod:         s.EmitEnumValidMethod,
 		EmitAllEnumValues:           s.EmitAllEnumValues,
+		EmitQueryNameInContext:      s.EmitQueryNameInContext,
 		JsonTagsCaseStyle:           s.JSONTagsCaseStyle,
 		Package:                     s.Package,
 		Out:                         s.Out,

internal/codegen/golang/gen.go

@@ -36,6 +36,7 @@ type tmplCtx struct {
 	EmitMethodsWithDBArgument bool
 	EmitEnumValidMethod       bool
 	EmitAllEnumValues         bool
+	EmitQueryNameInContext    bool
 	UsesCopyFrom              bool
 	UsesBatch                 bool
 }
@@ -57,6 +58,12 @@ func (t *tmplCtx) codegenEmitPreparedQueries() bool {
 	return t.EmitPreparedQueries
 }
 
+// Called as a global method since subtemplate queryCodeStdExec does not have
+// access to the toplevel tmplCtx
+func (t *tmplCtx) codegenEmitQueryNameInContext() bool {
+	return t.EmitQueryNameInContext
+}
+
 func (t *tmplCtx) codegenQueryMethod(q Query) string {
 	db := "q.db"
 	if t.EmitMethodsWithDBArgument {
@@ -135,6 +142,7 @@ func generate(req *plugin.CodeGenRequest, enums []Enum, structs []Struct, querie
 		EmitMethodsWithDBArgument: golang.EmitMethodsWithDbArgument,
 		EmitEnumValidMethod:       golang.EmitEnumValidMethod,
 		EmitAllEnumValues:         golang.EmitAllEnumValues,
+		EmitQueryNameInContext:    golang.EmitQueryNameInContext,
 		UsesCopyFrom:              usesCopyFrom(queries),
 		UsesBatch:                 usesBatch(queries),
 		SQLDriver:                 parseDriver(golang.SqlPackage),
@@ -169,10 +177,11 @@ func generate(req *plugin.CodeGenRequest, enums []Enum, structs []Struct, querie
 
 		// These methods are Go specific, they do not belong in the codegen package
 		// (as that is language independent)
-		"dbarg":               tctx.codegenDbarg,
-		"emitPreparedQueries": tctx.codegenEmitPreparedQueries,
-		"queryMethod":         tctx.codegenQueryMethod,
-		"queryRetval":         tctx.codegenQueryRetval,
+		"dbarg":                  tctx.codegenDbarg,
+		"emitPreparedQueries":    tctx.codegenEmitPreparedQueries,
+		"emitQueryNameInContext": tctx.codegenEmitQueryNameInContext,
+		"queryMethod":            tctx.codegenQueryMethod,
+		"queryRetval":            tctx.codegenQueryRetval,
 	}
 
 	tmpl := template.Must(

internal/codegen/golang/templates/stdlib/dbCode.tmpl

@@ -6,6 +6,11 @@ type DBTX interface {
 	QueryRowContext(context.Context, string, ...interface{}) *sql.Row
 }
 
+{{ if emitQueryNameInContext }}
+	type contextKey int8
+	var QueryNameContextKey contextKey = 0
+{{ end }}
+
 {{ if .EmitMethodsWithDBArgument}}
 func New() *Queries {
 	return &Queries{}

internal/codegen/golang/templates/stdlib/queryCode.tmpl

@@ -109,6 +109,9 @@ func (q *Queries) {{.MethodName}}(ctx context.Context, {{ dbarg }} {{.Arg.Pair}}
 {{end}}
 
 {{define "queryCodeStdExec"}}
+		{{- if emitQueryNameInContext }}
+        ctx = context.WithValue(ctx, QueryNameContextKey, "{{.MethodName}}")
+		{{ end }}
     {{- if .Arg.HasSqlcSlices }}
         query := {{.ConstantName}}
         var queryParams []interface{}

internal/config/config.go

@@ -143,6 +143,7 @@ type SQLGo struct {
 	EmitPointersForNullTypes    bool              `json:"emit_pointers_for_null_types" yaml:"emit_pointers_for_null_types"`
 	EmitEnumValidMethod         bool              `json:"emit_enum_valid_method,omitempty" yaml:"emit_enum_valid_method"`
 	EmitAllEnumValues           bool              `json:"emit_all_enum_values,omitempty" yaml:"emit_all_enum_values"`
+	EmitQueryNameInContext      bool              `json:"emit_query_name_in_context,omitempty" yaml:"emit_query_name_in_context"`
 	JSONTagsCaseStyle           string            `json:"json_tags_case_style,omitempty" yaml:"json_tags_case_style"`
 	Package                     string            `json:"package" yaml:"package"`
 	Out                         string            `json:"out" yaml:"out"`

internal/config/v_one.go

@@ -39,6 +39,7 @@ type v1PackageSettings struct {
 	EmitPointersForNullTypes  bool       `json:"emit_pointers_for_null_types" yaml:"emit_pointers_for_null_types"`
 	EmitEnumValidMethod       bool       `json:"emit_enum_valid_method,omitempty" yaml:"emit_enum_valid_method"`
 	EmitAllEnumValues         bool       `json:"emit_all_enum_values,omitempty" yaml:"emit_all_enum_values"`
+	EmitQueryNameInContext    bool       `json:"emit_query_name_in_context,omitempty" yaml:"emit_query_name_in_context"`
 	JSONTagsCaseStyle         string     `json:"json_tags_case_style,omitempty" yaml:"json_tags_case_style"`
 	SQLPackage                string     `json:"sql_package" yaml:"sql_package"`
 	SQLDriver                 string     `json:"sql_driver" yaml:"sql_driver"`
@@ -163,6 +164,7 @@ func (c *V1GenerateSettings) Translate() Config {
 					EmitPointersForNullTypes:  pkg.EmitPointersForNullTypes,
 					EmitEnumValidMethod:       pkg.EmitEnumValidMethod,
 					EmitAllEnumValues:         pkg.EmitAllEnumValues,
+					EmitQueryNameInContext:    pkg.EmitQueryNameInContext,
 					Package:                   pkg.Name,
 					Out:                       pkg.Path,
 					SQLPackage:                pkg.SQLPackage,

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

@@ -42,7 +42,8 @@
       "query_parameter_limit": 1,
       "output_batch_file_name": "",
       "json_tags_id_uppercase": false,
-      "omit_unused_structs": false
+      "omit_unused_structs": false,
+      "emit_query_name_in_context": false
     },
     "json": {
       "out": "gen",

internal/endtoend/testdata/emit_query_name_in_context/mysql/query.sql

@@ -0,0 +1,4 @@
+CREATE TABLE foo (id int not null);
+
+-- name: FuncParamIdent :many
+SELECT id FROM foo WHERE id = ?;

internal/endtoend/testdata/emit_query_name_in_context/mysql/sqlc.json

@@ -0,0 +1,19 @@
+{
+  "version": "1",
+  "packages": [
+    {
+      "engine": "mysql",
+      "path": "go",
+      "name": "querytest",
+      "schema": "query.sql",
+      "queries": "query.sql",
+			"emit_query_name_in_context": true
+    }
+  ],
+  "overrides": [
+    {
+    "column": "foo.mystr",
+    "go_type": "github.com/sqlc-dev/sqlc-testdata/mysql.ID"
+    }
+  ]
+}

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

@@ -42,7 +42,8 @@
       "query_parameter_limit": 1,
       "output_batch_file_name": "",
       "json_tags_id_uppercase": false,
-      "omit_unused_structs": false
+      "omit_unused_structs": false,
+      "emit_query_name_in_context": false
     },
     "json": {
       "out": "",

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

@@ -42,7 +42,8 @@
       "query_parameter_limit": 1,
       "output_batch_file_name": "",
       "json_tags_id_uppercase": false,
-      "omit_unused_structs": false
+      "omit_unused_structs": false,
+      "emit_query_name_in_context": false
     },
     "json": {
       "out": "",