Add long-alias support for aggregations. (#7244)

* Add long-alias support for aggregations.
firebase · MarkDuckworth · Apr 25, 2023 · Mar 29, 2023 · Mar 29, 2023 · Mar 29, 2023
commit 51e40595975758e52c5bf46d81818e357739244c
@@ -20,7 +20,6 @@ import { AggregateImpl } from '../core/aggregate';
 import { firestoreClientRunAggregateQuery } from '../core/firestore_client';
 import { count } from '../lite-api/aggregate';
 import { AggregateQuerySnapshot } from '../lite-api/aggregate_types';
-import { AggregateAlias } from '../model/aggregate_alias';
 import { ApiClientObjectMap, Value } from '../protos/firestore_proto_api';
 import { cast } from '../util/input_validation';
 import { mapToArray } from '../util/obj';
@@ -110,7 +109,7 @@ export function getAggregateFromServer<T extends AggregateSpec>(
 
   const internalAggregates = mapToArray(aggregateSpec, (aggregate, alias) => {
     return new AggregateImpl(
-      new AggregateAlias(alias),
+      alias,
       aggregate._aggregateType,
       aggregate._internalFieldPath
     );

@@ -15,7 +15,6 @@
  * limitations under the License.
  */
 
-import { AggregateAlias } from '../model/aggregate_alias';
 import { FieldPath } from '../model/path';
 
 /**
@@ -29,7 +28,7 @@ export type AggregateType = 'count' | 'avg' | 'sum';
  */
 export interface Aggregate {
   readonly fieldPath?: FieldPath;
-  readonly alias: AggregateAlias;
+  readonly alias: string;
   readonly aggregateType: AggregateType;
 }
 
@@ -38,7 +37,7 @@ export interface Aggregate {
  */
 export class AggregateImpl implements Aggregate {
   constructor(
-    readonly alias: AggregateAlias,
+    readonly alias: string,
     readonly aggregateType: AggregateType,
     readonly fieldPath?: FieldPath
   ) {}

@@ -18,7 +18,6 @@
 import { deepEqual } from '@firebase/util';
 
 import { AggregateImpl } from '../core/aggregate';
-import { AggregateAlias } from '../model/aggregate_alias';
 import { ApiClientObjectMap, Value } from '../protos/firestore_proto_api';
 import { invokeRunAggregationQueryRpc } from '../remote/datastore';
 import { cast } from '../util/input_validation';
@@ -96,7 +95,7 @@ export function getAggregate<T extends AggregateSpec>(
 
   const internalAggregates = mapToArray(aggregateSpec, (aggregate, alias) => {
     return new AggregateImpl(
-      new AggregateAlias(alias),
+      alias,
       aggregate._aggregateType,
       aggregate._internalFieldPath
     );

@@ -246,7 +246,7 @@ export async function invokeRunAggregationQueryRpc(
   aggregates: Aggregate[]
 ): Promise<ApiClientObjectMap<Value>> {
   const datastoreImpl = debugCast(datastore, DatastoreImpl);
-  const request = toRunAggregationQueryRequest(
+  const { request, aliasMap } = toRunAggregationQueryRequest(
     datastoreImpl.serializer,
     queryToTarget(query),
     aggregates
@@ -277,7 +277,22 @@ export async function invokeRunAggregationQueryRpc(
     'aggregationQueryResponse.result.aggregateFields'
   );
 
-  return filteredResult[0].result.aggregateFields;
+  // Remap the short-form aliases that were sent to the server
+  // to the client-side aliases. Users will access the results
+  // using the client-side alias.
+  const unmappedAggregateFields = filteredResult[0].result?.aggregateFields;
+  const remappedFields = Object.keys(unmappedAggregateFields).reduce<
+    ApiClientObjectMap<Value>
+  >((accumulator, key) => {
+    debugAssert(
+      !isNullOrUndefined(aliasMap[key]),
+      `'${key}' not present in aliasMap result`
+    );
+    accumulator[aliasMap[key]] = unmappedAggregateFields[key]!;
+    return accumulator;
+  }, {});
+
+  return remappedFields;
 }
 
 export function newPersistentWriteStream(

@@ -891,26 +891,38 @@ export function toRunAggregationQueryRequest(
   serializer: JsonProtoSerializer,
   target: Target,
   aggregates: Aggregate[]
-): ProtoRunAggregationQueryRequest {
+): {
+  request: ProtoRunAggregationQueryRequest;
+  aliasMap: Record<string, string>;
+} {
   const queryTarget = toQueryTarget(serializer, target);
+  const aliasMap: Record<string, string> = {};
 
   const aggregations: ProtoAggregation[] = [];
+  let aggregationNum = 0;
+
   aggregates.forEach(aggregate => {
+    // Map all client-side aliases to a unique short-form
+    // alias. This avoids issues with client-side aliases that
+    // exceed the 1500-byte string size limit.
+    const serverAlias = `aggregate_${aggregationNum++}`;
+    aliasMap[serverAlias] = aggregate.alias;
+
     if (aggregate.aggregateType === 'count') {
       aggregations.push({
-        alias: aggregate.alias.canonicalString(),
+        alias: serverAlias,
         count: {}
       });
     } else if (aggregate.aggregateType === 'avg') {
       aggregations.push({
-        alias: aggregate.alias.canonicalString(),
+        alias: serverAlias,
         avg: {
           field: toFieldPathReference(aggregate.fieldPath!)
         }
       });
     } else if (aggregate.aggregateType === 'sum') {
       aggregations.push({
-        alias: aggregate.alias.canonicalString(),
+        alias: serverAlias,
         sum: {
           field: toFieldPathReference(aggregate.fieldPath!)
         }
@@ -919,11 +931,14 @@ export function toRunAggregationQueryRequest(
   });
 
   return {
-    structuredAggregationQuery: {
-      aggregations,
-      structuredQuery: queryTarget.structuredQuery
+    request: {
+      structuredAggregationQuery: {
+        aggregations,
+        structuredQuery: queryTarget.structuredQuery
+      },
+      parent: queryTarget.parent
     },
-    parent: queryTarget.parent
+    aliasMap
   };
 }
 

@@ -218,6 +218,30 @@ apiDescribe('Aggregation queries', (persistence: boolean) => {
     });
   });
 
+  it('allows aliases with length greater than 1500 bytes', () => {
+    // Alias string length is bytes of UTF-8 encoded alias + 1;
+    let longAlias = '';
+    for (let i = 0; i < 150; i++) {
+      longAlias += '0123456789';
+    }
+
+    const longerAlias = longAlias + longAlias;
+
+    const testDocs = {
+      a: { num: 3 },
+      b: { num: 5 }
+    };
+    return withTestCollection(persistence, testDocs, async coll => {
+      const snapshot = await getAggregateFromServer(coll, {
+        [longAlias]: count(),
+        [longerAlias]: count()
+      });
+
+      expect(snapshot.data()[longAlias]).to.equal(2);
+      expect(snapshot.data()[longerAlias]).to.equal(2);
+    });
+  });
+
   it('can get duplicate aggregations using getAggregationFromServer', () => {
     const testDocs = {
       a: { author: 'authorA', title: 'titleA' },
@@ -432,7 +456,7 @@ apiDescribe.skip(
         });
 
         await expect(promise).to.eventually.be.rejectedWith(
-          /INVALID_ARGUMENT.*maximum number of aggregations/
+          /maximum number of aggregations/
         );
       });
     });