-
Notifications
You must be signed in to change notification settings - Fork 3.8k
/
type.go
286 lines (258 loc) · 8.6 KB
/
type.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
// Copyright 2021 The Cockroach Authors.
//
// Use of this software is governed by the Business Source License
// included in the file licenses/BSL.txt.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0, included in the file
// licenses/APL.txt.
package randgen
import (
"context"
"math/rand"
"sort"
clustersettings "github.com/cockroachdb/cockroach/pkg/settings/cluster"
"github.com/cockroachdb/cockroach/pkg/sql/catalog/colinfo"
"github.com/cockroachdb/cockroach/pkg/sql/rowenc/valueside"
"github.com/cockroachdb/cockroach/pkg/sql/types"
"github.com/lib/pq/oid"
)
var (
// SeedTypes includes the following types that form the basis of randomly
// generated types:
// - All scalar types, except UNKNOWN and ANY
// - ARRAY of ANY, where the ANY will be replaced with one of the legal
// array element types in RandType
// - OIDVECTOR and INT2VECTOR types
SeedTypes []*types.T
// arrayContentsTypes contains all of the types that are valid to store within
// an array.
arrayContentsTypes []*types.T
collationLocales = [...]string{"da", "de", "en"}
)
func init() {
for _, typ := range types.OidToType {
switch typ.Oid() {
case oid.T_regnamespace:
// Temporarily don't include this.
// TODO(msirek): Remove this exclusion once
// https://github.com/cockroachdb/cockroach/issues/55791 is fixed.
case oid.T_unknown, oid.T_anyelement:
// Don't include these.
case oid.T_anyarray, oid.T_oidvector, oid.T_int2vector:
// Include these.
SeedTypes = append(SeedTypes, typ)
default:
// Only include scalar types.
if typ.Family() != types.ArrayFamily {
SeedTypes = append(SeedTypes, typ)
}
}
}
for _, typ := range types.OidToType {
if IsAllowedForArray(typ) {
arrayContentsTypes = append(arrayContentsTypes, typ)
}
}
// Add a collated string separately (since it shares the oid with the STRING
// type and, thus, wasn't included above).
collatedStringType := types.MakeCollatedString(types.String, "en" /* locale */)
SeedTypes = append(SeedTypes, collatedStringType)
if IsAllowedForArray(collatedStringType) {
arrayContentsTypes = append(arrayContentsTypes, collatedStringType)
}
// Sort these so randomly chosen indexes always point to the same element.
sort.Slice(SeedTypes, func(i, j int) bool {
return SeedTypes[i].String() < SeedTypes[j].String()
})
sort.Slice(arrayContentsTypes, func(i, j int) bool {
return arrayContentsTypes[i].String() < arrayContentsTypes[j].String()
})
}
// IsAllowedForArray returns true iff the passed in type can be a valid ArrayContents()
func IsAllowedForArray(typ *types.T) bool {
// Don't include un-encodable types.
encTyp, err := valueside.DatumTypeToArrayElementEncodingType(typ)
if err != nil || encTyp == 0 {
return false
}
// Don't include reg types, since parser currently doesn't allow them to
// be declared as array element types.
if typ.Family() == types.OidFamily && typ.Oid() != oid.T_oid {
return false
}
return true
}
// RandType returns a random type value.
func RandType(rng *rand.Rand) *types.T {
return RandTypeFromSlice(rng, SeedTypes)
}
// RandArrayContentsType returns a random type that's guaranteed to be valid to
// use as the contents of an array.
func RandArrayContentsType(rng *rand.Rand) *types.T {
return RandTypeFromSlice(rng, arrayContentsTypes)
}
// RandTypeFromSlice returns a random type from the input slice of types.
func RandTypeFromSlice(rng *rand.Rand, typs []*types.T) *types.T {
typ := typs[rng.Intn(len(typs))]
switch typ.Family() {
case types.BitFamily:
return types.MakeBit(int32(rng.Intn(50)))
case types.CollatedStringFamily:
return types.MakeCollatedString(types.String, *RandCollationLocale(rng))
case types.ArrayFamily:
if typ.ArrayContents().Family() == types.AnyFamily {
inner := RandArrayContentsType(rng)
if inner.Family() == types.CollatedStringFamily {
// TODO(justin): change this when collated arrays are supported.
inner = types.String
}
return types.MakeArray(inner)
}
if typ.ArrayContents().Family() == types.TupleFamily {
// Generate tuples between 0 and 4 datums in length
len := rng.Intn(5)
contents := make([]*types.T, len)
for i := range contents {
contents[i] = RandTypeFromSlice(rng, typs)
}
return types.MakeArray(types.MakeTuple(contents))
}
case types.TupleFamily:
// Generate tuples between 0 and 4 datums in length
len := rng.Intn(5)
contents := make([]*types.T, len)
for i := range contents {
contents[i] = RandTypeFromSlice(rng, typs)
}
return types.MakeTuple(contents)
}
return typ
}
// RandColumnType returns a random type that is a legal column type (e.g. no
// nested arrays or tuples).
func RandColumnType(rng *rand.Rand) *types.T {
for {
typ := RandType(rng)
if IsLegalColumnType(typ) {
return typ
}
}
}
// IsLegalColumnType returns true if the given type can be
// given to a column in a user-created table.
func IsLegalColumnType(typ *types.T) bool {
switch typ.Oid() {
case oid.T_int2vector, oid.T_oidvector:
// OIDVECTOR and INT2VECTOR are not valid column types for
// user-created tables.
return false
case oid.T_regproc, oid.T_regprocedure:
// REGPROC and REGPROCEDURE columns hit an edge case. Customers are very
// unlikely to use these types of columns, so disabling their generation
// is low risk.
// TODO(#95641): Remove this once we correctly handle this edge case.
return false
}
ctx := context.Background()
version := clustersettings.MakeTestingClusterSettings().Version
return colinfo.ValidateColumnDefType(ctx, version, typ) == nil
}
// RandArrayType generates a random array type.
func RandArrayType(rng *rand.Rand) *types.T {
ctx := context.Background()
version := clustersettings.MakeTestingClusterSettings().Version
for {
typ := RandColumnType(rng)
resTyp := types.MakeArray(typ)
if err := colinfo.ValidateColumnDefType(ctx, version, resTyp); err == nil {
return resTyp
}
}
}
// RandColumnTypes returns a slice of numCols random types. These types must be
// legal table column types.
func RandColumnTypes(rng *rand.Rand, numCols int) []*types.T {
types := make([]*types.T, numCols)
for i := range types {
types[i] = RandColumnType(rng)
}
return types
}
// RandSortingType returns a column type which can be key-encoded.
func RandSortingType(rng *rand.Rand) *types.T {
typ := RandType(rng)
for colinfo.MustBeValueEncoded(typ) || typ == types.Void {
typ = RandType(rng)
}
return typ
}
// RandSortingTypes returns a slice of numCols random ColumnType values
// which are key-encodable.
func RandSortingTypes(rng *rand.Rand, numCols int) []*types.T {
types := make([]*types.T, numCols)
for i := range types {
types[i] = RandSortingType(rng)
}
return types
}
// RandCollationLocale returns a random element of collationLocales.
func RandCollationLocale(rng *rand.Rand) *string {
return &collationLocales[rng.Intn(len(collationLocales))]
}
// RandEncodableType wraps RandType in order to workaround #36736, which fails
// when name[] (or other type using DTypeWrapper) is encoded.
//
// TODO(andyk): Remove this workaround once #36736 is resolved. Also, RandDatum
// really should be extended to create DTypeWrapper datums with alternate OIDs
// like oid.T_varchar for better testing.
func RandEncodableType(rng *rand.Rand) *types.T {
var isEncodableType func(t *types.T) bool
isEncodableType = func(t *types.T) bool {
switch t.Family() {
case types.ArrayFamily:
// Due to #36736, any type returned by RandType that gets turned into
// a DTypeWrapper random datum will not work. Currently, that's just
// types.Name.
if t.ArrayContents().Oid() == oid.T_name {
return false
}
return isEncodableType(t.ArrayContents())
case types.TupleFamily:
for i := range t.TupleContents() {
if !isEncodableType(t.TupleContents()[i]) {
return false
}
}
case types.VoidFamily:
return false
}
return true
}
for {
typ := RandType(rng)
if isEncodableType(typ) {
return typ
}
}
}
// RandEncodableColumnTypes works around #36736, which fails when name[] (or
// other type using DTypeWrapper) is encoded.
//
// TODO(andyk): Remove this workaround once #36736 is resolved. Replace calls to
// it with calls to RandColumnTypes.
func RandEncodableColumnTypes(rng *rand.Rand, numCols int) []*types.T {
ctx := context.Background()
version := clustersettings.MakeTestingClusterSettings().Version
types := make([]*types.T, numCols)
for i := range types {
for {
types[i] = RandEncodableType(rng)
if err := colinfo.ValidateColumnDefType(ctx, version, types[i]); err == nil {
break
}
}
}
return types
}