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