Fixing memoization issue with queryToAggregateTarget.

Author: MarkDuckworth

packages/firestore/src/core/query.ts

@@ -269,62 +269,71 @@ export function queryNormalizedOrderBy(query: Query): OrderBy[] {
  * Converts this `Query` instance to it's corresponding `Target` representation.
  */
 export function queryToTarget(query: Query): Target {
-  return _queryToTarget(query, queryNormalizedOrderBy(query));
+  const queryImpl = debugCast(query, QueryImpl);
+  if (!queryImpl.memoizedTarget) {
+    queryImpl.memoizedTarget = _queryToTarget(
+      queryImpl,
+      queryNormalizedOrderBy(query)
+    );
+  }
+
+  return queryImpl.memoizedTarget;
 }
 
 /**
  * Converts this `Query` instance to it's corresponding `Target` representation,
  * for use within an aggregate query.
  */
-export function aggregateQueryToTarget(query: Query): Target {
-  // Do not include implicit order-bys for aggregate queries.
-  return _queryToTarget(query, query.explicitOrderBy);
-}
-
-export function _queryToTarget(query: Query, orderBys: OrderBy[]): Target {
+export function queryToAggregateTarget(query: Query): Target {
   const queryImpl = debugCast(query, QueryImpl);
-  if (!queryImpl.memoizedTarget) {
-    if (queryImpl.limitType === LimitType.First) {
-      queryImpl.memoizedTarget = newTarget(
-        queryImpl.path,
-        queryImpl.collectionGroup,
-        orderBys,
-        queryImpl.filters,
-        queryImpl.limit,
-        queryImpl.startAt,
-        queryImpl.endAt
-      );
-    } else {
-      // Flip the orderBy directions since we want the last results
-      orderBys = orderBys.map(orderBy => {
-        const dir =
-          orderBy.dir === Direction.DESCENDING
-            ? Direction.ASCENDING
-            : Direction.DESCENDING;
-        return new OrderBy(orderBy.field, dir);
-      });
-
-      // We need to swap the cursors to match the now-flipped query ordering.
-      const startAt = queryImpl.endAt
-        ? new Bound(queryImpl.endAt.position, queryImpl.endAt.inclusive)
-        : null;
-      const endAt = queryImpl.startAt
-        ? new Bound(queryImpl.startAt.position, queryImpl.startAt.inclusive)
-        : null;
-
-      // Now return as a LimitType.First query.
-      queryImpl.memoizedTarget = newTarget(
-        queryImpl.path,
-        queryImpl.collectionGroup,
-        orderBys,
-        queryImpl.filters,
-        queryImpl.limit,
-        startAt,
-        endAt
-      );
-    }
+
+  // Do not include implicit order-bys for aggregate queries.
+  return _queryToTarget(queryImpl, query.explicitOrderBy);
+}
+
+export function _queryToTarget(
+  queryImpl: QueryImpl,
+  orderBys: OrderBy[]
+): Target {
+  if (queryImpl.limitType === LimitType.First) {
+    return newTarget(
+      queryImpl.path,
+      queryImpl.collectionGroup,
+      orderBys,
+      queryImpl.filters,
+      queryImpl.limit,
+      queryImpl.startAt,
+      queryImpl.endAt
+    );
+  } else {
+    // Flip the orderBy directions since we want the last results
+    orderBys = orderBys.map(orderBy => {
+      const dir =
+        orderBy.dir === Direction.DESCENDING
+          ? Direction.ASCENDING
+          : Direction.DESCENDING;
+      return new OrderBy(orderBy.field, dir);
+    });
+
+    // We need to swap the cursors to match the now-flipped query ordering.
+    const startAt = queryImpl.endAt
+      ? new Bound(queryImpl.endAt.position, queryImpl.endAt.inclusive)
+      : null;
+    const endAt = queryImpl.startAt
+      ? new Bound(queryImpl.startAt.position, queryImpl.startAt.inclusive)
+      : null;
+
+    // Now return as a LimitType.First query.
+    return newTarget(
+      queryImpl.path,
+      queryImpl.collectionGroup,
+      orderBys,
+      queryImpl.filters,
+      queryImpl.limit,
+      startAt,
+      endAt
+    );
   }
-  return queryImpl.memoizedTarget!;
 }
 
 export function queryWithAddedFilter(query: Query, filter: Filter): Query {

packages/firestore/src/remote/datastore.ts

@@ -18,7 +18,7 @@
 import { CredentialsProvider } from '../api/credentials';
 import { User } from '../auth/user';
 import { Aggregate } from '../core/aggregate';
-import { aggregateQueryToTarget, Query, queryToTarget } from '../core/query';
+import { queryToAggregateTarget, Query, queryToTarget } from '../core/query';
 import { Document } from '../model/document';
 import { DocumentKey } from '../model/document_key';
 import { Mutation } from '../model/mutation';
@@ -248,7 +248,7 @@ export async function invokeRunAggregationQueryRpc(
   const datastoreImpl = debugCast(datastore, DatastoreImpl);
   const { request, aliasMap } = toRunAggregationQueryRequest(
     datastoreImpl.serializer,
-    aggregateQueryToTarget(query),
+    queryToAggregateTarget(query),
     aggregates
   );
 

packages/firestore/test/unit/core/query.test.ts

@@ -33,6 +33,7 @@ import {
   queryWithLimit,
   queryWithStartAt,
   stringifyQuery,
+  queryToAggregateTarget,
   queryToTarget,
   QueryImpl,
   queryEquals,
@@ -852,6 +853,57 @@ describe('Query', () => {
     assertQueryMatches(query5, [doc3], [doc1, doc2, doc4, doc5]);
   });
 
+  it('generates appropriate order-bys for aggregate and non-aggregate targets', () => {
+    const col = newQueryForPath(ResourcePath.fromString('collection'));
+
+    // Build two identical queries
+    const query1 = queryWithAddedFilter(col, filter('foo', '>', 1));
+    const query2 = queryWithAddedFilter(col, filter('foo', '>', 1));
+
+    // Compute an aggregate and non-aggregate target from the queries
+    const aggregateTarget = queryToAggregateTarget(query1);
+    const target = queryToTarget(query2);
+
+    expect(aggregateTarget.orderBy.length).to.equal(0);
+    expect(target.orderBy.length).to.equal(2);
+    expect(target.orderBy[0].dir).to.equal('asc');
+    expect(target.orderBy[0].field.canonicalString()).to.equal('foo');
+    expect(target.orderBy[1].dir).to.equal('asc');
+    expect(target.orderBy[1].field.canonicalString()).to.equal('__name__');
+  });
+
+  it('generated order-bys are not affected by previously memoized targets', () => {
+    const col = newQueryForPath(ResourcePath.fromString('collection'));
+
+    // Build two identical queries
+    const query1 = queryWithAddedFilter(col, filter('foo', '>', 1));
+    const query2 = queryWithAddedFilter(col, filter('foo', '>', 1));
+
+    // query1 - first to aggregate target, then to non-aggregate target
+    const aggregateTarget1 = queryToAggregateTarget(query1);
+    const target1 = queryToTarget(query1);
+
+    // query2 - first to non-aggregate target, then to aggregate target
+    const target2 = queryToTarget(query2);
+    const aggregateTarget2 = queryToAggregateTarget(query2);
+
+    expect(aggregateTarget1.orderBy.length).to.equal(0);
+
+    expect(aggregateTarget2.orderBy.length).to.equal(0);
+
+    expect(target1.orderBy.length).to.equal(2);
+    expect(target1.orderBy[0].dir).to.equal('asc');
+    expect(target1.orderBy[0].field.canonicalString()).to.equal('foo');
+    expect(target1.orderBy[1].dir).to.equal('asc');
+    expect(target1.orderBy[1].field.canonicalString()).to.equal('__name__');
+
+    expect(target2.orderBy.length).to.equal(2);
+    expect(target2.orderBy[0].dir).to.equal('asc');
+    expect(target2.orderBy[0].field.canonicalString()).to.equal('foo');
+    expect(target2.orderBy[1].dir).to.equal('asc');
+    expect(target2.orderBy[1].field.canonicalString()).to.equal('__name__');
+  });
+
   function assertQueryMatches(
     query: Query,
     matching: MutableDocument[],