sql/*: add hint injection

pkg/sql/conn_executor_exec.go

@@ -565,6 +565,7 @@ func (ex *connExecutor) execStmtInOpenState(
 		stmt.Hints = ps.Hints
 		stmt.HintIDs = ps.HintIDs
 		stmt.HintsGeneration = ps.HintsGeneration
+		stmt.ASTWithInjectedHints = ps.ASTWithInjectedHints
 		stmt.ReloadHintsIfStale(ctx, stmtFingerprintFmtMask, statementHintsCache)
 		res.ResetStmtType(ps.AST)
 
@@ -1465,6 +1466,7 @@ func (ex *connExecutor) execStmtInOpenStateWithPausablePortal(
 		vars.stmt.Hints = ps.Hints
 		vars.stmt.HintIDs = ps.HintIDs
 		vars.stmt.HintsGeneration = ps.HintsGeneration
+		vars.stmt.ASTWithInjectedHints = ps.ASTWithInjectedHints
 		vars.stmt.ReloadHintsIfStale(ctx, stmtFingerprintFmtMask, statementHintsCache)
 		res.ResetStmtType(ps.AST)
 

pkg/sql/conn_executor_prepare.go

@@ -247,6 +247,7 @@ func (ex *connExecutor) prepare(
 		prepared.Hints = stmt.Hints
 		prepared.HintIDs = stmt.HintIDs
 		prepared.HintsGeneration = stmt.HintsGeneration
+		prepared.ASTWithInjectedHints = stmt.ASTWithInjectedHints
 
 		// Point to the prepared state, which can be further populated during query
 		// preparation.

pkg/sql/faketreeeval/evalctx.go

@@ -603,6 +603,16 @@ func (ep *DummyEvalPlanner) InsertStatementHint(
 	return 0, nil
 }
 
+// UsingHintInjection is part of the eval.Planner interface.
+func (ep *DummyEvalPlanner) UsingHintInjection() bool {
+	return false
+}
+
+// GetHintIDs is part of the eval.Planner interface.
+func (ep *DummyEvalPlanner) GetHintIDs() []int64 {
+	return nil
+}
+
 // DummyPrivilegedAccessor implements the tree.PrivilegedAccessor interface by returning errors.
 type DummyPrivilegedAccessor struct{}
 

pkg/sql/hints/BUILD.bazel

@@ -24,12 +24,14 @@ go_library(
         "//pkg/sql/catalog/systemschema",
         "//pkg/sql/hintpb",
         "//pkg/sql/isql",
+        "//pkg/sql/parser",
         "//pkg/sql/rowenc",
         "//pkg/sql/sem/tree",
         "//pkg/sql/sessiondata",
         "//pkg/sql/types",
         "//pkg/util/buildutil",
         "//pkg/util/cache",
+        "//pkg/util/errorutil",
         "//pkg/util/hlc",
         "//pkg/util/log",
         "//pkg/util/metamorphic",

pkg/sql/hints/hint_cache.go

@@ -25,7 +25,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descs"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/systemschema"
-	"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/types"
@@ -406,8 +405,8 @@ func (c *StatementHintsCache) GetGeneration() int64 {
 // plans). It returns nil if the statement has no hints, or there was an error
 // retrieving them.
 func (c *StatementHintsCache) MaybeGetStatementHints(
-	ctx context.Context, statementFingerprint string,
-) (hints []hintpb.StatementHintUnion, ids []int64) {
+	ctx context.Context, statementFingerprint string, fingerprintFlags tree.FmtFlags,
+) (hints []Hint, ids []int64) {
 	hash := fnv.New64()
 	_, err := hash.Write([]byte(statementFingerprint))
 	if err != nil {
@@ -430,7 +429,7 @@ func (c *StatementHintsCache) MaybeGetStatementHints(
 	if !ok {
 		// The plan hints were evicted from the cache. Retrieve them from the
 		// database and add them to the cache.
-		return c.addCacheEntryLocked(ctx, statementHash, statementFingerprint)
+		return c.addCacheEntryLocked(ctx, statementHash, statementFingerprint, fingerprintFlags)
 	}
 	entry := e.(*cacheEntry)
 	c.maybeWaitForRefreshLocked(ctx, entry, statementHash)
@@ -464,8 +463,11 @@ func (c *StatementHintsCache) maybeWaitForRefreshLocked(
 // other queries to wait for the result via sync.Cond. Note that the lock is
 // released while reading from the db, and then reacquired.
 func (c *StatementHintsCache) addCacheEntryLocked(
-	ctx context.Context, statementHash int64, statementFingerprint string,
-) (hints []hintpb.StatementHintUnion, ids []int64) {
+	ctx context.Context,
+	statementHash int64,
+	statementFingerprint string,
+	fingerprintFlags tree.FmtFlags,
+) (hints []Hint, ids []int64) {
 	c.mu.AssertHeld()
 
 	// Add a cache entry that other queries can find and wait on until we have the
@@ -483,7 +485,7 @@ func (c *StatementHintsCache) addCacheEntryLocked(
 		defer c.mu.Lock()
 		log.VEventf(ctx, 1, "reading hints for query %s", statementFingerprint)
 		entry.ids, entry.fingerprints, entry.hints, err =
-			GetStatementHintsFromDB(ctx, c.db.Executor(), statementHash)
+			GetStatementHintsFromDB(ctx, c.db.Executor(), statementHash, fingerprintFlags)
 		log.VEventf(ctx, 1, "finished reading hints for query %s", statementFingerprint)
 	}()
 
@@ -517,16 +519,14 @@ type cacheEntry struct {
 	// be duplicate entries in the fingerprints slice.
 	// TODO(drewk): consider de-duplicating the fingerprint strings to reduce
 	// memory usage.
-	hints        []hintpb.StatementHintUnion
+	hints        []Hint
 	fingerprints []string
 	ids          []int64
 }
 
 // getMatchingHints returns the plan hints and row IDs for the given
 // fingerprint, or nil if they don't exist. The results are in order of row ID.
-func (entry *cacheEntry) getMatchingHints(
-	statementFingerprint string,
-) (hints []hintpb.StatementHintUnion, ids []int64) {
+func (entry *cacheEntry) getMatchingHints(statementFingerprint string) (hints []Hint, ids []int64) {
 	for i := range entry.hints {
 		if entry.fingerprints[i] == statementFingerprint {
 			hints = append(hints, entry.hints[i])

pkg/sql/hints/hint_table.go

@@ -10,11 +10,23 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/sql/hintpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/isql"
+	"github.com/cockroachdb/cockroach/pkg/sql/parser"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sessiondata"
+	"github.com/cockroachdb/cockroach/pkg/util/errorutil"
+	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/errors"
 )
 
+// Hint represents an unmarshaled hint that is ready to apply to statements.
+type Hint struct {
+	hintpb.StatementHintUnion
+
+	// HintInjectionDonor is the fully parsed donor statement fingerprint used for
+	// hint injection.
+	HintInjectionDonor *tree.HintInjectionDonor
+}
+
 // 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.
@@ -45,9 +57,13 @@ func CheckForStatementHintsInDB(
 // 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.
+//
+// GetStatementHintsFromDB will return an error if the query returns an
+// error. If one of the hints cannot be unmarshalled, the hint (and associated
+// fingerprint and ID) will be skipped but no error will be returned.
 func GetStatementHintsFromDB(
-	ctx context.Context, ex isql.Executor, statementHash int64,
-) (hintIDs []int64, fingerprints []string, hints []hintpb.StatementHintUnion, retErr error) {
+	ctx context.Context, ex isql.Executor, statementHash int64, fingerprintFlags tree.FmtFlags,
+) (hintIDs []int64, fingerprints []string, hints []Hint, retErr error) {
 	const opName = "get-plan-hints"
 	const getHintsStmt = `
     SELECT "row_id", "fingerprint", "hint"
@@ -72,18 +88,58 @@ func GetStatementHintsFromDB(
 		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])))
+		hintID, fingerprint, hint, err := parseHint(it.Cur(), fingerprintFlags)
 		if err != nil {
-			return nil, nil, nil, err
+			log.Dev.Warningf(
+				ctx, "could not decode hint ID %v for statement hash %v fingerprint %v: %v",
+				hintID, statementHash, fingerprint, err,
+			)
+			continue
 		}
+		hintIDs = append(hintIDs, hintID)
+		fingerprints = append(fingerprints, fingerprint)
 		hints = append(hints, hint)
 	}
 	return hintIDs, fingerprints, hints, nil
 }
 
+func parseHint(
+	datums tree.Datums, fingerprintFlags tree.FmtFlags,
+) (hintID int64, fingerprint string, hint Hint, err error) {
+	defer func() {
+		if r := recover(); r != nil {
+			// In the event of a "safe" panic, we only want to log the error and
+			// continue executing the query without this hint. This is only possible
+			// because the code does not update shared state and does not manipulate
+			// locks.
+			if ok, e := errorutil.ShouldCatch(r); ok {
+				err = e
+			} else {
+				// Other panic objects can't be considered "safe" and thus are
+				// propagated as crashes that terminate the session.
+				panic(r)
+			}
+		}
+	}()
+	hintID = int64(tree.MustBeDInt(datums[0]))
+	fingerprint = string(tree.MustBeDString(datums[1]))
+	hint.StatementHintUnion, err = hintpb.FromBytes([]byte(tree.MustBeDBytes(datums[2])))
+	if err != nil {
+		return hintID, fingerprint, Hint{}, err
+	}
+	if hint.InjectHints != nil && hint.InjectHints.DonorSQL != "" {
+		donorStmt, err := parser.ParseOne(hint.InjectHints.DonorSQL)
+		if err != nil {
+			return hintID, fingerprint, Hint{}, err
+		}
+		hint.HintInjectionDonor, err = tree.NewHintInjectionDonor(donorStmt.AST, fingerprintFlags)
+		if err != nil {
+			return hintID, fingerprint, Hint{}, err
+		}
+	}
+	return hintID, fingerprint, hint, 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(
@@ -107,3 +163,8 @@ func InsertHintIntoDB(
 	// local node's cache.
 	return int64(tree.MustBeDInt(row[0])), nil
 }
+
+func (hint *Hint) Size() int {
+	// TODO(michae2): add size of HintInjectionDonor
+	return hint.StatementHintUnion.Size()
+}