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

     1  // Copyright 2020 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 schemaexpr
    12  
    13  import (
    14  	"context"
    15  	"testing"
    16  
    17  	"github.com/cockroachdb/cockroach/pkg/sql/parser"
    18  	"github.com/cockroachdb/cockroach/pkg/sql/sem/builtins"
    19  	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
    20  	"github.com/cockroachdb/cockroach/pkg/sql/types"
    21  )
    22  
    23  func TestIndexPredicateValidator_Validate(t *testing.T) {
    24  	ctx := context.Background()
    25  	semaCtx := tree.MakeSemaContext()
    26  
    27  	// Trick to get the init() for the builtins package to run.
    28  	_ = builtins.AllBuiltinNames
    29  
    30  	database := tree.Name("foo")
    31  	table := tree.Name("bar")
    32  	tn := tree.MakeTableName(database, table)
    33  
    34  	desc := testTableDesc(
    35  		string(table),
    36  		[]testCol{{"a", types.Bool}, {"b", types.Int}},
    37  		[]testCol{{"c", types.String}},
    38  	)
    39  
    40  	validator := NewIndexPredicateValidator(ctx, tn, &desc, &semaCtx)
    41  
    42  	testData := []struct {
    43  		expr          string
    44  		expectedValid bool
    45  		expectedExpr  string
    46  	}{
    47  		// Allow expressions that result in a bool.
    48  		{"a", true, "a"},
    49  		{"b = 0", true, "b = 0"},
    50  		{"a AND b = 0", true, "a AND (b = 0)"},
    51  		{"a IS NULL", true, "a IS NULL"},
    52  		{"b IN (1, 2)", true, "b IN (1:::INT8, 2:::INT8)"},
    53  
    54  		// Allow immutable functions.
    55  		{"abs(b) > 0", true, "abs(b) > 0"},
    56  		{"c || c = 'foofoo'", true, "(c || c) = 'foofoo'"},
    57  		{"lower(c) = 'bar'", true, "lower(c) = 'bar'"},
    58  
    59  		// Disallow references to columns not in the table.
    60  		{"d", false, ""},
    61  		{"t.a", false, ""},
    62  
    63  		// Disallow expressions that do not result in a bool.
    64  		{"b", false, ""},
    65  		{"abs(b)", false, ""},
    66  		{"lower(c)", false, ""},
    67  
    68  		// Disallow subqueries.
    69  		{"exists(select 1)", false, ""},
    70  		{"b IN (select 1)", false, ""},
    71  
    72  		// Disallow mutable, aggregate, window, and set returning functions.
    73  		{"b > random()", false, ""},
    74  		{"sum(b) > 10", false, ""},
    75  		{"row_number() OVER () > 1", false, ""},
    76  		{"generate_series(1, 1) > 2", false, ""},
    77  
    78  		// Dequalify column names.
    79  		{"bar.a", true, "a"},
    80  		{"foo.bar.a", true, "a"},
    81  		{"bar.b = 0", true, "b = 0"},
    82  		{"foo.bar.b = 0", true, "b = 0"},
    83  		{"bar.a AND foo.bar.b = 0", true, "a AND (b = 0)"},
    84  	}
    85  
    86  	for _, d := range testData {
    87  		t.Run(d.expr, func(t *testing.T) {
    88  			expr, err := parser.ParseExpr(d.expr)
    89  			if err != nil {
    90  				t.Fatalf("%s: unexpected error: %s", d.expr, err)
    91  			}
    92  
    93  			r, err := validator.Validate(expr)
    94  
    95  			if !d.expectedValid {
    96  				if err == nil {
    97  					t.Fatalf("%s: expected invalid expression, but was valid", d.expr)
    98  				}
    99  				// The input expression is invalid so there is no need to check
   100  				// the output expression r.
   101  				return
   102  			}
   103  
   104  			if err != nil {
   105  				t.Fatalf("%s: expected valid expression, but found error: %s", d.expr, err)
   106  			}
   107  
   108  			s := tree.Serialize(r)
   109  			if s != d.expectedExpr {
   110  				t.Errorf("%s: expected %q, got %q", d.expr, d.expectedExpr, s)
   111  			}
   112  		})
   113  	}
   114  }