Skip to content

Commit 9f5ccfe

Browse files
kylemcckylemcc
committed
feat: allow setting a query comment through a context value
Signed-off-by: Kyle McCullough <[email protected]>
1 parent 187743b commit 9f5ccfe

14 files changed

+78
-0
lines changed

‎query_column_add.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -137,6 +137,9 @@ func (q *AddColumnQuery) Exec(ctx context.Context, dest ...interface{}) (sql.Res
137137
return nil, feature.NewNotSupportError(feature.AlterColumnExists)
138138
}
139139

140+
// if a comment is propagated via the context, use it
141+
setCommentFromContext(ctx, q)
142+
140143
queryBytes, err := q.AppendQuery(q.db.fmter, q.db.makeQueryBytes())
141144
if err != nil {
142145
return nil, err

‎query_column_drop.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -129,6 +129,9 @@ func (q *DropColumnQuery) AppendQuery(fmter schema.Formatter, b []byte) (_ []byt
129129
//------------------------------------------------------------------------------
130130

131131
func (q *DropColumnQuery) Exec(ctx context.Context, dest ...interface{}) (sql.Result, error) {
132+
// if a comment is propagated via the context, use it
133+
setCommentFromContext(ctx, q)
134+
132135
queryBytes, err := q.AppendQuery(q.db.fmter, q.db.makeQueryBytes())
133136
if err != nil {
134137
return nil, err

‎query_delete.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -321,6 +321,9 @@ func (q *DeleteQuery) scanOrExec(
321321
return nil, err
322322
}
323323

324+
// if a comment is propagated via the context, use it
325+
setCommentFromContext(ctx, q)
326+
324327
// Generate the query before checking hasReturning.
325328
queryBytes, err := q.AppendQuery(q.db.fmter, q.db.makeQueryBytes())
326329
if err != nil {

‎query_index_create.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -248,6 +248,9 @@ func (q *CreateIndexQuery) AppendQuery(fmter schema.Formatter, b []byte) (_ []by
248248
//------------------------------------------------------------------------------
249249

250250
func (q *CreateIndexQuery) Exec(ctx context.Context, dest ...interface{}) (sql.Result, error) {
251+
// if a comment is propagated via the context, use it
252+
setCommentFromContext(ctx, q)
253+
251254
queryBytes, err := q.AppendQuery(q.db.fmter, q.db.makeQueryBytes())
252255
if err != nil {
253256
return nil, err

‎query_index_drop.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -115,6 +115,9 @@ func (q *DropIndexQuery) AppendQuery(fmter schema.Formatter, b []byte) (_ []byte
115115
//------------------------------------------------------------------------------
116116

117117
func (q *DropIndexQuery) Exec(ctx context.Context, dest ...interface{}) (sql.Result, error) {
118+
// if a comment is propagated via the context, use it
119+
setCommentFromContext(ctx, q)
120+
118121
queryBytes, err := q.AppendQuery(q.db.fmter, q.db.makeQueryBytes())
119122
if err != nil {
120123
return nil, err

‎query_insert.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -586,6 +586,9 @@ func (q *InsertQuery) scanOrExec(
586586
return nil, err
587587
}
588588

589+
// if a comment is propagated via the context, use it
590+
setCommentFromContext(ctx, q)
591+
589592
// Generate the query before checking hasReturning.
590593
queryBytes, err := q.AppendQuery(q.db.fmter, q.db.makeQueryBytes())
591594
if err != nil {

‎query_merge.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -243,6 +243,9 @@ func (q *MergeQuery) scanOrExec(
243243
return nil, err
244244
}
245245

246+
// if a comment is propagated via the context, use it
247+
setCommentFromContext(ctx, q)
248+
246249
// Generate the query before checking hasReturning.
247250
queryBytes, err := q.AppendQuery(q.db.fmter, q.db.makeQueryBytes())
248251
if err != nil {

‎query_raw.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,9 @@ func (q *RawQuery) scanOrExec(
6767
}
6868
}
6969

70+
// if a comment is propagated via the context, use it
71+
setCommentFromContext(ctx, q)
72+
7073
query := q.db.format(q.query, q.args)
7174
var res sql.Result
7275

‎query_select.go

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -791,6 +791,9 @@ func (q *SelectQuery) Rows(ctx context.Context) (*sql.Rows, error) {
791791
return nil, err
792792
}
793793

794+
// if a comment is propagated via the context, use it
795+
setCommentFromContext(ctx, q)
796+
794797
queryBytes, err := q.AppendQuery(q.db.fmter, q.db.makeQueryBytes())
795798
if err != nil {
796799
return nil, err
@@ -812,6 +815,9 @@ func (q *SelectQuery) Exec(ctx context.Context, dest ...interface{}) (res sql.Re
812815
return nil, err
813816
}
814817

818+
// if a comment is propagated via the context, use it
819+
setCommentFromContext(ctx, q)
820+
815821
queryBytes, err := q.AppendQuery(q.db.fmter, q.db.makeQueryBytes())
816822
if err != nil {
817823
return nil, err
@@ -872,6 +878,9 @@ func (q *SelectQuery) scanResult(ctx context.Context, dest ...interface{}) (sql.
872878
return nil, err
873879
}
874880

881+
// if a comment is propagated via the context, use it
882+
setCommentFromContext(ctx, q)
883+
875884
queryBytes, err := q.AppendQuery(q.db.fmter, q.db.makeQueryBytes())
876885
if err != nil {
877886
return nil, err
@@ -924,6 +933,9 @@ func (q *SelectQuery) Count(ctx context.Context) (int, error) {
924933
return 0, q.err
925934
}
926935

936+
// if a comment is propagated via the context, use it
937+
setCommentFromContext(ctx, q)
938+
927939
qq := countQuery{q}
928940

929941
queryBytes, err := qq.AppendQuery(q.db.fmter, nil)
@@ -1028,6 +1040,9 @@ func (q *SelectQuery) Exists(ctx context.Context) (bool, error) {
10281040
}
10291041

10301042
func (q *SelectQuery) selectExists(ctx context.Context) (bool, error) {
1043+
// if a comment is propagated via the context, use it
1044+
setCommentFromContext(ctx, q)
1045+
10311046
qq := selectExistsQuery{q}
10321047

10331048
queryBytes, err := qq.AppendQuery(q.db.fmter, nil)
@@ -1047,6 +1062,9 @@ func (q *SelectQuery) selectExists(ctx context.Context) (bool, error) {
10471062
}
10481063

10491064
func (q *SelectQuery) whereExists(ctx context.Context) (bool, error) {
1065+
// if a comment is propagated via the context, use it
1066+
setCommentFromContext(ctx, q)
1067+
10501068
qq := whereExistsQuery{q}
10511069

10521070
queryBytes, err := qq.AppendQuery(q.db.fmter, nil)

‎query_table_create.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -358,6 +358,9 @@ func (q *CreateTableQuery) Exec(ctx context.Context, dest ...interface{}) (sql.R
358358
return nil, err
359359
}
360360

361+
// if a comment is propagated via the context, use it
362+
setCommentFromContext(ctx, q)
363+
361364
queryBytes, err := q.AppendQuery(q.db.fmter, q.db.makeQueryBytes())
362365
if err != nil {
363366
return nil, err

‎query_table_drop.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -123,6 +123,9 @@ func (q *DropTableQuery) Exec(ctx context.Context, dest ...interface{}) (sql.Res
123123
}
124124
}
125125

126+
// if a comment is propagated via the context, use it
127+
setCommentFromContext(ctx, q)
128+
126129
queryBytes, err := q.AppendQuery(q.db.fmter, q.db.makeQueryBytes())
127130
if err != nil {
128131
return nil, err

‎query_table_truncate.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -136,6 +136,9 @@ func (q *TruncateTableQuery) AppendQuery(
136136
//------------------------------------------------------------------------------
137137

138138
func (q *TruncateTableQuery) Exec(ctx context.Context, dest ...interface{}) (sql.Result, error) {
139+
// if a comment is propagated via the context, use it
140+
setCommentFromContext(ctx, q)
141+
139142
queryBytes, err := q.AppendQuery(q.db.fmter, q.db.makeQueryBytes())
140143
if err != nil {
141144
return nil, err

‎query_update.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -556,6 +556,9 @@ func (q *UpdateQuery) scanOrExec(
556556
return nil, err
557557
}
558558

559+
// if a comment is propagated via the context, use it
560+
setCommentFromContext(ctx, q)
561+
559562
// Generate the query before checking hasReturning.
560563
queryBytes, err := q.AppendQuery(q.db.fmter, q.db.makeQueryBytes())
561564
if err != nil {

‎util.go

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package bun
22

33
import (
4+
"context"
45
"fmt"
56
"reflect"
67
"strings"
@@ -86,3 +87,26 @@ func appendComment(b []byte, name string) []byte {
8687
name = strings.ReplaceAll(name, `*/`, `*\/`)
8788
return append(b, fmt.Sprintf("/* %s */ ", name)...)
8889
}
90+
91+
// queryCommentCtxKey is a context key for setting a query comment on a context instead of calling the Comment("...") API directly
92+
type queryCommentCtxKey struct{}
93+
94+
// WithComment returns a context that includes a comment that may be included in a query for debugging
95+
//
96+
// If a context with an attached query is used, a comment set by the Comment("...") API will be overwritten.
97+
func WithComment(ctx context.Context, comment string) context.Context {
98+
return context.WithValue(ctx, queryCommentCtxKey{}, comment)
99+
}
100+
101+
// commenter describes the Comment interface implemented by all of the query types
102+
type commenter[T any] interface {
103+
Comment(string) T
104+
}
105+
106+
// setCommentFromContext sets the comment on the given query from the supplied context if one is set using the Comment(...) method.
107+
func setCommentFromContext[T any](ctx context.Context, q commenter[T]) {
108+
s, _ := ctx.Value(queryCommentCtxKey{}).(string)
109+
if s != "" {
110+
q.Comment(s)
111+
}
112+
}

0 commit comments

Comments
 (0)