Skip to content

Commit d90b48f

Browse files
craig[bot]michae2
craig[bot]
and
michae2
committed
Merge #157107
157107: tree: handle statement fingerprints as HintInjectionDonors r=mgartner,DrewKimball a=michae2 Make changes to HintInjectionDonor in order to support using statement fingerprints as the donor. (See individual commits for details.) Informs: #153633 Release note: None Co-authored-by: Michael Erickson <[email protected]>
2 parents 5c4f3ee + 9578f83 commit d90b48f

File tree

4 files changed

+496
-104
lines changed

4 files changed

+496
-104
lines changed

‎pkg/sql/sem/tree/inject_hints.go‎

Lines changed: 216 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -8,12 +8,12 @@ package tree
88
import (
99
"reflect"
1010

11-
"github.com/cockroachdb/cockroach/pkg/settings"
1211
"github.com/cockroachdb/errors"
1312
)
1413

1514
// HintInjectionDonor holds the donor statement that provides the hints to be
16-
// injected into a target statement.
15+
// injected into a target statement. The donor statement could be a regular SQL
16+
// statement or a statement fingerprint.
1717
type HintInjectionDonor struct {
1818
// ast is the donor statement AST after parsing.
1919
ast Statement
@@ -23,63 +23,177 @@ type HintInjectionDonor struct {
2323
// statements match.
2424
validationSQL string
2525

26-
// walk is a pre-order traversal of the AST, holding Expr and TableExpr. As we
27-
// walk the target AST we expect to find matching nodes in the same order.
26+
// walk is a pre-order traversal of the AST, holding Expr, TableExpr, and
27+
// Statement nodes. As we walk the target AST we expect to find matching nodes
28+
// in the same order. Some nodes are ignored, such as ParenExpr and nodes
29+
// after collapse of a Tuple or Array or ValuesClause.
2830
walk []any
2931
}
3032

31-
// validationFmtFlags returns the FmtFlags used to check that the donor
32-
// statement and the target statement match. These should be the same FmtFlags
33-
// used for statement fingerprints (including
34-
// sql.stats.statement_fingerprint.format_mask) with FmtHideHints added.
35-
func validationFmtFlags(sv *settings.Values) FmtFlags {
36-
stmtFingerprintFmtMask := FmtHideConstants | FmtFlags(QueryFormattingForFingerprintsMask.Get(sv))
37-
return stmtFingerprintFmtMask | FmtHideHints
33+
// NewHintInjectionDonor creates a HintInjectionDonor from a parsed AST. The
34+
// parsed donor statement could be a regular SQL statement or a statement
35+
// fingerprint.
36+
//
37+
// After NewHintInjectionDonor returns a HintInjectionDonor, the donor becomes
38+
// read-only to allow concurrent use from multiple goroutines.
39+
func NewHintInjectionDonor(ast Statement, fingerprintFlags FmtFlags) (*HintInjectionDonor, error) {
40+
hd := &HintInjectionDonor{
41+
ast: ast,
42+
validationSQL: FormatStatementHideConstants(ast, fingerprintFlags, FmtHideHints),
43+
}
44+
v := newDonorVisitor{hd: hd, collapsed: []bool{false}}
45+
WalkStmt(&v, ast)
46+
return hd, nil
3847
}
3948

40-
// NewHintInjectionDonor creates a HintInjectionDonor from a parsed AST.
41-
func NewHintInjectionDonor(ast Statement, sv *settings.Values) (*HintInjectionDonor, error) {
42-
donor := &HintInjectionDonor{
43-
ast: ast,
44-
validationSQL: AsStringWithFlags(ast, validationFmtFlags(sv)),
49+
// newDonorVisitor is an ExtendedVisitor used to walk the donor AST and build
50+
// the ahead-of-time pre-order traversal at donor creation time.
51+
type newDonorVisitor struct {
52+
// hd is the hint donor being initialized.
53+
hd *HintInjectionDonor
54+
55+
// collapsed is a stack indicating whether we've collapsed the rest of the
56+
// current ValuesClause, Tuple, or Array list in the donor AST.
57+
collapsed []bool
58+
}
59+
60+
var _ ExtendedVisitor = &newDonorVisitor{}
61+
62+
func (v *newDonorVisitor) VisitPre(expr Expr) (recurse bool, newExpr Expr) {
63+
if v.collapsed[len(v.collapsed)-1] {
64+
return false, expr
65+
}
66+
switch expr.(type) {
67+
case *ParenExpr:
68+
// We sometimes wrap scalar expressions in extra ParenExpr when
69+
// printing. Skip over ParenExpr to match in more cases.
70+
return true, expr
71+
case *Tuple, *Array:
72+
// If the donor was a statement fingerprint, we might have collapse a Tuple
73+
// or Array into __more__. Push on a boolean tracking whether we have
74+
// collapsed it.
75+
v.collapsed = append(v.collapsed, false)
76+
case *UnresolvedName:
77+
if isArityIndicatorString(expr.String()) {
78+
// If we find __more__ or another arity indicator, treat the rest of the
79+
// current Tuple or Array or ValuesClause as collapsed, and skip forward
80+
// until we VisitPost it.
81+
v.collapsed[len(v.collapsed)-1] = true
82+
v.hd.walk = append(v.hd.walk, expr)
83+
return false, expr
84+
}
85+
}
86+
v.hd.walk = append(v.hd.walk, expr)
87+
return true, expr
88+
}
89+
90+
func (v *newDonorVisitor) VisitPost(expr Expr) (newNode Expr) {
91+
switch expr.(type) {
92+
case *Tuple, *Array:
93+
// Pop off the boolean tracking whether we collapsed the Tuple or Array.
94+
v.collapsed = v.collapsed[:len(v.collapsed)-1]
95+
}
96+
return expr
97+
}
98+
99+
func (v *newDonorVisitor) VisitTablePre(expr TableExpr) (recurse bool, newExpr TableExpr) {
100+
if v.collapsed[len(v.collapsed)-1] {
101+
return false, expr
45102
}
46-
if _, err := ExtendedSimpleStmtVisit(
47-
ast,
48-
func(expr Expr) (bool, Expr, error) {
49-
donor.walk = append(donor.walk, expr)
50-
return true, expr, nil
51-
},
52-
func(expr TableExpr) (bool, TableExpr, error) {
53-
donor.walk = append(donor.walk, expr)
54-
return true, expr, nil
55-
},
56-
); err != nil {
57-
return nil, err
103+
v.hd.walk = append(v.hd.walk, expr)
104+
return true, expr
105+
}
106+
107+
func (v *newDonorVisitor) VisitTablePost(expr TableExpr) (newNode TableExpr) {
108+
return expr
109+
}
110+
111+
func (v *newDonorVisitor) VisitStatementPre(expr Statement) (recurse bool, newNode Statement) {
112+
if v.collapsed[len(v.collapsed)-1] {
113+
return false, expr
114+
}
115+
switch expr.(type) {
116+
case *ValuesClause:
117+
// If the donor was a statement fingerprint, we might have collapse a
118+
// ValuesClause into __more__. Push on a boolean tracking whether we have
119+
// collapsed it.
120+
v.collapsed = append(v.collapsed, false)
121+
}
122+
v.hd.walk = append(v.hd.walk, expr)
123+
return true, expr
124+
}
125+
126+
func (v *newDonorVisitor) VisitStatementPost(expr Statement) (newNode Statement) {
127+
switch expr.(type) {
128+
case *ValuesClause:
129+
// Pop off the boolean tracking whether we collapsed the ValuesClause.
130+
v.collapsed = v.collapsed[:len(v.collapsed)-1]
58131
}
59-
return donor, nil
132+
return expr
60133
}
61134

135+
// hintInjectionVisitor is an ExtendedVisitor used to walk the target AST at
136+
// injection time. It performs the rewrite of the target AST.
62137
type hintInjectionVisitor struct {
63-
donor *HintInjectionDonor
138+
// hd is the hint donor. Note that multiple hintInjectionVisitors could share
139+
// the same donor concurrently from different sessions, so this is read-only.
140+
hd *HintInjectionDonor
141+
142+
// walkIdx is the current position within the donor AST walk.
64143
walkIdx int
65-
err error
144+
145+
// collapsed is a stack indicating whether we've collapsed the rest of the
146+
// current ValuesClause, Tuple, or Array list in the donor AST.
147+
collapsed []bool
148+
149+
// err is set if we detect a mismatch while walking the target AST.
150+
err error
66151
}
67152

153+
var _ ExtendedVisitor = &hintInjectionVisitor{}
154+
68155
func (v *hintInjectionVisitor) VisitPre(expr Expr) (recurse bool, newExpr Expr) {
69-
if v.err != nil {
156+
if v.err != nil || v.collapsed[len(v.collapsed)-1] {
70157
return false, expr
71158
}
72159

73-
if v.walkIdx >= len(v.donor.walk) {
160+
switch expr.(type) {
161+
case *ParenExpr:
162+
// Skip ParenExpr to match the donor walk.
163+
return true, expr
164+
}
165+
166+
if v.walkIdx >= len(v.hd.walk) {
74167
v.err = errors.Newf(
75168
"hint injection donor statement missing AST node corresponding to AST node %v", expr,
76169
)
77170
return false, expr
78171
}
79172

80-
donorExpr := v.donor.walk[v.walkIdx]
173+
donorExpr := v.hd.walk[v.walkIdx]
81174
v.walkIdx += 1
82175

176+
switch donor := donorExpr.(type) {
177+
case *Tuple, *Array:
178+
// If the donor was a statement fingerprint, we might have collapse a Tuple
179+
// or Array into __more__. Push on a boolean tracking whether we have
180+
// collapsed it.
181+
v.collapsed = append(v.collapsed, false)
182+
case *UnresolvedName:
183+
// If the donor was a fingerprint, then we might not exactly match _ or
184+
// __more__ or __more10_100__ etc in the donor. Treat these as wildcards
185+
// matching any subtree and don't recurse any further into the target AST.
186+
if donor.String() == string(StmtFingerprintPlaceholder) {
187+
return false, expr
188+
}
189+
if isArityIndicatorString(donor.String()) {
190+
// And skip the rest of the current Tuple or Array or ValuesClause if it
191+
// has been collapsed to __more__ or one of the other arity indicators.
192+
v.collapsed[len(v.collapsed)-1] = true
193+
return false, expr
194+
}
195+
}
196+
83197
// Check that the corresponding donor Expr matches this node.
84198
if reflect.TypeOf(expr) != reflect.TypeOf(donorExpr) {
85199
v.err = errors.Newf(
@@ -90,27 +204,37 @@ func (v *hintInjectionVisitor) VisitPre(expr Expr) (recurse bool, newExpr Expr)
90204
return true, expr
91205
}
92206

207+
func (v *hintInjectionVisitor) VisitPost(expr Expr) Expr {
208+
switch expr.(type) {
209+
case *Tuple, *Array:
210+
// Pop off the boolean tracking whether we collapsed the Tuple or Array.
211+
v.collapsed = v.collapsed[:len(v.collapsed)-1]
212+
}
213+
return expr
214+
}
215+
93216
func (v *hintInjectionVisitor) VisitTablePre(expr TableExpr) (recurse bool, newExpr TableExpr) {
94-
if v.err != nil {
217+
if v.err != nil || v.collapsed[len(v.collapsed)-1] {
95218
return false, expr
96219
}
97220

98-
if v.walkIdx >= len(v.donor.walk) {
221+
if v.walkIdx >= len(v.hd.walk) {
99222
v.err = errors.Newf(
100223
"hint injection donor statement missing AST node corresponding to AST node", expr,
101224
)
102225
return false, expr
103226
}
104227

105-
donorExpr := v.donor.walk[v.walkIdx]
228+
donorExpr := v.hd.walk[v.walkIdx]
106229
v.walkIdx += 1
107230

108-
// Copy hints from the corresponding donor TableExpr.
231+
// Remove any existing hints from the original TableExpr, and copy hints from
232+
// the corresponding donor TableExpr.
109233
switch donor := donorExpr.(type) {
110234
case *AliasedTableExpr:
111235
switch orig := expr.(type) {
112236
case *AliasedTableExpr:
113-
if donor.IndexFlags != nil {
237+
if !donor.IndexFlags.Equal(orig.IndexFlags) {
114238
// Create a new node with any pre-existing inline hints replaced with
115239
// hints from the donor.
116240
newNode := *orig
@@ -132,10 +256,11 @@ func (v *hintInjectionVisitor) VisitTablePre(expr TableExpr) (recurse bool, newE
132256
case *JoinTableExpr:
133257
switch orig := expr.(type) {
134258
case *JoinTableExpr:
135-
if donor.Hint != "" {
259+
if donor.JoinType != orig.JoinType || donor.Hint != orig.Hint {
136260
// Create a new node with any pre-existing inline hints replaced with
137261
// hints from the donor.
138262
newNode := *orig
263+
newNode.JoinType = donor.JoinType
139264
newNode.Hint = donor.Hint
140265
return true, &newNode
141266
}
@@ -153,15 +278,58 @@ func (v *hintInjectionVisitor) VisitTablePre(expr TableExpr) (recurse bool, newE
153278
return true, expr
154279
}
155280

156-
func (v *hintInjectionVisitor) VisitPost(expr Expr) Expr { return expr }
157-
func (v *hintInjectionVisitor) VisitTablePost(expr TableExpr) TableExpr { return expr }
281+
func (v *hintInjectionVisitor) VisitTablePost(expr TableExpr) TableExpr {
282+
return expr
283+
}
158284

159-
var _ ExtendedVisitor = &hintInjectionVisitor{}
285+
func (v *hintInjectionVisitor) VisitStatementPre(expr Statement) (recurse bool, newExpr Statement) {
286+
if v.err != nil || v.collapsed[len(v.collapsed)-1] {
287+
return false, expr
288+
}
289+
290+
if v.walkIdx >= len(v.hd.walk) {
291+
v.err = errors.Newf(
292+
"hint injection donor statement missing AST node corresponding to AST node %v", expr,
293+
)
294+
return false, expr
295+
}
296+
297+
donorExpr := v.hd.walk[v.walkIdx]
298+
v.walkIdx += 1
299+
300+
switch donorExpr.(type) {
301+
case *ValuesClause:
302+
// If the donor was a statement fingerprint, we might have collapse a
303+
// ValuesClause into __more__. Push on a boolean tracking whether we have
304+
// collapsed it.
305+
v.collapsed = append(v.collapsed, false)
306+
}
307+
308+
// Check that the corresponding donor Expr matches this node.
309+
if reflect.TypeOf(expr) != reflect.TypeOf(donorExpr) {
310+
v.err = errors.Newf(
311+
"hint injection donor statement AST node %v did not match AST node %v", donorExpr, expr,
312+
)
313+
return false, expr
314+
}
315+
return true, expr
316+
}
317+
318+
func (v *hintInjectionVisitor) VisitStatementPost(expr Statement) Statement {
319+
switch expr.(type) {
320+
case *ValuesClause:
321+
// Pop off the boolean tracking whether we collapsed the ValuesClause.
322+
v.collapsed = v.collapsed[:len(v.collapsed)-1]
323+
}
324+
return expr
325+
}
160326

161327
// Validate checks that the target statement exactly matches the donor (except
162328
// for hints).
163-
func (hd *HintInjectionDonor) Validate(stmt Statement, sv *settings.Values) error {
164-
sql := AsStringWithFlags(stmt, validationFmtFlags(sv))
329+
//
330+
// It is safe to call Validate concurrently from multiple goroutines.
331+
func (hd *HintInjectionDonor) Validate(stmt Statement, fingerprintFlags FmtFlags) error {
332+
sql := FormatStatementHideConstants(stmt, fingerprintFlags, FmtHideHints)
165333
if sql != hd.validationSQL {
166334
return errors.Newf(
167335
"statement does not match hint donor statement: %v vs %v", sql, hd.validationSQL,
@@ -183,9 +351,11 @@ func (hd *HintInjectionDonor) Validate(stmt Statement, sv *settings.Values) erro
183351
//
184352
// No semantic checks of the donor hints are performed. It is assumed that
185353
// semantic checks happen elsewhere.
354+
//
355+
// It is safe to call InjectHints concurrently from multiple goroutines.
186356
func (hd *HintInjectionDonor) InjectHints(stmt Statement) (Statement, bool, error) {
187357
// Walk the given statement, copying hints from the hints donor.
188-
v := hintInjectionVisitor{donor: hd}
358+
v := hintInjectionVisitor{hd: hd, collapsed: []bool{false}}
189359
newStmt, changed := WalkStmt(&v, stmt)
190360
if v.err != nil {
191361
return stmt, false, v.err

0 commit comments

Comments
 (0)