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  }