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

     1  // Copyright 2016 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 tree
    12  
    13  import (
    14  	"context"
    15  	"fmt"
    16  	"testing"
    17  
    18  	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
    19  	"github.com/cockroachdb/cockroach/pkg/sql/types"
    20  	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
    21  )
    22  
    23  type testVarContainer []Datum
    24  
    25  func (d testVarContainer) IndexedVarEval(idx int, ctx *EvalContext) (Datum, error) {
    26  	return d[idx].Eval(ctx)
    27  }
    28  
    29  func (d testVarContainer) IndexedVarResolvedType(idx int) *types.T {
    30  	return d[idx].ResolvedType()
    31  }
    32  
    33  func (d testVarContainer) IndexedVarNodeFormatter(idx int) NodeFormatter {
    34  	n := Name(fmt.Sprintf("var%d", idx))
    35  	return &n
    36  }
    37  
    38  func TestIndexedVars(t *testing.T) {
    39  	defer leaktest.AfterTest(t)()
    40  	c := make(testVarContainer, 4)
    41  	c[0] = NewDInt(3)
    42  	c[1] = NewDInt(5)
    43  	c[2] = NewDInt(6)
    44  	c[3] = NewDInt(0)
    45  
    46  	h := MakeIndexedVarHelper(c, 4)
    47  
    48  	// We use only the first three variables.
    49  	v0 := h.IndexedVar(0)
    50  	v1 := h.IndexedVar(1)
    51  	v2 := h.IndexedVar(2)
    52  
    53  	if !h.IndexedVarUsed(0) || !h.IndexedVarUsed(1) || !h.IndexedVarUsed(2) || h.IndexedVarUsed(3) {
    54  		t.Errorf("invalid IndexedVarUsed results %t %t %t %t (expected false false false true)",
    55  			h.IndexedVarUsed(0), h.IndexedVarUsed(1), h.IndexedVarUsed(2), h.IndexedVarUsed(3))
    56  	}
    57  
    58  	binary := func(op BinaryOperator, left, right Expr) Expr {
    59  		return &BinaryExpr{Operator: op, Left: left, Right: right}
    60  	}
    61  	expr := binary(Plus, v0, binary(Mult, v1, v2))
    62  
    63  	// Verify the expression evaluates correctly.
    64  	ctx := context.Background()
    65  	semaContext := MakeSemaContext()
    66  	semaContext.IVarContainer = c
    67  	typedExpr, err := expr.TypeCheck(ctx, &semaContext, types.Any)
    68  	if err != nil {
    69  		t.Fatal(err)
    70  	}
    71  
    72  	str := typedExpr.String()
    73  	expectedStr := "var0 + (var1 * var2)"
    74  	if str != expectedStr {
    75  		t.Errorf("invalid expression string '%s', expected '%s'", str, expectedStr)
    76  	}
    77  
    78  	// Test formatting using the indexed var format interceptor.
    79  	f := NewFmtCtx(FmtSimple)
    80  	f.SetIndexedVarFormat(func(ctx *FmtCtx, idx int) {
    81  		ctx.Printf("customVar%d", idx)
    82  	})
    83  	f.FormatNode(typedExpr)
    84  	str = f.CloseAndGetString()
    85  
    86  	expectedStr = "customVar0 + (customVar1 * customVar2)"
    87  	if str != expectedStr {
    88  		t.Errorf("invalid expression string '%s', expected '%s'", str, expectedStr)
    89  	}
    90  
    91  	typ := typedExpr.ResolvedType()
    92  	if !typ.Equivalent(types.Int) {
    93  		t.Errorf("invalid expression type %s", typ)
    94  	}
    95  	evalCtx := NewTestingEvalContext(cluster.MakeTestingClusterSettings())
    96  	defer evalCtx.Stop(context.Background())
    97  	evalCtx.IVarContainer = c
    98  	d, err := typedExpr.Eval(evalCtx)
    99  	if err != nil {
   100  		t.Fatal(err)
   101  	}
   102  	if d.Compare(evalCtx, NewDInt(3+5*6)) != 0 {
   103  		t.Errorf("invalid result %s (expected %d)", d, 3+5*6)
   104  	}
   105  }