Skip to content

Commit

Permalink
Add system tests and comments.
Browse files Browse the repository at this point in the history
  • Loading branch information
wu-hui committed Jul 5, 2020
1 parent 81ea916 commit 76e5416
Show file tree
Hide file tree
Showing 5 changed files with 260 additions and 80 deletions.
Expand Up @@ -17,28 +17,30 @@
package com.google.cloud.firestore;

import com.google.cloud.Timestamp;
import com.google.cloud.firestore.Query.LimitType;
import com.google.firestore.proto.BundleElement;
import com.google.firestore.proto.BundleMetadata;
import com.google.firestore.proto.BundledDocumentMetadata;
import com.google.firestore.proto.BundledQuery;
import com.google.firestore.proto.NamedQuery;
import com.google.firestore.v1.Document;
import com.google.firestore.v1.RunQueryRequest;
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.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;

// 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<>();
Expand All @@ -53,14 +55,21 @@ 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 documentSnap A document snapshot to add.
* @returns This instance.
*/
public Builder add(DocumentSnapshot documentSnap) {
BundledDocumentMetadata metadata =
BundledDocumentMetadata.newBuilder()
.setName(documentSnap.getReference().getName())
.setReadTime(documentSnap.getReadTime().toProto())
.setExists(documentSnap.exists())
.build();
Document document = documentSnap.toDocumentPb().build();
Document document = documentSnap.exists() ? documentSnap.toDocumentPb().build() : null;
documents.put(metadata.getName(), new BundledDocument(metadata, document));

if (documentSnap.getReadTime().compareTo(latestReadTime) > 0) {
Expand All @@ -70,21 +79,21 @@ public Builder add(DocumentSnapshot documentSnap) {
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) {
RunQueryRequest queryProto = querySnap.getQuery().toProto();
LimitType limitType = querySnap.getQuery().options.getLimitType();
BundledQuery query = querySnap.getQuery().toBundledQuery();
NamedQuery namedQuery =
NamedQuery.newBuilder()
.setName(queryName)
.setReadTime(querySnap.getReadTime().toProto())
.setBundledQuery(
BundledQuery.newBuilder()
.setParent(queryProto.getParent())
.setStructuredQuery(queryProto.getStructuredQuery())
.setLimitType(
limitType.equals(LimitType.Last)
? BundledQuery.LimitType.LAST
: BundledQuery.LimitType.FIRST))
.setBundledQuery(query)
.build();
namedQueries.put(queryName, namedQuery);

Expand Down Expand Up @@ -150,11 +159,15 @@ private FirestoreBundle(byte[] data) {
bundleData = data;
}

/** Returns the bundle content as a readonly {@link ByteBuffer}. */
public ByteBuffer toByteBuffer() {
return ByteBuffer.wrap(bundleData) /*.asReadOnlyBuffer()*/;
return ByteBuffer.wrap(bundleData).asReadOnlyBuffer();
}
}

/**
* Convenient class to hold both the metadata and the actual content of a document to be bundled.
*/
class BundledDocument {
private BundledDocumentMetadata metadata;
private Document document;
Expand Down
Expand Up @@ -36,6 +36,8 @@
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.firestore.proto.BundledQuery;
import com.google.firestore.proto.BundledQuery.LimitType;
import com.google.firestore.v1.Cursor;
import com.google.firestore.v1.Document;
import com.google.firestore.v1.RunQueryRequest;
Expand Down Expand Up @@ -1071,6 +1073,71 @@ public Query endAt(@Nonnull DocumentSnapshot snapshot) {

/** Build the final Firestore query. */
StructuredQuery.Builder buildQuery() {
StructuredQuery.Builder structuredQuery = buildWithoutClientTranslation();
if (options.getLimitType().equals(LimitType.Last)) {
// Apply client translation for limitToLast.

if (!options.getFieldOrders().isEmpty()) {
structuredQuery.clearOrderBy();
for (FieldOrder order : options.getFieldOrders()) {
// Flip the orderBy directions since we want the last results
order =
new FieldOrder(
order.fieldReference,
order.direction.equals(Direction.ASCENDING)
? Direction.DESCENDING
: Direction.ASCENDING);
structuredQuery.addOrderBy(order.toProto());
}
}

if (options.getStartCursor() != null) {
structuredQuery.clearEndAt();
// Swap the cursors to match the flipped query ordering.
Cursor cursor =
options
.getStartCursor()
.toBuilder()
.setBefore(!options.getStartCursor().getBefore())
.build();
structuredQuery.setEndAt(cursor);
}

if (options.getEndCursor() != null) {
structuredQuery.clearStartAt();
// Swap the cursors to match the flipped query ordering.
Cursor cursor =
options
.getEndCursor()
.toBuilder()
.setBefore(!options.getEndCursor().getBefore())
.build();
structuredQuery.setStartAt(cursor);
}
}

return structuredQuery;
}

/**
* Builds a query to saved in a bundle file.
*
* <p>This will not do client translation (limitToLast order flip, for example),
*/
BundledQuery toBundledQuery() {
StructuredQuery.Builder structuredQuery = buildWithoutClientTranslation();

return BundledQuery.newBuilder()
.setStructuredQuery(structuredQuery)
.setParent(options.getParentPath().toString())
.setLimitType(
options.getLimitType().equals(LimitType.Last)
? BundledQuery.LimitType.LAST
: BundledQuery.LimitType.FIRST)
.build();
}

private StructuredQuery.Builder buildWithoutClientTranslation() {
StructuredQuery.Builder structuredQuery = StructuredQuery.newBuilder();
CollectionSelector.Builder collectionSelector = CollectionSelector.newBuilder();
collectionSelector.setCollectionId(options.getCollectionId());
Expand Down Expand Up @@ -1100,21 +1167,7 @@ StructuredQuery.Builder buildQuery() {

if (!options.getFieldOrders().isEmpty()) {
for (FieldOrder order : options.getFieldOrders()) {
switch (options.getLimitType()) {
case First:
structuredQuery.addOrderBy(order.toProto());
break;
case Last:
// Flip the orderBy directions since we want the last results
order =
new FieldOrder(
order.fieldReference,
order.direction.equals(Direction.ASCENDING)
? Direction.DESCENDING
: Direction.ASCENDING);
structuredQuery.addOrderBy(order.toProto());
break;
}
structuredQuery.addOrderBy(order.toProto());
}
} else if (LimitType.Last.equals(options.getLimitType())) {
throw new IllegalStateException(
Expand All @@ -1134,39 +1187,11 @@ StructuredQuery.Builder buildQuery() {
}

if (options.getStartCursor() != null) {
switch (options.getLimitType()) {
case First:
structuredQuery.setStartAt(options.getStartCursor());
break;
case Last:
// Swap the cursors to match the flipped query ordering.
Cursor cursor =
options
.getStartCursor()
.toBuilder()
.setBefore(!options.getStartCursor().getBefore())
.build();
structuredQuery.setEndAt(cursor);
break;
}
structuredQuery.setStartAt(options.getStartCursor());
}

if (options.getEndCursor() != null) {
switch (options.getLimitType()) {
case First:
structuredQuery.setEndAt(options.getEndCursor());
break;
case Last:
// Swap the cursors to match the flipped query ordering.
Cursor cursor =
options
.getEndCursor()
.toBuilder()
.setBefore(!options.getEndCursor().getBefore())
.build();
structuredQuery.setStartAt(cursor);
break;
}
structuredQuery.setEndAt(options.getEndCursor());
}

return structuredQuery;
Expand Down

0 comments on commit 76e5416

Please sign in to comment.