sql/hints: extract hints system table interfacing logic

This commit extracts helper functions used by the `hints` package to
interface with the `system.statement_hints` table and adds tests for
the helpers.

Epic: None

Release note: None

Author: DrewKimball

pkg/sql/hints/BUILD.bazel

@@ -2,7 +2,10 @@ load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 
 go_library(
     name = "hints",
-    srcs = ["hint_cache.go"],
+    srcs = [
+        "hint_cache.go",
+        "hint_table.go",
+    ],
     importpath = "github.com/cockroachdb/cockroach/pkg/sql/hints",
     visibility = ["//visibility:public"],
     deps = [
@@ -20,6 +23,7 @@ go_library(
         "//pkg/sql/catalog/descs",
         "//pkg/sql/catalog/systemschema",
         "//pkg/sql/hintpb",
+        "//pkg/sql/isql",
         "//pkg/sql/rowenc",
         "//pkg/sql/sem/tree",
         "//pkg/sql/sessiondata",
@@ -42,6 +46,7 @@ go_test(
     size = "medium",
     srcs = [
         "hint_cache_test.go",
+        "hint_table_test.go",
         "main_test.go",
     ],
     exec_properties = select({
@@ -58,6 +63,7 @@ go_test(
         "//pkg/sql/catalog",
         "//pkg/sql/catalog/descs",
         "//pkg/sql/hintpb",
+        "//pkg/sql/isql",
         "//pkg/sql/randgen",
         "//pkg/sql/stats",
         "//pkg/testutils",

pkg/sql/hints/hint_cache.go

@@ -28,7 +28,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/hintpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/rowenc"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
-	"github.com/cockroachdb/cockroach/pkg/sql/sessiondata"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
 	"github.com/cockroachdb/cockroach/pkg/util/buildutil"
 	"github.com/cockroachdb/cockroach/pkg/util/cache"
@@ -291,7 +290,7 @@ func (c *StatementHintsCache) checkHashHasHintsAsync(
 		})
 		for {
 			log.VEventf(ctx, 1, "checking hints for fingerprint hash %d @ %v", hash, refreshTS)
-			hasHints, err := c.checkForStatementHintsInDB(ctx, hash)
+			hasHints, err := CheckForStatementHintsInDB(ctx, c.db.Executor(), hash)
 			if err != nil {
 				log.Dev.Warningf(ctx, "failed to check hints for hash %d: %v", hash, err)
 				if retryOnError.Next() {
@@ -323,27 +322,6 @@ func (c *StatementHintsCache) checkHashHasHintsAsync(
 	})
 }
 
-// checkForStatementHintsInDB queries the system.statement_hints table to
-// determine if there are any hints for the given fingerprint hash. The caller
-// must be able to retry if an error is returned.
-func (c *StatementHintsCache) checkForStatementHintsInDB(
-	ctx context.Context, statementHash int64,
-) (hasHints bool, retErr error) {
-	const opName = "get-plan-hints"
-	const getHintsStmt = `SELECT hash FROM system.statement_hints WHERE "hash" = $1 LIMIT 1`
-	it, err := c.db.Executor().QueryIteratorEx(
-		ctx, opName, nil /* txn */, sessiondata.NodeUserSessionDataOverride,
-		getHintsStmt, statementHash,
-	)
-	if err != nil {
-		return false, err
-	}
-	defer func() {
-		retErr = errors.CombineErrors(retErr, it.Close())
-	}()
-	return it.Next(ctx)
-}
-
 // decodeHashFromStatementHintsKey decodes the query hash from a range feed
 // event on system.statement_hints.
 func decodeHashFromStatementHintsKey(
@@ -466,7 +444,8 @@ func (c *StatementHintsCache) addCacheEntryLocked(
 		c.mu.Unlock()
 		defer c.mu.Lock()
 		log.VEventf(ctx, 1, "reading hints for query %s", statementFingerprint)
-		err = c.getStatementHintsFromDB(ctx, statementHash, entry)
+		entry.ids, entry.fingerprints, entry.hints, err =
+			GetStatementHintsFromDB(ctx, c.db.Executor(), statementHash)
 		log.VEventf(ctx, 1, "finished reading hints for query %s", statementFingerprint)
 	}()
 
@@ -483,47 +462,6 @@ func (c *StatementHintsCache) addCacheEntryLocked(
 	return entry.getMatchingHints(statementFingerprint)
 }
 
-// getStatementHintsFromDB queries the system.statement_hints table for hints
-// matching the given fingerprint hash. It is able to handle the case when
-// multiple fingerprints match the hash, as well as the case when there are no
-// hints for the fingerprint. Results are ordered by row ID.
-func (c *StatementHintsCache) getStatementHintsFromDB(
-	ctx context.Context, statementHash int64, entry *cacheEntry,
-) (retErr error) {
-	const opName = "get-plan-hints"
-	const getHintsStmt = `
-    SELECT "row_id", "fingerprint", "hint"
-    FROM system.statement_hints
-    WHERE "hash" = $1
-    ORDER BY "row_id" ASC`
-	it, err := c.db.Executor().QueryIteratorEx(
-		ctx, opName, nil /* txn */, sessiondata.NodeUserSessionDataOverride,
-		getHintsStmt, statementHash,
-	)
-	if err != nil {
-		return err
-	}
-	defer func() {
-		retErr = errors.CombineErrors(retErr, it.Close())
-	}()
-	for {
-		ok, err := it.Next(ctx)
-		if !ok || err != nil {
-			return err
-		}
-		datums := it.Cur()
-		rowID := int64(tree.MustBeDInt(datums[0]))
-		fingerprint := string(tree.MustBeDString(datums[1]))
-		hint, err := hintpb.FromBytes([]byte(tree.MustBeDBytes(datums[2])))
-		if err != nil {
-			return err
-		}
-		entry.ids = append(entry.ids, rowID)
-		entry.fingerprints = append(entry.fingerprints, fingerprint)
-		entry.hints = append(entry.hints, hint)
-	}
-}
-
 type cacheEntry struct {
 	// If mustWait is true, we do not have any hints for this hash, and we are in
 	// the process of fetching them from the database. Other callers can wait on

pkg/sql/hints/hint_table.go

@@ -0,0 +1,106 @@
+// Copyright 2025 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package hints
+
+import (
+	"context"
+
+	"github.com/cockroachdb/cockroach/pkg/sql/hintpb"
+	"github.com/cockroachdb/cockroach/pkg/sql/isql"
+	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
+	"github.com/cockroachdb/cockroach/pkg/sql/sessiondata"
+	"github.com/cockroachdb/errors"
+)
+
+// CheckForStatementHintsInDB queries the system.statement_hints table to
+// determine if there are any hints for the given fingerprint hash. The caller
+// must be able to retry if an error is returned.
+func CheckForStatementHintsInDB(
+	ctx context.Context, ex isql.Executor, statementHash int64,
+) (hasHints bool, retErr error) {
+	const opName = "get-plan-hints"
+	const getHintsStmt = `SELECT hash FROM system.statement_hints WHERE "hash" = $1 LIMIT 1`
+	it, err := ex.QueryIteratorEx(
+		ctx, opName, nil /* txn */, sessiondata.NodeUserSessionDataOverride,
+		getHintsStmt, statementHash,
+	)
+	if err != nil {
+		return false, err
+	}
+	defer func() {
+		retErr = errors.CombineErrors(retErr, it.Close())
+	}()
+	return it.Next(ctx)
+}
+
+// GetStatementHintsFromDB queries the system.statement_hints table for hints
+// matching the given fingerprint hash. It is able to handle the case when
+// multiple fingerprints match the hash, as well as the case when there are no
+// hints for the fingerprint.
+//
+// The returned slices (hints, fingerprints, and hintIDs) have the same length.
+// fingerprints[i] is the statement fingerprint to which hints[i] applies, while
+// hintIDs[i] uniquely identifies a hint in the system table. The results are in
+// order of hint ID.
+func GetStatementHintsFromDB(
+	ctx context.Context, ex isql.Executor, statementHash int64,
+) (hintIDs []int64, fingerprints []string, hints []hintpb.StatementHintUnion, retErr error) {
+	const opName = "get-plan-hints"
+	const getHintsStmt = `
+    SELECT "row_id", "fingerprint", "hint"
+    FROM system.statement_hints
+    WHERE "hash" = $1
+    ORDER BY "row_id" ASC`
+	it, err := ex.QueryIteratorEx(
+		ctx, opName, nil /* txn */, sessiondata.NodeUserSessionDataOverride,
+		getHintsStmt, statementHash,
+	)
+	if err != nil {
+		return nil, nil, nil, err
+	}
+	defer func() {
+		retErr = errors.CombineErrors(retErr, it.Close())
+	}()
+	for {
+		ok, err := it.Next(ctx)
+		if err != nil {
+			return nil, nil, nil, err
+		}
+		if !ok {
+			break
+		}
+		datums := it.Cur()
+		hintIDs = append(hintIDs, int64(tree.MustBeDInt(datums[0])))
+		fingerprints = append(fingerprints, string(tree.MustBeDString(datums[1])))
+		hint, err := hintpb.FromBytes([]byte(tree.MustBeDBytes(datums[2])))
+		if err != nil {
+			return nil, nil, nil, err
+		}
+		hints = append(hints, hint)
+	}
+	return hintIDs, fingerprints, hints, nil
+}
+
+// InsertHintIntoDB inserts a statement hint into the system.statement_hints
+// table. It returns the hint ID of the newly inserted hint if successful.
+func InsertHintIntoDB(
+	ctx context.Context, txn isql.Txn, fingerprint string, hint hintpb.StatementHintUnion,
+) (int64, error) {
+	const opName = "insert-statement-hint"
+	hintBytes, err := hintpb.ToBytes(hint)
+	if err != nil {
+		return 0, err
+	}
+	const insertStmt = `INSERT INTO system.statement_hints ("fingerprint", "hint") VALUES ($1, $2) RETURNING "row_id"`
+	row, err := txn.QueryRowEx(
+		ctx, opName, txn.KV(), sessiondata.NodeUserSessionDataOverride,
+		insertStmt, fingerprint, hintBytes,
+	)
+	if err != nil {
+		return 0, err
+	}
+	return int64(tree.MustBeDInt(row[0])), nil
+}

pkg/sql/hints/hint_table_test.go

@@ -0,0 +1,130 @@
+// Copyright 2025 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package hints_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/cockroachdb/cockroach/pkg/base"
+	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descs"
+	"github.com/cockroachdb/cockroach/pkg/sql/hintpb"
+	"github.com/cockroachdb/cockroach/pkg/sql/hints"
+	"github.com/cockroachdb/cockroach/pkg/sql/isql"
+	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
+	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
+	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/stretchr/testify/require"
+)
+
+// TestHintTableOperations tests the DB-interfacing functions in hint_table.go:
+// CheckForStatementHintsInDB, GetStatementHintsFromDB, and InsertHintIntoDB.
+func TestHintTableOperations(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+	srv := serverutils.StartServerOnly(t, base.TestServerArgs{})
+	defer srv.Stopper().Stop(ctx)
+	ts := srv.ApplicationLayer()
+	db := ts.InternalDB().(descs.DB)
+	ex := db.Executor()
+
+	// Create test hints.
+	fingerprint1 := "SELECT a FROM t WHERE b = $1"
+	fingerprint2 := "SELECT c FROM t WHERE d = $2"
+	hash1 := computeHash(t, fingerprint1)
+	hash2 := computeHash(t, fingerprint2)
+
+	var hint1, hint2 hintpb.StatementHintUnion
+	hint1.SetValue(&hintpb.InjectHints{DonorSQL: "SELECT a FROM t@t_b_idx WHERE b = $1"})
+	hint2.SetValue(&hintpb.InjectHints{DonorSQL: "SELECT c FROM t@{NO_FULL_SCAN} WHERE d = $2"})
+
+	// Check for nonexistent hints.
+	hasHints, err := hints.CheckForStatementHintsInDB(ctx, ex, hash1)
+	require.NoError(t, err)
+	require.False(t, hasHints)
+
+	// Retrieve nonexistent hints.
+	hintIDs, fingerprints, hintsFromDB, err := hints.GetStatementHintsFromDB(ctx, ex, hash1)
+	require.NoError(t, err)
+	require.Empty(t, hintIDs)
+	require.Empty(t, fingerprints)
+	require.Empty(t, hintsFromDB)
+
+	// Insert a hint.
+	var insertedHintID1 int64
+	err = db.Txn(ctx, func(ctx context.Context, txn isql.Txn) error {
+		insertedHintID1, err = hints.InsertHintIntoDB(ctx, txn, fingerprint1, hint1)
+		return err
+	})
+	require.NoError(t, err)
+	require.Greater(t, insertedHintID1, int64(0)) // Should return a valid ID.
+
+	// Check for the inserted hint.
+	hasHints, err = hints.CheckForStatementHintsInDB(ctx, ex, hash1)
+	require.NoError(t, err)
+	require.True(t, hasHints)
+
+	// Fetch the inserted hint.
+	hintIDs, fingerprints, hintsFromDB, err = hints.GetStatementHintsFromDB(ctx, ex, hash1)
+	require.NoError(t, err)
+	require.Len(t, hintIDs, 1)
+	require.Len(t, fingerprints, 1)
+	require.Len(t, hintsFromDB, 1)
+	require.Equal(t, fingerprint1, fingerprints[0])
+	require.Equal(t, hint1, hintsFromDB[0])
+	require.Equal(t, insertedHintID1, hintIDs[0])
+
+	// Insert multiple hints for the same fingerprint.
+	var insertedHintID2 int64
+	err = db.Txn(ctx, func(ctx context.Context, txn isql.Txn) error {
+		insertedHintID2, err = hints.InsertHintIntoDB(ctx, txn, fingerprint1, hint1)
+		return err
+	})
+	require.NoError(t, err)
+	require.Greater(t, insertedHintID2, int64(0))
+	require.NotEqual(t, insertedHintID1, insertedHintID2)
+
+	// Fetch all hints for the fingerprint.
+	hintIDs, fingerprints, hintsFromDB, err = hints.GetStatementHintsFromDB(ctx, ex, hash1)
+	require.NoError(t, err)
+	require.Len(t, hintIDs, 2)
+	require.Len(t, fingerprints, 2)
+	require.Len(t, hintsFromDB, 2)
+	require.Equal(t, fingerprint1, fingerprints[0])
+	require.Equal(t, fingerprint1, fingerprints[1])
+	require.Less(t, hintIDs[0], hintIDs[1])
+
+	// Insert hint for different fingerprint.
+	var insertedHintID3 int64
+	err = db.Txn(ctx, func(ctx context.Context, txn isql.Txn) error {
+		insertedHintID3, err = hints.InsertHintIntoDB(ctx, txn, fingerprint2, hint2)
+		return err
+	})
+	require.NoError(t, err)
+	require.Greater(t, insertedHintID3, int64(0))
+
+	// Retrieve hint for the new fingerprint.
+	hintIDs2, fingerprints2, hintsFromDB2, err := hints.GetStatementHintsFromDB(ctx, ex, hash2)
+	require.NoError(t, err)
+	require.Len(t, hintIDs2, 1)
+	require.Len(t, fingerprints2, 1)
+	require.Len(t, hintsFromDB2, 1)
+	require.Equal(t, fingerprint2, fingerprints2[0])
+	require.Equal(t, insertedHintID3, hintIDs2[0])
+
+	// Test InsertHintIntoDB with empty fingerprint and hint.
+	var emptyFingerprintHintID int64
+	var hintEmpty hintpb.StatementHintUnion
+	hintEmpty.SetValue(&hintpb.InjectHints{})
+	err = db.Txn(ctx, func(ctx context.Context, txn isql.Txn) error {
+		emptyFingerprintHintID, err = hints.InsertHintIntoDB(ctx, txn, "", hintEmpty)
+		return err
+	})
+	require.NoError(t, err)
+	require.Greater(t, emptyFingerprintHintID, int64(0))
+}