initial import

core/stores/mongo/query.go

@@ -0,0 +1,285 @@
+package mongo
+
+import (
+	"time"
+
+	"zero/core/breaker"
+
+	"github.com/globalsign/mgo"
+)
+
+type (
+	Query interface {
+		All(result interface{}) error
+		Apply(change mgo.Change, result interface{}) (*mgo.ChangeInfo, error)
+		Batch(n int) Query
+		Collation(collation *mgo.Collation) Query
+		Comment(comment string) Query
+		Count() (int, error)
+		Distinct(key string, result interface{}) error
+		Explain(result interface{}) error
+		For(result interface{}, f func() error) error
+		Hint(indexKey ...string) Query
+		Iter() Iter
+		Limit(n int) Query
+		LogReplay() Query
+		MapReduce(job *mgo.MapReduce, result interface{}) (*mgo.MapReduceInfo, error)
+		One(result interface{}) error
+		Prefetch(p float64) Query
+		Select(selector interface{}) Query
+		SetMaxScan(n int) Query
+		SetMaxTime(d time.Duration) Query
+		Skip(n int) Query
+		Snapshot() Query
+		Sort(fields ...string) Query
+		Tail(timeout time.Duration) Iter
+	}
+
+	promisedQuery struct {
+		*mgo.Query
+		promise keepablePromise
+	}
+
+	rejectedQuery struct{}
+)
+
+func (q promisedQuery) All(result interface{}) error {
+	return q.promise.keep(q.Query.All(result))
+}
+
+func (q promisedQuery) Apply(change mgo.Change, result interface{}) (*mgo.ChangeInfo, error) {
+	info, err := q.Query.Apply(change, result)
+	return info, q.promise.keep(err)
+}
+
+func (q promisedQuery) Batch(n int) Query {
+	return promisedQuery{
+		Query:   q.Query.Batch(n),
+		promise: q.promise,
+	}
+}
+
+func (q promisedQuery) Collation(collation *mgo.Collation) Query {
+	return promisedQuery{
+		Query:   q.Query.Collation(collation),
+		promise: q.promise,
+	}
+}
+
+func (q promisedQuery) Comment(comment string) Query {
+	return promisedQuery{
+		Query:   q.Query.Comment(comment),
+		promise: q.promise,
+	}
+}
+
+func (q promisedQuery) Count() (int, error) {
+	v, err := q.Query.Count()
+	return v, q.promise.keep(err)
+}
+
+func (q promisedQuery) Distinct(key string, result interface{}) error {
+	return q.promise.keep(q.Query.Distinct(key, result))
+}
+
+func (q promisedQuery) Explain(result interface{}) error {
+	return q.promise.keep(q.Query.Explain(result))
+}
+
+func (q promisedQuery) For(result interface{}, f func() error) error {
+	var ferr error
+	err := q.Query.For(result, func() error {
+		ferr = f()
+		return ferr
+	})
+	if ferr == err {
+		return q.promise.accept(err)
+	}
+
+	return q.promise.keep(err)
+}
+
+func (q promisedQuery) Hint(indexKey ...string) Query {
+	return promisedQuery{
+		Query:   q.Query.Hint(indexKey...),
+		promise: q.promise,
+	}
+}
+
+func (q promisedQuery) Iter() Iter {
+	return promisedIter{
+		Iter:    q.Query.Iter(),
+		promise: q.promise,
+	}
+}
+
+func (q promisedQuery) Limit(n int) Query {
+	return promisedQuery{
+		Query:   q.Query.Limit(n),
+		promise: q.promise,
+	}
+}
+
+func (q promisedQuery) LogReplay() Query {
+	return promisedQuery{
+		Query:   q.Query.LogReplay(),
+		promise: q.promise,
+	}
+}
+
+func (q promisedQuery) MapReduce(job *mgo.MapReduce, result interface{}) (*mgo.MapReduceInfo, error) {
+	info, err := q.Query.MapReduce(job, result)
+	return info, q.promise.keep(err)
+}
+
+func (q promisedQuery) One(result interface{}) error {
+	return q.promise.keep(q.Query.One(result))
+}
+
+func (q promisedQuery) Prefetch(p float64) Query {
+	return promisedQuery{
+		Query:   q.Query.Prefetch(p),
+		promise: q.promise,
+	}
+}
+
+func (q promisedQuery) Select(selector interface{}) Query {
+	return promisedQuery{
+		Query:   q.Query.Select(selector),
+		promise: q.promise,
+	}
+}
+
+func (q promisedQuery) SetMaxScan(n int) Query {
+	return promisedQuery{
+		Query:   q.Query.SetMaxScan(n),
+		promise: q.promise,
+	}
+}
+
+func (q promisedQuery) SetMaxTime(d time.Duration) Query {
+	return promisedQuery{
+		Query:   q.Query.SetMaxTime(d),
+		promise: q.promise,
+	}
+}
+
+func (q promisedQuery) Skip(n int) Query {
+	return promisedQuery{
+		Query:   q.Query.Skip(n),
+		promise: q.promise,
+	}
+}
+
+func (q promisedQuery) Snapshot() Query {
+	return promisedQuery{
+		Query:   q.Query.Snapshot(),
+		promise: q.promise,
+	}
+}
+
+func (q promisedQuery) Sort(fields ...string) Query {
+	return promisedQuery{
+		Query:   q.Query.Sort(fields...),
+		promise: q.promise,
+	}
+}
+
+func (q promisedQuery) Tail(timeout time.Duration) Iter {
+	return promisedIter{
+		Iter:    q.Query.Tail(timeout),
+		promise: q.promise,
+	}
+}
+
+func (q rejectedQuery) All(result interface{}) error {
+	return breaker.ErrServiceUnavailable
+}
+
+func (q rejectedQuery) Apply(change mgo.Change, result interface{}) (*mgo.ChangeInfo, error) {
+	return nil, breaker.ErrServiceUnavailable
+}
+
+func (q rejectedQuery) Batch(n int) Query {
+	return q
+}
+
+func (q rejectedQuery) Collation(collation *mgo.Collation) Query {
+	return q
+}
+
+func (q rejectedQuery) Comment(comment string) Query {
+	return q
+}
+
+func (q rejectedQuery) Count() (int, error) {
+	return 0, breaker.ErrServiceUnavailable
+}
+
+func (q rejectedQuery) Distinct(key string, result interface{}) error {
+	return breaker.ErrServiceUnavailable
+}
+
+func (q rejectedQuery) Explain(result interface{}) error {
+	return breaker.ErrServiceUnavailable
+}
+
+func (q rejectedQuery) For(result interface{}, f func() error) error {
+	return breaker.ErrServiceUnavailable
+}
+
+func (q rejectedQuery) Hint(indexKey ...string) Query {
+	return q
+}
+
+func (q rejectedQuery) Iter() Iter {
+	return rejectedIter{}
+}
+
+func (q rejectedQuery) Limit(n int) Query {
+	return q
+}
+
+func (q rejectedQuery) LogReplay() Query {
+	return q
+}
+
+func (q rejectedQuery) MapReduce(job *mgo.MapReduce, result interface{}) (*mgo.MapReduceInfo, error) {
+	return nil, breaker.ErrServiceUnavailable
+}
+
+func (q rejectedQuery) One(result interface{}) error {
+	return breaker.ErrServiceUnavailable
+}
+
+func (q rejectedQuery) Prefetch(p float64) Query {
+	return q
+}
+
+func (q rejectedQuery) Select(selector interface{}) Query {
+	return q
+}
+
+func (q rejectedQuery) SetMaxScan(n int) Query {
+	return q
+}
+
+func (q rejectedQuery) SetMaxTime(d time.Duration) Query {
+	return q
+}
+
+func (q rejectedQuery) Skip(n int) Query {
+	return q
+}
+
+func (q rejectedQuery) Snapshot() Query {
+	return q
+}
+
+func (q rejectedQuery) Sort(fields ...string) Query {
+	return q
+}
+
+func (q rejectedQuery) Tail(timeout time.Duration) Iter {
+	return rejectedIter{}
+}