From 3f627ce1dd098e371bb410b03db24b32bf5f87cb Mon Sep 17 00:00:00 2001 From: Kyle McCullough Date: Thu, 25 Apr 2024 08:01:30 -0700 Subject: [PATCH] feat: add mechanism for prefixing queries with a comment string Signed-off-by: Kyle McCullough --- query_base.go | 43 +++++++++++++++++++++++++++++++++++++++++++ query_delete.go | 18 +++++++++++++++++- query_insert.go | 11 ++++++++++- query_merge.go | 11 ++++++++++- query_select.go | 38 ++++++++++++++++++++++++++++++++------ query_update.go | 18 +++++++++++++++++- 6 files changed, 129 insertions(+), 10 deletions(-) diff --git a/query_base.go b/query_base.go index 2321a7537..0605ce3d5 100644 --- a/query_base.go +++ b/query_base.go @@ -74,6 +74,7 @@ var ( // QueryBuilder is used for common query methods type QueryBuilder interface { Query + Comment(c string) QueryBuilder Where(query string, args ...interface{}) QueryBuilder WhereGroup(sep string, fn func(QueryBuilder) QueryBuilder) QueryBuilder WhereOr(query string, args ...interface{}) QueryBuilder @@ -104,6 +105,8 @@ type baseQuery struct { tables []schema.QueryWithArgs columns []schema.QueryWithArgs + comment string + flags internal.Flag } @@ -146,6 +149,32 @@ func (q *baseQuery) GetTableName() string { return "" } +func (q *baseQuery) setComment(c string) { + q.comment = c +} + +func (q *baseQuery) appendComment(b []byte) []byte { + if len(q.comment) < 1 { + return b + } + + lastIdx := len(q.comment) - 1 + b = append(b, "/*"...) + for i := range q.comment { + // if the comment contains the closing sequence, escape it + if q.comment[i] == '*' && i < lastIdx && q.comment[i+1] == '/' { + b = append(b, '\\') + b = append(b, q.comment[i]) + b = append(b, '\\') + } else { + b = append(b, q.comment[i]) + } + } + b = append(b, "*/ "...) + + return b +} + func (q *baseQuery) setConn(db IConn) { // Unwrap Bun wrappers to not call query hooks twice. switch db := db.(type) { @@ -1347,3 +1376,17 @@ func (ih *idxHintsQuery) bufIndexHint( b = append(b, ")"...) return b, nil } + +//------------------------------------------------------------------------------ + +type queryCommentCtxKey struct{} + +// CtxWithComment returns a context that includes a comment that may be included in a query for debugging +func ContextWithComment(ctx context.Context, comment string) context.Context { + return context.WithValue(ctx, queryCommentCtxKey{}, comment) +} + +func commentFromContext(ctx context.Context) string { + c, _ := ctx.Value(queryCommentCtxKey{}).(string) + return c +} diff --git a/query_delete.go b/query_delete.go index 49a750cc8..d50df8b07 100644 --- a/query_delete.go +++ b/query_delete.go @@ -44,6 +44,13 @@ func (q *DeleteQuery) Err(err error) *DeleteQuery { return q } +// Comment prepends a comment to the beginning of the query. +// This is useful for debugging or identifying queries in database logs. +func (q *DeleteQuery) Comment(c string) *DeleteQuery { + q.setComment(c) + return q +} + // Apply calls the fn passing the DeleteQuery as an argument. func (q *DeleteQuery) Apply(fn func(*DeleteQuery) *DeleteQuery) *DeleteQuery { if fn != nil { @@ -259,8 +266,10 @@ func (q *DeleteQuery) scanOrExec( return nil, err } + q.setComment(commentFromContext(ctx)) + // Generate the query before checking hasReturning. - queryBytes, err := q.AppendQuery(q.db.fmter, q.db.makeQueryBytes()) + queryBytes, err := q.AppendQuery(q.db.fmter, q.appendComment(q.db.makeQueryBytes())) if err != nil { return nil, err } @@ -342,6 +351,13 @@ type deleteQueryBuilder struct { *DeleteQuery } +// Comment prepends a comment to the beginning of the query. +// This is useful for debugging or identifying queries in database logs. +func (q *deleteQueryBuilder) Comment(c string) QueryBuilder { + q.setComment(c) + return q +} + func (q *deleteQueryBuilder) WhereGroup( sep string, fn func(QueryBuilder) QueryBuilder, ) QueryBuilder { diff --git a/query_insert.go b/query_insert.go index 6d38a4efe..089f0bede 100644 --- a/query_insert.go +++ b/query_insert.go @@ -53,6 +53,13 @@ func (q *InsertQuery) Err(err error) *InsertQuery { return q } +// Comment prepends a comment to the beginning of the query. +// This is useful for debugging or identifying queries in database logs. +func (q *InsertQuery) Comment(c string) *InsertQuery { + q.setComment(c) + return q +} + // Apply calls the fn passing the SelectQuery as an argument. func (q *InsertQuery) Apply(fn func(*InsertQuery) *InsertQuery) *InsertQuery { if fn != nil { @@ -574,8 +581,10 @@ func (q *InsertQuery) scanOrExec( return nil, err } + q.setComment(commentFromContext(ctx)) + // Generate the query before checking hasReturning. - queryBytes, err := q.AppendQuery(q.db.fmter, q.db.makeQueryBytes()) + queryBytes, err := q.AppendQuery(q.db.fmter, q.appendComment(q.db.makeQueryBytes())) if err != nil { return nil, err } diff --git a/query_merge.go b/query_merge.go index 626752b8a..9e27df338 100644 --- a/query_merge.go +++ b/query_merge.go @@ -50,6 +50,13 @@ func (q *MergeQuery) Err(err error) *MergeQuery { return q } +// Comment prepends a comment to the beginning of the query. +// This is useful for debugging or identifying queries in database logs. +func (q *MergeQuery) Comment(c string) *MergeQuery { + q.setComment(c) + return q +} + // Apply calls the fn passing the MergeQuery as an argument. func (q *MergeQuery) Apply(fn func(*MergeQuery) *MergeQuery) *MergeQuery { if fn != nil { @@ -231,8 +238,10 @@ func (q *MergeQuery) scanOrExec( return nil, err } + q.setComment(commentFromContext(ctx)) + // Generate the query before checking hasReturning. - queryBytes, err := q.AppendQuery(q.db.fmter, q.db.makeQueryBytes()) + queryBytes, err := q.AppendQuery(q.db.fmter, q.appendComment(q.db.makeQueryBytes())) if err != nil { return nil, err } diff --git a/query_select.go b/query_select.go index 932cd48be..b7e4278ed 100644 --- a/query_select.go +++ b/query_select.go @@ -66,6 +66,13 @@ func (q *SelectQuery) Err(err error) *SelectQuery { return q } +// Comment prepends a comment to the beginning of the query. +// This is useful for debugging or identifying queries in database logs. +func (q *SelectQuery) Comment(c string) *SelectQuery { + q.setComment(c) + return q +} + // Apply calls the fn passing the SelectQuery as an argument. func (q *SelectQuery) Apply(fn func(*SelectQuery) *SelectQuery) *SelectQuery { if fn != nil { @@ -807,7 +814,9 @@ func (q *SelectQuery) Rows(ctx context.Context) (*sql.Rows, error) { return nil, err } - queryBytes, err := q.AppendQuery(q.db.fmter, q.db.makeQueryBytes()) + q.setComment(commentFromContext(ctx)) + + queryBytes, err := q.AppendQuery(q.db.fmter, q.appendComment(q.db.makeQueryBytes())) if err != nil { return nil, err } @@ -828,7 +837,9 @@ func (q *SelectQuery) Exec(ctx context.Context, dest ...interface{}) (res sql.Re return nil, err } - queryBytes, err := q.AppendQuery(q.db.fmter, q.db.makeQueryBytes()) + q.setComment(commentFromContext(ctx)) + + queryBytes, err := q.AppendQuery(q.db.fmter, q.appendComment(q.db.makeQueryBytes())) if err != nil { return nil, err } @@ -875,7 +886,9 @@ func (q *SelectQuery) Scan(ctx context.Context, dest ...interface{}) error { return err } - queryBytes, err := q.AppendQuery(q.db.fmter, q.db.makeQueryBytes()) + q.setComment(commentFromContext(ctx)) + + queryBytes, err := q.AppendQuery(q.db.fmter, q.appendComment(q.db.makeQueryBytes())) if err != nil { return err } @@ -929,7 +942,9 @@ func (q *SelectQuery) Count(ctx context.Context) (int, error) { qq := countQuery{q} - queryBytes, err := qq.AppendQuery(q.db.fmter, nil) + q.setComment(commentFromContext(ctx)) + + queryBytes, err := qq.AppendQuery(q.db.fmter, q.appendComment(nil)) if err != nil { return 0, err } @@ -1021,7 +1036,9 @@ func (q *SelectQuery) Exists(ctx context.Context) (bool, error) { func (q *SelectQuery) selectExists(ctx context.Context) (bool, error) { qq := selectExistsQuery{q} - queryBytes, err := qq.AppendQuery(q.db.fmter, nil) + q.setComment(commentFromContext(ctx)) + + queryBytes, err := qq.AppendQuery(q.db.fmter, q.appendComment(nil)) if err != nil { return false, err } @@ -1040,7 +1057,9 @@ func (q *SelectQuery) selectExists(ctx context.Context) (bool, error) { func (q *SelectQuery) whereExists(ctx context.Context) (bool, error) { qq := whereExistsQuery{q} - queryBytes, err := qq.AppendQuery(q.db.fmter, nil) + q.setComment(commentFromContext(ctx)) + + queryBytes, err := qq.AppendQuery(q.db.fmter, q.appendComment(nil)) if err != nil { return false, err } @@ -1082,6 +1101,13 @@ type selectQueryBuilder struct { *SelectQuery } +// Comment prepends a comment to the beginning of the query. +// This is useful for debugging or identifying queries in database logs. +func (q *selectQueryBuilder) Comment(c string) QueryBuilder { + q.setComment(c) + return q +} + func (q *selectQueryBuilder) WhereGroup( sep string, fn func(QueryBuilder) QueryBuilder, ) QueryBuilder { diff --git a/query_update.go b/query_update.go index e56ba20d1..3494a0400 100644 --- a/query_update.go +++ b/query_update.go @@ -53,6 +53,13 @@ func (q *UpdateQuery) Err(err error) *UpdateQuery { return q } +// Comment prepends a comment to the beginning of the query. +// This is useful for debugging or identifying queries in database logs. +func (q *UpdateQuery) Comment(c string) *UpdateQuery { + q.setComment(c) + return q +} + // Apply calls the fn passing the SelectQuery as an argument. func (q *UpdateQuery) Apply(fn func(*UpdateQuery) *UpdateQuery) *UpdateQuery { if fn != nil { @@ -505,8 +512,10 @@ func (q *UpdateQuery) scanOrExec( return nil, err } + q.setComment(commentFromContext(ctx)) + // Generate the query before checking hasReturning. - queryBytes, err := q.AppendQuery(q.db.fmter, q.db.makeQueryBytes()) + queryBytes, err := q.AppendQuery(q.db.fmter, q.appendComment(q.db.makeQueryBytes())) if err != nil { return nil, err } @@ -604,6 +613,13 @@ type updateQueryBuilder struct { *UpdateQuery } +// Comment prepends a comment to the beginning of the query. +// This is useful for debugging or identifying queries in database logs. +func (q *updateQueryBuilder) Comment(c string) QueryBuilder { + q.setComment(c) + return q +} + func (q *updateQueryBuilder) WhereGroup( sep string, fn func(QueryBuilder) QueryBuilder, ) QueryBuilder {