vitess.io/vitess@v0.16.2/go/vt/vtgate/planbuilder/fallback_planner_test.go (about) 1 /* 2 Copyright 2021 The Vitess Authors. 3 4 Licensed under the Apache License, Version 2.0 (the "License"); 5 you may not use this file except in compliance with the License. 6 You may obtain a copy of the License at 7 8 http://www.apache.org/licenses/LICENSE-2.0 9 10 Unless required by applicable law or agreed to in writing, software 11 distributed under the License is distributed on an "AS IS" BASIS, 12 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 See the License for the specific language governing permissions and 14 limitations under the License. 15 */ 16 17 package planbuilder 18 19 import ( 20 "fmt" 21 "testing" 22 23 "vitess.io/vitess/go/vt/vtgate/planbuilder/plancontext" 24 25 "github.com/stretchr/testify/assert" 26 27 "vitess.io/vitess/go/vt/sqlparser" 28 29 "vitess.io/vitess/go/vt/vtgate/engine" 30 ) 31 32 type testPlanner struct { 33 panic any 34 err error 35 res engine.Primitive 36 messWithAST func(sqlparser.Statement) 37 called bool 38 } 39 40 var _ stmtPlanner = (*testPlanner)(nil).plan 41 42 func (tp *testPlanner) plan(statement sqlparser.Statement, vars *sqlparser.ReservedVars, schema plancontext.VSchema) (*planResult, error) { 43 tp.called = true 44 if tp.panic != nil { 45 panic(tp.panic) 46 } 47 if tp.messWithAST != nil { 48 tp.messWithAST(statement) 49 } 50 return newPlanResult(tp.res), tp.err 51 } 52 53 func TestFallbackPlanner(t *testing.T) { 54 a := &testPlanner{} 55 b := &testPlanner{} 56 fb := &fallbackPlanner{ 57 primary: a.plan, 58 fallback: b.plan, 59 } 60 61 stmt := &sqlparser.Select{} 62 var vschema plancontext.VSchema 63 64 // first planner succeeds 65 _, _ = fb.plan(stmt, nil, vschema) 66 assert.True(t, a.called) 67 assert.False(t, b.called) 68 a.called = false 69 70 // first planner errors 71 a.err = fmt.Errorf("fail") 72 _, _ = fb.plan(stmt, nil, vschema) 73 assert.True(t, a.called) 74 assert.True(t, b.called) 75 76 a.called = false 77 b.called = false 78 79 // first planner panics 80 a.panic = "oh noes" 81 _, _ = fb.plan(stmt, nil, vschema) 82 assert.True(t, a.called) 83 assert.True(t, b.called) 84 } 85 86 func TestFallbackClonesBeforePlanning(t *testing.T) { 87 a := &testPlanner{ 88 messWithAST: func(statement sqlparser.Statement) { 89 sel := statement.(*sqlparser.Select) 90 sel.SelectExprs = nil 91 }, 92 } 93 b := &testPlanner{} 94 fb := &fallbackPlanner{ 95 primary: a.plan, 96 fallback: b.plan, 97 } 98 99 stmt := &sqlparser.Select{ 100 SelectExprs: sqlparser.SelectExprs{&sqlparser.StarExpr{}}, 101 } 102 var vschema plancontext.VSchema 103 104 // first planner succeeds 105 _, _ = fb.plan(stmt, nil, vschema) 106 107 assert.NotNilf(t, stmt.SelectExprs, "should not have changed") 108 }