feat: add mechanism for prefixing queries with a comment string

Signed-off-by: Kyle McCullough <[email protected]>
uptrace · Jul 1, 2024 · 3f627ce · 3f627ce
1 parent 5be346f
commit 3f627ce
Show file tree

Hide file tree

Showing 6 changed files with 129 additions and 10 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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 {