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

     1  // Copyright 2019 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 memo_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"
    21  	"github.com/cockroachdb/cockroach/pkg/sql/opt/memo"
    22  	"github.com/cockroachdb/cockroach/pkg/sql/opt/optbuilder"
    23  	"github.com/cockroachdb/cockroach/pkg/sql/opt/optgen/exprgen"
    24  	"github.com/cockroachdb/cockroach/pkg/sql/opt/testutils/opttester"
    25  	"github.com/cockroachdb/cockroach/pkg/sql/opt/testutils/testcat"
    26  	"github.com/cockroachdb/cockroach/pkg/sql/opt/xform"
    27  	"github.com/cockroachdb/cockroach/pkg/sql/parser"
    28  	_ "github.com/cockroachdb/cockroach/pkg/sql/sem/builtins"
    29  	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
    30  	"github.com/cockroachdb/cockroach/pkg/sql/types"
    31  	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
    32  	"github.com/cockroachdb/datadriven"
    33  )
    34  
    35  // TestExprIsNeverNull runs data-driven testcases of the form
    36  //   <command> [<args>]...
    37  //   <SQL statement or expression>
    38  //   ----
    39  //   <expected results>
    40  //
    41  // See OptTester.Handle for supported commands. In addition to those, we
    42  // support:
    43  //
    44  //  - scalar-is-not-nullable [args]
    45  //
    46  //    Builds a scalar expression using the input and performs a best-effort
    47  //    check to see if the scalar expression is nullable. It outputs this
    48  //    result as a boolean.
    49  //
    50  //    The supported args (in addition to the ones supported by OptTester):
    51  //
    52  //      - vars=(type1,type2,...)
    53  //
    54  //      Adding a !null suffix on a var type is used to mark that var as
    55  //      non-nullable.
    56  func TestExprIsNeverNull(t *testing.T) {
    57  	defer leaktest.AfterTest(t)()
    58  
    59  	datadriven.Walk(t, "testdata/expr", func(t *testing.T, path string) {
    60  		catalog := testcat.New()
    61  
    62  		datadriven.RunTest(t, path, func(t *testing.T, d *datadriven.TestData) string {
    63  			var varTypes []*types.T
    64  			var err error
    65  
    66  			tester := opttester.New(catalog, d.Input)
    67  			switch d.Cmd {
    68  			case "scalar-is-not-nullable":
    69  				var notNullCols opt.ColSet
    70  				for _, arg := range d.CmdArgs {
    71  					key, vals := arg.Key, arg.Vals
    72  					switch key {
    73  					case "vars":
    74  						for i := 0; i < len(vals); i++ {
    75  							if strings.HasSuffix(strings.ToLower(vals[i]), "!null") {
    76  								vals[i] = strings.TrimSuffix(strings.ToLower(vals[i]), "!null")
    77  								notNullCols.Add(opt.ColumnID(i + 1))
    78  							}
    79  						}
    80  						varTypes, err = exprgen.ParseTypes(vals)
    81  						if err != nil {
    82  							d.Fatalf(t, "%v", err)
    83  						}
    84  
    85  					default:
    86  						if err := tester.Flags.Set(arg); err != nil {
    87  							d.Fatalf(t, "%s", err)
    88  						}
    89  					}
    90  				}
    91  
    92  				expr, err := parser.ParseExpr(d.Input)
    93  				if err != nil {
    94  					d.Fatalf(t, "%v", err)
    95  				}
    96  
    97  				ctx := context.Background()
    98  				semaCtx := tree.MakeSemaContext()
    99  				evalCtx := tree.MakeTestingEvalContext(cluster.MakeTestingClusterSettings())
   100  
   101  				var o xform.Optimizer
   102  				o.Init(&evalCtx, catalog)
   103  				for i, typ := range varTypes {
   104  					o.Memo().Metadata().AddColumn(fmt.Sprintf("@%d", i+1), typ)
   105  				}
   106  				b := optbuilder.NewScalar(ctx, &semaCtx, &evalCtx, o.Factory())
   107  				err = b.Build(expr)
   108  				if err != nil {
   109  					return fmt.Sprintf("error: %s\n", strings.TrimSpace(err.Error()))
   110  				}
   111  				return fmt.Sprintf("%t\n", memo.ExprIsNeverNull(o.Memo().RootExpr().(opt.ScalarExpr), notNullCols))
   112  
   113  			default:
   114  				return tester.RunCommand(t, d)
   115  			}
   116  		})
   117  	})
   118  }