package uodatastore
import (
"errors"
"fmt"
)
type KeyRangeState int
const (
KeyRangeError KeyRangeState = iota
KeyRangeCollision
KeyRangeContention
KeyRangeEmpty
)
func AllocateIdRange(c context.Context, kind string, parent *datastore.Key, start int64, end int64) (KeyRangeState, error) {
if kind == "" {
return KeyRangeError, errors.New("datastore: AllocateIdRange given an empty kind")
}
if start < 1 {
return KeyRangeError, fmt.Errorf("datastore: AllocateIdRange Illegal start %d: less than 1", start)
}
if end < start {
return KeyRangeError, fmt.Errorf("datastore: AllocateIdRange Illegal end %d: less than start %d", end, start);
}
if parent != nil && parent.Incomplete() {
return KeyRangeError, errors.New("datastore: AllocateIdRange parent must be complete");
}
req := &pb.AllocateIdsRequest{
ModelKey: keyToProto("", datastore.NewIncompleteKey(c, kind, parent)),
Max: proto.Int64(end),
}
res := &pb.AllocateIdsResponse{}
if err := internal.Call(c, "datastore_v3", "AllocateIds", req, res); err != nil {
return KeyRangeError, err
}
q := datastore.NewQuery(kind).KeysOnly().Limit(1)
q = q.Filter("__key__ >=", datastore.NewKey(c, kind, "", start, parent))
q = q.Filter("__key__ <=", datastore.NewKey(c, kind, "", end, parent))
t := q.Run(c)
if _, err := t.Next(nil); err == nil {
return KeyRangeCollision, nil
} else if err != datastore.Done {
return KeyRangeError, err
}
if raceCondition := start < res.GetStart(); raceCondition {
return KeyRangeContention, nil
} else {
return KeyRangeEmpty, nil
}
}
// keyToProto converts a *Key to a Reference proto.
func keyToProto(defaultAppID string, k *datastore.Key) *pb.Reference {
appID := k.AppID()
if appID == "" {
appID = defaultAppID
}
n := 0
for i := k; i != nil; i = i.Parent() {
n++
}
e := make([]*pb.Path_Element, n)
for i := k; i != nil; i = i.Parent() {
n--
iKind := i.Kind()
e[n] = &pb.Path_Element{
Type: &iKind,
}
// At most one of {Name,Id} should be set.
// Neither will be set for incomplete keys.
if iStringID := i.StringID(); iStringID != "" {
e[n].Name = &iStringID
} else if iIntID := i.IntID(); iIntID != 0 {
e[n].Id = &iIntID
}
}
var namespace *string
if k.Namespace() != "" {
namespace = proto.String(k.Namespace())
}
return &pb.Reference{
App: proto.String(appID),
NameSpace: namespace,
Path: &pb.Path{
Element: e,
},
}
}