Ideas to support converting bundle protos to Proto3 JSON compatible format

Author: MarkDuckworth

packages/firestore/src/lite-api/user_data_reader.ts

@@ -778,7 +778,7 @@ export function parseData(
   }
 }
 
-function parseObject(
+export function parseObject(
   obj: Dict<unknown>,
   context: ParseContextImpl
 ): { mapValue: ProtoMapValue } {

packages/firestore/src/util/bundle_builder_impl.ts

@@ -15,37 +15,51 @@
  * limitations under the License.
  */
 
-import { queryToTarget } from '../../src/core/query';
+import {queryToTarget} from '../../src/core/query';
 import {
   JsonProtoSerializer,
   toDocument,
   toName,
   toQueryTarget,
+  toTimestamp,
 } from '../../src/remote/serializer';
-import { Firestore } from '../api/database';
-import { DatabaseId } from '../core/database_info';
-import { DocumentSnapshot, QuerySnapshot } from '../lite-api/snapshot';
-import { Timestamp } from '../lite-api/timestamp';
+import {Firestore} from '../api/database';
+import {DatabaseId} from '../core/database_info';
+import {DocumentSnapshot, QuerySnapshot} from '../lite-api/snapshot';
+import {Timestamp} from '../lite-api/timestamp';
 import {
+  BundledDocumentMetadata as ProtoBundledDocumentMetadata,
   BundleElement as ProtoBundleElement,
   BundleMetadata as ProtoBundleMetadata,
-  BundledDocumentMetadata as ProtoBundledDocumentMetadata,
   NamedQuery as ProtoNamedQuery,
 } from '../protos/firestore_bundle_proto';
-import { Document } from '../protos/firestore_proto_api';
+import {
+  Document as ProtoDocument,
+  Document
+} from '../protos/firestore_proto_api';
 
 import {
   invalidArgumentMessage,
   validateString,
 } from './bundle_builder_validation_utils';
+import {encoder} from "../../test/unit/util/bundle_data";
+import {
+  parseData, parseObject,
+  UserDataReader,
+  UserDataSource
+} from "../lite-api/user_data_reader";
+import {AbstractUserDataWriter} from "../lite-api/user_data_writer";
+import {ExpUserDataWriter} from "../api/reference_impl";
+import {MutableDocument} from "../model/document";
+import {debugAssert} from "./assert";
 
 const BUNDLE_VERSION = 1;
 
 /**
  * Builds a Firestore data bundle with results from the given document and query snapshots.
  */
 export class BundleBuilder {
-  
+
   // Resulting documents for the bundle, keyed by full document path.
   private documents: Map<string, BundledDocument> = new Map();
   // Named queries saved in the bundle, keyed by query name.
@@ -56,10 +70,21 @@ export class BundleBuilder {
 
   private databaseId: DatabaseId;
 
+  private readonly serializer: JsonProtoSerializer;
+  private readonly userDataReader: UserDataReader;
+  private readonly userDataWriter: AbstractUserDataWriter;
+
   constructor(private firestore: Firestore, readonly bundleId: string) {
     this.databaseId = firestore._databaseId;
+
+    // useProto3Json is true because the objects will be serialized to JSON string
+    // before being written to the bundle buffer.
+    this.serializer = new JsonProtoSerializer(this.databaseId, /*useProto3Json=*/ true);
+
+    this.userDataWriter = new ExpUserDataWriter(firestore);
+    this.userDataReader = new UserDataReader(this.databaseId, true, this.serializer);
   }
-  
+
   /**
    * Adds a Firestore document snapshot or query snapshot to the bundle.
    * Both the documents data and the query read time will be included in the bundle.
@@ -97,7 +122,34 @@ export class BundleBuilder {
     }
     return this;
   }
-  
+
+  toBundleDocument(
+    document: MutableDocument
+  ): ProtoDocument {
+    // TODO handle documents that have mutations
+    debugAssert(
+      !document.hasLocalMutations,
+      "Can't serialize documents with mutations."
+    );
+
+    // Convert document fields proto to DocumentData and then back
+    // to Proto3 JSON objects. This is the same approach used in
+    // bundling in the nodejs-firestore SDK. It may not be the most
+    // performant approach.
+    const documentData = this.userDataWriter.convertObjectMap(document.data.value.mapValue.fields, 'previous');
+    // a parse context is typically used for validating and parsing user data, but in this
+    // case we are using it internally to convert DocumentData to Proto3 JSON
+    const context = this.userDataReader.createContext(UserDataSource.ArrayArgument, 'internal toBundledDocument');
+    const proto3Fields = parseObject(documentData, context);
+
+    return {
+      name: toName(this.serializer, document.key),
+      fields: proto3Fields.mapValue.fields,
+      updateTime: toTimestamp(this.serializer, document.version.toTimestamp()),
+      createTime: toTimestamp(this.serializer, document.createTime.toTimestamp())
+    };
+  }
+
   private addBundledDocument(snap: DocumentSnapshot, queryName?: string): void {
     // TODO:  is this a valid shortcircuit?
     if(!snap._document || !snap._document.isValidDocument()) {
@@ -114,15 +166,11 @@ export class BundleBuilder {
         (snapReadTime && originalDocument.metadata.readTime! < snapReadTime)
     ) {
 
-      // TODO: Should I create on serializer for the bundler instance, or just created one adhoc
-      // like this?
-      const serializer = new JsonProtoSerializer(this.databaseId, /*useProto3Json=*/ false);
-      
       this.documents.set(snap.ref.path, {
-        document: snap._document.isFoundDocument() ? toDocument(serializer, mutableCopy) : undefined,
+        document: snap._document.isFoundDocument() ? this.toBundleDocument(mutableCopy) : undefined,
         metadata: {
-          name: toName(serializer, mutableCopy.key),
-          readTime: snapReadTime,
+          name: toName(this.serializer, mutableCopy.key),
+          readTime: !!snapReadTime ? toTimestamp(this.serializer, snapReadTime) : undefined,
           exists: snap.exists(),
         },
       });
@@ -145,10 +193,8 @@ export class BundleBuilder {
     if (this.namedQueries.has(name)) {
       throw new Error(`Query name conflict: ${name} has already been added.`);
     }
+    const queryTarget = toQueryTarget(this.serializer, queryToTarget(querySnap.query._query));
 
-    const serializer = new JsonProtoSerializer(this.databaseId, /*useProto3Json=*/ false);
-    const queryTarget = toQueryTarget(serializer, queryToTarget(querySnap.query._query));
-    
     // TODO: if we can't resolve the query's readTime then can we set it to the latest
     // of the document collection?
     let latestReadTime = new Timestamp(0, 0);
@@ -169,7 +215,7 @@ export class BundleBuilder {
     this.namedQueries.set(name, {
       name,
       bundledQuery,
-      readTime: latestReadTime
+      readTime: toTimestamp(this.serializer, latestReadTime)
     });
   }
 
@@ -178,65 +224,54 @@ export class BundleBuilder {
    * of the element.
    * @private
    * @internal
+   * @param bundleElement A ProtoBundleElement that is expected to be Proto3 JSON compatible.
    */
-  private elementToLengthPrefixedBuffer(
+  private lengthPrefixedString(
     bundleElement: ProtoBundleElement
-  ): Buffer {
-    // Convert to a valid proto message object then take its JSON representation.
-    // This take cares of stuff like converting internal byte array fields
-    // to Base64 encodings.
-    
-    // TODO: This fails. BundleElement doesn't have a toJSON method.
-    const message = require('../protos/firestore_v1_proto_api')
-      .firestore.BundleElement.fromObject(bundleElement)
-      .toJSON();
-    const buffer = Buffer.from(JSON.stringify(message), 'utf-8');
-    const lengthBuffer = Buffer.from(buffer.length.toString());
-    return Buffer.concat([lengthBuffer, buffer]);
+  ): string {
+    const str = JSON.stringify(bundleElement);
+    // TODO: it's not ideal to have to re-encode all of these strings multiple times
+    //       It may be more performant to return a UInt8Array that is concatenated to other
+    //       UInt8Arrays instead of returning and concatenating strings and then
+    //       converting the full string to UInt8Array.
+    const l = encoder.encode(str).byteLength;
+    return `${l}${str}`;
   }
 
-  build(): Buffer {
-    
-    let bundleBuffer = Buffer.alloc(0);
+  build(): Uint8Array {
+    let bundleString = '';
 
     for (const namedQuery of this.namedQueries.values()) {
-      bundleBuffer = Buffer.concat([
-        bundleBuffer,
-        this.elementToLengthPrefixedBuffer({namedQuery}),
-      ]);
+      bundleString += this.lengthPrefixedString({namedQuery});
     }
 
     for (const bundledDocument of this.documents.values()) {
       const documentMetadata: ProtoBundledDocumentMetadata =
         bundledDocument.metadata;
 
-      bundleBuffer = Buffer.concat([
-        bundleBuffer,
-        this.elementToLengthPrefixedBuffer({documentMetadata}),
-      ]);
+      bundleString += this.lengthPrefixedString({documentMetadata});
       // Write to the bundle if document exists.
       const document = bundledDocument.document;
       if (document) {
-        bundleBuffer = Buffer.concat([
-          bundleBuffer,
-          this.elementToLengthPrefixedBuffer({document}),
-        ]);
+        bundleString += this.lengthPrefixedString({document});
       }
     }
 
     const metadata: ProtoBundleMetadata = {
       id: this.bundleId,
-      createTime: this.latestReadTime,
+      createTime: toTimestamp(this.serializer, this.latestReadTime),
       version: BUNDLE_VERSION,
       totalDocuments: this.documents.size,
-      totalBytes: bundleBuffer.length,
+      // TODO: it's not ideal to have to re-encode all of these strings multiple times
+      totalBytes: encoder.encode(bundleString).length,
     };
     // Prepends the metadata element to the bundleBuffer: `bundleBuffer` is the second argument to `Buffer.concat`.
-    bundleBuffer = Buffer.concat([
-      this.elementToLengthPrefixedBuffer({metadata}),
-      bundleBuffer,
-    ]);
-    return bundleBuffer;
+    bundleString = this.lengthPrefixedString({metadata}) + bundleString;
+
+    // TODO: it's not ideal to have to re-encode all of these strings multiple times
+    //       the implementation in nodejs-firestore concatenates Buffers instead of
+    //       concatenating strings.
+    return encoder.encode(bundleString);
   }
 }
 
@@ -278,4 +313,4 @@ function validateQuerySnapshot(arg: string | number, value: unknown): void {
   if (!(value instanceof QuerySnapshot)) {
     throw new Error(invalidArgumentMessage(arg, 'QuerySnapshot'));
   }
-}
+}