Skip to content

Commit

Permalink
fix(firestore): Correct the cursors when LimitToLast is used (googlea…
Browse files Browse the repository at this point in the history 
…pis#9413)
  • Loading branch information
@bhshkh
bhshkh authored Feb 15, 2024
1 parent 9fe6061 commit 2090651
Show file tree
Hide file tree
Showing 2 changed files with 104 additions and 47 deletions.
84 changes: 58 additions & 26 deletions firestore/integration_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1020,13 +1020,22 @@ func TestIntegration_QueryDocuments_WhereEntity(t *testing.T) {
})
}

func reverseSlice(s []map[string]interface{}) []map[string]interface{} {
reversed := make([]map[string]interface{}, len(s))
for i, j := 0, len(s)-1; i <= j; i, j = i+1, j-1 {
reversed[i] = s[j]
reversed[j] = s[i]
}
return reversed
}

func TestIntegration_QueryDocuments(t *testing.T) {
ctx := context.Background()
coll := integrationColl(t)
h := testHelper{t}
var wants []map[string]interface{}
var createdDocRefs []*DocumentRef
for i := 0; i < 3; i++ {
for i := 0; i < 8; i++ {
doc := coll.NewDoc()
createdDocRefs = append(createdDocRefs, doc)

Expand All @@ -1037,46 +1046,69 @@ func TestIntegration_QueryDocuments(t *testing.T) {
}
q := coll.Select("q")
for i, test := range []struct {
q Query
want []map[string]interface{}
orderBy bool // Some query types do not allow ordering.
desc string
q Query
want []map[string]interface{}
orderBy bool // Some query types do not allow ordering.
orderByDir Direction
}{
{q, wants, true},
{q.Where("q", ">", 1), wants[2:], true},
{q.Where("q", "<", 1), wants[:1], true},
{q.Where("q", "==", 1), wants[1:2], false},
{q.Where("q", "!=", 0), wants[1:], true},
{q.Where("q", ">=", 1), wants[1:], true},
{q.Where("q", "<=", 1), wants[:2], true},
{q.Where("q", "in", []int{0}), wants[:1], false},
{q.Where("q", "not-in", []int{0, 1}), wants[2:], true},
{q.WherePath([]string{"q"}, ">", 1), wants[2:], true},
{q.Offset(1).Limit(1), wants[1:2], true},
{q.StartAt(1), wants[1:], true},
{q.StartAfter(1), wants[2:], true},
{q.EndAt(1), wants[:2], true},
{q.EndBefore(1), wants[:1], true},
{q.LimitToLast(2), wants[1:], true},
{q.EndBefore(2).LimitToLast(2), wants[:2], true},
{q.StartAt(1).EndBefore(2).LimitToLast(3), wants[1:2], true},
{"Without filters", q, wants, true, 0},
{"> filter", q.Where("q", ">", 1), wants[2:], true, Asc},
{"< filter", q.Where("q", "<", 1), wants[:1], true, Asc},
{"== filter", q.Where("q", "==", 1), wants[1:2], false, 0},
{"!= filter", q.Where("q", "!=", 0), wants[1:], true, Asc},
{">= filter", q.Where("q", ">=", 1), wants[1:], true, Asc},
{"<= filter", q.Where("q", "<=", 1), wants[:2], true, Asc},
{"in filter", q.Where("q", "in", []int{0}), wants[:1], false, 0},
{"not-in filter", q.Where("q", "not-in", []int{0, 1}), wants[2:], true, Asc},
{"WherePath", q.WherePath([]string{"q"}, ">", 1), wants[2:], true, Asc},
{"Offset with Limit", q.Offset(1).Limit(1), wants[1:2], true, Asc},
{"StartAt", q.StartAt(1), wants[1:], true, Asc},
{"StartAfter", q.StartAfter(1), wants[2:], true, Asc},
{"EndAt", q.EndAt(1), wants[:2], true, Asc},
{"EndBefore", q.EndBefore(1), wants[:1], true, Asc},
{"Open range with DESC order", q.StartAfter(6).EndBefore(2), reverseSlice(wants[3:6]), true, Desc},
{"LimitToLast", q.LimitToLast(2), wants[len(wants)-2:], true, Asc},
{"StartAfter with LimitToLast", q.StartAfter(2).LimitToLast(2), wants[len(wants)-2:], true, Asc},
{"StartAt with LimitToLast", q.StartAt(2).LimitToLast(2), wants[len(wants)-2:], true, Asc},
{"EndBefore with LimitToLast", q.EndBefore(7).LimitToLast(2), wants[5:7], true, Asc},
{"EndAt with LimitToLast", q.EndAt(7).LimitToLast(2), wants[6:8], true, Asc},
{"LimitToLast greater than no. of results", q.StartAt(1).EndBefore(2).LimitToLast(3), wants[1:2], true, Asc},
{"Closed range with LimitToLast ASC order", q.StartAt(2).EndAt(6).LimitToLast(2), wants[5:7], true, Asc},
{"Left closed right open range with LimitToLast ASC order", q.StartAt(2).EndBefore(6).LimitToLast(2), wants[4:6], true, Asc},
{"Left open right closed with LimitToLast ASC order", q.StartAfter(2).EndAt(6).LimitToLast(2), wants[5:7], true, Asc},
{"Open range with LimitToLast ASC order", q.StartAfter(2).EndBefore(6).LimitToLast(2), wants[4:6], true, Asc},
{"Closed range with LimitToLast DESC order", q.StartAt(6).EndAt(2).LimitToLast(2), reverseSlice(wants[2:4]), true, Desc},
{"Left closed right open range with LimitToLast DESC order", q.StartAt(6).EndBefore(2).LimitToLast(2), reverseSlice(wants[3:5]), true, Desc},
{"Left open right closed with LimitToLast DESC order", q.StartAfter(6).EndAt(2).LimitToLast(2), reverseSlice(wants[2:4]), true, Desc},
{"Open range with LimitToLast DESC order", q.StartAfter(6).EndBefore(2).LimitToLast(2), reverseSlice(wants[3:5]), true, Desc},
} {
if test.orderBy {
test.q = test.q.OrderBy("q", Asc)
test.q = test.q.OrderBy("q", test.orderByDir)
}
gotDocs, err := test.q.Documents(ctx).GetAll()
if err != nil {
t.Errorf("#%d: %+v: %v", i, test.q, err)
t.Errorf("#%d %v: %+v: %v", i, test.desc, test.q, err)
continue
}
if len(gotDocs) != len(test.want) {
t.Errorf("#%d: %+v: got %d docs, want %d", i, test.q, len(gotDocs), len(test.want))
t.Errorf("#%d %v: %+v: got %d docs, want %d", i, test.desc, test.q, len(gotDocs), len(test.want))
continue
}

fmt.Printf("test.want: %+v\n", test.want)

docsEqual := true
docsNotEqualErr := ""
for j, g := range gotDocs {
if got, want := g.Data(), test.want[j]; !testEqual(got, want) {
t.Errorf("#%d: %+v, #%d: got\n%+v\nwant\n%+v", i, test.q, j, got, want)
docsNotEqualErr += fmt.Sprintf("\n\t#%d: got %+v want %+v", j, got, want)
docsEqual = false
}
}
if !docsEqual {
t.Errorf("#%d %v: %+v %v", i, test.desc, test.q, docsNotEqualErr)
}
}
_, err := coll.Select("q").Where("x", "==", 1).OrderBy("q", Asc).Documents(ctx).GetAll()
codeEq(t, "Where and OrderBy on different fields without an index", codes.FailedPrecondition, err)
Expand Down
67 changes: 46 additions & 21 deletions firestore/query.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,18 +36,21 @@ import (
// Query values are immutable. Each Query method creates
// a new Query; it does not modify the old.
type Query struct {
c *Client
path string // path to query (collection)
parentPath string // path of the collection's parent (document)
collectionID string
selection []*pb.StructuredQuery_FieldReference
filters []*pb.StructuredQuery_Filter
orders []order
offset int32
limit *wrappers.Int32Value
limitToLast bool
startVals, endVals []interface{}
startDoc, endDoc *DocumentSnapshot
c *Client
path string // path to query (collection)
parentPath string // path of the collection's parent (document)
collectionID string
selection []*pb.StructuredQuery_FieldReference
filters []*pb.StructuredQuery_Filter
orders []order
offset int32
limit *wrappers.Int32Value
limitToLast bool
startVals, endVals []interface{}
startDoc, endDoc *DocumentSnapshot

// Set startBefore to true when doc in startVals needs to be included in result
// Set endBefore to false when doc in endVals needs to be included in result
startBefore, endBefore bool
err error

Expand Down Expand Up @@ -293,23 +296,37 @@ func (q *Query) processCursorArg(name string, docSnapshotOrFieldValues []interfa

func (q *Query) processLimitToLast() {
if q.limitToLast {
// Flip order statements before posting a request.
// Firestore service does not provide limit to last behaviour out of the box. This is a client-side concept
// So, flip order statements and cursors before posting a request. The response is flipped by other methods before returning to user
// E.g.
// If id of documents is 1, 2, 3, 4, 5, 6, 7 and query is (OrderBy(id, ASC), StartAt(2), EndAt(6), LimitToLast(3))
// request sent to server is (OrderBy(id, DESC), StartAt(6), EndAt(2), Limit(3))
for i := range q.orders {
if q.orders[i].dir == Asc {
q.orders[i].dir = Desc
} else {
q.orders[i].dir = Asc
}
}
// Swap cursors.
q.startVals, q.endVals = q.endVals, q.startVals
q.startDoc, q.endDoc = q.endDoc, q.startDoc
if q.endBefore {

if q.startBefore == q.endBefore && q.startCursorSpecified() && q.endCursorSpecified() {
// E.g. query.StartAt(2).EndBefore(6).LimitToLast(3).OrderBy(Asc) i.e. cursors are [2, 6)
// E.g. query.StartAfter(2).EndAt(6).LimitToLast(3).OrderBy(Asc) i.e. cursors are (2, 6]
q.startBefore, q.endBefore = !q.startBefore, !q.endBefore
} else if !q.startCursorSpecified() && q.endCursorSpecified() {
// E.g. query.EndAt(6).LimitToLast(3).OrderBy(Asc) i.e. cursors are (-inf, 6]
q.startBefore = !q.endBefore
q.endBefore = false
} else if q.startCursorSpecified() && !q.endCursorSpecified() {
// E.g. query.StartAt(2).LimitToLast(3).OrderBy(Asc) i.e. cursors are [2, inf)
q.endBefore = !q.startBefore
q.startBefore = false
} else {
q.startBefore, q.endBefore = q.endBefore, q.startBefore
}

// Swap cursors.
q.startVals, q.endVals = q.endVals, q.startVals
q.startDoc, q.endDoc = q.endDoc, q.startDoc

q.limitToLast = false
}
}
Expand Down Expand Up @@ -463,6 +480,14 @@ func (q Query) fromProto(pbQuery *pb.RunQueryRequest) (Query, error) {
return q, q.err
}

func (q Query) startCursorSpecified() bool {
return len(q.startVals) != 0 || q.startDoc != nil
}

func (q Query) endCursorSpecified() bool {
return len(q.endVals) != 0 || q.endDoc != nil
}

func (q Query) toProto() (*pb.StructuredQuery, error) {
if q.err != nil {
return nil, q.err
Expand All @@ -471,12 +496,12 @@ func (q Query) toProto() (*pb.StructuredQuery, error) {
return nil, errors.New("firestore: query created without CollectionRef")
}
if q.startBefore {
if len(q.startVals) == 0 && q.startDoc == nil {
if !q.startCursorSpecified() {
return nil, errors.New("firestore: StartAt/StartAfter must be called with at least one value")
}
}
if q.endBefore {
if len(q.endVals) == 0 && q.endDoc == nil {
if !q.endCursorSpecified() {
return nil, errors.New("firestore: EndAt/EndBefore must be called with at least one value")
}
}
Expand Down

0 comments on commit 2090651

Please sign in to comment.