From 25ecf26a2ced9d8febe805bc4c4173bfd440564c Mon Sep 17 00:00:00 2001 From: Sebastian Schmidt Date: Thu, 27 May 2021 10:08:08 -0600 Subject: [PATCH 1/3] fix: return results from getPartitions() in order --- dev/src/collection-group.ts | 32 +++++++++++---------- dev/src/order.ts | 2 +- dev/test/partition-query.ts | 55 ++++++++++++++++++++++++++++++++++++- 3 files changed, 73 insertions(+), 16 deletions(-) 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..e487d05eb 100644 --- a/dev/test/partition-query.ts +++ b/dev/test/partition-query.ts @@ -24,7 +24,7 @@ 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, @@ -102,6 +102,28 @@ 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 = { @@ -255,4 +277,35 @@ describe('Partition Query', () => { return result[0].toQuery().get(); }); }); + + it('sorts partitions', () => { + const desiredPartitionsCount = 3; + + const overrides: ApiOverride = { + partitionQueryStream: () => { + const doc1 = 'projects/p1/databases/d1/documents/coll/doc1'; + const doc2 = 'projects/p1/databases/d1/documents/coll/doc2'; + + 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); + }); + }); }); From 6b05270fad51e2ef8b1619cd85e1cd6534db3e05 Mon Sep 17 00:00:00 2001 From: Sebastian Schmidt Date: Thu, 27 May 2021 10:11:00 -0600 Subject: [PATCH 2/3] Inline --- dev/test/partition-query.ts | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/dev/test/partition-query.ts b/dev/test/partition-query.ts index e487d05eb..983c57cc5 100644 --- a/dev/test/partition-query.ts +++ b/dev/test/partition-query.ts @@ -18,7 +18,7 @@ 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'; @@ -283,15 +283,16 @@ describe('Partition Query', () => { const overrides: ApiOverride = { partitionQueryStream: () => { - const doc1 = 'projects/p1/databases/d1/documents/coll/doc1'; - const doc2 = 'projects/p1/databases/d1/documents/coll/doc2'; - return stream( { - values: [{referenceValue: doc2}], + values: [ + {referenceValue: 'projects/p1/databases/d1/documents/coll/doc2'}, + ], }, { - values: [{referenceValue: doc1}], + values: [ + {referenceValue: 'projects/p1/databases/d1/documents/coll/doc1'}, + ], } ); }, From 9577774e149b322e4c1c042d80f440b2c618caee Mon Sep 17 00:00:00 2001 From: Sebastian Schmidt Date: Thu, 27 May 2021 10:24:24 -0600 Subject: [PATCH 3/3] Fix tests --- dev/test/partition-query.ts | 24 +++++++++++------------- 1 file changed, 11 insertions(+), 13 deletions(-) diff --git a/dev/test/partition-query.ts b/dev/test/partition-query.ts index 983c57cc5..70a338971 100644 --- a/dev/test/partition-query.ts +++ b/dev/test/partition-query.ts @@ -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, @@ -127,7 +129,7 @@ describe('Partition Query', () => { 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 = { @@ -178,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, ]; @@ -196,10 +198,10 @@ describe('Partition Query', () => { return stream( { - values: [{referenceValue: 'coll/doc1'}], + values: [{referenceValue: DOC1}], }, { - values: [{referenceValue: 'coll/doc2'}], + values: [{referenceValue: DOC2}], } ); }, @@ -285,14 +287,10 @@ describe('Partition Query', () => { partitionQueryStream: () => { return stream( { - values: [ - {referenceValue: 'projects/p1/databases/d1/documents/coll/doc2'}, - ], + values: [{referenceValue: DOC2}], }, { - values: [ - {referenceValue: 'projects/p1/databases/d1/documents/coll/doc1'}, - ], + values: [{referenceValue: DOC1}], } ); },