Refactor keys and queries in datastore service and implementation.

service/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 datastore
+
+import (
+        "bytes"
+        "fmt"
+        "sort"
+        "strings"
+
+        "github.com/luci/luci-go/common/errors"
+        "github.com/luci/luci-go/common/stringset"
+)
+
+var ErrMultipleInequalityFilter = errors.New(
+        "inequality filters on multiple properties in the same Query is not allowed")
+
+var ErrNullQuery = errors.New(
+        "the query is overconstraind and can never have results")
+
+type Query struct {
+        kind string
+
+        eventualConsistency bool
+        keysOnly            bool
+        distinct            bool
+
+        limit  *int32
+        offset *int32
+
+        order   []IndexColumn
+        project stringset.Set
+
+        eqFilts map[string]PropertySlice
+
+        ineqFiltProp     string
+        ineqFiltLow      Property
+        ineqFiltLowIncl  bool
+        ineqFiltLowSet   bool
+        ineqFiltHigh     Property
+        ineqFiltHighIncl bool
+        ineqFiltHighSet  bool
+
+        start Cursor
+        end   Cursor
+
+        // These are set by Finalize as a way to cache the 1-1 correspondence of
+        // a Query to its FinalizedQuery form. err may also be set by intermediate
+        // Query functions if there's a problem before finalization.
+        finalized *FinalizedQuery
+        err       error
+}
+
+func NewQuery(kind string) *Query {
+        return &Query{kind: kind}
+}
+
+func (q *Query) mod(cb func(*Query)) *Query {
+        if q.err != nil {
+                return q
+        }
+
+        ret := *q
+        ret.finalized = nil
+        if len(q.order) > 0 {
+                ret.order = make([]IndexColumn, len(q.order))
+                copy(ret.order, q.order)
+        }
+        if q.project != nil {
+                ret.project = q.project.Dup()
+        }
+        if len(q.eqFilts) > 0 {
+                ret.eqFilts = make(map[string]PropertySlice, len(q.eqFilts))
+                for k, v := range q.eqFilts {
+                        newV := make(PropertySlice, len(v))
+                        copy(newV, v)
+                        ret.eqFilts[k] = newV
+                }
+        }
+        cb(&ret)
+        return &ret
+}
+
+func (q *Query) Kind(kind string) *Query {
+        return q.mod(func(q *Query) {
+                q.kind = kind
+        })
+}
+
+func (q *Query) Ancestor(ancestor *Key) *Query {
+        return q.mod(func(q *Query) {
+                if q.eqFilts == nil {
+                        q.eqFilts = map[string]PropertySlice{}
+                }
+                if ancestor == nil {
+                        delete(q.eqFilts, "__ancestor__")
+                        if len(q.eqFilts) == 0 {
+                                q.eqFilts = nil
+                        }
+                } else {
+                        q.eqFilts["__ancestor__"] = PropertySlice{MkProperty(ancestor)}
+                }
+        })
+}
+
+func (q *Query) EventualConsistency(on bool) *Query {
+        return q.mod(func(q *Query) {
+                q.eventualConsistency = on
+        })
+}
+
+func (q *Query) Limit(limit int32) *Query {
+        return q.mod(func(q *Query) {
+                if limit < 0 {
+                        q.limit = nil
+                } else {
+                        q.limit = &limit
+                }
+        })
+}
+
+func (q *Query) Offset(offset int32) *Query {
+        return q.mod(func(q *Query) {
+                if offset < 0 {
+                        q.offset = nil
+                } else {
+                        q.offset = &offset
+                }
+        })
+}
+
+func (q *Query) KeysOnly(on bool) *Query {
+        return q.mod(func(q *Query) {
+                q.keysOnly = on
+        })
+}
+
+func (q *Query) Order(fieldNames ...string) *Query {
+        if len(fieldNames) == 0 {
+                return q
+        }
+        return q.mod(func(q *Query) {
+                for _, fn := range fieldNames {
+                        ic, err := ParseIndexColumn(fn)
+                        if err != nil {
+                                q.err = err
+                                return
+                        }
+                        if q.reserved(ic.Property) {
+                                return
+                        }
+                        q.order = append(q.order, ic)
+                }
+        })
+}
+
+func (q *Query) ClearOrder() *Query {
+        return q.mod(func(q *Query) {
+                q.order = nil
+        })
+}
+
+func (q *Query) Project(fieldNames ...string) *Query {
+        if len(fieldNames) == 0 {
+                return q
+        }
+        return q.mod(func(q *Query) {
+                for _, f := range fieldNames {
+                        if q.reserved(f) {
+                                return
+                        }
+                        if f == "__key__" {
+                                q.err = fmt.Errorf("cannot project on %q", f)
+                                return
+                        }
+                        if q.project == nil {
+                                q.project = stringset.New(1)
+                        }
+                        q.project.Add(f)
+                }
+        })
+}
+
+func (q *Query) Distinct(on bool) *Query {
+        return q.mod(func(q *Query) {
+                q.distinct = on
+        })
+}
+
+func (q *Query) ClearProject() *Query {
+        return q.mod(func(q *Query) {
+                q.project = nil
+        })
+}
+
+func (q *Query) Start(c Cursor) *Query {
+        return q.mod(func(q *Query) {
+                q.start = c
+        })
+}
+
+func (q *Query) End(c Cursor) *Query {
+        return q.mod(func(q *Query) {
+                q.end = c
+        })
+}
+
+func (q *Query) Eq(field string, values ...interface{}) *Query {
+        return q.mod(func(q *Query) {
+                if !q.reserved(field) {
+                        if q.eqFilts == nil {
+                                q.eqFilts = make(map[string]PropertySlice, 1)
+                        }
+                        s := q.eqFilts[field]
+                        for _, value := range values {
+                                p := Property{}
+                                if q.err = p.SetValue(value, ShouldIndex); q.err != nil {
+                                        return
+                                }
+                                idx := sort.Search(len(s), func(i int) bool {
+                                        return s[i].Equal(&p)
+                                })
+                                if idx == len(s) {
+                                        s = append(s, p)
+                                        sort.Sort(s)
+                                }
+                        }
+                        q.eqFilts[field] = s
+                }
+        })
+}
+
+func (q *Query) reserved(field string) bool {
+        if field == "__key__" {
+                return false
+        }
+        if field == "" {
+                q.err = fmt.Errorf(
+                        "cannot filter/project on: %q", field)
+                return true
+        }
+        if strings.HasPrefix(field, "__") && strings.HasSuffix(field, "__") {
+                q.err = fmt.Errorf(
+                        "cannot filter/project on reserved property: %q", field)
+                return true
+        }
+        return false
+}
+
+func (q *Query) ineqOK(field string, value Property) bool {
+        if q.reserved(field) {
+                return false
+        }
+        if field == "__key__" && value.Type() != PTKey {
+                q.err = fmt.Errorf(
+                        "filters on %q must have type *Key (got %s)", field, value.Type())
+                return false
+        }
+        if q.ineqFiltProp != "" && q.ineqFiltProp != field {
+                q.err = ErrMultipleInequalityFilter
+                return false
+        }
+        return true
+}
+
+func (q *Query) Lt(field string, value interface{}) *Query {
+        p := Property{}
+        err := p.SetValue(value, ShouldIndex)
+
+        if err == nil && q.ineqFiltHighSet {
+                if q.ineqFiltHigh.Less(&p) {
+                        return q
+                } else if q.ineqFiltHigh.Equal(&p) && !q.ineqFiltHighIncl {
+                        return q
+                }
+        }
+
+        return q.mod(func(q *Query) {
+                if q.err = err; err != nil {
+                        return
+                }
+                if q.ineqOK(field, p) {
+                        q.ineqFiltProp = field
+                        q.ineqFiltHighSet = true
+                        q.ineqFiltHigh = p
+                        q.ineqFiltHighIncl = false
+                }
+        })
+}
+
+func (q *Query) Lte(field string, value interface{}) *Query {
+        p := Property{}
+        err := p.SetValue(value, ShouldIndex)
+
+        if err == nil && q.ineqFiltHighSet {
+                if q.ineqFiltHigh.Less(&p) {
+                        return q
+                } else if q.ineqFiltHigh.Equal(&p) {
+                        return q
+                }
+        }
+
+        return q.mod(func(q *Query) {
+                if q.err = err; err != nil {
+                        return
+                }
+                if q.ineqOK(field, p) {
+                        q.ineqFiltProp = field
+                        q.ineqFiltHighSet = true
+                        q.ineqFiltHigh = p
+                        q.ineqFiltHighIncl = true
+                }
+        })
+}
+
+func (q *Query) Gt(field string, value interface{}) *Query {
+        p := Property{}
+        err := p.SetValue(value, ShouldIndex)
+
+        if err == nil && q.ineqFiltLowSet {
+                if p.Less(&q.ineqFiltLow) {
+                        return q
+                } else if p.Equal(&q.ineqFiltLow) && !q.ineqFiltLowIncl {
+                        return q
+                }
+        }
+
+        return q.mod(func(q *Query) {
+                if q.err = err; err != nil {
+                        return
+                }
+                if q.ineqOK(field, p) {
+                        q.ineqFiltProp = field
+                        q.ineqFiltLowSet = true
+                        q.ineqFiltLow = p
+                        q.ineqFiltLowIncl = false
+                }
+        })
+}
+
+func (q *Query) Gte(field string, value interface{}) *Query {
+        p := Property{}
+        err := p.SetValue(value, ShouldIndex)
+
+        if err == nil && q.ineqFiltLowSet {
+                if p.Less(&q.ineqFiltLow) {
+                        return q
+                } else if p.Equal(&q.ineqFiltLow) {
+                        return q
+                }
+        }
+
+        return q.mod(func(q *Query) {
+                if q.err = err; err != nil {
+                        return
+                }
+                if q.ineqOK(field, p) {
+                        q.ineqFiltProp = field
+                        q.ineqFiltLowSet = true
+                        q.ineqFiltLow = p
+                        q.ineqFiltLowIncl = true
+                }
+        })
+}
+
+func (q *Query) ClearFilters() *Query {
+        return q.mod(func(q *Query) {
+                anc := q.eqFilts["__ancestor__"]
+                if anc != nil {
+                        q.eqFilts = map[string]PropertySlice{"__ancestor__": anc}
+                } else {
+                        q.eqFilts = nil
+                }
+                q.ineqFiltLowSet = false
+                q.ineqFiltHighSet = false
+        })
+}
+
+func (q *Query) Finalize() (*FinalizedQuery, error) {
+        if q.err != nil || q.finalized != nil {
+                return q.finalized, q.err
+        }
+
+        err := func() error {
+                if q.kind == "" { // kindless query checks
+                        if q.ineqFiltProp != "" && q.ineqFiltProp != "__key__" {
+                                return fmt.Errorf(
+                                        "kindless queries can only filter on __key__, got %q", q.ineqFiltProp)
+                        }
+                        if len(q.eqFilts) > 0 {
+                                return fmt.Errorf("kindless queries not have any equality filters")
+                        }
+                        for _, o := range q.order {
+                                if o.Property != "__key__" || o.Direction != ASCENDING {
+                                        return fmt.Errorf("invalid order for kindless query: %#v", o)
+                                }
+                        }
+                }
+
+                if q.keysOnly && q.project != nil && q.project.Len() > 0 {
+                        return errors.New("cannot project a keysOnly query")
+                }
+
+                if q.ineqFiltProp != "" {
+                        if len(q.order) > 0 && q.order[0].Property != q.ineqFiltProp {
+                                return fmt.Errorf(
+                                        "first sort order must match inequality filter: %q v %q",
+                                        q.order[0].Property, q.ineqFiltProp)
+                        }
+                        if q.ineqFiltLowSet && q.ineqFiltHighSet {
+                                if q.ineqFiltHigh.Less(&q.ineqFiltLow) &&
+                                        (q.ineqFiltHigh.Equal(&q.ineqFiltLow) &&
+                                                (!q.ineqFiltLowIncl || !q.ineqFiltHighIncl)) {
+                                        return ErrNullQuery
+                                }
+                        }
+                }
+
+                err := error(nil)
+                if q.project != nil {
+                        q.project.Iter(func(p string) bool {
+                                if _, iseq := q.eqFilts[p]; iseq {
+                                        err = fmt.Errorf("cannot project on equality filter field: %s", p)
+                                        return false
+                                }
+                                return true
+                        })
+                }
+                return err
+        }()
+        if err != nil {
+                q.err = err
+                return nil, err
+        }
+
+        ret := &FinalizedQuery{
+                original: q,
+                kind:     q.kind,
+
+                keysOnly:             q.keysOnly,
+                eventuallyConsistent: q.eventualConsistency || q.eqFilts["__ancestor__"] == nil,
+                limit:                q.limit,
+                offset:               q.offset,
+                start:                q.start,
+                end:                  q.end,
+
+                eqFilts: q.eqFilts,
+
+                ineqFiltProp:     q.ineqFiltProp,
+                ineqFiltLow:      q.ineqFiltLow,
+                ineqFiltLowIncl:  q.ineqFiltLowIncl,
+                ineqFiltLowSet:   q.ineqFiltLowSet,
+                ineqFiltHigh:     q.ineqFiltHigh,
+                ineqFiltHighIncl: q.ineqFiltHighIncl,
+                ineqFiltHighSet:  q.ineqFiltHighSet,
+        }
+
+        if q.project != nil {
+                ret.project = q.project.ToSlice()
+                ret.distinct = q.distinct && q.project.Len() > 0
+        }
+
+        // if len(q.order) > 0, we already enforce that the first order
+        // is the same as the inequality above. Otherwise we need to add it.
+        if len(q.order) == 0 && q.ineqFiltProp != "" {
+                ret.orders = []IndexColumn{{q.ineqFiltProp, ASCENDING}}
+        }
+
+        seenOrders := stringset.New(len(q.order))
+
+        // drop orders where there's an equality filter
+        //   https://cloud.google.com/appengine/docs/go/datastore/queries#sort_orders_are_ignored_on_properties_with_equality_filters
+        // Deduplicate orders
+        for _, o := range q.order {
+                if _, iseq := q.eqFilts[o.Property]; !iseq {
+                        if seenOrders.Add(o.Property) {
+                                ret.orders = append(ret.orders, o)
+                        }
+                }
+        }
+
+        // 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 suffix entry for it.
+        //
+        // NOTE: if you want to use an index that sorts by -__key__, you MUST
+        // include all of the projected fields for that index in the order explicitly.
+        // Otherwise the generated orders will be wacky. So:
+        //   Query("Foo").Project("A", "B").Order("A").Order("-__key__")
+        //
+        // will turn into a orders of:
+        //   A, ASCENDING
+        //   __key__, DESCENDING
+        //   B, ASCENDING
+        //   __key__, ASCENDING
+        //
+        // To prevent this, your query should have another Order("B") clause before
+        // the -__key__ clause.
+        if len(ret.project) > 0 {
+                sort.Strings(ret.project)
+                for _, p := range ret.project {
+                        if !seenOrders.Has(p) {
+                                ret.orders = append(ret.orders, IndexColumn{p, ASCENDING})
+                        }
+                }
+        }
+
+        // If the suffix format ends with __key__ already (e.g. .Order("__key__")),
+        // then we're good to go. Otherwise we need to add it as the last bit of the
+        // suffix, since all indexes implicitly have it as the last column.
+        if len(ret.orders) == 0 || ret.orders[len(ret.orders)-1].Property != "__key__" {
+                ret.orders = append(ret.orders, IndexColumn{"__key__", ASCENDING})
+        }
+
+        q.finalized = ret
+        return ret, nil
+}
+
+func (q *Query) String() string {
+        ret := &bytes.Buffer{}
+        needComma := false
+        p := func(fmtStr string, stuff ...interface{}) {
+                if needComma {
+                        ret.WriteString(", ")
+                }
+                needComma = true
+                fmt.Fprintf(ret, fmtStr, stuff...)
+        }
+        ret.WriteString("Query(")
+        if q.err != nil {
+                p("ERROR=%q", q.err.Error())
+        }
+
+        // Filters
+        if q.kind != "" {
+                p("Kind=%q", q.kind)
+        }
+        if q.eqFilts["__ancestor__"] != nil {
+                p("Ancestor=%s", q.eqFilts["__ancestor__"][0].Value().(*Key).String())
+        }
+        for prop, vals := range q.eqFilts {
+                if prop == "__ancestor__" {
+                        continue
+                }
+                for _, v := range vals {
+                        p("Filter(%q == %s)", prop, v.GQL())
+                }
+        }
+        if q.ineqFiltProp != "" {
+                if q.ineqFiltLowSet {
+                        op := ">"
+                        if q.ineqFiltLowIncl {
+                                op = ">="
+                        }
+                        p("Filter(%q %s %s)", q.ineqFiltProp, op, q.ineqFiltLow.GQL())
+                }
+                if q.ineqFiltHighSet {
+                        op := "<"
+                        if q.ineqFiltHighIncl {
+                                op = "<="
+                        }
+                        p("Filter(%q %s %s)", q.ineqFiltProp, op, q.ineqFiltHigh.GQL())
+                }
+        }
+
+        // Order
+        if len(q.order) > 0 {
+                orders := make([]string, len(q.order))
+                for i, o := range q.order {
+                        orders[i] = o.String()
+                }
+                p("Order(%s)", strings.Join(orders, ", "))
+        }
+
+        // Projection
+        if q.project != nil && q.project.Len() > 0 {
+                f := "Project(%s)"
+                if q.distinct {
+                        f = "Project[DISTINCT](%s)"
+                }
+                p(f, strings.Join(q.project.ToSlice(), ", "))
+        }
+
+        // Cursors
+        if q.start != nil {
+                p("Start(%q)", q.start.String())
+        }
+        if q.end != nil {
+                p("End(%q)", q.end.String())
+        }
+
+        // Modifiiers
+        if q.limit != nil {
+                p("Limit=%d", *q.limit)
+        }
+        if q.offset != nil {
+                p("Offset=%d", *q.offset)
+        }
+        if q.eventualConsistency {
+                p("EventualConsistency")
+        }
+        if q.keysOnly {
+                p("KeysOnly")
+        }
+
+        ret.WriteRune(')')
+
+        return ret.String()
+}