impl/memory: Implement Queries

impl/memory/datastore_query.go

 // Copyright 2015 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 package memory
 
 import (
+        "bytes"
+        "encoding/base64"
         "errors"
         "fmt"
         "math"
         "strings"
 
         ds "github.com/luci/gae/service/datastore"
         "github.com/luci/gae/service/datastore/serialize"
+        "github.com/luci/luci-go/common/cmpbin"
 )
 
 // MaxQueryComponents was lifted from a hard-coded constant in dev_appserver.
 // No idea if it's a real limit or just a convenience in the current dev
 // appserver implementation.
 const MaxQueryComponents = 100
 
 var errQueryDone = errors.New("query is done")
 
 type queryOp int
@@ skipping 29 matching lines @@
                 op = queryOpMap[toks[1]]
                 if op == qInvalid {
                         err = fmt.Errorf("datastore: invalid operator %q in filter %q", toks[1], f)
                 } else {
                         prop = toks[0]
                 }
         }
         return
 }
 
-type queryCursor string
+// A queryCursor is:
+//   {#orders} ++ IndexColumn* ++ RawRowData
+type queryCursor []byte
 
-func (q queryCursor) String() string { return string(q) }
-func (q queryCursor) Valid() bool    { return q != "" }
+func (q queryCursor) String() string {
+        return base64.URLEncoding.EncodeToString([]byte(q))
+}
+
+func (q queryCursor) decode() (cols []ds.IndexColumn, data []byte, err error) {
+        buf := bytes.NewBuffer([]byte(q))
+        count, _, err := cmpbin.ReadUint(buf)
+        if err != nil {
+                err = fmt.Errorf("invalid cursor")
+                return
+        }
+
+        if count > ds.MaxIndexColumns {
+                err = fmt.Errorf("invalid cursor")
+                return
+        }
+
+        cols = make([]ds.IndexColumn, count)
+        for i := range cols {
+                if cols[i], err = serialize.ReadIndexColumn(buf); err != nil {
+                        err = fmt.Errorf("invalid cursor")
+                        return
+                }
+        }
+
+        data = buf.Bytes()
+        return
+}
 
 type queryIneqFilter struct {
         prop string
 
-        low  *string
-        high *string
+        low  []byte
+        high []byte
 }
 
