vitess.io/vitess@v0.16.2/go/vt/vtgate/planbuilder/simplifier_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/vterrors"
    24  
    25  	"github.com/stretchr/testify/assert"
    26  
    27  	"vitess.io/vitess/go/vt/vtgate/simplifier"
    28  
    29  	"github.com/stretchr/testify/require"
    30  
    31  	"vitess.io/vitess/go/vt/log"
    32  
    33  	"vitess.io/vitess/go/vt/sqlparser"
    34  )
    35  
    36  // TestSimplifyBuggyQuery should be used to whenever we get a planner bug reported
    37  // It will try to minimize the query to make it easier to understand and work with the bug.
    38  func TestSimplifyBuggyQuery(t *testing.T) {
    39  	query := "(select id from unsharded union select id from unsharded_auto) union (select id from user union select name from unsharded)"
    40  	vschema := &vschemaWrapper{
    41  		v:       loadSchema(t, "vschemas/schema.json", true),
    42  		version: Gen4,
    43  	}
    44  	stmt, reserved, err := sqlparser.Parse2(query)
    45  	require.NoError(t, err)
    46  	rewritten, _ := sqlparser.RewriteAST(sqlparser.CloneStatement(stmt), vschema.currentDb(), sqlparser.SQLSelectLimitUnset, "", nil, nil)
    47  	reservedVars := sqlparser.NewReservedVars("vtg", reserved)
    48  
    49  	simplified := simplifier.SimplifyStatement(
    50  		stmt.(sqlparser.SelectStatement),
    51  		vschema.currentDb(),
    52  		vschema,
    53  		keepSameError(query, reservedVars, vschema, rewritten.BindVarNeeds),
    54  	)
    55  
    56  	fmt.Println(sqlparser.String(simplified))
    57  }
    58  
    59  func TestSimplifyPanic(t *testing.T) {
    60  	t.Skip("not needed to run")
    61  	query := "(select id from unsharded union select id from unsharded_auto) union (select id from unsharded_auto union select name from unsharded)"
    62  	vschema := &vschemaWrapper{
    63  		v:       loadSchema(t, "vschemas/schema.json", true),
    64  		version: Gen4,
    65  	}
    66  	stmt, reserved, err := sqlparser.Parse2(query)
    67  	require.NoError(t, err)
    68  	rewritten, _ := sqlparser.RewriteAST(sqlparser.CloneStatement(stmt), vschema.currentDb(), sqlparser.SQLSelectLimitUnset, "", nil, nil)
    69  	reservedVars := sqlparser.NewReservedVars("vtg", reserved)
    70  
    71  	simplified := simplifier.SimplifyStatement(
    72  		stmt.(sqlparser.SelectStatement),
    73  		vschema.currentDb(),
    74  		vschema,
    75  		keepPanicking(query, reservedVars, vschema, rewritten.BindVarNeeds),
    76  	)
    77  
    78  	fmt.Println(sqlparser.String(simplified))
    79  }
    80  
    81  func TestUnsupportedFile(t *testing.T) {
    82  	t.Skip("run manually to see if any queries can be simplified")
    83  	vschema := &vschemaWrapper{
    84  		v:       loadSchema(t, "vschemas/schema.json", true),
    85  		version: Gen4,
    86  	}
    87  	fmt.Println(vschema)
    88  	for _, tcase := range readJSONTests("unsupported_cases.txt") {
    89  		t.Run(tcase.Query, func(t *testing.T) {
    90  			log.Errorf("unsupported_cases.txt - %s", tcase.Query)
    91  			stmt, reserved, err := sqlparser.Parse2(tcase.Query)
    92  			require.NoError(t, err)
    93  			_, ok := stmt.(sqlparser.SelectStatement)
    94  			if !ok {
    95  				t.Skip()
    96  				return
    97  			}
    98  			rewritten, err := sqlparser.RewriteAST(stmt, vschema.currentDb(), sqlparser.SQLSelectLimitUnset, "", nil, nil)
    99  			if err != nil {
   100  				t.Skip()
   101  			}
   102  			vschema.currentDb()
   103  
   104  			reservedVars := sqlparser.NewReservedVars("vtg", reserved)
   105  			ast := rewritten.AST
   106  			origQuery := sqlparser.String(ast)
   107  			stmt, _, _ = sqlparser.Parse2(tcase.Query)
   108  			simplified := simplifier.SimplifyStatement(
   109  				stmt.(sqlparser.SelectStatement),
   110  				vschema.currentDb(),
   111  				vschema,
   112  				keepSameError(tcase.Query, reservedVars, vschema, rewritten.BindVarNeeds),
   113  			)
   114  
   115  			if simplified == nil {
   116  				t.Skip()
   117  			}
   118  
   119  			simpleQuery := sqlparser.String(simplified)
   120  			fmt.Println(simpleQuery)
   121  
   122  			assert.Equal(t, origQuery, simpleQuery)
   123  		})
   124  	}
   125  }
   126  
   127  func keepSameError(query string, reservedVars *sqlparser.ReservedVars, vschema *vschemaWrapper, needs *sqlparser.BindVarNeeds) func(statement sqlparser.SelectStatement) bool {
   128  	stmt, _, err := sqlparser.Parse2(query)
   129  	if err != nil {
   130  		panic(err)
   131  	}
   132  	rewritten, _ := sqlparser.RewriteAST(stmt, vschema.currentDb(), sqlparser.SQLSelectLimitUnset, "", nil, nil)
   133  	ast := rewritten.AST
   134  	_, expected := BuildFromStmt(query, ast, reservedVars, vschema, rewritten.BindVarNeeds, true, true)
   135  	if expected == nil {
   136  		panic("query does not fail to plan")
   137  	}
   138  	return func(statement sqlparser.SelectStatement) bool {
   139  		_, myErr := BuildFromStmt(query, statement, reservedVars, vschema, needs, true, true)
   140  		if myErr == nil {
   141  			return false
   142  		}
   143  		state := vterrors.ErrState(expected)
   144  		if state == vterrors.Undefined {
   145  			return expected.Error() == myErr.Error()
   146  		}
   147  		return vterrors.ErrState(myErr) == state
   148  	}
   149  }
   150  
   151  func keepPanicking(query string, reservedVars *sqlparser.ReservedVars, vschema *vschemaWrapper, needs *sqlparser.BindVarNeeds) func(statement sqlparser.SelectStatement) bool {
   152  	cmp := func(statement sqlparser.SelectStatement) (res bool) {
   153  		defer func() {
   154  			r := recover()
   155  			if r != nil {
   156  				log.Errorf("panicked with %v", r)
   157  				res = true
   158  			}
   159  		}()
   160  		log.Errorf("trying %s", sqlparser.String(statement))
   161  		_, _ = BuildFromStmt(query, statement, reservedVars, vschema, needs, true, true)
   162  		log.Errorf("did not panic")
   163  
   164  		return false
   165  	}
   166  
   167  	stmt, _, err := sqlparser.Parse2(query)
   168  	if err != nil {
   169  		panic(err.Error())
   170  	}
   171  	if !cmp(stmt.(sqlparser.SelectStatement)) {
   172  		panic("query is not panicking")
   173  	}
   174  
   175  	return cmp
   176  }