Merge #158096

158096: sql/*: add hint injection r=DrewKimball a=michae2

**sql: add parseHint step when loading hint into hints cache**

When loading an external statement hint into the statement hints cache,
we might need to call some function to get the hint ready for use. (For
hint injections, this function is `tree.NewHintInjectionDonor` which
parses and walks the donor statement fingerprint.) This function could
fail, in which case we want to skip over the hint but not return an
error from `GetStatementHintsFromDB`. This function could succeed but
create some extra state which we need to save.

This commit adds a new `parseHint` step which calls any functions needed
to get the hint ready, and creates a new `hints.Hint` struct which holds
the object(s) created when parsing hints. (These are analogous to
`parseStats` and `TableStatistic` from the stats cache.)

Informs: #153633

Release note: None

---

**sql/\*: add hint injection**

1. During `ReloadHintsIfStale` we now call `Validate` and `InjectHints`
   using the donor to perform the AST rewrite. We save the rewritten AST
   in the statement separately from the original AST.

2. We wrap `prepareUsingOptimizer` and `makeOptimizerPlan` with
   functions that first try preparing / planning with injected hints,
   and then try again without injected hints in case the injected hints
   are invalid.

With these two pieces we can now actually perform hint injection.

Fixes: #153633

Release note (sql change): A new "hint injection" ability has been
added, which allows operators to dynamically inject inline hints into
statements, without modifying the text of those statements. Hints can be
injected using the builtin function `crdb_internal.inject_hint` with the
target statement fingerprint to rewrite. For example, to add an index
hint to the statement `SELECT * FROM my_table WHERE col = 3`, use:

```
SELECT crdb_internal.inject_hint(
  'SELECT * FROM my_table WHERE col = _',
  'SELECT * FROM my_table@my_table_col_idx WHERE col = _'
);
```

Whenever a statement is executed matching statement fingerprint
`SELECT * FROM my_table WHERE col = _`, it will first be rewritten
to include the injected index hint.

---

**sql/\*: invalidate cached memos after hint injection changes**

If we build a memo with hint injection, and then later we realize that
memo won't work (maybe because we discover the hint is unsatisfiable
during execution of a prepared statement) we need to invalidate the
cached memo.

To do this, add a usingHintInjection field which tells the memo
staleness check whether we're trying with or without hint injection.

Also, in a related but separate change, this commit adds all matching
HintIDs to the optimizer metadata so that we don't invalidate cached
memos if the hintsGeneration changed due to some unrelated statement
hints changing.

Informs: #153633

Release note: None

Co-authored-by: Michael Erickson <michae2@cockroachlabs.com>

Author: craig[bot]

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/parserutils",
         "//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",
@@ -59,11 +61,14 @@ go_test(
         "//pkg/security/securityassets",
         "//pkg/security/securitytest",
         "//pkg/server",
+        "//pkg/settings/cluster",
         "//pkg/sql",
         "//pkg/sql/catalog/descs",
         "//pkg/sql/hintpb",
         "//pkg/sql/isql",
+        "//pkg/sql/parserutils",
         "//pkg/sql/randgen",
+        "//pkg/sql/sem/tree",
         "//pkg/sql/stats",
         "//pkg/storage/fs",
         "//pkg/testutils",

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])