hints: add injection hint and refactor usage of hint union

This commit adds the protobuf struct for injecting "internal" hints
like (index or join hints) into a statement's AST. This can be serialized
and inserted into the `system.statement_hints` table. This commit also
refactors the way hints are represented in memory.

Informs #153633

Release note: None

Author: DrewKimball

pkg/BUILD.bazel

@@ -505,6 +505,7 @@ ALL_TESTS = [
     "//pkg/sql/gcjob/gcjobnotifier:gcjobnotifier_test",
     "//pkg/sql/gcjob:gcjob_test",
     "//pkg/sql/gcjob_test:gcjob_test_test",
+    "//pkg/sql/hintpb:hintpb_test",
     "//pkg/sql/hints:hints_test",
     "//pkg/sql/idxrecommendations:idxrecommendations_test",
     "//pkg/sql/idxusage:idxusage_test",
@@ -2071,6 +2072,7 @@ GO_TARGETS = [
     "//pkg/sql/gcjob:gcjob_test",
     "//pkg/sql/gcjob_test:gcjob_test_test",
     "//pkg/sql/hintpb:hintpb",
+    "//pkg/sql/hintpb:hintpb_test",
     "//pkg/sql/hints:hints",
     "//pkg/sql/hints:hints_test",
     "//pkg/sql/idxrecommendations:idxrecommendations",

pkg/sql/hintpb/BUILD.bazel

@@ -1,5 +1,5 @@
 load("@rules_proto//proto:defs.bzl", "proto_library")
-load("@io_bazel_rules_go//go:def.bzl", "go_library")
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 load("@io_bazel_rules_go//proto:def.bzl", "go_proto_library")
 
 proto_library(
@@ -25,5 +25,15 @@ go_library(
     embed = [":hintpb_go_proto"],
     importpath = "github.com/cockroachdb/cockroach/pkg/sql/hintpb",
     visibility = ["//visibility:public"],
-    deps = ["//pkg/util/protoutil"],
+    deps = [
+        "//pkg/util/protoutil",
+        "@com_github_cockroachdb_errors//:errors",
+    ],
+)
+
+go_test(
+    name = "hintpb_test",
+    srcs = ["statement_hint_test.go"],
+    embed = [":hintpb"],
+    deps = ["@com_github_stretchr_testify//require"],
 )

pkg/sql/hintpb/statement_hint.go

@@ -5,24 +5,29 @@
 
 package hintpb
 
-import "github.com/cockroachdb/cockroach/pkg/util/protoutil"
+import (
+	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
+	"github.com/cockroachdb/errors"
+)
 
-type StatementHint interface {
-	protoutil.Message
-}
-
-// NewStatementHint converts the raw bytes from system.statement_hints into a
+// FromBytes converts the raw bytes from system.statement_hints into a
 // StatementHintUnion object.
-func NewStatementHint(bytes []byte) (*StatementHintUnion, error) {
-	res := &StatementHintUnion{}
-	if err := protoutil.Unmarshal(bytes, res); err != nil {
-		return nil, err
+func FromBytes(bytes []byte) (StatementHintUnion, error) {
+	res := StatementHintUnion{}
+	if err := protoutil.Unmarshal(bytes, &res); err != nil {
+		return StatementHintUnion{}, err
+	}
+	if res.GetValue() == nil {
+		return StatementHintUnion{}, errors.New("invalid hint bytes: no value set")
 	}
 	return res, nil
 }
 
 // ToBytes converts the StatementHintUnion to a raw bytes representation that
 // can be inserted into the system.statement_hints table.
-func (hint *StatementHintUnion) ToBytes() ([]byte, error) {
-	return protoutil.Marshal(hint)
+func ToBytes(hint StatementHintUnion) ([]byte, error) {
+	if hint.GetValue() == nil {
+		return nil, errors.New("cannot convert empty hint to bytes")
+	}
+	return protoutil.Marshal(&hint)
 }

pkg/sql/hintpb/statement_hint.proto

@@ -12,6 +12,15 @@ import "gogoproto/gogo.proto";
 // that can be applied to statements with a given fingerprint using the
 // system.statement_hints table.
 message StatementHintUnion {
-  // TODO: add a field for each type of statement hint.
-  // oneof value {}
+  option (gogoproto.onlyone) = true;
+
+  InjectHints inject_hints = 1;
+}
+
+// InjectHints applies inline query plan hints (join and index hints) from the
+// DonorSQL field to the hinted statement. The DonorSQL string must have
+// the same syntactical structure as the hinted statement, with only the
+// addition of inline hints.
+message InjectHints {
+  string donor_sql = 1 [(gogoproto.customname) = "DonorSQL"];
 }

pkg/sql/hintpb/statement_hint_test.go

@@ -0,0 +1,42 @@
+// Copyright 2025 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package hintpb
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestFromToBytes(t *testing.T) {
+	// Test writing empty hint.
+	_, err := ToBytes(StatementHintUnion{})
+	require.EqualError(t, err, "cannot convert empty hint to bytes")
+
+	// Test reading empty bytes.
+	_, err = FromBytes(nil)
+	require.EqualError(t, err, "invalid hint bytes: no value set")
+	_, err = FromBytes([]byte{})
+	require.EqualError(t, err, "invalid hint bytes: no value set")
+
+	// Test reading invalid bytes.
+	_, err = FromBytes([]byte{0xFF, 0xFF, 0xFF})
+	require.Error(t, err)
+
+	// Test that a valid hint round trips.
+	testRT := func(hint interface{}) {
+		var hintUnion StatementHintUnion
+		hintUnion.SetValue(hint)
+		bytes, err := ToBytes(hintUnion)
+		require.NoError(t, err)
+		require.NotEmpty(t, bytes)
+		decodedHintUnion, err := FromBytes(bytes)
+		require.NoError(t, err)
+		require.Equal(t, hint, decodedHintUnion.GetValue())
+	}
+	testRT(&InjectHints{})
+	testRT(&InjectHints{DonorSQL: "SELECT * FROM t"})
+}

pkg/sql/hints/hint_cache.go

@@ -147,8 +147,6 @@ func NewStatementHintsCache(
 func (c *StatementHintsCache) Start(
 	ctx context.Context, sysTableResolver catalog.SystemTableIDResolver,
 ) error {
-	// TODO(drewk): in a follow-up commit, block until the initial scan is done
-	// so that statements without hints don't have to check the LRU cache.
 	return c.startRangefeedInternal(ctx, sysTableResolver)
 }
 
@@ -393,7 +391,7 @@ func (c *StatementHintsCache) GetGeneration() int64 {
 // retrieving them.
 func (c *StatementHintsCache) MaybeGetStatementHints(
 	ctx context.Context, statementFingerprint string,
-) (hints []hintpb.StatementHint, ids []int64) {
+) (hints []hintpb.StatementHintUnion, ids []int64) {
 	hash := fnv.New64()
 	_, err := hash.Write([]byte(statementFingerprint))
 	if err != nil {
@@ -451,7 +449,7 @@ func (c *StatementHintsCache) maybeWaitForRefreshLocked(
 // released while reading from the db, and then reacquired.
 func (c *StatementHintsCache) addCacheEntryLocked(
 	ctx context.Context, statementHash int64, statementFingerprint string,
-) (hints []hintpb.StatementHint, ids []int64) {
+) (hints []hintpb.StatementHintUnion, ids []int64) {
 	c.mu.AssertHeld()
 
 	// Add a cache entry that other queries can find and wait on until we have the
@@ -516,7 +514,7 @@ func (c *StatementHintsCache) getStatementHintsFromDB(
 		datums := it.Cur()
 		rowID := int64(tree.MustBeDInt(datums[0]))
 		fingerprint := string(tree.MustBeDString(datums[1]))
-		hint, err := hintpb.NewStatementHint([]byte(tree.MustBeDBytes(datums[2])))
+		hint, err := hintpb.FromBytes([]byte(tree.MustBeDBytes(datums[2])))
 		if err != nil {
 			return err
 		}
@@ -543,7 +541,7 @@ type cacheEntry struct {
 	// be duplicate entries in the fingerprints slice.
 	// TODO(drewk): consider de-duplicating the fingerprint strings to reduce
 	// memory usage.
-	hints        []hintpb.StatementHint
+	hints        []hintpb.StatementHintUnion
 	fingerprints []string
 	ids          []int64
 }
@@ -552,7 +550,7 @@ type cacheEntry struct {
 // fingerprint, or nil if they don't exist. The results are in order of row ID.
 func (entry *cacheEntry) getMatchingHints(
 	statementFingerprint string,
-) (hints []hintpb.StatementHint, ids []int64) {
+) (hints []hintpb.StatementHintUnion, ids []int64) {
 	for i := range entry.hints {
 		if entry.fingerprints[i] == statementFingerprint {
 			hints = append(hints, entry.hints[i])

pkg/sql/hints/hint_cache_test.go

@@ -539,11 +539,13 @@ func checkIDOrder(t *testing.T, ids []int64) {
 	}
 }
 
-// insertStatementHint inserts an empty statement hint into the
+// insertStatementHint inserts a random statement hint into the
 // system.statement_hints table.
 func insertStatementHint(t *testing.T, r *sqlutils.SQLRunner, fingerprint string) {
-	emptyHint := &hintpb.StatementHintUnion{}
-	hintBytes, err := emptyHint.ToBytes()
+	// TODO(drewk,michae2): randomly choose the hint type once we support others.
+	var hint hintpb.StatementHintUnion
+	hint.SetValue(&hintpb.InjectHints{})
+	hintBytes, err := hintpb.ToBytes(hint)
 	require.NoError(t, err)
 	const insertStmt = `INSERT INTO system.statement_hints ("fingerprint", "hint") VALUES ($1, $2)`
 	r.Exec(t, insertStmt, fingerprint, hintBytes)