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

     1  // Copyright 2017 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  	"fmt"
    15  	"time"
    16  
    17  	"github.com/cockroachdb/cockroach/pkg/geo"
    18  	"github.com/cockroachdb/cockroach/pkg/sql/types"
    19  	"github.com/cockroachdb/cockroach/pkg/util/timeofday"
    20  	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
    21  	"github.com/cockroachdb/cockroach/pkg/util/timeutil/pgdate"
    22  )
    23  
    24  // presetTypesForTesting is a mapping of qualified names to types that can be mocked out
    25  // for tests to allow the qualified names to be type checked without throwing an error.
    26  var presetTypesForTesting map[string]*types.T
    27  
    28  // MockNameTypes populates presetTypesForTesting for a test.
    29  func MockNameTypes(types map[string]*types.T) func() {
    30  	presetTypesForTesting = types
    31  	return func() {
    32  		presetTypesForTesting = nil
    33  	}
    34  }
    35  
    36  // SampleDatum is intended to be a more lightweight version of RandDatum for
    37  // when you just need one consistent example of a datum.
    38  func SampleDatum(t *types.T) Datum {
    39  	switch t.Family() {
    40  	case types.BitFamily:
    41  		a, _ := NewDBitArrayFromInt(123, 40)
    42  		return a
    43  	case types.BoolFamily:
    44  		return MakeDBool(true)
    45  	case types.IntFamily:
    46  		return NewDInt(123)
    47  	case types.FloatFamily:
    48  		f := DFloat(123.456)
    49  		return &f
    50  	case types.DecimalFamily:
    51  		d := &DDecimal{}
    52  		// int64(rng.Uint64()) to get negative numbers, too
    53  		d.Decimal.SetFinite(3, 6)
    54  		return d
    55  	case types.StringFamily:
    56  		return NewDString("Carl")
    57  	case types.BytesFamily:
    58  		return NewDBytes("Princess")
    59  	case types.DateFamily:
    60  		return NewDDate(pgdate.MakeCompatibleDateFromDisk(123123))
    61  	case types.TimeFamily:
    62  		return MakeDTime(timeofday.FromInt(789))
    63  	case types.TimeTZFamily:
    64  		return NewDTimeTZFromOffset(timeofday.FromInt(345), 5*60*60 /* OffsetSecs */)
    65  	case types.TimestampFamily:
    66  		return MustMakeDTimestamp(timeutil.Unix(123, 123), time.Second)
    67  	case types.TimestampTZFamily:
    68  		return MustMakeDTimestampTZ(timeutil.Unix(123, 123), time.Second)
    69  	case types.IntervalFamily:
    70  		i, _ := ParseDInterval("1h1m1s")
    71  		return i
    72  	case types.UuidFamily:
    73  		u, _ := ParseDUuidFromString("3189ad07-52f2-4d60-83e8-4a8347fef718")
    74  		return u
    75  	case types.INetFamily:
    76  		i, _ := ParseDIPAddrFromINetString("127.0.0.1")
    77  		return i
    78  	case types.JsonFamily:
    79  		j, _ := ParseDJSON(`{"a": "b"}`)
    80  		return j
    81  	case types.OidFamily:
    82  		return NewDOid(DInt(1009))
    83  	case types.GeographyFamily:
    84  		return NewDGeography(geo.MustParseGeographyFromEWKBRaw([]byte("\x01\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\xf0\x3f\x00\x00\x00\x00\x00\x00\xf0\x3f")))
    85  	case types.GeometryFamily:
    86  		return NewDGeometry(geo.MustParseGeometryFromEWKBRaw([]byte("\x01\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\xf0\x3f\x00\x00\x00\x00\x00\x00\xf0\x3f")))
    87  	default:
    88  		panic(fmt.Sprintf("SampleDatum not implemented for %s", t))
    89  	}
    90  }