add select hints for postgresql

dialect/psql/dialect/hints.go

@@ -0,0 +1,183 @@
+package dialect
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"strings"
+
+	"github.com/stephenafamo/bob"
+)
+
+type hints struct {
+	hints []string
+}
+
+func (h *hints) AppendHint(hint string) {
+	h.hints = append(h.hints, hint)
+}
+
+func (h hints) WriteSQL(ctx context.Context, w io.StringWriter, d bob.Dialect, start int) ([]any, error) {
+	return bob.ExpressSlice(ctx, w, d, start, h.hints, "/*+ ", " ", " */ ")
+}
+
+type hintable interface{ AppendHint(string) }
+
+// Scan hints
+
+func SeqScan[Q hintable](table string) bob.Mod[Q] {
+	hint := fmt.Sprintf("SeqScan(%s)", table)
+	return bob.ModFunc[Q](func (q Q) {
+		q.AppendHint(hint)
+	})
+}
+
+func NoSeqScan[Q hintable](table string) bob.Mod[Q] {
+	hint := fmt.Sprintf("NoSeqScan(%s)", table)
+	return bob.ModFunc[Q](func(q Q) {
+		q.AppendHint(hint)
+	})
+}
+
+func IndexScan[Q hintable](table string, indexes ...string) bob.Mod[Q] {
+	hint := fmt.Sprintf("IndexScan(%s)", joinWithTable(table, indexes))
+	return bob.ModFunc[Q](func(q Q) {
+		q.AppendHint(hint)
+	})
+}
+
+func NoIndexScan[Q hintable](table string) bob.Mod[Q] {
+	hint := fmt.Sprintf("NoIndexScan(%s)", table)
+	return bob.ModFunc[Q](func(q Q) {
+		q.AppendHint(hint)
+	})
+}
+
+func IndexOnlyScan[Q hintable](table string, indexes ...string) bob.Mod[Q] {
+	hint := fmt.Sprintf("IndexOnlyScan(%s)", joinWithTable(table, indexes))
+	return bob.ModFunc[Q](func(q Q) {
+		q.AppendHint(hint)
+	})
+}
+
+func NoIndexOnlyScan[Q hintable](table string) bob.Mod[Q] {
+	hint := fmt.Sprintf("NoIndexOnlyScan(%s)", table)
+	return bob.ModFunc[Q](func(q Q) {
+		q.AppendHint(hint)
+	})
+}
+
+func BitmapScan[Q hintable](table string, indexes ...string) bob.Mod[Q] {
+	hint := fmt.Sprintf("BitmapScan(%s)", joinWithTable(table, indexes))
+	return bob.ModFunc[Q](func(q Q) {
+		q.AppendHint(hint)
+	})
+}
+
+func NoBitmapScan[Q hintable](table string) bob.Mod[Q] {
+	hint := fmt.Sprintf("NoBitmapScan(%s)", table)
+	return bob.ModFunc[Q](func(q Q) {
+		q.AppendHint(hint)
+	})
+}
+
+func TidScan[Q hintable](table string) bob.Mod[Q] {
+	hint := fmt.Sprintf("TidScan(%s)", table)
+	return bob.ModFunc[Q](func(q Q) {
+		q.AppendHint(hint)
+	})
+}
+
+func NoTidScan[Q hintable](table string) bob.Mod[Q] {
+	hint := fmt.Sprintf("NoTidScan(%s)", table)
+	return bob.ModFunc[Q](func(q Q) {
+		q.AppendHint(hint)
+	})
+}
+
+// Join hints
+
+func NestLoop[Q hintable](tables ...string) bob.Mod[Q] {
+	hint := fmt.Sprintf("NestLoop(%s)", strings.Join(tables, " "))
+	return bob.ModFunc[Q](func(q Q) {
+		q.AppendHint(hint)
+	})
+}
+
+func NoNestLoop[Q hintable](tables ...string) bob.Mod[Q] {
+	hint := fmt.Sprintf("NoNestLoop(%s)", strings.Join(tables, " "))
+	return bob.ModFunc[Q](func(q Q) {
+		q.AppendHint(hint)
+	})
+}
+
+func HashJoin[Q hintable](tables ...string) bob.Mod[Q] {
+	hint := fmt.Sprintf("HashJoin(%s)", strings.Join(tables, " "))
+	return bob.ModFunc[Q](func(q Q) {
+		q.AppendHint(hint)
+	})
+}
+
+func NoHashJoin[Q hintable](tables ...string) bob.Mod[Q] {
+	hint := fmt.Sprintf("NoHashJoin(%s)", strings.Join(tables, " "))
+	return bob.ModFunc[Q](func(q Q) {
+		q.AppendHint(hint)
+	})
+}
+
+func MergeJoin[Q hintable](tables ...string) bob.Mod[Q] {
+	hint := fmt.Sprintf("MergeJoin(%s)", strings.Join(tables, " "))
+	return bob.ModFunc[Q](func(q Q) {
+		q.AppendHint(hint)
+	})
+}
+
+func NoMergeJoin[Q hintable](tables ...string) bob.Mod[Q] {
+	hint := fmt.Sprintf("NoMergeJoin(%s)", strings.Join(tables, " "))
+	return bob.ModFunc[Q](func(q Q) {
+		q.AppendHint(hint)
+	})
+}
+
+// Join order hint
+
+func Leading[Q hintable](spec string) bob.Mod[Q] {
+	hint := fmt.Sprintf("Leading(%s)", spec)
+	return bob.ModFunc[Q](func(q Q) {
+		q.AppendHint(hint)
+	})
+}
+
+// Row estimation hint
+
+func Rows[Q hintable](spec string) bob.Mod[Q] {
+	hint := fmt.Sprintf("Rows(%s)", spec)
+	return bob.ModFunc[Q](func(q Q) {
+		q.AppendHint(hint)
+	})
+}
+
+// Parallel hint
+
+func Parallel[Q hintable](table string, nworkers int, strength string) bob.Mod[Q] {
+	hint := fmt.Sprintf("Parallel(%s %d %s)", table, nworkers, strength)
+	return bob.ModFunc[Q](func(q Q) {
+		q.AppendHint(hint)
+	})
+}
+
+// GUC hint
+
+func Set[Q hintable](variable string, value string) bob.Mod[Q] {
+	hint := fmt.Sprintf("Set(%s %s)", variable, value)
+	return bob.ModFunc[Q](func(q Q) {
+		q.AppendHint(hint)
+	})
+}
+
+func joinWithTable(table string, extras []string) string {
+	if len(extras) == 0 {
+		return table
+	}
+	return table + " " + strings.Join(extras, " ")
+}

