feat: add support for Partition API

dev/src/collection-group.ts

@@ -0,0 +1,191 @@
+/*
+ * 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 api = protos.google.firestore.v1;
+
+/**
+ * A `CollectionGroup` refers to all documents that are contained in a
+ * collection or subcollection with a specific collection ID.
+ */
+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.
+   *
+   * @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 *getPartitionsAsync(
+    desiredPartitionCount: number
+  ): AsyncIterable<QueryPartition<T>> {
+    const tag = requestTag();
+    await this.firestore.initializeIfNeeded(tag);
+
+    // 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 en 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();
+
+    let lastValues: api.IValue[] | undefined = undefined;
+    let partitionCount = 0;
+
+    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.getPartitionsAsync',
+      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
+    );
+  }
+
+  /**
+   * 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.
+   *
+   * @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 {Promise<QueryPartition[]>} A Promise with the `QueryPartition`s
+   * returned as an array.
+   */
+  async getPartitions(
+    desiredPartitionCount: number
+  ): Promise<QueryPartition<T>[]> {
+    const result: QueryPartition<T>[] = [];
+    for await (const partition of this.getPartitionsAsync(
+      desiredPartitionCount
+    )) {
+      result.push(partition);
+    }
+    return result;
+  }
+
+  /**
+   * 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 postSnap = await Firestore()
+   *   .collectionGroup('posts')
+   *   .withConverter(postConverter)
+   *   .doc().get();
+   * const post = postSnap.data();
+   * if (post !== undefined) {
+   *   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/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);
   }
 
   /**

dev/src/query-partition.ts

@@ -0,0 +1,116 @@
+/*
+ * 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 {FieldOrder, Query, QueryOptions} from './reference';
+import {FieldPath} from './path';
+import {Serializer} from './serializer';
+import {Firestore} from './index';
+
+import api = protos.google.firestore.v1;
+
+/**
+ * A split point that can be used in a query as a starting and/or end point for
+ * the query results. The cursors returned by {@link #startAt} and {@link
+ * #endBefore} can only be used in a query that matches the constraint of query
+ * that produced this partition.
+ */
+export class QueryPartition<T = firestore.DocumentData>
+  implements firestore.QueryPartition<T> {
+  private readonly _serializer: Serializer;
+  private _memoizedStartAt: unknown[] | undefined;
+  private _memoizedEndBefore: unknown[] | undefined;
+
+  constructor(
+    private readonly _firestore: Firestore,
+    private readonly _collectionId: string,
+    private readonly _converter:
+      | firestore.FirestoreDataConverter<T>
+      | undefined,
+    private readonly _startAt: api.IValue[] | undefined,
+    private readonly _endBefore: api.IValue[] | undefined
+  ) {
+    this._serializer = new Serializer(_firestore);
+  }
+
+  /**
+   * The cursor that defines the first result for this partition or `undefined`
+   * if this is the first partition.
+   *
+   * @return a cursor value that can be used with {@link Query#startAt} or
+   * `undefined` if this is the first partition. The returned array must be
+   * destructured when passed to `startAt`.
+   */
+  get startAt(): unknown[] | undefined {
+    if (this._startAt && !this._memoizedStartAt) {
+      this._memoizedStartAt = this._startAt.map(v =>
+        this._serializer.decodeValue(v)
+      );
+    }
+
+    return this._memoizedStartAt;
+  }
+
+  /**
+   * The cursor that defines the first result after this partition or
+   * `undefined` if this is the last partition.
+   *
+   * @return a cursor value that can be used with {@link Query#endBefore} or
+   * `undefined` if this is the last partition. The returned array must be
+   * destructured when passed to `endBefore`.
+   */
+  get endBefore(): unknown[] | undefined {
+    if (this._endBefore && !this._memoizedEndBefore) {
+      this._memoizedEndBefore = this._endBefore.map(v =>
+        this._serializer.decodeValue(v)
+      );
+    }
+
+    return this._memoizedEndBefore;
+  }
+
+  /**
+   * Returns a query that only returns the documents for this partition.
+   *
+   * @return a query partitioned by a {@link Query#startAt} and {@link
+   *     Query#endBefore} cursor.
+   */
+  toQuery(): Query<T> {
+    // Since the api.Value to JavaScript type conversion can be lossy (unless
+    // `useBigInt` is used), we pass the original protobuf representaion to the
+    // created query.
+    let queryOptions = QueryOptions.forCollectionGroupQuery(
+      this._collectionId,
+      this._converter
+    );
+    queryOptions = queryOptions.with({
+      fieldOrders: [new FieldOrder(FieldPath.documentId())],
+    });
+    if (this._startAt !== undefined) {
+      queryOptions = queryOptions.with({
+        startAt: {before: true, values: this._startAt},
+      });
+    }
+    if (this._endBefore !== undefined) {
+      queryOptions = queryOptions.with({
+        endAt: {before: true, values: this._endBefore},
+      });
+    }
+    return new Query(this._firestore, queryOptions);
+  }
+}