vendor/google.golang.org/appengine/datastore/key.go

   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
   5 package datastore
   6
   7 import (
   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
  23 type KeyRangeCollisionError struct {
  24         start int64
  25         end   int64
  26 }
  27
  28 func (e *KeyRangeCollisionError) Error() string {
  29         return fmt.Sprintf("datastore: Collision when attempting to allocate range [%d, %d]",
  30                 e.start, e.end)
  31 }
  32
  33 type KeyRangeContentionError struct {
  34         start int64
  35         end   int64
  36 }
  37
  38 func (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.
  44 type 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).
  54 func (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 "".
  60 func (k *Key) StringID() string {
  61         return k.stringID
  62 }
  63
  64 // IntID returns the key's integer ID, which may be 0.
  65 func (k *Key) IntID() int64 {
  66         return k.intID
  67 }
  68
  69 // Parent returns the key's parent key, which may be nil.
  70 func (k *Key) Parent() *Key {
  71         return k.parent
  72 }
  73
  74 // AppID returns the key's application ID.
  75 func (k *Key) AppID() string {
  76         return k.appID
  77 }
  78
  79 // Namespace returns the key's namespace.
  80 func (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.
  86 func (k *Key) Incomplete() bool {
  87         return k.stringID == "" && k.intID == 0
  88 }
  89
  90 // valid returns whether the key is valid.
  91 func (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.
 115 func (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.
 126 func (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.
 134 func (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.
 149 func (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
 158 type gobKey struct {
 159         Kind      string
 160         StringID  string
 161         IntID     int64
 162         Parent    *gobKey
 163         AppID     string
 164         Namespace string
 165 }
 166
 167 func 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
 181 func 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
 195 func (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
 203 func (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
 212 func (k *Key) MarshalJSON() ([]byte, error) {
 213         return []byte(`"` + k.Encode() + `"`), nil
 214 }
 215
 216 func (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.
 231 func (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.
 244 func 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                 return nil, err
 258         }
 259
 260         return protoToKey(ref)
 261 }
 262
 263 // NewIncompleteKey creates a new incomplete key.
 264 // kind cannot be empty.
 265 func NewIncompleteKey(c context.Context, kind string, parent *Key) *Key {
 266         return NewKey(c, kind, "", 0, parent)
 267 }
 268
 269 // NewKey creates a new key.
 270 // kind cannot be empty.
 271 // Either one or both of stringID and intID must be zero. If both are zero,
 272 // the key returned is incomplete.
 273 // parent must either be a complete key or nil.
 274 func NewKey(c context.Context, kind, stringID string, intID int64, parent *Key) *Key {
 275         // If there's a parent key, use its namespace.
 276         // Otherwise, use any namespace attached to the context.
 277         var namespace string
 278         if parent != nil {
 279                 namespace = parent.namespace
 280         } else {
 281                 namespace = internal.NamespaceFromContext(c)
 282         }
 283
 284         return &Key{
 285                 kind:      kind,
 286                 stringID:  stringID,
 287                 intID:     intID,
 288                 parent:    parent,
 289                 appID:     internal.FullyQualifiedAppID(c),
 290                 namespace: namespace,
 291         }
 292 }
 293
 294 // AllocateIDs returns a range of n integer IDs with the given kind and parent
 295 // combination. kind cannot be empty; parent may be nil. The IDs in the range
 296 // returned will not be used by the datastore's automatic ID sequence generator
 297 // and may be used with NewKey without conflict.
 298 //
 299 // The range is inclusive at the low end and exclusive at the high end. In
 300 // other words, valid intIDs x satisfy low <= x && x < high.
 301 //
 302 // If no error is returned, low + n == high.
 303 func AllocateIDs(c context.Context, kind string, parent *Key, n int) (low, high int64, err error) {
 304         if kind == "" {
 305                 return 0, 0, errors.New("datastore: AllocateIDs given an empty kind")
 306         }
 307         if n < 0 {
 308                 return 0, 0, fmt.Errorf("datastore: AllocateIDs given a negative count: %d", n)
 309         }
 310         if n == 0 {
 311                 return 0, 0, nil
 312         }
 313         req := &pb.AllocateIdsRequest{
 314                 ModelKey: keyToProto("", NewIncompleteKey(c, kind, parent)),
 315                 Size:     proto.Int64(int64(n)),
 316         }
 317         res := &pb.AllocateIdsResponse{}
 318         if err := internal.Call(c, "datastore_v3", "AllocateIds", req, res); err != nil {
 319                 return 0, 0, err
 320         }
 321         // The protobuf is inclusive at both ends. Idiomatic Go (e.g. slices, for loops)
 322         // is inclusive at the low end and exclusive at the high end, so we add 1.
 323         low = res.GetStart()
 324         high = res.GetEnd() + 1
 325         if low+int64(n) != high {
 326                 return 0, 0, fmt.Errorf("datastore: internal error: could not allocate %d IDs", n)
 327         }
 328         return low, high, nil
 329 }
 330
 331 // AllocateIDRange allocates a range of IDs with specific endpoints.
 332 // The range is inclusive at both the low and high end. Once these IDs have been
 333 // allocated, you can manually assign them to newly created entities.
 334 //
 335 // The Datastore's automatic ID allocator never assigns a key that has already
 336 // been allocated (either through automatic ID allocation or through an explicit
 337 // AllocateIDs call). As a result, entities written to the given key range will
 338 // never be overwritten. However, writing entities with manually assigned keys in
 339 // this range may overwrite existing entities (or new entities written by a separate
 340 // request), depending on the error returned.
 341 //
 342 // Use this only if you have an existing numeric ID range that you want to reserve
 343 // (for example, bulk loading entities that already have IDs). If you don't care
 344 // about which IDs you receive, use AllocateIDs instead.
 345 //
 346 // AllocateIDRange returns nil if the range is successfully allocated. If one or more
 347 // entities with an ID in the given range already exist, it returns a KeyRangeCollisionError.
 348 // If the Datastore has already cached IDs in this range (e.g. from a previous call to
 349 // AllocateIDRange), it returns a KeyRangeContentionError. Errors of other types indicate
 350 // problems with arguments or an error returned directly from the Datastore.
 351 func AllocateIDRange(c context.Context, kind string, parent *Key, start, end int64) (err error) {
 352         if kind == "" {
 353                 return errors.New("datastore: AllocateIDRange given an empty kind")
 354         }
 355
 356         if start < 1 || end < 1 {
 357                 return errors.New("datastore: AllocateIDRange start and end must both be greater than 0")
 358         }
 359
 360         if start > end {
 361                 return errors.New("datastore: AllocateIDRange start must be before end")
 362         }
 363
 364         req := &pb.AllocateIdsRequest{
 365                 ModelKey: keyToProto("", NewIncompleteKey(c, kind, parent)),
 366                 Max:      proto.Int64(end),
 367         }
 368         res := &pb.AllocateIdsResponse{}
 369         if err := internal.Call(c, "datastore_v3", "AllocateIds", req, res); err != nil {
 370                 return err
 371         }
 372
 373         // Check for collisions, i.e. existing entities with IDs in this range.
 374         // We could do this before the allocation, but we'd still have to do it
 375         // afterward as well to catch the race condition where an entity is inserted
 376         // after that initial check but before the allocation. Skip the up-front check
 377         // and just do it once.
 378         q := NewQuery(kind).Filter("__key__ >=", NewKey(c, kind, "", start, parent)).
 379                 Filter("__key__ <=", NewKey(c, kind, "", end, parent)).KeysOnly().Limit(1)
 380
 381         keys, err := q.GetAll(c, nil)
 382         if err != nil {
 383                 return err
 384         }
 385         if len(keys) != 0 {
 386                 return &KeyRangeCollisionError{start: start, end: end}
 387         }
 388
 389         // Check for a race condition, i.e. cases where the datastore may have
 390         // cached ID batches that contain IDs in this range.
 391         if start < res.GetStart() {
 392                 return &KeyRangeContentionError{start: start, end: end}
 393         }
 394
 395         return nil
 396 }