Skip to content

Commit 8cc65a7

Browse files
michae2michae2
committed
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.
1 parent 6d09b87 commit 8cc65a7

File tree

8 files changed

+490
-29
lines changed

8 files changed

+490
-29
lines changed

‎pkg/sql/conn_executor_exec.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -565,6 +565,7 @@ func (ex *connExecutor) execStmtInOpenState(
565565
stmt.Hints = ps.Hints
566566
stmt.HintIDs = ps.HintIDs
567567
stmt.HintsGeneration = ps.HintsGeneration
568+
stmt.ASTWithInjectedHints = ps.ASTWithInjectedHints
568569
stmt.ReloadHintsIfStale(ctx, stmtFingerprintFmtMask, statementHintsCache)
569570
res.ResetStmtType(ps.AST)
570571

@@ -1465,6 +1466,7 @@ func (ex *connExecutor) execStmtInOpenStateWithPausablePortal(
14651466
vars.stmt.Hints = ps.Hints
14661467
vars.stmt.HintIDs = ps.HintIDs
14671468
vars.stmt.HintsGeneration = ps.HintsGeneration
1469+
vars.stmt.ASTWithInjectedHints = ps.ASTWithInjectedHints
14681470
vars.stmt.ReloadHintsIfStale(ctx, stmtFingerprintFmtMask, statementHintsCache)
14691471
res.ResetStmtType(ps.AST)
14701472

‎pkg/sql/conn_executor_prepare.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -247,6 +247,7 @@ func (ex *connExecutor) prepare(
247247
prepared.Hints = stmt.Hints
248248
prepared.HintIDs = stmt.HintIDs
249249
prepared.HintsGeneration = stmt.HintsGeneration
250+
prepared.ASTWithInjectedHints = stmt.ASTWithInjectedHints
250251

251252
// Point to the prepared state, which can be further populated during query
252253
// preparation.

‎pkg/sql/logictest/testdata/logic_test/statement_hint_builtins‎

Lines changed: 302 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -179,3 +179,305 @@ FROM crdb_internal.feature_usage
179179
WHERE feature_name = 'sql.session.statement-hints'
180180
----
181181
6
182+
183+
statement ok
184+
DEALLOCATE ALL
185+
186+
# Testcases for hint injection.
187+
188+
statement ok
189+
CREATE TABLE abc (a INT PRIMARY KEY, b INT, c INT, INDEX (b))
190+
191+
# Temporary, until next commit.
192+
statement ok
193+
SET CLUSTER SETTING sql.query_cache.enabled = off
194+
195+
# Try some simple hint injections. First, an index hint.
196+
197+
query T
198+
EXPLAIN SELECT a FROM abc WHERE a = 10
199+
----
200+
distribution: local
201+
vectorized: true
202+
·
203+
• scan
204+
missing stats
205+
table: abc@abc_pkey
206+
spans: [/10 - /10]
207+
208+
statement ok
209+
SELECT crdb_internal.inject_hint(
210+
'SELECT a FROM abc WHERE a = _',
211+
'SELECT a FROM abc@abc_b_idx WHERE a = _'
212+
)
213+
214+
statement ok
215+
SELECT crdb_internal.await_statement_hints_cache()
216+
217+
statement ok
218+
SET tracing = on
219+
220+
statement ok
221+
SELECT a FROM abc WHERE a = 5
222+
223+
statement ok
224+
SET tracing = off
225+
226+
query T
227+
SELECT regexp_replace(split_part(message, ': SELECT', 1), E'\\d+', 'x')
228+
FROM [SHOW TRACE FOR SESSION]
229+
WHERE message LIKE '%injected hints%'
230+
----
231+
injected hints from external statement hint x
232+
trying planning with injected hints
233+
234+
query T
235+
EXPLAIN SELECT a FROM abc WHERE a = 10
236+
----
237+
distribution: local
238+
vectorized: true
239+
statement hints count: 1
240+
·
241+
• filter
242+
│ filter: a = 10
243+
│
244+
└── • scan
245+
missing stats
246+
table: abc@abc_b_idx
247+
spans: FULL SCAN
248+
249+
# Try injecting a join hint.
250+
251+
query T
252+
EXPLAIN SELECT a, x FROM abc JOIN xy ON y = b WHERE a = 10
253+
----
254+
distribution: local
255+
vectorized: true
256+
·
257+
• lookup join
258+
│ table: xy@xy_y_idx
259+
│ equality: (b) = (y)
260+
│
261+
└── • scan
262+
missing stats
263+
table: abc@abc_pkey
264+
spans: [/10 - /10]
265+
266+
statement ok
267+
SELECT crdb_internal.inject_hint(
268+
'SELECT a, x FROM abc JOIN xy ON y = b WHERE a = _',
269+
'SELECT a, x FROM abc INNER HASH JOIN xy ON y = b WHERE a = _'
270+
)
271+
272+
statement ok
273+
SELECT crdb_internal.await_statement_hints_cache()
274+
275+
statement ok
276+
SET tracing = on
277+
278+
statement ok
279+
SELECT a, x FROM abc JOIN xy ON y = b WHERE a = 5
280+
281+
statement ok
282+
SET tracing = off
283+
284+
query T
285+
SELECT regexp_replace(split_part(message, ': SELECT', 1), E'\\d+', 'x')
286+
FROM [SHOW TRACE FOR SESSION]
287+
WHERE message LIKE '%injected hints%'
288+
----
289+
injected hints from external statement hint x
290+
trying planning with injected hints
291+
292+
query T
293+
EXPLAIN SELECT a, x FROM abc JOIN xy ON y = b WHERE a = 10
294+
----
295+
distribution: local
296+
vectorized: true
297+
statement hints count: 1
298+
·
299+
• hash join
300+
│ equality: (b) = (y)
301+
│ left cols are key
302+
│
303+
├── • scan
304+
│ missing stats
305+
│ table: abc@abc_pkey
306+
│ spans: [/10 - /10]
307+
│
308+
└── • scan
309+
missing stats
310+
table: xy@xy_pkey
311+
spans: FULL SCAN
312+
313+
# Try removing a hint.
314+
315+
statement ok
316+
SELECT crdb_internal.inject_hint(
317+
'SELECT a FROM abc@abc_pkey WHERE b = _',
318+
'SELECT a FROM abc WHERE b = _'
319+
)
320+
321+
statement ok
322+
SELECT crdb_internal.await_statement_hints_cache()
323+
324+
statement ok
325+
SET tracing = on
326+
327+
statement ok
328+
SELECT a FROM abc@abc_pkey WHERE b = 5
329+
330+
statement ok
331+
SET tracing = off
332+
333+
query T
334+
SELECT regexp_replace(split_part(message, ': SELECT', 1), E'\\d+', 'x')
335+
FROM [SHOW TRACE FOR SESSION]
336+
WHERE message LIKE '%injected hints%'
337+
----
338+
injected hints from external statement hint x
339+
trying planning with injected hints
340+
341+
query T
342+
EXPLAIN SELECT a FROM abc@abc_pkey WHERE b = 10
343+
----
344+
distribution: local
345+
vectorized: true
346+
statement hints count: 1
347+
·
348+
• scan
349+
missing stats
350+
table: abc@abc_b_idx
351+
spans: [/10 - /10]
352+
353+
# Check that we do not use an invalid injected index hint.
354+
355+
statement ok
356+
SELECT crdb_internal.inject_hint(
357+
'SELECT a + _ FROM abc WHERE a = _',
358+
'SELECT a + _ FROM abc@foo WHERE a = _'
359+
)
360+
361+
statement ok
362+
SELECT crdb_internal.await_statement_hints_cache()
363+
364+
statement ok
365+
SET tracing = on
366+
367+
statement ok
368+
SELECT a + 1 FROM abc WHERE a = 5
369+
370+
statement ok
371+
SET tracing = off
372+
373+
query T
374+
SELECT regexp_replace(split_part(message, ': SELECT', 1), E'\\d+', 'x')
375+
FROM [SHOW TRACE FOR SESSION]
376+
WHERE message LIKE '%injected hints%'
377+
----
378+
injected hints from external statement hint x
379+
trying planning with injected hints
380+
planning with injected hints failed with: index "foo" not found
381+
falling back to planning without injected hints
382+
383+
query T
384+
EXPLAIN SELECT a + 1 FROM abc WHERE a = 10
385+
----
386+
distribution: local
387+
vectorized: true
388+
statement hints count: 1
389+
·
390+
• render
391+
│
392+
└── • scan
393+
missing stats
394+
table: abc@abc_pkey
395+
spans: [/10 - /10]
396+
397+
# Check that we do not use an unsatisfiable injected hint.
398+
399+
statement ok
400+
SELECT crdb_internal.inject_hint(
401+
'SELECT c FROM xy JOIN abc ON c = y WHERE x = _',
402+
'SELECT c FROM xy INNER LOOKUP JOIN abc ON c = y WHERE x = _'
403+
)
404+
405+
statement ok
406+
SELECT crdb_internal.await_statement_hints_cache()
407+
408+
statement ok
409+
SET tracing = on
410+
411+
statement ok
412+
SELECT c FROM xy JOIN abc ON c = y WHERE x = 5
413+
414+
statement ok
415+
SET tracing = off
416+
417+
query T
418+
SELECT regexp_replace(split_part(message, ': SELECT', 1), E'\\d+', 'x')
419+
FROM [SHOW TRACE FOR SESSION]
420+
WHERE message LIKE '%injected hints%'
421+
----
422+
injected hints from external statement hint x
423+
trying planning with injected hints
424+
planning with injected hints failed with: could not produce a query plan conforming to the LOOKUP JOIN hint
425+
falling back to planning without injected hints
426+
427+
query T
428+
EXPLAIN SELECT c FROM xy JOIN abc ON c = y WHERE x = 10
429+
----
430+
distribution: local
431+
vectorized: true
432+
statement hints count: 1
433+
·
434+
• hash join
435+
│ equality: (c) = (y)
436+
│ right cols are key
437+
│
438+
├── • scan
439+
│ missing stats
440+
│ table: abc@abc_pkey
441+
│ spans: FULL SCAN
442+
│
443+
└── • scan
444+
missing stats
445+
table: xy@xy_pkey
446+
spans: [/10 - /10]
447+
448+
# Try a prepared statement with an injected hint.
449+
450+
statement ok
451+
SELECT crdb_internal.inject_hint(
452+
'SELECT c FROM abc WHERE b > _',
453+
'SELECT c FROM abc@{NO_INDEX_JOIN} WHERE b > _'
454+
)
455+
456+
statement ok
457+
SELECT crdb_internal.await_statement_hints_cache()
458+
459+
statement ok
460+
SET tracing = on
461+
462+
statement ok
463+
PREPARE p AS SELECT c FROM abc WHERE b > $1
464+
465+
statement ok
466+
EXECUTE p (5)
467+
468+
statement ok
469+
SET tracing = off
470+
471+
query T
472+
SELECT regexp_replace(split_part(message, ': SELECT', 1), E'\\d+', 'x')
473+
FROM [SHOW TRACE FOR SESSION]
474+
WHERE message LIKE '%injected hints%'
475+
----
476+
injected hints from external statement hint x
477+
injected hints from external statement hint x
478+
trying preparing with injected hints
479+
trying planning with injected hints
480+
481+
# Temporary, until next commit.
482+
statement ok
483+
RESET CLUSTER SETTING sql.query_cache.enabled

‎pkg/sql/opt/exec/execbuilder/testdata/explain‎

Lines changed: 17 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2594,10 +2594,13 @@ distribution: local
25942594
vectorized: true
25952595
statement hints count: 1
25962596
·
2597-
• scan
2598-
missing stats
2599-
table: t_hints@t_hints_pkey
2600-
spans: [/100 - ]
2597+
• filter
2598+
│ filter: k >= 100
2599+
│
2600+
└── • scan
2601+
missing stats
2602+
table: t_hints@t_hints_v_idx
2603+
spans: FULL SCAN
26012604

26022605
query T
26032606
EXPLAIN (VERBOSE) SELECT k FROM t_hints WHERE k >= 100
@@ -2606,8 +2609,13 @@ distribution: local
26062609
vectorized: true
26072610
statement hints count: 1
26082611
·
2609-
• scan
2610-
columns: (k)
2611-
estimated row count: 333 (missing stats)
2612-
table: t_hints@t_hints_pkey
2613-
spans: /100-
2612+
• filter
2613+
│ columns: (k)
2614+
│ estimated row count: 333 (missing stats)
2615+
│ filter: k >= 100
2616+
│
2617+
└── • scan
2618+
columns: (k)
2619+
estimated row count: 1,000 (missing stats)
2620+
table: t_hints@t_hints_v_idx
2621+
spans: FULL SCAN

0 commit comments

Comments
 (0)