feat: Add support for header comments in ORM-generated queries (#2011)

* feat: Added comment for public functions

feat: Add query comments to headers for improved debugging

This feature enhances debugging by including query comments in headers, allowing developers to trace the query's origin. By matching the query comment in the header, it's easier to identify the specific location in the code where the query was executed.

* feat: Fixed query buffer append

* feat: Changed query comment template

* Reverted changes in test

* Revised query comment implementation

* Fixed appendComment func doc

* Kept if cond only inside appendComment

* Implemented tests for appendComment to verify proper escaping of comment symbols

Author: wwoytenko

orm/composite_create.go

@@ -53,6 +53,7 @@ func (q *CreateCompositeQuery) AppendTemplate(b []byte) ([]byte, error) {
 }
 
 func (q *CreateCompositeQuery) AppendQuery(fmter QueryFormatter, b []byte) ([]byte, error) {
+	b = appendComment(b, q.q.comment)
 	if q.q.stickyErr != nil {
 		return nil, q.q.stickyErr
 	}

orm/composite_drop.go

@@ -50,6 +50,7 @@ func (q *DropCompositeQuery) AppendTemplate(b []byte) ([]byte, error) {
 }
 
 func (q *DropCompositeQuery) AppendQuery(fmter QueryFormatter, b []byte) ([]byte, error) {
+	b = appendComment(b, q.q.comment)
 	if q.q.stickyErr != nil {
 		return nil, q.q.stickyErr
 	}

orm/delete.go

@@ -52,6 +52,7 @@ func (q *DeleteQuery) AppendTemplate(b []byte) ([]byte, error) {
 }
 
 func (q *DeleteQuery) AppendQuery(fmter QueryFormatter, b []byte) (_ []byte, err error) {
+	b = appendComment(b, q.q.comment)
 	if q.q.stickyErr != nil {
 		return nil, q.q.stickyErr
 	}

orm/insert.go

@@ -56,6 +56,7 @@ func (q *InsertQuery) AppendTemplate(b []byte) ([]byte, error) {
 var _ QueryAppender = (*InsertQuery)(nil)
 
 func (q *InsertQuery) AppendQuery(fmter QueryFormatter, b []byte) (_ []byte, err error) {
+	b = appendComment(b, q.q.comment)
 	if q.q.stickyErr != nil {
 		return nil, q.q.stickyErr
 	}

orm/query.go

@@ -81,6 +81,7 @@ type Query struct {
 
 	onConflict *SafeQueryAppender
 	returning  []*SafeQueryAppender
+	comment string
 }
 
 func NewQuery(db DB, model ...interface{}) *Query {
@@ -779,6 +780,12 @@ func (q *Query) Apply(fn func(*Query) (*Query, error)) *Query {
 	return qq
 }
 
+// Comment adds a comment to the query, wrapped by /* ... */.
+func (q *Query) Comment(c string) *Query {
+	q.comment = c
+	return q
+}
+
 // Count returns number of rows matching the query using count aggregate function.
 func (q *Query) Count() (int, error) {
 	if q.stickyErr != nil {

orm/select.go

@@ -53,6 +53,7 @@ func (q *SelectQuery) AppendTemplate(b []byte) ([]byte, error) {
 }
 
 func (q *SelectQuery) AppendQuery(fmter QueryFormatter, b []byte) (_ []byte, err error) { //nolint:gocyclo
+	b = appendComment(b, q.q.comment)
 	if q.q.stickyErr != nil {
 		return nil, q.q.stickyErr
 	}

orm/table_create.go

@@ -64,6 +64,7 @@ func (q *CreateTableQuery) AppendTemplate(b []byte) ([]byte, error) {
 }
 
 func (q *CreateTableQuery) AppendQuery(fmter QueryFormatter, b []byte) (_ []byte, err error) {
+	b = appendComment(b, q.q.comment)
 	if q.q.stickyErr != nil {
 		return nil, q.q.stickyErr
 	}

orm/table_drop.go

@@ -50,6 +50,7 @@ func (q *DropTableQuery) AppendTemplate(b []byte) ([]byte, error) {
 }
 
 func (q *DropTableQuery) AppendQuery(fmter QueryFormatter, b []byte) (_ []byte, err error) {
+	b = appendComment(b, q.q.comment)
 	if q.q.stickyErr != nil {
 		return nil, q.q.stickyErr
 	}

orm/update.go

@@ -57,6 +57,7 @@ func (q *UpdateQuery) AppendTemplate(b []byte) ([]byte, error) {
 }
 
 func (q *UpdateQuery) AppendQuery(fmter QueryFormatter, b []byte) (_ []byte, err error) {
+	b = appendComment(b, q.q.comment)
 	if q.q.stickyErr != nil {
 		return nil, q.q.stickyErr
 	}

orm/util.go

@@ -1,7 +1,9 @@
 package orm
 
 import (
+	"fmt"
 	"reflect"
+	"strings"
 
 	"github.com/go-pg/pg/v10/types"
 )
@@ -149,3 +151,13 @@ func appendColumns(b []byte, table types.Safe, fields []*Field) []byte {
 	}
 	return b
 }
+
+// appendComment adds comment in the header of the query into buffer
+func appendComment(b []byte, name string) []byte {
+	if name == "" {
+		return b
+	}
+	name = strings.ReplaceAll(name, `/*`, `/\*`)
+	name = strings.ReplaceAll(name, `*/`, `*\/`)
+	return append(b, fmt.Sprintf("/* %s */ ", name)...)
+}

orm/util_test.go

@@ -0,0 +1,30 @@
+package orm
+
+import (
+	. "github.com/onsi/ginkgo"
+	. "github.com/onsi/gomega"
+)
+
+var _ = Describe("Comment - escape comment symbols", func() {
+	It("only open sequence", func() {
+		var res []byte
+		c := "/* comment"
+
+		s := appendComment(res, c)
+		Expect(s).To(Equal([]byte("/* /\\* comment */ ")))
+	})
+	It("only close sequence", func() {
+		var res []byte
+		c := "comment */"
+
+		s := appendComment(res, c)
+		Expect(s).To(Equal([]byte("/* comment *\\/ */ ")))
+	})
+	It("open and close sequences", func() {
+		var res []byte
+		c := "/* comment */"
+
+		s := appendComment(res, c)
+		Expect(s).To(Equal([]byte("/* /\\* comment *\\/ */ ")))
+	})
+})