Refactor keys and queries in datastore service and implementation.

service/datastore/query.go

Index: service/datastore/query.go
diff --git a/service/datastore/query.go b/service/datastore/query.go
new file mode 100644
index 0000000000000000000000000000000000000000..83396fe8ec63bb863e696b31ac17d857239759fe
--- /dev/null
+++ b/service/datastore/query.go
@@ -0,0 +1,603 @@
+// 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 q.eqFilts != nil {

[dnj, 2015/09/18 16:47:58]

len(q.eqFilts) > 0

[iannucci, 2015/09/18 22:25:49]

done

+		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 {

[dnj, 2015/09/18 16:47:58]

Worth doing pre-mod identity checks for these? e.g.,
if kind != q.kind {
  ...
}

[iannucci, 2015/09/18 22:25:48]

Not really. If you have allocation bottlenecks in your query construction,
you're doing it really, really, really wrong. This is already over-optimized.

+	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 {
+	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) RevertOrder() *Query {

[dnj, 2015/09/18 16:47:59]

ClearOrder()

+	return q.mod(func(q *Query) {
+		q.order = nil
+	})
+}
+
+func (q *Query) Project(fieldNames ...string) *Query {
+	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) RevertProject() *Query {

[dnj, 2015/09/18 16:47:59]

ClearProject

+	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 = map[string]PropertySlice{}

[dnj, 2015/09/18 16:47:58]

make(..., 1)

[iannucci, 2015/09/18 22:25:48]

has anyone told you that you're the king of warrantless pre-micro-optimization
:)?

+			}
+			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, "__") {

[dnj, 2015/09/18 16:47:58]

Is "_ _ _" (without spaces) reserved? Granted it's a stupid field name...

[iannucci, 2015/09/18 22:25:49]

They're not actually reserved in the underlying thing, but I don't want
conflation with "__ancestor__", and I may need additional things later. If it
becomes a conflict for something, we can change this logic to a whitelist... or
something. don't overthink it :p

+		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 {

[dnj, 2015/09/18 16:47:59]

IMO the auto-adjustment is convenient, but it should be an error to set multiple
inequality filters on a Query.

[iannucci, 2015/09/18 22:25:49]

yep, it does in q.ineqOK (which modifies q.err as a side effect)

+		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) RevertFilters() *Query {

[dnj, 2015/09/18 16:47:58]

ClearFilters

[iannucci, 2015/09/18 22:25:48]

done for all

+	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())
+		}

[dnj, 2015/09/18 16:47:58]

ineqFiltHighSet?

[iannucci, 2015/09/18 22:25:48]

oops

+	}
+
+	// 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()
+}