Merge branch 'Pilatuz/master' into dev

Author: aarondl

queries/qm/query_mods.go

@@ -554,6 +554,22 @@ func For(clause string) QueryMod {
 	}
 }
 
+type commentQueryMod struct {
+	comment string
+}
+
+// Apply implements QueryMod.Apply.
+func (qm commentQueryMod) Apply(q *queries.Query) {
+	queries.SetComment(q, qm.comment)
+}
+
+// Comment inserts a custom comment at the begin of your query
+func Comment(comment string) QueryMod {
+	return commentQueryMod{
+		comment: comment,
+	}
+}
+
 // Rels is an alias for strings.Join to make it easier to use relationship name
 // constants in Load.
 func Rels(r ...string) string {

queries/query.go

@@ -44,6 +44,7 @@ type Query struct {
 	offset     int
 	forlock    string
 	distinct   string
+	comment    string
 }
 
 // Applicator exists only to allow
@@ -270,6 +271,11 @@ func SetFor(q *Query, clause string) {
 	q.forlock = clause
 }
 
+// SetComment on the query.
+func SetComment(q *Query, comment string) {
+	q.comment = comment
+}
+
 // SetUpdate on the query.
 func SetUpdate(q *Query, cols map[string]interface{}) {
 	q.update = cols

queries/query_builders.go

@@ -48,6 +48,7 @@ func buildSelectQuery(q *Query) (*bytes.Buffer, []interface{}) {
 	buf := strmangle.GetBuffer()
 	var args []interface{}
 
+	writeComment(q, buf)
 	writeCTEs(q, buf, &args)
 
 	buf.WriteString("SELECT ")
@@ -138,6 +139,7 @@ func buildDeleteQuery(q *Query) (*bytes.Buffer, []interface{}) {
 	var args []interface{}
 	buf := strmangle.GetBuffer()
 
+	writeComment(q, buf)
 	writeCTEs(q, buf, &args)
 
 	buf.WriteString("DELETE FROM ")
@@ -160,6 +162,7 @@ func buildUpdateQuery(q *Query) (*bytes.Buffer, []interface{}) {
 	buf := strmangle.GetBuffer()
 	var args []interface{}
 
+	writeComment(q, buf)
 	writeCTEs(q, buf, &args)
 
 	buf.WriteString("UPDATE ")
@@ -586,6 +589,19 @@ func parseFromClause(toks []string) (alias, name string, ok bool) {
 	return alias, name, ok
 }
 
+func writeComment(q *Query, buf *bytes.Buffer) {
+	if len(q.comment) == 0 {
+		return
+	}
+
+	lines := strings.Split(q.comment, "\n")
+	for _, line := range lines {
+		buf.WriteString("-- ")
+		buf.WriteString(line)
+		buf.WriteByte('\n')
+	}
+}
+
 func writeCTEs(q *Query, buf *bytes.Buffer, args *[]interface{}) {
 	if len(q.withs) == 0 {
 		return

queries/query_builders_test.go

@@ -629,3 +629,36 @@ func TestWriteAsStatements(t *testing.T) {
 		}
 	}
 }
+
+func TestWriteComment(t *testing.T) {
+	t.Parallel()
+
+	var buf bytes.Buffer
+	query := Query{
+		dialect: &drivers.Dialect{LQ: '"', RQ: '"', UseIndexPlaceholders: true},
+	}
+
+	// empty comment
+	buf.Reset()
+	query.comment = ""
+	writeComment(&query, &buf)
+	if got := buf.String(); got != "" {
+		t.Errorf(`bad empty comment, got: %s`, got)
+	}
+
+	// one line comment
+	buf.Reset()
+	query.comment = "comment"
+	writeComment(&query, &buf)
+	if got := buf.String(); got != "-- comment\n" {
+		t.Errorf(`bad one line comment, got: %s`, got)
+	}
+
+	// two lines comment
+	buf.Reset()
+	query.comment = "first\nsecond"
+	writeComment(&query, &buf)
+	if got := buf.String(); got != "-- first\n-- second\n" {
+		t.Errorf(`bad two lines comment, got: %s`, got)
+	}
+}

queries/query_test.go

@@ -640,3 +640,14 @@ func TestAppendWith(t *testing.T) {
 		t.Errorf("Got invalid with on string: %#v", q.withs)
 	}
 }
+
+func TestSetComment(t *testing.T) {
+	t.Parallel()
+
+	q := &Query{}
+	SetComment(q, "my comment")
+
+	if q.comment != "my comment" {
+		t.Errorf("Got invalid comment: %s", q.comment)
+	}
+}