A Sequence object using Appengine's unimplemented AllocateIDs RPC

sequence.go

/* Sequence: https://gist.github.com/1342205

 Example:

 var defSeq *sequence.Sequence

 func init() {
 	defSeq = sequence.Make("Default")
 	defSeq.BatchSize = 3 // for testing...
 }

 func NextHandler(w http.ResponseWriter, r *http.Request) {
 	c := appengine.NewContext(r)

 	next, err := defSeq.NextID(c)

 	fmt.Fprintf(w, "next: %d err: %v", next, err)
 }

*/

package sequence

import (
	"appengine"
	"appengine/datastore"
	"os"
	pb "appengine_internal/datastore"
	"sync"
)

type Sequence struct {
	sync.Mutex
	BatchSize   int64
	key         *datastore.Key
	kind        string
	next        int64
	end         int64
}

func Make(kind string) *Sequence {
	return &Sequence{
		BatchSize: 1000,
		kind: kind,
	}
}

func (s *Sequence) NextID(c appengine.Context) (next int64, err os.Error) {
	s.Lock()
	defer s.Unlock()

	if s.key == nil {
		s.key = datastore.NewIncompleteKey(c, s.kind, nil)
	}

	if s.next == s.end {
		s.next, s.end, err = allocateIds(c, s.key, s.BatchSize)
	} else {
		s.next++
	}
	return s.next, err
}

// Implements Appengine's missing AllocateIDs API.
func allocateIds(c appengine.Context, key *datastore.Key, size int64) (int64, int64, os.Error) {

	req := &pb.AllocateIdsRequest{ModelKey: keyToProto(c.FullyQualifiedAppID(), key)}
	req.Size = &size

	res := &pb.AllocateIdsResponse{}
	err := c.Call("datastore_v3", "AllocateIds", req, res, nil)

	if err != nil {
		return 0, 0, err
	}
	return *res.Start, *res.End, nil
}

// Private helper coppied from Appengine: 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--
		kk := i.Kind()
		ks := i.StringID()
		ki := i.IntID()
		e[n] = &pb.Path_Element{
			Type: &kk,
		}
		// Both Name and Id are optional proto fields, but the App Server expects
		// that exactly one of those fields are set.
		if i.StringID() != "" {
			e[n].Name = &ks
		} else {
			e[n].Id = &ki
		}
	}
	return &pb.Reference{
		App: &appID,
		Path: &pb.Path{
			Element: e,
		},
	}
}