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

     1  // Copyright 2018 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 constraint
    12  
    13  import (
    14  	"strconv"
    15  	"strings"
    16  	"time"
    17  
    18  	"github.com/cockroachdb/cockroach/pkg/sql/opt"
    19  	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
    20  	"github.com/cockroachdb/cockroach/pkg/sql/types"
    21  	"github.com/cockroachdb/errors"
    22  )
    23  
    24  // ParseConstraint parses a constraint in the format of Constraint.String, e.g:
    25  //   "/1/2/3: [/1 - /2]".
    26  func ParseConstraint(evalCtx *tree.EvalContext, str string) Constraint {
    27  	s := strings.SplitN(str, ": ", 2)
    28  	if len(s) != 2 {
    29  		panic(str)
    30  	}
    31  	var cols []opt.OrderingColumn
    32  	for _, v := range parseIntPath(s[0]) {
    33  		cols = append(cols, opt.OrderingColumn(v))
    34  	}
    35  	var c Constraint
    36  	c.Columns.Init(cols)
    37  	c.Spans = parseSpans(evalCtx, s[1])
    38  	return c
    39  }
    40  
    41  // parseSpans parses a list of spans with integer values like:
    42  //   "[/1 - /2] [/5 - /6]".
    43  func parseSpans(evalCtx *tree.EvalContext, str string) Spans {
    44  	if str == "" {
    45  		return Spans{}
    46  	}
    47  	s := strings.Split(str, " ")
    48  	// Each span has three pieces.
    49  	if len(s)%3 != 0 {
    50  		panic(str)
    51  	}
    52  	var result Spans
    53  	for i := 0; i < len(s)/3; i++ {
    54  		sp := ParseSpan(evalCtx, strings.Join(s[i*3:i*3+3], " "), types.IntFamily)
    55  		result.Append(&sp)
    56  	}
    57  	return result
    58  }
    59  
    60  // ParseSpan parses a span in the format of Span.String, e.g: [/1 - /2].
    61  func ParseSpan(evalCtx *tree.EvalContext, str string, typ types.Family) Span {
    62  	if len(str) < len("[ - ]") {
    63  		panic(str)
    64  	}
    65  	boundary := map[byte]SpanBoundary{
    66  		'[': IncludeBoundary,
    67  		']': IncludeBoundary,
    68  		'(': ExcludeBoundary,
    69  		')': ExcludeBoundary,
    70  	}
    71  	s, e := str[0], str[len(str)-1]
    72  	if (s != '[' && s != '(') || (e != ']' && e != ')') {
    73  		panic(str)
    74  	}
    75  	keys := strings.Split(str[1:len(str)-1], " - ")
    76  	if len(keys) != 2 {
    77  		panic(str)
    78  	}
    79  	var sp Span
    80  	startVals := parseDatumPath(evalCtx, keys[0], typ)
    81  	endVals := parseDatumPath(evalCtx, keys[1], typ)
    82  	sp.Init(
    83  		MakeCompositeKey(startVals...), boundary[s],
    84  		MakeCompositeKey(endVals...), boundary[e],
    85  	)
    86  	return sp
    87  }
    88  
    89  // parseIntPath parses a string like "/1/2/3" into a list of integers.
    90  func parseIntPath(str string) []int {
    91  	var res []int
    92  	for _, valStr := range parsePath(str) {
    93  		val, err := strconv.Atoi(valStr)
    94  		if err != nil {
    95  			panic(err)
    96  		}
    97  		res = append(res, val)
    98  	}
    99  	return res
   100  }
   101  
   102  // parseDatumPath parses a span key string like "/1/2/3".
   103  // Only NULL and a subset of types are currently supported.
   104  func parseDatumPath(evalCtx *tree.EvalContext, str string, typ types.Family) []tree.Datum {
   105  	var res []tree.Datum
   106  	for _, valStr := range parsePath(str) {
   107  		if valStr == "NULL" {
   108  			res = append(res, tree.DNull)
   109  			continue
   110  		}
   111  		var val tree.Datum
   112  		var err error
   113  		switch typ {
   114  		case types.BoolFamily:
   115  			val, err = tree.ParseDBool(valStr)
   116  		case types.IntFamily:
   117  			val, err = tree.ParseDInt(valStr)
   118  		case types.FloatFamily:
   119  			val, err = tree.ParseDFloat(valStr)
   120  		case types.DecimalFamily:
   121  			val, err = tree.ParseDDecimal(valStr)
   122  		case types.DateFamily:
   123  			val, err = tree.ParseDDate(evalCtx, valStr)
   124  		case types.TimestampFamily:
   125  			val, err = tree.ParseDTimestamp(evalCtx, valStr, time.Microsecond)
   126  		case types.TimestampTZFamily:
   127  			val, err = tree.ParseDTimestampTZ(evalCtx, valStr, time.Microsecond)
   128  		case types.StringFamily:
   129  			val = tree.NewDString(valStr)
   130  		default:
   131  			panic(errors.AssertionFailedf("type %s not supported", typ.String()))
   132  		}
   133  		if err != nil {
   134  			panic(err)
   135  		}
   136  		res = append(res, val)
   137  	}
   138  	return res
   139  }
   140  
   141  // parsePath splits a string of the form "/foo/bar" into strings ["foo", "bar"].
   142  // An empty string is allowed, otherwise the string must start with /.
   143  func parsePath(str string) []string {
   144  	if str == "" {
   145  		return nil
   146  	}
   147  	if str[0] != '/' {
   148  		panic(str)
   149  	}
   150  	return strings.Split(str, "/")[1:]
   151  }