1// Copyright 2011 Google Inc. All rights reserved. 2// Use of this source code is governed by the Apache 2.0 3// license that can be found in the LICENSE file. 4 5package datastore 6 7import ( 8 "bytes" 9 "encoding/base64" 10 "encoding/gob" 11 "errors" 12 "fmt" 13 "strconv" 14 "strings" 15 16 "github.com/golang/protobuf/proto" 17 "golang.org/x/net/context" 18 19 "google.golang.org/appengine/internal" 20 pb "google.golang.org/appengine/internal/datastore" 21) 22 23type KeyRangeCollisionError struct { 24 start int64 25 end int64 26} 27 28func (e *KeyRangeCollisionError) Error() string { 29 return fmt.Sprintf("datastore: Collision when attempting to allocate range [%d, %d]", 30 e.start, e.end) 31} 32 33type KeyRangeContentionError struct { 34 start int64 35 end int64 36} 37 38func (e *KeyRangeContentionError) Error() string { 39 return fmt.Sprintf("datastore: Contention when attempting to allocate range [%d, %d]", 40 e.start, e.end) 41} 42 43// Key represents the datastore key for a stored entity, and is immutable. 44type Key struct { 45 kind string 46 stringID string 47 intID int64 48 parent *Key 49 appID string 50 namespace string 51} 52 53// Kind returns the key's kind (also known as entity type). 54func (k *Key) Kind() string { 55 return k.kind 56} 57 58// StringID returns the key's string ID (also known as an entity name or key 59// name), which may be "". 60func (k *Key) StringID() string { 61 return k.stringID 62} 63 64// IntID returns the key's integer ID, which may be 0. 65func (k *Key) IntID() int64 { 66 return k.intID 67} 68 69// Parent returns the key's parent key, which may be nil. 70func (k *Key) Parent() *Key { 71 return k.parent 72} 73 74// AppID returns the key's application ID. 75func (k *Key) AppID() string { 76 return k.appID 77} 78 79// Namespace returns the key's namespace. 80func (k *Key) Namespace() string { 81 return k.namespace 82} 83 84// Incomplete returns whether the key does not refer to a stored entity. 85// In particular, whether the key has a zero StringID and a zero IntID. 86func (k *Key) Incomplete() bool { 87 return k.stringID == "" && k.intID == 0 88} 89 90// valid returns whether the key is valid. 91func (k *Key) valid() bool { 92 if k == nil { 93 return false 94 } 95 for ; k != nil; k = k.parent { 96 if k.kind == "" || k.appID == "" { 97 return false 98 } 99 if k.stringID != "" && k.intID != 0 { 100 return false 101 } 102 if k.parent != nil { 103 if k.parent.Incomplete() { 104 return false 105 } 106 if k.parent.appID != k.appID || k.parent.namespace != k.namespace { 107 return false 108 } 109 } 110 } 111 return true 112} 113 114// Equal returns whether two keys are equal. 115func (k *Key) Equal(o *Key) bool { 116 for k != nil && o != nil { 117 if k.kind != o.kind || k.stringID != o.stringID || k.intID != o.intID || k.appID != o.appID || k.namespace != o.namespace { 118 return false 119 } 120 k, o = k.parent, o.parent 121 } 122 return k == o 123} 124 125// root returns the furthest ancestor of a key, which may be itself. 126func (k *Key) root() *Key { 127 for k.parent != nil { 128 k = k.parent 129 } 130 return k 131} 132 133// marshal marshals the key's string representation to the buffer. 134func (k *Key) marshal(b *bytes.Buffer) { 135 if k.parent != nil { 136 k.parent.marshal(b) 137 } 138 b.WriteByte('/') 139 b.WriteString(k.kind) 140 b.WriteByte(',') 141 if k.stringID != "" { 142 b.WriteString(k.stringID) 143 } else { 144 b.WriteString(strconv.FormatInt(k.intID, 10)) 145 } 146} 147 148// String returns a string representation of the key. 149func (k *Key) String() string { 150 if k == nil { 151 return "" 152 } 153 b := bytes.NewBuffer(make([]byte, 0, 512)) 154 k.marshal(b) 155 return b.String() 156} 157 158type gobKey struct { 159 Kind string 160 StringID string 161 IntID int64 162 Parent *gobKey 163 AppID string 164 Namespace string 165} 166 167func keyToGobKey(k *Key) *gobKey { 168 if k == nil { 169 return nil 170 } 171 return &gobKey{ 172 Kind: k.kind, 173 StringID: k.stringID, 174 IntID: k.intID, 175 Parent: keyToGobKey(k.parent), 176 AppID: k.appID, 177 Namespace: k.namespace, 178 } 179} 180 181func gobKeyToKey(gk *gobKey) *Key { 182 if gk == nil { 183 return nil 184 } 185 return &Key{ 186 kind: gk.Kind, 187 stringID: gk.StringID, 188 intID: gk.IntID, 189 parent: gobKeyToKey(gk.Parent), 190 appID: gk.AppID, 191 namespace: gk.Namespace, 192 } 193} 194 195func (k *Key) GobEncode() ([]byte, error) { 196 buf := new(bytes.Buffer) 197 if err := gob.NewEncoder(buf).Encode(keyToGobKey(k)); err != nil { 198 return nil, err 199 } 200 return buf.Bytes(), nil 201} 202 203func (k *Key) GobDecode(buf []byte) error { 204 gk := new(gobKey) 205 if err := gob.NewDecoder(bytes.NewBuffer(buf)).Decode(gk); err != nil { 206 return err 207 } 208 *k = *gobKeyToKey(gk) 209 return nil 210} 211 212func (k *Key) MarshalJSON() ([]byte, error) { 213 return []byte(`"` + k.Encode() + `"`), nil 214} 215 216func (k *Key) UnmarshalJSON(buf []byte) error { 217 if len(buf) < 2 || buf[0] != '"' || buf[len(buf)-1] != '"' { 218 return errors.New("datastore: bad JSON key") 219 } 220 k2, err := DecodeKey(string(buf[1 : len(buf)-1])) 221 if err != nil { 222 return err 223 } 224 *k = *k2 225 return nil 226} 227 228// Encode returns an opaque representation of the key 229// suitable for use in HTML and URLs. 230// This is compatible with the Python and Java runtimes. 231func (k *Key) Encode() string { 232 ref := keyToProto("", k) 233 234 b, err := proto.Marshal(ref) 235 if err != nil { 236 panic(err) 237 } 238 239 // Trailing padding is stripped. 240 return strings.TrimRight(base64.URLEncoding.EncodeToString(b), "=") 241} 242 243// DecodeKey decodes a key from the opaque representation returned by Encode. 244func DecodeKey(encoded string) (*Key, error) { 245 // Re-add padding. 246 if m := len(encoded) % 4; m != 0 { 247 encoded += strings.Repeat("=", 4-m) 248 } 249 250 b, err := base64.URLEncoding.DecodeString(encoded) 251 if err != nil { 252 return nil, err 253 } 254 255 ref := new(pb.Reference) 256 if err := proto.Unmarshal(b, ref); err != nil { 257 // Couldn't decode it as an App Engine key, try decoding it as a key encoded by cloud.google.com/go/datastore. 258 if k := decodeCloudKey(encoded); k != nil { 259 return k, nil 260 } 261 return nil, err 262 } 263 264 return protoToKey(ref) 265} 266 267// NewIncompleteKey creates a new incomplete key. 268// kind cannot be empty. 269func NewIncompleteKey(c context.Context, kind string, parent *Key) *Key { 270 return NewKey(c, kind, "", 0, parent) 271} 272 273// NewKey creates a new key. 274// kind cannot be empty. 275// Either one or both of stringID and intID must be zero. If both are zero, 276// the key returned is incomplete. 277// parent must either be a complete key or nil. 278func NewKey(c context.Context, kind, stringID string, intID int64, parent *Key) *Key { 279 // If there's a parent key, use its namespace. 280 // Otherwise, use any namespace attached to the context. 281 var namespace string 282 if parent != nil { 283 namespace = parent.namespace 284 } else { 285 namespace = internal.NamespaceFromContext(c) 286 } 287 288 return &Key{ 289 kind: kind, 290 stringID: stringID, 291 intID: intID, 292 parent: parent, 293 appID: internal.FullyQualifiedAppID(c), 294 namespace: namespace, 295 } 296} 297 298// AllocateIDs returns a range of n integer IDs with the given kind and parent 299// combination. kind cannot be empty; parent may be nil. The IDs in the range 300// returned will not be used by the datastore's automatic ID sequence generator 301// and may be used with NewKey without conflict. 302// 303// The range is inclusive at the low end and exclusive at the high end. In 304// other words, valid intIDs x satisfy low <= x && x < high. 305// 306// If no error is returned, low + n == high. 307func AllocateIDs(c context.Context, kind string, parent *Key, n int) (low, high int64, err error) { 308 if kind == "" { 309 return 0, 0, errors.New("datastore: AllocateIDs given an empty kind") 310 } 311 if n < 0 { 312 return 0, 0, fmt.Errorf("datastore: AllocateIDs given a negative count: %d", n) 313 } 314 if n == 0 { 315 return 0, 0, nil 316 } 317 req := &pb.AllocateIdsRequest{ 318 ModelKey: keyToProto("", NewIncompleteKey(c, kind, parent)), 319 Size: proto.Int64(int64(n)), 320 } 321 res := &pb.AllocateIdsResponse{} 322 if err := internal.Call(c, "datastore_v3", "AllocateIds", req, res); err != nil { 323 return 0, 0, err 324 } 325 // The protobuf is inclusive at both ends. Idiomatic Go (e.g. slices, for loops) 326 // is inclusive at the low end and exclusive at the high end, so we add 1. 327 low = res.GetStart() 328 high = res.GetEnd() + 1 329 if low+int64(n) != high { 330 return 0, 0, fmt.Errorf("datastore: internal error: could not allocate %d IDs", n) 331 } 332 return low, high, nil 333} 334 335// AllocateIDRange allocates a range of IDs with specific endpoints. 336// The range is inclusive at both the low and high end. Once these IDs have been 337// allocated, you can manually assign them to newly created entities. 338// 339// The Datastore's automatic ID allocator never assigns a key that has already 340// been allocated (either through automatic ID allocation or through an explicit 341// AllocateIDs call). As a result, entities written to the given key range will 342// never be overwritten. However, writing entities with manually assigned keys in 343// this range may overwrite existing entities (or new entities written by a separate 344// request), depending on the error returned. 345// 346// Use this only if you have an existing numeric ID range that you want to reserve 347// (for example, bulk loading entities that already have IDs). If you don't care 348// about which IDs you receive, use AllocateIDs instead. 349// 350// AllocateIDRange returns nil if the range is successfully allocated. If one or more 351// entities with an ID in the given range already exist, it returns a KeyRangeCollisionError. 352// If the Datastore has already cached IDs in this range (e.g. from a previous call to 353// AllocateIDRange), it returns a KeyRangeContentionError. Errors of other types indicate 354// problems with arguments or an error returned directly from the Datastore. 355func AllocateIDRange(c context.Context, kind string, parent *Key, start, end int64) (err error) { 356 if kind == "" { 357 return errors.New("datastore: AllocateIDRange given an empty kind") 358 } 359 360 if start < 1 || end < 1 { 361 return errors.New("datastore: AllocateIDRange start and end must both be greater than 0") 362 } 363 364 if start > end { 365 return errors.New("datastore: AllocateIDRange start must be before end") 366 } 367 368 req := &pb.AllocateIdsRequest{ 369 ModelKey: keyToProto("", NewIncompleteKey(c, kind, parent)), 370 Max: proto.Int64(end), 371 } 372 res := &pb.AllocateIdsResponse{} 373 if err := internal.Call(c, "datastore_v3", "AllocateIds", req, res); err != nil { 374 return err 375 } 376 377 // Check for collisions, i.e. existing entities with IDs in this range. 378 // We could do this before the allocation, but we'd still have to do it 379 // afterward as well to catch the race condition where an entity is inserted 380 // after that initial check but before the allocation. Skip the up-front check 381 // and just do it once. 382 q := NewQuery(kind).Filter("__key__ >=", NewKey(c, kind, "", start, parent)). 383 Filter("__key__ <=", NewKey(c, kind, "", end, parent)).KeysOnly().Limit(1) 384 385 keys, err := q.GetAll(c, nil) 386 if err != nil { 387 return err 388 } 389 if len(keys) != 0 { 390 return &KeyRangeCollisionError{start: start, end: end} 391 } 392 393 // Check for a race condition, i.e. cases where the datastore may have 394 // cached ID batches that contain IDs in this range. 395 if start < res.GetStart() { 396 return &KeyRangeContentionError{start: start, end: end} 397 } 398 399 return nil 400} 401