feat: add implicit ordering for startAt(DocumentReference) calls (#1328)

googleapis · Oct 20, 2020 · e9afa38 · e9afa38
1 parent fdf5462
commit e9afa38
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 21 deletions.
diff --git a/dev/src/reference.ts b/dev/src/reference.ts
@@ -1115,23 +1115,6 @@ export class Query<T = firestore.DocumentData> implements firestore.Query<T> {
       .ignoreUndefinedProperties;
   }
 
-  /**
-   * Detects the argument type for Firestore cursors.
-   *
-   * @private
-   * @param fieldValuesOrDocumentSnapshot A snapshot of the document or a set
-   * of field values.
-   * @returns 'true' if the input is a single DocumentSnapshot..
-   */
-  static _isDocumentSnapshot(
-    fieldValuesOrDocumentSnapshot: Array<DocumentSnapshot<unknown> | unknown>
-  ): boolean {
-    return (
-      fieldValuesOrDocumentSnapshot.length === 1 &&
-      fieldValuesOrDocumentSnapshot[0] instanceof DocumentSnapshot
-    );
-  }
-
   /**
    * Extracts field values from the DocumentSnapshot based on the provided
    * field order.
@@ -1471,7 +1454,15 @@ export class Query<T = firestore.DocumentData> implements firestore.Query<T> {
   private createImplicitOrderBy(
     cursorValuesOrDocumentSnapshot: Array<DocumentSnapshot<unknown> | unknown>
   ): FieldOrder[] {
-    if (!Query._isDocumentSnapshot(cursorValuesOrDocumentSnapshot)) {
+    // Add an implicit orderBy if the only cursor value is a DocumentSnapshot
+    // or a DocumentReference.
+    if (
+      cursorValuesOrDocumentSnapshot.length !== 1 ||
+      !(
+        cursorValuesOrDocumentSnapshot[0] instanceof DocumentSnapshot ||
+        cursorValuesOrDocumentSnapshot[0] instanceof DocumentReference
+      )
+    ) {
       return this._queryOptions.fieldOrders;
     }
 
@@ -1527,7 +1518,10 @@ export class Query<T = firestore.DocumentData> implements firestore.Query<T> {
   ): QueryCursor {
     let fieldValues;
 
-    if (Query._isDocumentSnapshot(cursorValuesOrDocumentSnapshot)) {
+    if (
+      cursorValuesOrDocumentSnapshot.length === 1 &&
+      cursorValuesOrDocumentSnapshot[0] instanceof DocumentSnapshot
+    ) {
       fieldValues = Query._extractFieldValues(
         cursorValuesOrDocumentSnapshot[0] as DocumentSnapshot,
         fieldOrders

diff --git a/dev/system-test/firestore.ts b/dev/system-test/firestore.ts
@@ -271,8 +271,7 @@ describe('CollectionGroup class', () => {
 
     const documents: QueryDocumentSnapshot<DocumentData>[] = [];
     for (const partition of partitions) {
-      // TODO(mrschmidt); Remove the need to add an `orderBy` here
-      let partitionedQuery = collectionGroup.orderBy(FieldPath.documentId());
+      let partitionedQuery: Query = collectionGroup;
       if (partition.startAt) {
         partitionedQuery = partitionedQuery.startAt(...partition.startAt);
       }
@@ -1623,6 +1622,12 @@ describe('Query class', () => {
     expectDocs(res, {foo: 'b'});
   });
 
+  it('startAt() adds implicit order by for DocumentReference', async () => {
+    const references = await addDocs({foo: 'a'}, {foo: 'b'});
+    const res = await randomCol.startAt(references[1]).get();
+    expectDocs(res, {foo: 'b'});
+  });
+
   it('has startAfter() method', async () => {
     await addDocs({foo: 'a'}, {foo: 'b'});
     const res = await randomCol.orderBy('foo').startAfter('a').get();

diff --git a/dev/test/query.ts b/dev/test/query.ts
@@ -1852,6 +1852,32 @@ describe('startAt() interface', () => {
     });
   });
 
+  it('appends orderBy for DocumentReference cursors', () => {
+    const overrides: ApiOverride = {
+      runQuery: request => {
+        queryEquals(
+          request,
+          orderBy('__name__', 'ASCENDING'),
+          startAt(true, {
+            referenceValue:
+              `projects/${PROJECT_ID}/databases/(default)/` +
+              'documents/collectionId/doc',
+          })
+        );
+
+        return stream();
+      },
+    };
+
+    return createInstance(overrides).then(firestore => {
+      return snapshot('collectionId/doc', {foo: 'bar'}).then(doc => {
+        let query: Query = firestore.collection('collectionId');
+        query = query.startAt(doc.ref);
+        return query.get();
+      });
+    });
+  });
+
   it('can extract implicit direction for document snapshot', () => {
     const overrides: ApiOverride = {
       runQuery: request => {