-func increment(bstr string, positive bool) string {
+func increment(bstr []byte) []byte {
+        // Copy bstr
+        bstr = bjoin(bstr)
+
         lastIdx := len(bstr) - 1
-        last := bstr[lastIdx]
-        if positive {
-                if last == 0xFF {
-                        return bstr + "\x00"
-                }
-                return bstr[:lastIdx] + string([]byte{last + 1})
-        } else {
-                if last == 0 {
-                        return bstr[:lastIdx-1]
-                }
-                return bstr[:lastIdx] + string([]byte{last - 1})
+        if bstr[lastIdx] == 0xFF {
+                return append(bstr, 0)
         }
+        bstr[lastIdx] += 1
+        return bstr
 }
 
 // constrain 'folds' a new inequality into the current inequality filter.
 //
 // It will bump the high bound down, or the low bound up, assuming the incoming
 // constraint does so.
 //
 // It returns true iff the filter is overconstrained (i.e.  low > high)
-func (q *queryIneqFilter) constrain(op queryOp, val string) bool {
+func (q *queryIneqFilter) constrain(op queryOp, val []byte) bool {
         switch op {
+        case qLessEq:
+                val = increment(val)
+                fallthrough
         case qLessThan:
-                val = increment(val, true)
-                fallthrough
-        case qLessEq:
                 // adjust upper bound downwards
-                if q.high == nil || *q.high > val {
-                        q.high = &val
+                if q.high == nil || bytes.Compare(q.high, val) > 0 {
+                        q.high = val
                 }
 
         case qGreaterThan:
-                val = increment(val, false)
+                val = increment(val)
                 fallthrough
         case qGreaterEq:
                 // adjust lower bound upwards
-                if q.low == nil || *q.low < val {
-                        q.low = &val
+                if q.low == nil || bytes.Compare(q.low, val) < 0 {
+                        q.low = val
                 }
 
         default:
                 panic(fmt.Errorf("constrain cannot handle filter op %d", op))
         }
 
         if q.low != nil && q.high != nil {
-                return *q.low > *q.high
+                return bytes.Compare(q.low, q.high) >= 0
         }
         return false
 }
 
 type queryImpl struct {
         ns string
 
         kind     string
         ancestor ds.Key
 
         // prop -> encoded values
         eqFilters  map[string]map[string]struct{}
         ineqFilter queryIneqFilter
         order      []ds.IndexColumn
-        project    map[string]struct{}
+        project    []string
 
         distinct            bool
         eventualConsistency bool
         keysOnly            bool
         limit               int32
         offset              int32
 
         start queryCursor
         end   queryCursor
 
         err error
 }
 
 var _ ds.Query = (*queryImpl)(nil)
 
-func (q *queryImpl) valid(ns string, isTxn bool) (done bool, err error) {
-        if q.err == errQueryDone {
-                done = true
-        } else if q.err != nil {
-                err = q.err
-        } else if ns != q.ns {
-                err = errors.New(
+func sortOrdersEqual(as, bs []ds.IndexColumn) bool {
+        if len(as) != len(bs) {
+                return false
+        }
+        for i, a := range as {
+                if a != bs[i] {
+                        return false
+                }
+        }
+        return true
+}
+
+func (q *queryImpl) prep(ns string, isTxn bool) (*reducedQuery, error) {
+        if q.err != nil {
+                return nil, q.err
+        }
+        if ns != q.ns {
+                return nil, errors.New(
                         "gae/memory: Namespace mismatched. Query and Datastore don't agree " +
                                 "on the current namespace")
         } else if isTxn && q.ancestor == nil {
-                err = errors.New(
+                return nil, errors.New(
                         "gae/memory: Only ancestor queries are allowed inside transactions")
         } else if q.numComponents() > MaxQueryComponents {
-                err = fmt.Errorf(
+                return nil, fmt.Errorf(
                         "gae/memory: query is too large. may not have more than "+
                                 "%d filters + sort orders + ancestor total: had %d",
                         MaxQueryComponents, q.numComponents())
         } else if len(q.project) == 0 && q.distinct {
                 // This must be delayed, because q.Distinct().Project("foo") is a valid
                 // construction. If we checked this in Distinct, it could be too early, and
                 // checking it in Project doesn't matter.
-                err = errors.New(
+                return nil, errors.New(
                         "gae/memory: Distinct() only makes sense on projection queries.")
         }
-        return
+        ret := &reducedQuery{
+                ns:           q.ns,
+                kind:         q.kind,
+                eqFilters:    q.eqFilters,
+                suffixFormat: q.order,
+        }
+        if len(ret.suffixFormat) == 0 && q.ineqFilter.prop != "" {
+                ret.suffixFormat = []ds.IndexColumn{{Property: q.ineqFilter.prop}}
+        }
+        if q.ancestor != nil {
+                ret.ancestor = serialize.ToBytes(q.ancestor)
+        }
+
+        // The inequality is specified in natural (ascending) order in the query, but
+        // the order information may indicate to use an descending index column.  If
+        // that's the case, then we must invert, swap and increment the inequality
+        // endpoints.
+        //
+        // Invert so that the desired numbers are represented correctly in the index.
+        // Swap so that our iterators still go from >= low to < high.
+        // Increment so that >= and < get correctly bounded (since the iterator is
+        //   still using natrual ordering)
+        shouldInvertIneq := false
+        if len(ret.suffixFormat) > 0 && ret.suffixFormat[0].Direction == ds.DESCENDING {
+                shouldInvertIneq = true
+        }
+
+        err := error(nil)
+        cols := []ds.IndexColumn(nil)
+        if q.start != nil {
+                cols, ret.low, err = q.start.decode()
+                if err != nil {
+                        return nil, err
+                }
+                if !sortOrdersEqual(cols, ret.suffixFormat) {
+                        return nil, errors.New("gae/memory: start cursor is invalid.")
+                }
+        } else {
+                if shouldInvertIneq {
+                        ret.high = increment(invert(q.ineqFilter.low))
+                } else {
+                        ret.low = q.ineqFilter.low
+                }
+        }
+
+        if q.end != nil {
+                cols, ret.high, err = q.end.decode()
+                if err != nil {
+                        return nil, err
+                }
+                if !sortOrdersEqual(cols, ret.suffixFormat) {
+                        return nil, errors.New("gae/memory: end cursor is invalid.")
+                }
+        } else {
+                if shouldInvertIneq {
+                        ret.low = increment(invert(q.ineqFilter.high))
+                } else {
+                        ret.high = q.ineqFilter.high
+                }
+        }
+
+        // Add any projection columns not mentioned in the user-defined order as
+        // ASCENDING orders. Technically we could be smart and automatically use
+        // a DESCENDING ordered index, if it fit, but the logic gets insane, since all
+        // suffixes of all used indexes need to be PRECISELY equal (and so you'd have
+        // to hunt/invalidate/something to find the combination of indexes that are
+        // compatible with each other as well as the query). If you want to use
+        // a DESCENDING column, just add it to the user sort order, and this loop will
+        // not synthesize a new sort order for it.
+        originalStop := len(ret.suffixFormat)
+        for _, p := range q.project {
+                needAdd := true
+                // originalStop prevents this loop from getting longer every time we add
+                // a projected property.
+                for _, col := range ret.suffixFormat[:originalStop] {
+                        if col.Property == p {
+                                needAdd = false
+                                break
+                        }
+                }
+                if needAdd {
+                        ret.suffixFormat = append(ret.suffixFormat, ds.IndexColumn{Property: p})
+                }
+        }
+
+        return ret, nil
+
 }
 
 func (q *queryImpl) numComponents() int {
         numComponents := len(q.order)
         if q.ineqFilter.prop != "" {
                 if q.ineqFilter.low != nil {
                         numComponents++
                 }
                 if q.ineqFilter.high != nil {
                         numComponents++
                 }
         }
         for _, v := range q.eqFilters {
                 numComponents += len(v)
         }
         if q.ancestor != nil {
                 numComponents++
         }
         return numComponents
 }
 
-func (q *queryImpl) calculateIndex() *ds.IndexDefinition {
-        // as a nod to simplicity in this code, we'll require that a single index
-        // is able to service the entire query. E.g. no zigzag merge joins or
-        // multiqueries. This will mean that the user will need to rely on
-        // dev_appserver to tell them what indicies they need for real, and for thier
-        // tests they'll need to specify the missing composite indices manually.
-        //
-        // This COULD lead to an exploding indicies problem, but we can fix that when
-        // we get to it.
-
-        //sortOrders := []qSortBy{}
-
-        return nil
-}
-
 // checkMutateClone sees if the query has an error. If not, it clones the query,
 // and assigns the output of `check` to the query error slot. If check returns
 // nil, it calls `mutate` on the cloned query. The (possibly new) query is then
 // returned.
 func (q *queryImpl) checkMutateClone(check func() error, mutate func(*queryImpl)) *queryImpl {
         if q.err != nil {
                 return q
         }
         nq := *q
         nq.eqFilters = make(map[string]map[string]struct{}, len(q.eqFilters))
         for prop, vals := range q.eqFilters {
                 nq.eqFilters[prop] = make(map[string]struct{}, len(vals))
                 for v := range vals {
                         nq.eqFilters[prop][v] = struct{}{}
                 }
         }
         nq.order = make([]ds.IndexColumn, len(q.order))
         copy(nq.order, q.order)
-        nq.project = make(map[string]struct{}, len(q.project))
-        for f := range q.project {
-                nq.project[f] = struct{}{}
-        }
+        nq.project = make([]string, len(q.project))
+        copy(nq.project, q.project)
         if check != nil {
                 nq.err = check()
         }
         if nq.err == nil {
                 mutate(&nq)
         }
         return &nq
 }
 
 func (q *queryImpl) Ancestor(k ds.Key) ds.Query {
         return q.checkMutateClone(
                 func() error {
                         if k == nil {
                                 // SDK has an explicit nil-check
                                 return errors.New("datastore: nil query ancestor")
                         }
+                        if k.Namespace() != q.ns {
+                                return fmt.Errorf("bad namespace: %q (expected %q)", k.Namespace(), q.ns)
+                        }
                         if !k.Valid(false, globalAppID, q.ns) {
                                 // technically the SDK implementation does a Weird Thing (tm) if both the
                                 // stringID and intID are set on a key; it only serializes the stringID in
                                 // the proto. This means that if you set the Ancestor to an invalid key,
                                 // you'll never actually hear about it. Instead of doing that insanity, we
                                 // just swap to an error here.
                                 return ds.ErrInvalidKey
                         }
-                        if k.Namespace() != q.ns {
-                                return fmt.Errorf("bad namespace: %q (expected %q)", k.Namespace(), q.ns)
-                        }
                         if q.ancestor != nil {
                                 return errors.New("cannot have more than one ancestor")
                         }
                         return nil
                 },
                 func(q *queryImpl) {
                         q.ancestor = k
                 })
 }
 
 func (q *queryImpl) Distinct() ds.Query {
         return q.checkMutateClone(nil, func(q *queryImpl) {
                 q.distinct = true
         })
 }
 
 func (q *queryImpl) Filter(fStr string, val interface{}) ds.Query {
         prop := ""
         op := qInvalid
-        binVal := ""
+        binVal := []byte(nil)
         return q.checkMutateClone(
                 func() error {
                         var err error
                         prop, op, err = parseFilter(fStr)
                         if err != nil {
                                 return err
                         }
 
                         if q.kind == "" && prop != "__key__" {
                                 // https://cloud.google.com/appengine/docs/go/datastore/queries#Go_Kindless_queries
@@ skipping 27 matching lines @@
                                 if q.ineqFilter.prop != "" && q.ineqFilter.prop != prop {
                                         return fmt.Errorf(
                                                 "inequality filters on multiple properties: %q and %q",
                                                 q.ineqFilter.prop, prop)
                                 }
                                 if len(q.order) > 0 && q.order[0].Property != prop {
                                         return fmt.Errorf(
                                                 "first sort order must match inequality filter: %q v %q",
                                                 q.order[0].Property, prop)
                                 }
-                        } else if _, ok := q.project[prop]; ok {
-                                return fmt.Errorf(
-                                        "cannot project on field which is used in an equality filter: %q",
-                                        prop)
+                        } else {
+                                for _, p := range q.project {
+                                        if p == prop {
+                                                return fmt.Errorf(
+                                                        "cannot project on field which is used in an equality filter: %q",
+                                                        prop)
+                                        }
+                                }
                         }
-                        binVal = string(serialize.ToBytes(p))
+                        binVal = serialize.ToBytes(p)
                         return err
                 },
                 func(q *queryImpl) {
                         if op == qEqual {
                                 // add it to eq filters
                                 if _, ok := q.eqFilters[prop]; !ok {
-                                        q.eqFilters[prop] = map[string]struct{}{binVal: {}}
+                                        q.eqFilters[prop] = map[string]struct{}{string(binVal): {}}
                                 } else {
-                                        q.eqFilters[prop][binVal] = struct{}{}
+                                        q.eqFilters[prop][string(binVal)] = struct{}{}
                                 }
 
                                 // remove it from sort orders.
                                 // https://cloud.google.com/appengine/docs/go/datastore/queries#sort_orders_are_ignored_on_properties_with_equality_filters
                                 toRm := -1
                                 for i, o := range q.order {
                                         if o.Property == prop {
                                                 toRm = i
                                                 break
                                         }
@@ skipping 56 matching lines @@
                         }
                 })
 }
 
 func (q *queryImpl) Project(fieldName ...string) ds.Query {
         return q.checkMutateClone(
                 func() error {
                         if q.keysOnly {
                                 return errors.New("cannot project a keysOnly query")
                         }
+                        dupCheck := map[string]struct{}{}
                         for _, f := range fieldName {
+                                if _, ok := dupCheck[f]; ok {
+                                        return fmt.Errorf("cannot project on the same field twice: %q", f)
+                                }
+                                dupCheck[f] = struct{}{}
                                 if f == "" {
                                         return errors.New("cannot project on an empty field name")
                                 }
-                                if strings.HasPrefix(f, "__") && strings.HasSuffix(f, "__") {
-                                        return fmt.Errorf("cannot project on %q", f)
+                                if f == "__key__" {
+                                        return fmt.Errorf("cannot project on __key__")
                                 }
                                 if _, ok := q.eqFilters[f]; ok {
                                         return fmt.Errorf(
                                                 "cannot project on field which is used in an equality filter: %q", f)
                                 }
+                                for _, p := range q.project {
+                                        if p == f {
+                                                return fmt.Errorf("cannot project on the same field twice: %q", f)
+                                        }
+                                }
                         }
                         return nil
                 },
                 func(q *queryImpl) {
-                        for _, f := range fieldName {
-                                q.project[f] = struct{}{}
-                        }
+                        q.project = append(q.project, fieldName...)
                 })
 }
 
 func (q *queryImpl) KeysOnly() ds.Query {
         return q.checkMutateClone(
                 func() error {
                         if len(q.project) != 0 {
                                 return errors.New("cannot project a keysOnly query")
                         }
                         return nil
@@ skipping 25 matching lines @@
                         if offset > math.MaxInt32 {
                                 return errors.New("datastore: query offset overflow")
                         }
                         return nil
                 },
                 func(q *queryImpl) {
                         q.offset = int32(offset)
                 })
 }
 
+func queryCursorCheck(ns, flavor string, current queryCursor, newCursor ds.Cursor) (queryCursor, error) {
+        if current != nil {
+                return nil, fmt.Errorf("%s cursor is multiply defined", flavor)
+        }
+        curs, ok := newCursor.(queryCursor)
+        if !ok {
+                return nil, fmt.Errorf("%s cursor is unknown type: %T", flavor, curs)
+        }
+        _, _, err := curs.decode()
+        return curs, err
+}
+
 func (q *queryImpl) Start(c ds.Cursor) ds.Query {
-        curs := queryCursor("")
+        curs := queryCursor(nil)
         return q.checkMutateClone(
-                func() error {
-                        ok := false
-                        if curs, ok = c.(queryCursor); !ok {
-                                return fmt.Errorf("start cursor is unknown type: %T", c)
-                        }
-                        if !curs.Valid() {
-                                return errors.New("datastore: invalid cursor")
-                        }
-                        return nil
+                func() (err error) {
+                        curs, err = queryCursorCheck(q.ns, "start", q.start, c)
+                        return
                 },
                 func(q *queryImpl) {
                         q.start = curs
                 })
 }
 
 func (q *queryImpl) End(c ds.Cursor) ds.Query {
-        curs := queryCursor("")
+        curs := queryCursor(nil)
         return q.checkMutateClone(
-                func() error {
-                        ok := false
-                        if curs, ok = c.(queryCursor); !ok {
-                                return fmt.Errorf("end cursor is unknown type: %T", c)
-                        }
-                        if !curs.Valid() {
-                                return errors.New("datastore: invalid cursor")
-                        }
-                        return nil
+                func() (err error) {
+                        curs, err = queryCursorCheck(q.ns, "end", q.end, c)
+                        return
                 },
                 func(q *queryImpl) {
                         q.end = curs
                 })
 }
 
 func (q *queryImpl) EventualConsistency() ds.Query {
         return q.checkMutateClone(
                 nil, func(q *queryImpl) {
                         q.eventualConsistency = true
                 })
 }