Interface update, per-method Contexts.

impl/memory/datastore_query_execution.go

 // Copyright 2015 The LUCI Authors. All rights reserved.
 // Use of this source code is governed under the Apache License, Version 2.0
 // that can be found in the LICENSE file.
 
 package memory
 
 import (
         "bytes"
         "errors"
         "fmt"
@@ skipping 79 matching lines @@
 func (s *keysOnlyStrategy) handle(rawData [][]byte, _ []ds.Property, key *ds.Key, gc func() (ds.Cursor, error)) error {
         if !s.dedup.Add(string(rawData[len(rawData)-1])) {
                 return nil
         }
         return s.cb(key, nil, gc)
 }
 
 type normalStrategy struct {
         cb ds.RawRunCB
 
-        aid   string
-        ns    string
+        kc    ds.KeyContext
         head  memCollection
         dedup stringset.Set
 }
 
-func newNormalStrategy(aid, ns string, cb ds.RawRunCB, head memStore) queryStrategy {
-        coll := head.GetCollection("ents:" + ns)
+func newNormalStrategy(kc ds.KeyContext, cb ds.RawRunCB, head memStore) queryStrategy {
+        coll := head.GetCollection("ents:" + kc.Namespace)
         if coll == nil {
                 return nil
         }
-        return &normalStrategy{cb, aid, ns, coll, stringset.New(0)}
+        return &normalStrategy{cb, kc, coll, stringset.New(0)}
 }
 
 func (s *normalStrategy) handle(rawData [][]byte, _ []ds.Property, key *ds.Key, gc func() (ds.Cursor, error)) error {
         rawKey := rawData[len(rawData)-1]
         if !s.dedup.Add(string(rawKey)) {
                 return nil
         }
 
         rawEnt := s.head.Get(rawKey)
         if rawEnt == nil {
                 // entity doesn't exist at head
                 return nil
         }
-        pm, err := serialize.ReadPropertyMap(bytes.NewBuffer(rawEnt), serialize.WithoutContext, s.aid, s.ns)
+        pm, err := serialize.ReadPropertyMap(bytes.NewBuffer(rawEnt), serialize.WithoutContext, s.kc)
         memoryCorruption(err)
 
         return s.cb(key, pm, gc)
 }
 
 func pickQueryStrategy(fq *ds.FinalizedQuery, rq *reducedQuery, cb ds.RawRunCB, head memStore) queryStrategy {
         if fq.KeysOnly() {
                 return &keysOnlyStrategy{cb, stringset.New(0)}
         }
         if len(fq.Project()) > 0 {
                 return newProjectionStrategy(fq, rq, cb)
         }
-        return newNormalStrategy(rq.aid, rq.ns, cb, head)
+        return newNormalStrategy(rq.kc, cb, head)
 }
 
 func parseSuffix(aid, ns string, suffixFormat []ds.IndexColumn, suffix []byte, count int) (raw [][]byte, decoded []ds.Property) {
         buf := serialize.Invertible(bytes.NewBuffer(suffix))
         decoded = make([]ds.Property, len(suffixFormat))
         raw = make([][]byte, len(suffixFormat))
 
         err := error(nil)
+        kc := ds.KeyContext{aid, ns}
         for i := range decoded {
                 if count >= 0 && i >= count {
                         break
                 }
                 needInvert := suffixFormat[i].Descending
 
                 buf.SetInvert(needInvert)
-                decoded[i], err = serialize.ReadProperty(buf, serialize.WithoutContext, aid, ns)
+                decoded[i], err = serialize.ReadProperty(buf, serialize.WithoutContext, kc)
                 memoryCorruption(err)
 
                 offset := len(suffix) - buf.Len()
                 raw[i] = suffix[:offset]
                 suffix = suffix[offset:]
                 if needInvert {
                         raw[i] = serialize.Invert(raw[i])
                 }
         }
 
         return
 }
 
-func countQuery(fq *ds.FinalizedQuery, aid, ns string, isTxn bool, idx, head memStore) (ret int64, err error) {
+func countQuery(fq *ds.FinalizedQuery, kc ds.KeyContext, isTxn bool, idx, head memStore) (ret int64, err error) {
         if len(fq.Project()) == 0 && !fq.KeysOnly() {
                 fq, err = fq.Original().KeysOnly(true).Finalize()
                 if err != nil {
                         return
                 }
         }
-        err = executeQuery(fq, aid, ns, isTxn, idx, head, func(_ *ds.Key, _ ds.PropertyMap, _ ds.CursorCB) error {
+        err = executeQuery(fq, kc, isTxn, idx, head, func(_ *ds.Key, _ ds.PropertyMap, _ ds.CursorCB) error {
                 ret++
                 return nil
         })
         return
 }
 
-func executeNamespaceQuery(fq *ds.FinalizedQuery, aid string, head memStore, cb ds.RawRunCB) error {
+func executeNamespaceQuery(fq *ds.FinalizedQuery, kc ds.KeyContext, head memStore, cb ds.RawRunCB) error {
         // these objects have no properties, so any filters on properties cause an
         // empty result.
         if len(fq.EqFilters()) > 0 || len(fq.Project()) > 0 || len(fq.Orders()) > 1 {
                 return nil
         }
         if !(fq.IneqFilterProp() == "" || fq.IneqFilterProp() == "__key__") {
                 return nil
         }
         limit, hasLimit := fq.Limit()
         offset, hasOffset := fq.Offset()
         start, end := fq.Bounds()
 
         cursErr := errors.New("cursors not supported for __namespace__ query")
         cursFn := func() (ds.Cursor, error) { return nil, cursErr }
         if !(start == nil && end == nil) {
                 return cursErr
         }
+
+        kc.Namespace = ""
         for _, ns := range namespaces(head) {
                 if hasOffset && offset > 0 {
                         offset--
                         continue
                 }
                 if hasLimit {
                         if limit <= 0 {
                                 return ds.Stop
                         }
                         limit--
                 }
                 k := (*ds.Key)(nil)
                 if ns == "" {
                         // Datastore uses an id of 1 to indicate the default namespace in its
                         // metadata API.
-                        k = ds.MakeKey(aid, "", "__namespace__", 1)
+                        k = kc.MakeKey("__namespace__", 1)
                 } else {
-                        k = ds.MakeKey(aid, "", "__namespace__", ns)
+                        k = kc.MakeKey("__namespace__", ns)
                 }
                 if err := cb(k, nil, cursFn); err != nil {
                         return err
                 }
         }
         return nil
 }
 
-func executeQuery(fq *ds.FinalizedQuery, aid, ns string, isTxn bool, idx, head memStore, cb ds.RawRunCB) error {
-        rq, err := reduce(fq, aid, ns, isTxn)
+func executeQuery(fq *ds.FinalizedQuery, kc ds.KeyContext, isTxn bool, idx, head memStore, cb ds.RawRunCB) error {
+        rq, err := reduce(fq, kc, isTxn)
         if err == ds.ErrNullQuery {
                 return nil
         }
         if err != nil {
                 return err
         }
 
         if rq.kind == "__namespace__" {
-                return executeNamespaceQuery(fq, aid, head, cb)
+                return executeNamespaceQuery(fq, kc, head, cb)
         }
 
         idxs, err := getIndexes(rq, idx)
         if err == ds.ErrNullQuery {
                 return nil
         }
         if err != nil {
                 return err
         }
 
@@ skipping 32 matching lines @@
                         offset--
                         return nil
                 }
                 if hasLimit {
                         if limit <= 0 {
                                 return ds.Stop
                         }
                         limit--
                 }
 
-                rawData, decodedProps := parseSuffix(aid, ns, rq.suffixFormat, suffix, -1)
+                rawData, decodedProps := parseSuffix(kc.AppID, kc.Namespace, rq.suffixFormat, suffix, -1)
 
                 keyProp := decodedProps[len(decodedProps)-1]
                 if keyProp.Type() != ds.PTKey {
                         impossible(fmt.Errorf("decoded index row doesn't end with a Key: %#v", keyProp))
                 }
 
                 return strategy.handle(
                         rawData, decodedProps, keyProp.Value().(*ds.Key),
                         getCursorFn(suffix))
         })
 }