dialect/psql/dialect/select.go

@@ -11,6 +11,7 @@ import (
 // Trying to represent the select query structure as documented in
 // https://www.postgresql.org/docs/current/sql-select.html
 type SelectQuery struct {
+	hints
 	clause.With
 	clause.SelectList
 	Distinct
@@ -44,6 +45,13 @@ func (s SelectQuery) WriteSQL(ctx context.Context, w io.StringWriter, d bob.Dial
 		return nil, err
 	}
 
+	// no optimizer hint args
+	_, err = bob.ExpressIf(ctx, w, d, start+len(args), s.hints,
+		len(s.hints.hints) > 0, "", "")
+	if err != nil {
+		return nil, err
+	}
+
 	withArgs, err := bob.ExpressIf(ctx, w, d, start+len(args), s.With,
 		len(s.With.CTEs) > 0, "\n", "")
 	if err != nil {

dialect/psql/select_test.go

@@ -290,6 +290,56 @@ ORDER BY id LIMIT 1000`,
 	testutils.RunTests(t, examples, formatter)
 }
 
+func TestSelectHints(t *testing.T) {
+	examples := testutils.Testcases{
+		"with optimizer hints": {
+			Query: psql.Select(
+				sm.SeqScan("t1"),
+				sm.Columns("id", "name"),
+				sm.From("t1"),
+			),
+			ExpectedSQL: `/*+ SeqScan(t1) */ SELECT id, name FROM t1`,
+		},
+		"with multiple optimizer hints": {
+			Query: psql.Select(
+				sm.SeqScan("t1"),
+				sm.NestLoop("t1", "t2"),
+				sm.Columns("id", "name"),
+				sm.From("t1"),
+				sm.InnerJoin("t2").On(psql.Quote("t1", "id").EQ(psql.Quote("t2", "id"))),
+			),
+			ExpectedSQL: `/*+ SeqScan(t1) NestLoop(t1 t2) */ SELECT id, name FROM t1 INNER JOIN t2 ON ("t1"."id" = "t2"."id")`,
+		},
+		"with index scan hint": {
+			Query: psql.Select(
+				sm.IndexScan("t1", "idx_t1_id", "idx_t2_id"),
+				sm.Columns("id"),
+				sm.From("t1"),
+			),
+			ExpectedSQL: `/*+ IndexScan(t1 idx_t1_id idx_t2_id) */ SELECT id FROM t1`,
+		},
+		"with leading hint": {
+			Query: psql.Select(
+				sm.Leading("t1 t2 t3"),
+				sm.Columns("*"),
+				sm.From("t1"),
+			),
+			ExpectedSQL: `/*+ Leading(t1 t2 t3) */ SELECT * FROM t1`,
+		},
+		"with set hint": {
+			Query: psql.Select(
+				sm.Set("random_page_cost", "2.0"),
+				sm.Columns("*"),
+				sm.From("t1"),
+			),
+			ExpectedSQL: `/*+ Set(random_page_cost 2.0) */ SELECT * FROM t1`,
+		},
+	}
+
+	// pgparse strips comments, so we bypass the formatter
+	testutils.RunTests(t, examples, nil)
+}
+
 func formatter(s string) (string, error) {
 	aTree, err := pgparse.Parse(s)
 	if err != nil {

dialect/psql/sm/hints.go

@@ -0,0 +1,86 @@
+package sm
+
+import (
+	"github.com/stephenafamo/bob"
+	"github.com/stephenafamo/bob/dialect/psql/dialect"
+)
+
+func SeqScan(table string) bob.Mod[*dialect.SelectQuery] {
+	return dialect.SeqScan[*dialect.SelectQuery](table)
+}
+
+func NoSeqScan(table string) bob.Mod[*dialect.SelectQuery] {
+	return dialect.NoSeqScan[*dialect.SelectQuery](table)
+}
+
+func IndexScan(table string, indexes ...string) bob.Mod[*dialect.SelectQuery] {
+	return dialect.IndexScan[*dialect.SelectQuery](table, indexes...)
+}
+
+func NoIndexScan(table string) bob.Mod[*dialect.SelectQuery] {
+	return dialect.NoIndexScan[*dialect.SelectQuery](table)
+}
+
+func IndexOnlyScan(table string, indexes ...string) bob.Mod[*dialect.SelectQuery] {
+	return dialect.IndexOnlyScan[*dialect.SelectQuery](table, indexes...)
+}
+
+func NoIndexOnlyScan(table string) bob.Mod[*dialect.SelectQuery] {
+	return dialect.NoIndexOnlyScan[*dialect.SelectQuery](table)
+}
+
+func BitmapScan(table string, indexes ...string) bob.Mod[*dialect.SelectQuery] {
+	return dialect.BitmapScan[*dialect.SelectQuery](table, indexes...)
+}
+
+func NoBitmapScan(table string) bob.Mod[*dialect.SelectQuery] {
+	return dialect.NoBitmapScan[*dialect.SelectQuery](table)
+}
+
+func TidScan(table string) bob.Mod[*dialect.SelectQuery] {
+	return dialect.TidScan[*dialect.SelectQuery](table)
+}
+
+func NoTidScan(table string) bob.Mod[*dialect.SelectQuery] {
+	return dialect.NoTidScan[*dialect.SelectQuery](table)
+}
+
+func NestLoop(tables ...string) bob.Mod[*dialect.SelectQuery] {
+	return dialect.NestLoop[*dialect.SelectQuery](tables...)
+}
+
+func NoNestLoop(tables ...string) bob.Mod[*dialect.SelectQuery] {
+	return dialect.NoNestLoop[*dialect.SelectQuery](tables...)
+}
+
+func HashJoin(tables ...string) bob.Mod[*dialect.SelectQuery] {
+	return dialect.HashJoin[*dialect.SelectQuery](tables...)
+}
+
+func NoHashJoin(tables ...string) bob.Mod[*dialect.SelectQuery] {
+	return dialect.NoHashJoin[*dialect.SelectQuery](tables...)
+}
+
+func MergeJoin(tables ...string) bob.Mod[*dialect.SelectQuery] {
+	return dialect.MergeJoin[*dialect.SelectQuery](tables...)
+}
+
+func NoMergeJoin(tables ...string) bob.Mod[*dialect.SelectQuery] {
+	return dialect.NoMergeJoin[*dialect.SelectQuery](tables...)
+}
+
+func Leading(spec string) bob.Mod[*dialect.SelectQuery] {
+	return dialect.Leading[*dialect.SelectQuery](spec)
+}
+
+func Rows(spec string) bob.Mod[*dialect.SelectQuery] {
+	return dialect.Rows[*dialect.SelectQuery](spec)
+}
+
+func Parallel(table string, nworkers int, strength string) bob.Mod[*dialect.SelectQuery] {
+	return dialect.Parallel[*dialect.SelectQuery](table, nworkers, strength)
+}
+
+func Set(variable string, value string) bob.Mod[*dialect.SelectQuery] {
+	return dialect.Set[*dialect.SelectQuery](variable, value)
+}