feat: add support for Partition API (#1320)

dev/src/bulk-writer.ts

@@ -330,7 +330,7 @@ class BulkCommitBatch {
  * A Firestore BulkWriter than can be used to perform a large number of writes
  * in parallel. Writes to the same document will be executed sequentially.
  *
- * @class
+ * @class BulkWriter
  */
 export class BulkWriter {
   /**

dev/src/collection-group.ts

@@ -0,0 +1,185 @@
+/*
+ * Copyright 2020 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import * as firestore from '@google-cloud/firestore';
+import * as protos from '../protos/firestore_v1_proto_api';
+
+import {QueryPartition} from './query-partition';
+import {requestTag} from './util';
+import {logger} from './logger';
+import {Query, QueryOptions} from './reference';
+import {FieldPath} from './path';
+import {Firestore} from './index';
+import {validateInteger} from './validate';
+
+import api = protos.google.firestore.v1;
+
+/**
+ * A `CollectionGroup` refers to all documents that are contained in a
+ * collection or subcollection with a specific collection ID.
+ *
+ * @class CollectionGroup
+ */
+export class CollectionGroup<T = firestore.DocumentData>
+  extends Query<T>
+  implements firestore.CollectionGroup<T> {
+  /** @hideconstructor */
+  constructor(
+    firestore: Firestore,
+    collectionId: string,
+    converter: firestore.FirestoreDataConverter<T> | undefined
+  ) {
+    super(
+      firestore,
+      QueryOptions.forCollectionGroupQuery(collectionId, converter)
+    );
+  }
+
+  /**
+   * Partitions a query by returning partition cursors that can be used to run
+   * the query in parallel. The returned cursors are split points that can be
+   * used as starting and end points for individual query invocations.
+   *
+   * @example
+   * const query = firestore.collectionGroup('collectionId');
+   * for await (const partition of query.getPartitions(42)) {
+   *   const partitionedQuery = partition.toQuery();
+   *   const querySnapshot = await partitionedQuery.get();
+   *   console.log(`Partition contained ${querySnapshot.length} documents`);
+   * }
+   *
+   * @param {number} desiredPartitionCount The desired maximum number of
+   * partition points. The number must be strictly positive. The actual number
+   * of partitions returned may be fewer.
+   * @return {AsyncIterable<QueryPartition>} An AsyncIterable of
+   * `QueryPartition`s.
+   */
+  async *getPartitions(
+    desiredPartitionCount: number
+  ): AsyncIterable<QueryPartition<T>> {
+    validateInteger('desiredPartitionCount', desiredPartitionCount, {
+      minValue: 1,
+    });
+
+    const tag = requestTag();
+    await this.firestore.initializeIfNeeded(tag);
+
+    let lastValues: api.IValue[] | undefined = undefined;
+    let partitionCount = 0;
+
+    if (desiredPartitionCount > 1) {
+      // Partition queries require explicit ordering by __name__.
+      const queryWithDefaultOrder = this.orderBy(FieldPath.documentId());
+      const request: api.IPartitionQueryRequest = queryWithDefaultOrder.toProto();
+
+      // Since we are always returning an extra partition (with an empty endBefore
+      // cursor), we reduce the desired partition count by one.
+      request.partitionCount = desiredPartitionCount - 1;
+
+      const stream = await this.firestore.requestStream(
+        'partitionQueryStream',
+        request,
+        tag
+      );
+      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;
+      }
+    }
+
+    logger(
+      'Firestore.getPartitions',
+      tag,
+      'Received %d partitions',
+      partitionCount
+    );
+
+    // Return the extra partition with the empty cursor.
+    yield new QueryPartition(
+      this._firestore,
+      this._queryOptions.collectionId,
+      this._queryOptions.converter,
+      lastValues,
+      undefined
+    );
+  }
+
+  /**
+   * Applies a custom data converter to this `CollectionGroup`, allowing you
+   * to use your own custom model objects with Firestore. When you call get()
+   * on the returned `CollectionGroup`, the provided converter will convert
+   * between Firestore data and your custom type U.
+   *
+   * Using the converter allows you to specify generic type arguments when
+   * storing and retrieving objects from Firestore.
+   *
+   * @example
+   * class Post {
+   *   constructor(readonly title: string, readonly author: string) {}
+   *
+   *   toString(): string {
+   *     return this.title + ', by ' + this.author;
+   *   }
+   * }
+   *
+   * const postConverter = {
+   *   toFirestore(post: Post): FirebaseFirestore.DocumentData {
+   *     return {title: post.title, author: post.author};
+   *   },
+   *   fromFirestore(
+   *     snapshot: FirebaseFirestore.QueryDocumentSnapshot
+   *   ): Post {
+   *     const data = snapshot.data();
+   *     return new Post(data.title, data.author);
+   *   }
+   * };
+   *
+   * const querySnapshot = await Firestore()
+   *   .collectionGroup('posts')
+   *   .withConverter(postConverter)
+   *   .get();
+   * for (const doc of querySnapshot.docs) {
+   *   const post = doc.data();
+   *   post.title; // string
+   *   post.toString(); // Should be defined
+   *   post.someNonExistentProperty; // TS error
+   * }
+   *
+   * @param {FirestoreDataConverter} converter Converts objects to and from
+   * Firestore.
+   * @return {CollectionGroup} A `CollectionGroup<U>` that uses the provided
+   * converter.
+   */
+  withConverter<U>(
+    converter: firestore.FirestoreDataConverter<U>
+  ): CollectionGroup<U> {
+    return new CollectionGroup<U>(
+      this.firestore,
+      this._queryOptions.collectionId,
+      converter
+    );
+  }
+}

dev/src/document-change.ts

@@ -24,7 +24,7 @@ export type DocumentChangeType = 'added' | 'removed' | 'modified';
  * A DocumentChange represents a change to the documents matching a query.
  * It contains the document affected and the type of change that occurred.
  *
- * @class
+ * @class DocumentChange
  */
 export class DocumentChange<T = firestore.DocumentData>
   implements firestore.DocumentChange {

dev/src/document.ts

@@ -94,7 +94,7 @@ export class DocumentSnapshotBuilder<T = firestore.DocumentData> {
  * [exists]{@link DocumentSnapshot#exists} property to explicitly verify a
  * document's existence.
  *
- * @class
+ * @class DocumentSnapshot
  */
 export class DocumentSnapshot<T = firestore.DocumentData>
   implements firestore.DocumentSnapshot<T> {
@@ -534,7 +534,7 @@ export class DocumentSnapshot<T = firestore.DocumentData>
  * always be true and [data()]{@link QueryDocumentSnapshot#data} will never
  * return 'undefined'.
  *
- * @class
+ * @class QueryDocumentSnapshot
  * @extends DocumentSnapshot
  */
 export class QueryDocumentSnapshot<T = firestore.DocumentData>

dev/src/field-value.ts

@@ -34,7 +34,7 @@ import api = proto.google.firestore.v1;
  * Sentinel values that can be used when writing documents with set(), create()
  * or update().
  *
- * @class
+ * @class FieldValue
  */
 export class FieldValue implements firestore.FieldValue {
   /**

dev/src/index.ts

@@ -40,7 +40,7 @@ import {
   validateResourcePath,
 } from './path';
 import {ClientPool} from './pool';
-import {CollectionReference, Query, QueryOptions} from './reference';
+import {CollectionReference} from './reference';
 import {DocumentReference} from './reference';
 import {Serializer} from './serializer';
 import {Timestamp} from './timestamp';
@@ -75,6 +75,7 @@ import {interfaces} from './v1/firestore_client_config.json';
 const serviceConfig = interfaces['google.firestore.v1.Firestore'];
 
 import api = google.firestore.v1;
+import {CollectionGroup} from './collection-group';
 
 export {
   CollectionReference,
@@ -91,6 +92,8 @@ export {Timestamp} from './timestamp';
 export {DocumentChange} from './document-change';
 export {FieldPath} from './path';
 export {GeoPoint} from './geo-point';
+export {CollectionGroup};
+export {QueryPartition} from './query-partition';
 export {setLogFunction} from './logger';
 export {Status as GrpcStatus} from 'google-gax';
 
@@ -632,7 +635,7 @@ export class Firestore implements firestore.Firestore {
    * @param {string} collectionId Identifies the collections to query over.
    * Every collection or subcollection with this ID as the last segment of its
    * path will be included. Cannot contain a slash.
-   * @returns {Query} The created Query.
+   * @returns {CollectionGroup} The created CollectionGroup.
    *
    * @example
    * let docA = firestore.doc('mygroup/docA').set({foo: 'bar'});
@@ -646,14 +649,14 @@ export class Firestore implements firestore.Firestore {
    *    });
    * });
    */
-  collectionGroup(collectionId: string): Query {
+  collectionGroup(collectionId: string): CollectionGroup {
     if (collectionId.indexOf('/') !== -1) {
       throw new Error(
         `Invalid collectionId '${collectionId}'. Collection IDs must not contain '/'.`
       );
     }
 
-    return new Query(this, QueryOptions.forCollectionGroupQuery(collectionId));
+    return new CollectionGroup(this, collectionId, /* converter= */ undefined);
   }
 
   /**