1/* 2Copyright 2017 Google LLC 3 4Licensed under the Apache License, Version 2.0 (the "License"); 5you may not use this file except in compliance with the License. 6You may obtain a copy of the License at 7 8 http://www.apache.org/licenses/LICENSE-2.0 9 10Unless required by applicable law or agreed to in writing, software 11distributed under the License is distributed on an "AS IS" BASIS, 12WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13See the License for the specific language governing permissions and 14limitations under the License. 15*/ 16 17package spanner 18 19import ( 20 "bytes" 21 "fmt" 22 "math/big" 23 "time" 24 25 "cloud.google.com/go/civil" 26 proto3 "github.com/golang/protobuf/ptypes/struct" 27 sppb "google.golang.org/genproto/googleapis/spanner/v1" 28 "google.golang.org/grpc/codes" 29) 30 31// A Key can be either a Cloud Spanner row's primary key or a secondary index 32// key. It is essentially an interface{} array, which represents a set of Cloud 33// Spanner columns. A Key can be used as: 34// 35// - A primary key which uniquely identifies a Cloud Spanner row. 36// - A secondary index key which maps to a set of Cloud Spanner rows indexed under it. 37// - An endpoint of primary key/secondary index ranges; see the KeyRange type. 38// 39// Rows that are identified by the Key type are outputs of read operation or 40// targets of delete operation in a mutation. Note that for 41// Insert/Update/InsertOrUpdate/Update mutation types, although they don't 42// require a primary key explicitly, the column list provided must contain 43// enough columns that can comprise a primary key. 44// 45// Keys are easy to construct. For example, suppose you have a table with a 46// primary key of username and product ID. To make a key for this table: 47// 48// key := spanner.Key{"john", 16} 49// 50// See the description of Row and Mutation types for how Go types are mapped to 51// Cloud Spanner types. For convenience, Key type supports a wide range of Go 52// types: 53// - int, int8, int16, int32, int64, and NullInt64 are mapped to Cloud Spanner's INT64 type. 54// - uint8, uint16 and uint32 are also mapped to Cloud Spanner's INT64 type. 55// - float32, float64, NullFloat64 are mapped to Cloud Spanner's FLOAT64 type. 56// - bool and NullBool are mapped to Cloud Spanner's BOOL type. 57// - []byte is mapped to Cloud Spanner's BYTES type. 58// - string and NullString are mapped to Cloud Spanner's STRING type. 59// - time.Time and NullTime are mapped to Cloud Spanner's TIMESTAMP type. 60// - civil.Date and NullDate are mapped to Cloud Spanner's DATE type. 61type Key []interface{} 62 63// errInvdKeyPartType returns error for unsupported key part type. 64func errInvdKeyPartType(part interface{}) error { 65 return spannerErrorf(codes.InvalidArgument, "key part has unsupported type %T", part) 66} 67 68// keyPartValue converts a part of the Key (which is a valid Cloud Spanner type) 69// into a proto3.Value. Used for encoding Key type into protobuf. 70func keyPartValue(part interface{}) (pb *proto3.Value, err error) { 71 switch v := part.(type) { 72 case int: 73 pb, _, err = encodeValue(int64(v)) 74 case int8: 75 pb, _, err = encodeValue(int64(v)) 76 case int16: 77 pb, _, err = encodeValue(int64(v)) 78 case int32: 79 pb, _, err = encodeValue(int64(v)) 80 case uint8: 81 pb, _, err = encodeValue(int64(v)) 82 case uint16: 83 pb, _, err = encodeValue(int64(v)) 84 case uint32: 85 pb, _, err = encodeValue(int64(v)) 86 case float32: 87 pb, _, err = encodeValue(float64(v)) 88 case int64, float64, NullInt64, NullFloat64, bool, NullBool, []byte, string, NullString, time.Time, civil.Date, NullTime, NullDate, big.Rat, NullNumeric: 89 pb, _, err = encodeValue(v) 90 case Encoder: 91 part, err = v.EncodeSpanner() 92 if err != nil { 93 return nil, err 94 } 95 pb, err = keyPartValue(part) 96 default: 97 return nil, errInvdKeyPartType(v) 98 } 99 return pb, err 100} 101 102// proto converts a spanner.Key into a proto3.ListValue. 103func (key Key) proto() (*proto3.ListValue, error) { 104 lv := &proto3.ListValue{} 105 lv.Values = make([]*proto3.Value, 0, len(key)) 106 for _, part := range key { 107 v, err := keyPartValue(part) 108 if err != nil { 109 return nil, err 110 } 111 lv.Values = append(lv.Values, v) 112 } 113 return lv, nil 114} 115 116// keySetProto lets a single Key act as a KeySet. 117func (key Key) keySetProto() (*sppb.KeySet, error) { 118 kp, err := key.proto() 119 if err != nil { 120 return nil, err 121 } 122 return &sppb.KeySet{Keys: []*proto3.ListValue{kp}}, nil 123} 124 125// String implements fmt.Stringer for Key. For string, []byte and NullString, it 126// prints the uninterpreted bytes of their contents, leaving caller with the 127// opportunity to escape the output. 128func (key Key) String() string { 129 b := &bytes.Buffer{} 130 fmt.Fprint(b, "(") 131 for i, part := range []interface{}(key) { 132 if i != 0 { 133 fmt.Fprint(b, ",") 134 } 135 key.elemString(b, part) 136 } 137 fmt.Fprint(b, ")") 138 return b.String() 139} 140 141func (key Key) elemString(b *bytes.Buffer, part interface{}) { 142 switch v := part.(type) { 143 case int, int8, int16, int32, int64, uint, uint8, uint16, uint32, float32, float64, bool: 144 // Use %v to print numeric types and bool. 145 fmt.Fprintf(b, "%v", v) 146 case string: 147 fmt.Fprintf(b, "%q", v) 148 case []byte: 149 if v != nil { 150 fmt.Fprintf(b, "%q", v) 151 } else { 152 fmt.Fprint(b, nullString) 153 } 154 case NullInt64, NullFloat64, NullBool, NullNumeric: 155 // The above types implement fmt.Stringer. 156 fmt.Fprintf(b, "%s", v) 157 case NullString, NullDate, NullTime: 158 // Quote the returned string if it is not null. 159 if v.(NullableValue).IsNull() { 160 fmt.Fprintf(b, "%s", nullString) 161 } else { 162 fmt.Fprintf(b, "%q", v) 163 } 164 case civil.Date: 165 fmt.Fprintf(b, "%q", v) 166 case time.Time: 167 fmt.Fprintf(b, "%q", v.Format(time.RFC3339Nano)) 168 case big.Rat: 169 fmt.Fprintf(b, "%v", NumericString(&v)) 170 case Encoder: 171 var err error 172 part, err = v.EncodeSpanner() 173 if err != nil { 174 fmt.Fprintf(b, "error") 175 } else { 176 key.elemString(b, part) 177 } 178 default: 179 fmt.Fprintf(b, "%v", v) 180 } 181} 182 183// AsPrefix returns a KeyRange for all keys where k is the prefix. 184func (key Key) AsPrefix() KeyRange { 185 return KeyRange{ 186 Start: key, 187 End: key, 188 Kind: ClosedClosed, 189 } 190} 191 192// KeyRangeKind describes the kind of interval represented by a KeyRange: 193// whether it is open or closed on the left and right. 194type KeyRangeKind int 195 196const ( 197 // ClosedOpen is closed on the left and open on the right: the Start 198 // key is included, the End key is excluded. 199 ClosedOpen KeyRangeKind = iota 200 201 // ClosedClosed is closed on the left and the right: both keys are included. 202 ClosedClosed 203 204 // OpenClosed is open on the left and closed on the right: the Start 205 // key is excluded, the End key is included. 206 OpenClosed 207 208 // OpenOpen is open on the left and the right: neither key is included. 209 OpenOpen 210) 211 212// A KeyRange represents a range of rows in a table or index. 213// 214// A range has a Start key and an End key. IncludeStart and IncludeEnd 215// indicate whether the Start and End keys are included in the range. 216// 217// For example, consider the following table definition: 218// 219// CREATE TABLE UserEvents ( 220// UserName STRING(MAX), 221// EventDate STRING(10), 222// ) PRIMARY KEY(UserName, EventDate); 223// 224// The following keys name rows in this table: 225// 226// spanner.Key{"Bob", "2014-09-23"} 227// spanner.Key{"Alfred", "2015-06-12"} 228// 229// Since the UserEvents table's PRIMARY KEY clause names two columns, each 230// UserEvents key has two elements; the first is the UserName, and the second 231// is the EventDate. 232// 233// Key ranges with multiple components are interpreted lexicographically by 234// component using the table or index key's declared sort order. For example, 235// the following range returns all events for user "Bob" that occurred in the 236// year 2015: 237// 238// spanner.KeyRange{ 239// Start: spanner.Key{"Bob", "2015-01-01"}, 240// End: spanner.Key{"Bob", "2015-12-31"}, 241// Kind: ClosedClosed, 242// } 243// 244// Start and end keys can omit trailing key components. This affects the 245// inclusion and exclusion of rows that exactly match the provided key 246// components: if IncludeStart is true, then rows that exactly match the 247// provided components of the Start key are included; if IncludeStart is false 248// then rows that exactly match are not included. IncludeEnd and End key 249// behave in the same fashion. 250// 251// For example, the following range includes all events for "Bob" that occurred 252// during and after the year 2000: 253// 254// spanner.KeyRange{ 255// Start: spanner.Key{"Bob", "2000-01-01"}, 256// End: spanner.Key{"Bob"}, 257// Kind: ClosedClosed, 258// } 259// 260// The next example retrieves all events for "Bob": 261// 262// spanner.Key{"Bob"}.AsPrefix() 263// 264// To retrieve events before the year 2000: 265// 266// spanner.KeyRange{ 267// Start: spanner.Key{"Bob"}, 268// End: spanner.Key{"Bob", "2000-01-01"}, 269// Kind: ClosedOpen, 270// } 271// 272// Although we specified a Kind for this KeyRange, we didn't need to, because 273// the default is ClosedOpen. In later examples we'll omit Kind if it is 274// ClosedOpen. 275// 276// The following range includes all rows in a table or under a 277// index: 278// 279// spanner.AllKeys() 280// 281// This range returns all users whose UserName begins with any 282// character from A to C: 283// 284// spanner.KeyRange{ 285// Start: spanner.Key{"A"}, 286// End: spanner.Key{"D"}, 287// } 288// 289// This range returns all users whose UserName begins with B: 290// 291// spanner.KeyRange{ 292// Start: spanner.Key{"B"}, 293// End: spanner.Key{"C"}, 294// } 295// 296// Key ranges honor column sort order. For example, suppose a table is defined 297// as follows: 298// 299// CREATE TABLE DescendingSortedTable { 300// Key INT64, 301// ... 302// ) PRIMARY KEY(Key DESC); 303// 304// The following range retrieves all rows with key values between 1 and 100 305// inclusive: 306// 307// spanner.KeyRange{ 308// Start: spanner.Key{100}, 309// End: spanner.Key{1}, 310// Kind: ClosedClosed, 311// } 312// 313// Note that 100 is passed as the start, and 1 is passed as the end, because 314// Key is a descending column in the schema. 315type KeyRange struct { 316 // Start specifies the left boundary of the key range; End specifies 317 // the right boundary of the key range. 318 Start, End Key 319 320 // Kind describes whether the boundaries of the key range include 321 // their keys. 322 Kind KeyRangeKind 323} 324 325// String implements fmt.Stringer for KeyRange type. 326func (r KeyRange) String() string { 327 var left, right string 328 switch r.Kind { 329 case ClosedClosed: 330 left, right = "[", "]" 331 case ClosedOpen: 332 left, right = "[", ")" 333 case OpenClosed: 334 left, right = "(", "]" 335 case OpenOpen: 336 left, right = "(", ")" 337 default: 338 left, right = "?", "?" 339 } 340 return fmt.Sprintf("%s%s,%s%s", left, r.Start, r.End, right) 341} 342 343// proto converts KeyRange into sppb.KeyRange. 344func (r KeyRange) proto() (*sppb.KeyRange, error) { 345 var err error 346 var start, end *proto3.ListValue 347 pb := &sppb.KeyRange{} 348 if start, err = r.Start.proto(); err != nil { 349 return nil, err 350 } 351 if end, err = r.End.proto(); err != nil { 352 return nil, err 353 } 354 if r.Kind == ClosedClosed || r.Kind == ClosedOpen { 355 pb.StartKeyType = &sppb.KeyRange_StartClosed{StartClosed: start} 356 } else { 357 pb.StartKeyType = &sppb.KeyRange_StartOpen{StartOpen: start} 358 } 359 if r.Kind == ClosedClosed || r.Kind == OpenClosed { 360 pb.EndKeyType = &sppb.KeyRange_EndClosed{EndClosed: end} 361 } else { 362 pb.EndKeyType = &sppb.KeyRange_EndOpen{EndOpen: end} 363 } 364 return pb, nil 365} 366 367// keySetProto lets a KeyRange act as a KeySet. 368func (r KeyRange) keySetProto() (*sppb.KeySet, error) { 369 rp, err := r.proto() 370 if err != nil { 371 return nil, err 372 } 373 return &sppb.KeySet{Ranges: []*sppb.KeyRange{rp}}, nil 374} 375 376// A KeySet defines a collection of Cloud Spanner keys and/or key ranges. All 377// the keys are expected to be in the same table or index. The keys need not be 378// sorted in any particular way. 379// 380// An individual Key can act as a KeySet, as can a KeyRange. Use the KeySets 381// function to create a KeySet consisting of multiple Keys and KeyRanges. To 382// obtain an empty KeySet, call KeySets with no arguments. 383// 384// If the same key is specified multiple times in the set (for example if two 385// ranges, two keys, or a key and a range overlap), the Cloud Spanner backend 386// behaves as if the key were only specified once. 387type KeySet interface { 388 keySetProto() (*sppb.KeySet, error) 389} 390 391// AllKeys returns a KeySet that represents all Keys of a table or a index. 392func AllKeys() KeySet { 393 return all{} 394} 395 396type all struct{} 397 398func (all) keySetProto() (*sppb.KeySet, error) { 399 return &sppb.KeySet{All: true}, nil 400} 401 402// KeySets returns the union of the KeySets. If any of the KeySets is AllKeys, 403// then the resulting KeySet will be equivalent to AllKeys. 404func KeySets(keySets ...KeySet) KeySet { 405 u := make(union, len(keySets)) 406 copy(u, keySets) 407 return u 408} 409 410// KeySetFromKeys returns a KeySet containing the given slice of keys. 411func KeySetFromKeys(keys ...Key) KeySet { 412 u := make(union, len(keys)) 413 for i, k := range keys { 414 u[i] = k 415 } 416 return u 417} 418 419type union []KeySet 420 421func (u union) keySetProto() (*sppb.KeySet, error) { 422 upb := &sppb.KeySet{} 423 for _, ks := range u { 424 pb, err := ks.keySetProto() 425 if err != nil { 426 return nil, err 427 } 428 if pb.All { 429 return pb, nil 430 } 431 upb.Keys = append(upb.Keys, pb.Keys...) 432 upb.Ranges = append(upb.Ranges, pb.Ranges...) 433 } 434 return upb, nil 435} 436