Skip to content

Commit a3e810d

Browse files
DrewKimballDrewKimball
committed
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
1 parent d78b3ad commit a3e810d

File tree

4 files changed

+246
-66
lines changed

4 files changed

+246
-66
lines changed

‎pkg/sql/hints/BUILD.bazel‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,10 @@ load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
22

33
go_library(
44
name = "hints",
5-
srcs = ["hint_cache.go"],
5+
srcs = [
6+
"hint_cache.go",
7+
"hint_table.go",
8+
],
69
importpath = "github.com/cockroachdb/cockroach/pkg/sql/hints",
710
visibility = ["//visibility:public"],
811
deps = [
@@ -20,6 +23,7 @@ go_library(
2023
"//pkg/sql/catalog/descs",
2124
"//pkg/sql/catalog/systemschema",
2225
"//pkg/sql/hintpb",
26+
"//pkg/sql/isql",
2327
"//pkg/sql/rowenc",
2428
"//pkg/sql/sem/tree",
2529
"//pkg/sql/sessiondata",
@@ -42,6 +46,7 @@ go_test(
4246
size = "medium",
4347
srcs = [
4448
"hint_cache_test.go",
49+
"hint_table_test.go",
4550
"main_test.go",
4651
],
4752
exec_properties = select({
@@ -58,6 +63,7 @@ go_test(
5863
"//pkg/sql/catalog",
5964
"//pkg/sql/catalog/descs",
6065
"//pkg/sql/hintpb",
66+
"//pkg/sql/isql",
6167
"//pkg/sql/randgen",
6268
"//pkg/sql/stats",
6369
"//pkg/testutils",

‎pkg/sql/hints/hint_cache.go‎

Lines changed: 3 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,6 @@ import (
2828
"github.com/cockroachdb/cockroach/pkg/sql/hintpb"
2929
"github.com/cockroachdb/cockroach/pkg/sql/rowenc"
3030
"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
31-
"github.com/cockroachdb/cockroach/pkg/sql/sessiondata"
3231
"github.com/cockroachdb/cockroach/pkg/sql/types"
3332
"github.com/cockroachdb/cockroach/pkg/util/buildutil"
3433
"github.com/cockroachdb/cockroach/pkg/util/cache"
@@ -291,7 +290,7 @@ func (c *StatementHintsCache) checkHashHasHintsAsync(
291290
})
292291
for {
293292
log.VEventf(ctx, 1, "checking hints for fingerprint hash %d @ %v", hash, refreshTS)
294-
hasHints, err := c.checkForStatementHintsInDB(ctx, hash)
293+
hasHints, err := CheckForStatementHintsInDB(ctx, c.db.Executor(), hash)
295294
if err != nil {
296295
log.Dev.Warningf(ctx, "failed to check hints for hash %d: %v", hash, err)
297296
if retryOnError.Next() {
@@ -323,27 +322,6 @@ func (c *StatementHintsCache) checkHashHasHintsAsync(
323322
})
324323
}
325324

326-
// checkForStatementHintsInDB queries the system.statement_hints table to
327-
// determine if there are any hints for the given fingerprint hash. The caller
328-
// must be able to retry if an error is returned.
329-
func (c *StatementHintsCache) checkForStatementHintsInDB(
330-
ctx context.Context, statementHash int64,
331-
) (hasHints bool, retErr error) {
332-
const opName = "get-plan-hints"
333-
const getHintsStmt = `SELECT hash FROM system.statement_hints WHERE "hash" = $1 LIMIT 1`
334-
it, err := c.db.Executor().QueryIteratorEx(
335-
ctx, opName, nil /* txn */, sessiondata.NodeUserSessionDataOverride,
336-
getHintsStmt, statementHash,
337-
)
338-
if err != nil {
339-
return false, err
340-
}
341-
defer func() {
342-
retErr = errors.CombineErrors(retErr, it.Close())
343-
}()
344-
return it.Next(ctx)
345-
}
346-
347325
// decodeHashFromStatementHintsKey decodes the query hash from a range feed
348326
// event on system.statement_hints.
349327
func decodeHashFromStatementHintsKey(
@@ -466,7 +444,8 @@ func (c *StatementHintsCache) addCacheEntryLocked(
466444
c.mu.Unlock()
467445
defer c.mu.Lock()
468446
log.VEventf(ctx, 1, "reading hints for query %s", statementFingerprint)
469-
err = c.getStatementHintsFromDB(ctx, statementHash, entry)
447+
entry.ids, entry.fingerprints, entry.hints, err =
448+
GetStatementHintsFromDB(ctx, c.db.Executor(), statementHash)
470449
log.VEventf(ctx, 1, "finished reading hints for query %s", statementFingerprint)
471450
}()
472451

@@ -483,47 +462,6 @@ func (c *StatementHintsCache) addCacheEntryLocked(
483462
return entry.getMatchingHints(statementFingerprint)
484463
}
485464

486-
// getStatementHintsFromDB queries the system.statement_hints table for hints
487-
// matching the given fingerprint hash. It is able to handle the case when
488-
// multiple fingerprints match the hash, as well as the case when there are no
489-
// hints for the fingerprint. Results are ordered by row ID.
490-
func (c *StatementHintsCache) getStatementHintsFromDB(
491-
ctx context.Context, statementHash int64, entry *cacheEntry,
492-
) (retErr error) {
493-
const opName = "get-plan-hints"
494-
const getHintsStmt = `
495-
SELECT "row_id", "fingerprint", "hint"
496-
FROM system.statement_hints
497-
WHERE "hash" = $1
498-
ORDER BY "row_id" ASC`
499-
it, err := c.db.Executor().QueryIteratorEx(
500-
ctx, opName, nil /* txn */, sessiondata.NodeUserSessionDataOverride,
501-
getHintsStmt, statementHash,
502-
)
503-
if err != nil {
504-
return err
505-
}
506-
defer func() {
507-
retErr = errors.CombineErrors(retErr, it.Close())
508-
}()
509-
for {
510-
ok, err := it.Next(ctx)
511-
if !ok || err != nil {
512-
return err
513-
}
514-
datums := it.Cur()
515-
rowID := int64(tree.MustBeDInt(datums[0]))
516-
fingerprint := string(tree.MustBeDString(datums[1]))
517-
hint, err := hintpb.FromBytes([]byte(tree.MustBeDBytes(datums[2])))
518-
if err != nil {
519-
return err
520-
}
521-
entry.ids = append(entry.ids, rowID)
522-
entry.fingerprints = append(entry.fingerprints, fingerprint)
523-
entry.hints = append(entry.hints, hint)
524-
}
525-
}
526-
527465
type cacheEntry struct {
528466
// If mustWait is true, we do not have any hints for this hash, and we are in
529467
// the process of fetching them from the database. Other callers can wait on

‎pkg/sql/hints/hint_table.go‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
1+
// Copyright 2025 The Cockroach Authors.
2+
//
3+
// Use of this software is governed by the CockroachDB Software License
4+
// included in the /LICENSE file.
5+
6+
package hints
7+
8+
import (
9+
"context"
10+
11+
"github.com/cockroachdb/cockroach/pkg/sql/hintpb"
12+
"github.com/cockroachdb/cockroach/pkg/sql/isql"
13+
"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
14+
"github.com/cockroachdb/cockroach/pkg/sql/sessiondata"
15+
"github.com/cockroachdb/errors"
16+
)
17+
18+
// CheckForStatementHintsInDB queries the system.statement_hints table to
19+
// determine if there are any hints for the given fingerprint hash. The caller
20+
// must be able to retry if an error is returned.
21+
func CheckForStatementHintsInDB(
22+
ctx context.Context, ex isql.Executor, statementHash int64,
23+
) (hasHints bool, retErr error) {
24+
const opName = "get-plan-hints"
25+
const getHintsStmt = `SELECT hash FROM system.statement_hints WHERE "hash" = $1 LIMIT 1`
26+
it, err := ex.QueryIteratorEx(
27+
ctx, opName, nil /* txn */, sessiondata.NodeUserSessionDataOverride,
28+
getHintsStmt, statementHash,
29+
)
30+
if err != nil {
31+
return false, err
32+
}
33+
defer func() {
34+
retErr = errors.CombineErrors(retErr, it.Close())
35+
}()
36+
return it.Next(ctx)
37+
}
38+
39+
// GetStatementHintsFromDB queries the system.statement_hints table for hints
40+
// matching the given fingerprint hash. It is able to handle the case when
41+
// multiple fingerprints match the hash, as well as the case when there are no
42+
// hints for the fingerprint.
43+
//
44+
// The returned slices (hints, fingerprints, and hintIDs) have the same length.
45+
// fingerprints[i] is the statement fingerprint to which hints[i] applies, while
46+
// hintIDs[i] uniquely identifies a hint in the system table. The results are in
47+
// order of hint ID.
48+
func GetStatementHintsFromDB(
49+
ctx context.Context, ex isql.Executor, statementHash int64,
50+
) (hintIDs []int64, fingerprints []string, hints []hintpb.StatementHintUnion, retErr error) {
51+
const opName = "get-plan-hints"
52+
const getHintsStmt = `
53+
SELECT "row_id", "fingerprint", "hint"
54+
FROM system.statement_hints
55+
WHERE "hash" = $1
56+
ORDER BY "row_id" ASC`
57+
it, err := ex.QueryIteratorEx(
58+
ctx, opName, nil /* txn */, sessiondata.NodeUserSessionDataOverride,
59+
getHintsStmt, statementHash,
60+
)
61+
if err != nil {
62+
return nil, nil, nil, err
63+
}
64+
defer func() {
65+
retErr = errors.CombineErrors(retErr, it.Close())
66+
}()
67+
for {
68+
ok, err := it.Next(ctx)
69+
if err != nil {
70+
return nil, nil, nil, err
71+
}
72+
if !ok {
73+
break
74+
}
75+
datums := it.Cur()
76+
hintIDs = append(hintIDs, int64(tree.MustBeDInt(datums[0])))
77+
fingerprints = append(fingerprints, string(tree.MustBeDString(datums[1])))
78+
hint, err := hintpb.FromBytes([]byte(tree.MustBeDBytes(datums[2])))
79+
if err != nil {
80+
return nil, nil, nil, err
81+
}
82+
hints = append(hints, hint)
83+
}
84+
return hintIDs, fingerprints, hints, nil
85+
}
86+
87+
// InsertHintIntoDB inserts a statement hint into the system.statement_hints
88+
// table. It returns the hint ID of the newly inserted hint if successful.
89+
func InsertHintIntoDB(
90+
ctx context.Context, txn isql.Txn, fingerprint string, hint hintpb.StatementHintUnion,
91+
) (int64, error) {
92+
const opName = "insert-statement-hint"
93+
hintBytes, err := hintpb.ToBytes(hint)
94+
if err != nil {
95+
return 0, err
96+
}
97+
const insertStmt = `INSERT INTO system.statement_hints ("fingerprint", "hint") VALUES ($1, $2) RETURNING "row_id"`
98+
row, err := txn.QueryRowEx(
99+
ctx, opName, txn.KV(), sessiondata.NodeUserSessionDataOverride,
100+
insertStmt, fingerprint, hintBytes,
101+
)
102+
if err != nil {
103+
return 0, err
104+
}
105+
return int64(tree.MustBeDInt(row[0])), nil
106+
}

‎pkg/sql/hints/hint_table_test.go‎

Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
// Copyright 2025 The Cockroach Authors.
2+
//
3+
// Use of this software is governed by the CockroachDB Software License
4+
// included in the /LICENSE file.
5+
6+
package hints_test
7+
8+
import (
9+
"context"
10+
"testing"
11+
12+
"github.com/cockroachdb/cockroach/pkg/base"
13+
"github.com/cockroachdb/cockroach/pkg/sql/catalog/descs"
14+
"github.com/cockroachdb/cockroach/pkg/sql/hintpb"
15+
"github.com/cockroachdb/cockroach/pkg/sql/hints"
16+
"github.com/cockroachdb/cockroach/pkg/sql/isql"
17+
"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
18+
"github.com/cockroachdb/cockroach/pkg/util/leaktest"
19+
"github.com/cockroachdb/cockroach/pkg/util/log"
20+
"github.com/stretchr/testify/require"
21+
)
22+
23+
// TestHintTableOperations tests the DB-interfacing functions in hint_table.go:
24+
// CheckForStatementHintsInDB, GetStatementHintsFromDB, and InsertHintIntoDB.
25+
func TestHintTableOperations(t *testing.T) {
26+
defer leaktest.AfterTest(t)()
27+
defer log.Scope(t).Close(t)
28+
29+
ctx := context.Background()
30+
srv := serverutils.StartServerOnly(t, base.TestServerArgs{})
31+
defer srv.Stopper().Stop(ctx)
32+
ts := srv.ApplicationLayer()
33+
db := ts.InternalDB().(descs.DB)
34+
ex := db.Executor()
35+
36+
// Create test hints.
37+
fingerprint1 := "SELECT a FROM t WHERE b = $1"
38+
fingerprint2 := "SELECT c FROM t WHERE d = $2"
39+
hash1 := computeHash(t, fingerprint1)
40+
hash2 := computeHash(t, fingerprint2)
41+
42+
var hint1, hint2 hintpb.StatementHintUnion
43+
hint1.SetValue(&hintpb.InjectHints{DonorSQL: "SELECT a FROM t@t_b_idx WHERE b = $1"})
44+
hint2.SetValue(&hintpb.InjectHints{DonorSQL: "SELECT c FROM t@{NO_FULL_SCAN} WHERE d = $2"})
45+
46+
// Check for nonexistent hints.
47+
hasHints, err := hints.CheckForStatementHintsInDB(ctx, ex, hash1)
48+
require.NoError(t, err)
49+
require.False(t, hasHints)
50+
51+
// Retrieve nonexistent hints.
52+
hintIDs, fingerprints, hintsFromDB, err := hints.GetStatementHintsFromDB(ctx, ex, hash1)
53+
require.NoError(t, err)
54+
require.Empty(t, hintIDs)
55+
require.Empty(t, fingerprints)
56+
require.Empty(t, hintsFromDB)
57+
58+
// Insert a hint.
59+
var insertedHintID1 int64
60+
err = db.Txn(ctx, func(ctx context.Context, txn isql.Txn) error {
61+
insertedHintID1, err = hints.InsertHintIntoDB(ctx, txn, fingerprint1, hint1)
62+
return err
63+
})
64+
require.NoError(t, err)
65+
require.Greater(t, insertedHintID1, int64(0)) // Should return a valid ID.
66+
67+
// Check for the inserted hint.
68+
hasHints, err = hints.CheckForStatementHintsInDB(ctx, ex, hash1)
69+
require.NoError(t, err)
70+
require.True(t, hasHints)
71+
72+
// Fetch the inserted hint.
73+
hintIDs, fingerprints, hintsFromDB, err = hints.GetStatementHintsFromDB(ctx, ex, hash1)
74+
require.NoError(t, err)
75+
require.Len(t, hintIDs, 1)
76+
require.Len(t, fingerprints, 1)
77+
require.Len(t, hintsFromDB, 1)
78+
require.Equal(t, fingerprint1, fingerprints[0])
79+
require.Equal(t, hint1, hintsFromDB[0])
80+
require.Equal(t, insertedHintID1, hintIDs[0])
81+
82+
// Insert multiple hints for the same fingerprint.
83+
var insertedHintID2 int64
84+
err = db.Txn(ctx, func(ctx context.Context, txn isql.Txn) error {
85+
insertedHintID2, err = hints.InsertHintIntoDB(ctx, txn, fingerprint1, hint1)
86+
return err
87+
})
88+
require.NoError(t, err)
89+
require.Greater(t, insertedHintID2, int64(0))
90+
require.NotEqual(t, insertedHintID1, insertedHintID2)
91+
92+
// Fetch all hints for the fingerprint.
93+
hintIDs, fingerprints, hintsFromDB, err = hints.GetStatementHintsFromDB(ctx, ex, hash1)
94+
require.NoError(t, err)
95+
require.Len(t, hintIDs, 2)
96+
require.Len(t, fingerprints, 2)
97+
require.Len(t, hintsFromDB, 2)
98+
require.Equal(t, fingerprint1, fingerprints[0])
99+
require.Equal(t, fingerprint1, fingerprints[1])
100+
require.Less(t, hintIDs[0], hintIDs[1])
101+
102+
// Insert hint for different fingerprint.
103+
var insertedHintID3 int64
104+
err = db.Txn(ctx, func(ctx context.Context, txn isql.Txn) error {
105+
insertedHintID3, err = hints.InsertHintIntoDB(ctx, txn, fingerprint2, hint2)
106+
return err
107+
})
108+
require.NoError(t, err)
109+
require.Greater(t, insertedHintID3, int64(0))
110+
111+
// Retrieve hint for the new fingerprint.
112+
hintIDs2, fingerprints2, hintsFromDB2, err := hints.GetStatementHintsFromDB(ctx, ex, hash2)
113+
require.NoError(t, err)
114+
require.Len(t, hintIDs2, 1)
115+
require.Len(t, fingerprints2, 1)
116+
require.Len(t, hintsFromDB2, 1)
117+
require.Equal(t, fingerprint2, fingerprints2[0])
118+
require.Equal(t, insertedHintID3, hintIDs2[0])
119+
120+
// Test InsertHintIntoDB with empty fingerprint and hint.
121+
var emptyFingerprintHintID int64
122+
var hintEmpty hintpb.StatementHintUnion
123+
hintEmpty.SetValue(&hintpb.InjectHints{})
124+
err = db.Txn(ctx, func(ctx context.Context, txn isql.Txn) error {
125+
emptyFingerprintHintID, err = hints.InsertHintIntoDB(ctx, txn, "", hintEmpty)
126+
return err
127+
})
128+
require.NoError(t, err)
129+
require.Greater(t, emptyFingerprintHintID, int64(0))
130+
}

0 commit comments

Comments
 (0)