diff --git a/dev/src/collection-group.ts b/dev/src/collection-group.ts index 70a371937..663c229b0 100644 --- a/dev/src/collection-group.ts +++ b/dev/src/collection-group.ts @@ -27,6 +27,7 @@ import {validateInteger} from './validate'; import api = protos.google.firestore.v1; import {defaultConverter} from './types'; +import {compareArrays} from './order'; /** * A `CollectionGroup` refers to all documents that are contained in a @@ -79,8 +80,7 @@ export class CollectionGroup const tag = requestTag(); await this.firestore.initializeIfNeeded(tag); - let lastValues: api.IValue[] | undefined = undefined; - let partitionCount = 0; + const partitions: Array[] = []; if (desiredPartitionCount > 1) { // Partition queries require explicit ordering by __name__. @@ -100,16 +100,7 @@ export class CollectionGroup stream.resume(); for await (const currentCursor of stream) { - ++partitionCount; - const currentValues = currentCursor.values ?? []; - yield new QueryPartition( - this._firestore, - this._queryOptions.collectionId, - this._queryOptions.converter, - lastValues, - currentValues - ); - lastValues = currentValues; + partitions.push(currentCursor.values ?? []); } } @@ -117,15 +108,28 @@ export class CollectionGroup 'Firestore.getPartitions', tag, 'Received %d partitions', - partitionCount + partitions.length ); + // Sort the partitions as they may not be ordered if responses are paged. + partitions.sort((l, r) => compareArrays(l, r)); + + for (let i = 0; i < partitions.length; ++i) { + yield new QueryPartition( + this._firestore, + this._queryOptions.collectionId, + this._queryOptions.converter, + i > 0 ? partitions[i - 1] : undefined, + partitions[i] + ); + } + // Return the extra partition with the empty cursor. yield new QueryPartition( this._firestore, this._queryOptions.collectionId, this._queryOptions.converter, - lastValues, + partitions.pop(), undefined ); } diff --git a/dev/src/order.ts b/dev/src/order.ts index 88c62c63b..5aa264b09 100644 --- a/dev/src/order.ts +++ b/dev/src/order.ts @@ -179,7 +179,7 @@ function compareGeoPoints( /*! * @private */ -function compareArrays(left: api.IValue[], right: api.IValue[]): number { +export function compareArrays(left: api.IValue[], right: api.IValue[]): number { for (let i = 0; i < left.length && i < right.length; i++) { const valueComparison = compare(left[i], right[i]); if (valueComparison !== 0) { diff --git a/dev/test/partition-query.ts b/dev/test/partition-query.ts index 647c3d391..70a338971 100644 --- a/dev/test/partition-query.ts +++ b/dev/test/partition-query.ts @@ -18,13 +18,13 @@ import { QueryPartition, } from '@google-cloud/firestore'; -import {describe, it, beforeEach, afterEach} from 'mocha'; +import {afterEach, beforeEach, describe, it} from 'mocha'; import {expect, use} from 'chai'; import * as chaiAsPromised from 'chai-as-promised'; import * as extend from 'extend'; import {google} from '../protos/firestore_v1_proto_api'; -import {Firestore} from '../src'; +import {DocumentReference, Firestore} from '../src'; import {setTimeoutHandler} from '../src/backoff'; import { ApiOverride, @@ -39,6 +39,8 @@ import api = google.firestore.v1; const PROJECT_ID = 'test-project'; const DATABASE_ROOT = `projects/${PROJECT_ID}/databases/(default)`; +const DOC1 = `${DATABASE_ROOT}/documents/coll/doc1`; +const DOC2 = `${DATABASE_ROOT}/documents/coll/doc2`; export function partitionQueryEquals( actual: api.IPartitionQueryRequest | undefined, @@ -102,10 +104,32 @@ describe('Partition Query', () => { return partitions; } + function verifyPartition( + partition: FirebaseFirestore.QueryPartition, + startAt: string | null, + endBefore: string | null + ) { + if (startAt) { + expect( + partition.startAt?.map(value => (value as DocumentReference).path) + ).to.have.members([startAt]); + } else { + expect(partition.startAt).to.be.undefined; + } + + if (endBefore) { + expect( + partition.endBefore?.map(value => (value as DocumentReference).path) + ).to.have.members([endBefore]); + } else { + expect(partition.endBefore).to.be.undefined; + } + } + it('requests one less than desired partitions', () => { const desiredPartitionsCount = 2; const cursorValue = { - values: [{referenceValue: 'projects/p1/databases/d1/documents/coll/doc'}], + values: [{referenceValue: DOC1}], }; const overrides: ApiOverride = { @@ -156,12 +180,12 @@ describe('Partition Query', () => { const expectedStartAt: Array = [ undefined, - {referenceValue: 'coll/doc1'}, - {referenceValue: 'coll/doc2'}, + {referenceValue: DOC1}, + {referenceValue: DOC2}, ]; const expectedEndBefore: Array = [ - {referenceValue: 'coll/doc1'}, - {referenceValue: 'coll/doc2'}, + {referenceValue: DOC1}, + {referenceValue: DOC2}, undefined, ]; @@ -174,10 +198,10 @@ describe('Partition Query', () => { return stream( { - values: [{referenceValue: 'coll/doc1'}], + values: [{referenceValue: DOC1}], }, { - values: [{referenceValue: 'coll/doc2'}], + values: [{referenceValue: DOC2}], } ); }, @@ -255,4 +279,32 @@ describe('Partition Query', () => { return result[0].toQuery().get(); }); }); + + it('sorts partitions', () => { + const desiredPartitionsCount = 3; + + const overrides: ApiOverride = { + partitionQueryStream: () => { + return stream( + { + values: [{referenceValue: DOC2}], + }, + { + values: [{referenceValue: DOC1}], + } + ); + }, + }; + + return createInstance(overrides).then(async firestore => { + const query = firestore.collectionGroup('collectionId'); + + const partitions = await getPartitions(query, desiredPartitionsCount); + expect(partitions.length).to.equal(3); + + verifyPartition(partitions[0], null, 'coll/doc1'); + verifyPartition(partitions[1], 'coll/doc1', 'coll/doc2'); + verifyPartition(partitions[2], 'coll/doc2', null); + }); + }); });