feat: Implementation of Firestore Bundle Builder

google-cloud-firestore/pom.xml

@@ -157,6 +157,12 @@
       <version>3.11</version>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>com.google.cloud</groupId>
+      <artifactId>proto-google-cloud-firestore-bundle-v1</artifactId>
+      <version>2.1.1-SNAPSHOT</version>
+      <scope>compile</scope>
+    </dependency>
   </dependencies>
 
   <reporting>

google-cloud-firestore/src/main/java/com/google/cloud/firestore/DocumentSnapshot.java

@@ -412,6 +412,16 @@ Write.Builder toPb() {
     return write;
   }
 
+  Document.Builder toDocumentPb() {
+    Preconditions.checkState(exists(), "Can't call toDocument() on a document that doesn't exist");
+    Document.Builder document = Document.newBuilder();
+    return document
+        .setName(docRef.getName())
+        .putAllFields(fields)
+        .setCreateTime(createTime.toProto())
+        .setUpdateTime(updateTime.toProto());
+  }
+
   /**
    * Returns true if the document's data and path in this DocumentSnapshot equals the provided
    * snapshot.

google-cloud-firestore/src/main/java/com/google/cloud/firestore/FirestoreBundle.java

@@ -0,0 +1,228 @@
+/*
+ * 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.
+ */
+
+package com.google.cloud.firestore;
+
+import com.google.api.client.util.Lists;
+import com.google.cloud.Timestamp;
+import com.google.common.base.Optional;
+import com.google.firestore.bundle.BundleElement;
+import com.google.firestore.bundle.BundleMetadata;
+import com.google.firestore.bundle.BundledDocumentMetadata;
+import com.google.firestore.bundle.BundledQuery;
+import com.google.firestore.bundle.NamedQuery;
+import com.google.firestore.v1.Document;
+import com.google.protobuf.InvalidProtocolBufferException;
+import com.google.protobuf.util.JsonFormat;
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+/** Represents a Firestore data bundle with results from the given document and query snapshots. */
+public final class FirestoreBundle {
+
+  static int BUNDLE_SCHEMA_VERSION = 1;
+  // Printer to encode protobuf objects into JSON string.
+  private static JsonFormat.Printer PRINTER = JsonFormat.printer();
+
+  // Raw byte array to hold the content of the bundle.
+  private byte[] bundleData;
+
+  /** Builds a Firestore data bundle with results from the given document and query snapshots. */
+  public static final class Builder {
+    // Id of the bundle.
+    private String id;
+    // Resulting documents for the bundle, keyed by full document path.
+    private Map<String, BundledDocument> documents = new HashMap<>();
+    // Named queries saved in the bundle, keyed by query name.
+    private Map<String, NamedQuery> namedQueries = new HashMap<>();
+    // The latest read time among all bundled documents and queries.
+    private Timestamp latestReadTime = Timestamp.MIN_VALUE;
+
+    public Builder(String id) {
+      this.id = id;
+    }
+
+    /**
+     * Adds a Firestore document snapshot to the bundle. Both the documents data and the document
+     * read time will be included in the bundle.
+     *
+     * @param documentSnapshot A document snapshot to add.
+     * @returns This instance.
+     */
+    public Builder add(DocumentSnapshot documentSnapshot) {
+      return add(documentSnapshot, Optional.<String>absent());
+    }
+
+    private Builder add(DocumentSnapshot documentSnapshot, Optional<String> queryName) {
+      String documentName = documentSnapshot.getReference().getName();
+      BundledDocument originalDocument = documents.get(documentSnapshot.getReference().getName());
+      List<String> queries =
+          originalDocument == null
+              ? Lists.<String>newArrayList()
+              : Lists.newArrayList(originalDocument.getMetadata().getQueriesList());
+
+      // Update with document built from `documentSnapshot` because it is newer.
+      if (originalDocument == null
+          || documentSnapshot
+                  .getReadTime()
+                  .compareTo(Timestamp.fromProto(originalDocument.getMetadata().getReadTime()))
+              > 0) {
+        BundledDocumentMetadata metadata =
+            BundledDocumentMetadata.newBuilder()
+                .setName(documentName)
+                .setReadTime(documentSnapshot.getReadTime().toProto())
+                .setExists(documentSnapshot.exists())
+                .build();
+        Document document =
+            documentSnapshot.exists() ? documentSnapshot.toDocumentPb().build() : null;
+        documents.put(documentName, new BundledDocument(metadata, document));
+      }
+
+      // Update queries to include all queries whose results include this document.
+      if (queryName.isPresent()) {
+        queries.add(queryName.get());
+      }
+      documents
+          .get(documentName)
+          .setMetadata(
+              documents
+                  .get(documentName)
+                  .getMetadata()
+                  .toBuilder()
+                  .clearQueries()
+                  .addAllQueries(queries)
+                  .build());
+
+      if (documentSnapshot.getReadTime().compareTo(latestReadTime) > 0) {
+        latestReadTime = documentSnapshot.getReadTime();
+      }
+
+      return this;
+    }
+
+    /**
+     * Adds a Firestore query snapshots to the bundle. Both the documents in the query snapshots and
+     * the query read time will be included in the bundle.
+     *
+     * @param queryName The name of the query to add.
+     * @param querySnap The query snapshot to add.
+     * @returns This instance.
+     */
+    public Builder add(String queryName, QuerySnapshot querySnap) {
+      BundledQuery query = querySnap.getQuery().toBundledQuery();
+      NamedQuery namedQuery =
+          NamedQuery.newBuilder()
+              .setName(queryName)
+              .setReadTime(querySnap.getReadTime().toProto())
+              .setBundledQuery(query)
+              .build();
+      namedQueries.put(queryName, namedQuery);
+
+      for (QueryDocumentSnapshot snapshot : querySnap.getDocuments()) {
+        add(snapshot, Optional.of(queryName));
+      }
+
+      if (querySnap.getReadTime().compareTo(latestReadTime) > 0) {
+        latestReadTime = querySnap.getReadTime();
+      }
+
+      return this;
+    }
+
+    public FirestoreBundle build() {
+      StringBuilder buffer = new StringBuilder();
+
+      for (NamedQuery namedQuery : namedQueries.values()) {
+        buffer.append(
+            elementToLengthPrefixedStringBuilder(
+                BundleElement.newBuilder().setNamedQuery(namedQuery).build()));
+      }
+
+      for (BundledDocument bundledDocument : documents.values()) {
+        buffer.append(
+            elementToLengthPrefixedStringBuilder(
+                BundleElement.newBuilder()
+                    .setDocumentMetadata(bundledDocument.getMetadata())
+                    .build()));
+        if (bundledDocument.getDocument() != null) {
+          buffer.append(
+              elementToLengthPrefixedStringBuilder(
+                  BundleElement.newBuilder().setDocument(bundledDocument.getDocument()).build()));
+        }
+      }
+
+      BundleMetadata metadata =
+          BundleMetadata.newBuilder()
+              .setId(id)
+              .setCreateTime(latestReadTime.toProto())
+              .setVersion(BUNDLE_SCHEMA_VERSION)
+              .setTotalDocuments(documents.size())
+              .setTotalBytes(buffer.toString().getBytes().length)
+              .build();
+      BundleElement element = BundleElement.newBuilder().setMetadata(metadata).build();
+      buffer.insert(0, elementToLengthPrefixedStringBuilder(element));
+
+      return new FirestoreBundle(buffer.toString().getBytes(StandardCharsets.UTF_8));
+    }
+
+    private StringBuilder elementToLengthPrefixedStringBuilder(BundleElement element) {
+      String elementJson = null;
+      try {
+        elementJson = PRINTER.print(element);
+      } catch (InvalidProtocolBufferException e) {
+        throw new RuntimeException(e);
+      }
+      return new StringBuilder().append(elementJson.getBytes().length).append(elementJson);
+    }
+  }
+
+  private FirestoreBundle(byte[] data) {
+    bundleData = data;
+  }
+
+  /** Returns the bundle content as a readonly {@link ByteBuffer}. */
+  public ByteBuffer toByteBuffer() {
+    return ByteBuffer.wrap(bundleData).asReadOnlyBuffer();
+  }
+
+  /**
+   * Convenient class to hold both the metadata and the actual content of a document to be bundled.
+   */
+  private static class BundledDocument {
+    private BundledDocumentMetadata metadata;
+    private final Document document;
+
+    BundledDocument(BundledDocumentMetadata metadata, Document document) {
+      this.metadata = metadata;
+      this.document = document;
+    }
+
+    public BundledDocumentMetadata getMetadata() {
+      return metadata;
+    }
+
+    void setMetadata(BundledDocumentMetadata metadata) {
+      this.metadata = metadata;
+    }
+
+    public Document getDocument() {
+      return document;
+    }
+  }
+}