github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/opt/optbuilder/builder_test.go (about)

     1  // Copyright 2018 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package optbuilder_test
    12  
    13  import (
    14  	"context"
    15  	"fmt"
    16  	"strings"
    17  	"testing"
    18  
    19  	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
    20  	"github.com/cockroachdb/cockroach/pkg/sql/opt/memo"
    21  	"github.com/cockroachdb/cockroach/pkg/sql/opt/optbuilder"
    22  	"github.com/cockroachdb/cockroach/pkg/sql/opt/optgen/exprgen"
    23  	"github.com/cockroachdb/cockroach/pkg/sql/opt/testutils/opttester"
    24  	"github.com/cockroachdb/cockroach/pkg/sql/opt/testutils/testcat"
    25  	"github.com/cockroachdb/cockroach/pkg/sql/opt/xform"
    26  	"github.com/cockroachdb/cockroach/pkg/sql/parser"
    27  	_ "github.com/cockroachdb/cockroach/pkg/sql/sem/builtins"
    28  	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
    29  	"github.com/cockroachdb/cockroach/pkg/sql/types"
    30  	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
    31  	"github.com/cockroachdb/datadriven"
    32  )
    33  
    34  // TestBuilder runs data-driven testcases of the form
    35  //   <command> [<args>]...
    36  //   <SQL statement or expression>
    37  //   ----
    38  //   <expected results>
    39  //
    40  // See OptTester.Handle for supported commands. In addition to those, we
    41  // support:
    42  //
    43  //  - build-scalar [args]
    44  //
    45  //    Builds a memo structure from a SQL scalar expression and outputs a
    46  //    representation of the "expression view" of the memo structure.
    47  //
    48  //    The supported args (in addition to the ones supported by OptTester):
    49  //
    50  //      - vars=(type1,type2,...)
    51  //
    52  //        Information about IndexedVar columns.
    53  //
    54  func TestBuilder(t *testing.T) {
    55  	defer leaktest.AfterTest(t)()
    56  
    57  	datadriven.Walk(t, "testdata", func(t *testing.T, path string) {
    58  		catalog := testcat.New()
    59  
    60  		datadriven.RunTest(t, path, func(t *testing.T, d *datadriven.TestData) string {
    61  			var varTypes []*types.T
    62  			var err error
    63  
    64  			tester := opttester.New(catalog, d.Input)
    65  			tester.Flags.ExprFormat = memo.ExprFmtHideMiscProps |
    66  				memo.ExprFmtHideConstraints |
    67  				memo.ExprFmtHideFuncDeps |
    68  				memo.ExprFmtHideRuleProps |
    69  				memo.ExprFmtHideStats |
    70  				memo.ExprFmtHideCost |
    71  				memo.ExprFmtHideQualifications |
    72  				memo.ExprFmtHideScalars |
    73  				memo.ExprFmtHideTypes
    74  
    75  			switch d.Cmd {
    76  			case "build-scalar":
    77  				// Remove the HideScalars, HideTypes flag for build-scalars.
    78  				tester.Flags.ExprFormat &= ^(memo.ExprFmtHideScalars | memo.ExprFmtHideTypes)
    79  				for _, arg := range d.CmdArgs {
    80  					key, vals := arg.Key, arg.Vals
    81  					switch key {
    82  					case "vars":
    83  						varTypes, err = exprgen.ParseTypes(vals)
    84  						if err != nil {
    85  							d.Fatalf(t, "%v", err)
    86  						}
    87  
    88  					default:
    89  						if err := tester.Flags.Set(arg); err != nil {
    90  							d.Fatalf(t, "%s", err)
    91  						}
    92  					}
    93  				}
    94  
    95  				expr, err := parser.ParseExpr(d.Input)
    96  				if err != nil {
    97  					d.Fatalf(t, "%v", err)
    98  				}
    99  
   100  				ctx := context.Background()
   101  				semaCtx := tree.MakeSemaContext()
   102  				evalCtx := tree.MakeTestingEvalContext(cluster.MakeTestingClusterSettings())
   103  				evalCtx.SessionData.OptimizerFKChecks = true
   104  				evalCtx.SessionData.OptimizerFKCascades = true
   105  				evalCtx.SessionData.OptimizerUseHistograms = true
   106  				evalCtx.SessionData.OptimizerUseMultiColStats = true
   107  
   108  				var o xform.Optimizer
   109  				o.Init(&evalCtx, catalog)
   110  				for i, typ := range varTypes {
   111  					o.Memo().Metadata().AddColumn(fmt.Sprintf("@%d", i+1), typ)
   112  				}
   113  				// Disable normalization rules: we want the tests to check the result
   114  				// of the build process.
   115  				o.DisableOptimizations()
   116  				b := optbuilder.NewScalar(ctx, &semaCtx, &evalCtx, o.Factory())
   117  				err = b.Build(expr)
   118  				if err != nil {
   119  					return fmt.Sprintf("error: %s\n", strings.TrimSpace(err.Error()))
   120  				}
   121  				f := memo.MakeExprFmtCtx(tester.Flags.ExprFormat, o.Memo(), catalog)
   122  				f.FormatExpr(o.Memo().RootExpr())
   123  				return f.Buffer.String()
   124  
   125  			default:
   126  				return tester.RunCommand(t, d)
   127  			}
   128  		})
   129  	})
   130  }