feature: WithOptimizerHint

Author: mickeyreiss

queries/qm/query_mods.go

@@ -227,6 +227,24 @@ func Select(columns ...string) QueryMod {
 	}
 }
 
+type optimizerHintQueryMod struct {
+	optimizerHint string
+	args          []interface{}
+}
+
+// Apply implements QueryMod.Apply.
+func (qm optimizerHintQueryMod) Apply(q *queries.Query) {
+	queries.AppendOptimizerHint(q, qm.optimizerHint, qm.args...)
+}
+
+// WithOptimizerHint adds an optimizer hint comment to the query
+func WithOptimizerHint(optimizerHint string, args ...interface{}) QueryMod {
+	return optimizerHintQueryMod{
+		optimizerHint: optimizerHint,
+		args:          args,
+	}
+}
+
 // Where allows you to specify a where clause for your statement. If multiple
 // Where statements are used they are combined with 'and'
 func Where(clause string, args ...interface{}) QueryMod {

queries/query.go

@@ -30,22 +30,23 @@ type Query struct {
 	load     []string
 	loadMods map[string]Applicator
 
-	delete     bool
-	update     map[string]interface{}
-	withs      []argClause
-	selectCols []string
-	count      bool
-	from       []string
-	joins      []join
-	where      []where
-	groupBy    []string
-	orderBy    []argClause
-	having     []argClause
-	limit      *int
-	offset     int
-	forlock    string
-	distinct   string
-	comment    string
+	delete         bool
+	update         map[string]interface{}
+	withs          []argClause
+	optimizerHints []string
+	selectCols     []string
+	count          bool
+	from           []string
+	joins          []join
+	where          []where
+	groupBy        []string
+	orderBy        []argClause
+	having         []argClause
+	limit          *int
+	offset         int
+	forlock        string
+	distinct       string
+	comment        string
 
 	// This field is a hack to allow a query to strip out the reference
 	// to deleted at is null.
@@ -246,6 +247,16 @@ func GetSelect(q *Query) []string {
 	return q.selectCols
 }
 
+// SetOptimizerHints on the query.
+func SetOptimizerHints(q *Query, optimizerHints []string) {
+	q.optimizerHints = optimizerHints
+}
+
+// GetOptimizerHints from the query
+func GetOptimizerHints(q *Query) []string {
+	return q.optimizerHints
+}
+
 // SetDistinct on the query.
 func SetDistinct(q *Query, distinct string) {
 	q.distinct = distinct
@@ -291,6 +302,11 @@ func AppendSelect(q *Query, columns ...string) {
 	q.selectCols = append(q.selectCols, columns...)
 }
 
+// AppendOptimizerHint on the query.
+func AppendOptimizerHint(q *Query, optimizerHint string, args ...interface{}) {
+	q.optimizerHints = append(q.optimizerHints, fmt.Sprintf(optimizerHint, args...))
+}
+
 // AppendFrom on the query.
 func AppendFrom(q *Query, from ...string) {
 	q.from = append(q.from, from...)

queries/query_builders.go

@@ -54,6 +54,7 @@ func buildSelectQuery(q *Query) (*bytes.Buffer, []interface{}) {
 	writeCTEs(q, buf, &args)
 
 	buf.WriteString("SELECT ")
+	writeOptimizerHints(q, buf)
 
 	if q.dialect.UseTopClause {
 		if q.limit != nil && q.offset == 0 {
@@ -632,3 +633,16 @@ func writeCTEs(q *Query, buf *bytes.Buffer, args *[]interface{}) {
 	fmt.Fprintf(buf, resp)
 	strmangle.PutBuffer(withBuf)
 }
+
+func writeOptimizerHints(q *Query, buf *bytes.Buffer) {
+	if len(q.optimizerHints) == 0 {
+		return
+	}
+
+	buf.WriteString("/*+ ")
+	for _, hint := range q.optimizerHints {
+		buf.WriteString(hint)
+		buf.WriteString(" ")
+	}
+	buf.WriteString("*/ ")
+}