Skip to content

Commit f8b9386

Browse files
tom-andersentom-andersen
committed
Mirror Query class pattern
1 parent 77decb5 commit f8b9386

File tree

4 files changed

+72
-105
lines changed

4 files changed

+72
-105
lines changed

‎google-cloud-firestore/src/main/java/com/google/cloud/firestore/Pipeline.java

Lines changed: 60 additions & 92 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,6 @@
3535
import com.google.firestore.v1.ExecutePipelineRequest;
3636
import com.google.firestore.v1.ExecutePipelineResponse;
3737
import com.google.firestore.v1.StructuredPipeline;
38-
import com.google.firestore.v1.Value;
3938
import io.opencensus.trace.AttributeValue;
4039
import io.opencensus.trace.Tracing;
4140
import java.util.ArrayList;
@@ -98,20 +97,20 @@
9897
public final class Pipeline {
9998
private static Logger logger = Logger.getLogger(Pipeline.class.getName());
10099
private final FluentIterable<Stage> stages;
101-
private final Firestore db;
100+
private final FirestoreRpcContext<?> rpcContext;
102101

103-
private Pipeline(Firestore db, FluentIterable<Stage> stages) {
104-
this.db = db;
102+
private Pipeline(FirestoreRpcContext<?> rpcContext, FluentIterable<Stage> stages) {
103+
this.rpcContext = rpcContext;
105104
this.stages = stages;
106105
}
107106

108107
@InternalApi
109-
Pipeline(Firestore db, Stage stage) {
110-
this(db, FluentIterable.of(stage));
108+
Pipeline(FirestoreRpcContext<?> rpcContext, Stage stage) {
109+
this(rpcContext, FluentIterable.of(stage));
111110
}
112111

113112
private Pipeline append(Stage stage) {
114-
return new Pipeline(this.db, stages.append(stage));
113+
return new Pipeline(this.rpcContext, stages.append(stage));
115114
}
116115

117116
/**
@@ -607,47 +606,29 @@ public Pipeline genericStage(String name, List<Object> params) {
607606
*/
608607
@BetaApi
609608
public ApiFuture<List<PipelineResult>> execute() {
610-
if (db instanceof FirestoreImpl) {
611-
FirestoreImpl firestoreImpl = (FirestoreImpl) db;
612-
Value pipelineValue = toProto();
613-
ExecutePipelineRequest request =
614-
ExecutePipelineRequest.newBuilder()
615-
.setDatabase(firestoreImpl.getResourcePath().getDatabaseName().toString())
616-
.setStructuredPipeline(
617-
StructuredPipeline.newBuilder()
618-
.setPipeline(pipelineValue.getPipelineValue())
619-
.build())
620-
.build();
621-
622-
SettableApiFuture<List<PipelineResult>> futureResult = SettableApiFuture.create();
623-
624-
pipelineInternalStream( // Assuming you have this method
625-
firestoreImpl,
626-
request,
627-
new PipelineResultObserver() {
628-
final List<PipelineResult> results = new ArrayList<>();
629-
630-
@Override
631-
public void onCompleted() {
632-
futureResult.set(results);
633-
}
609+
SettableApiFuture<List<PipelineResult>> futureResult = SettableApiFuture.create();
634610

635-
@Override
636-
public void onNext(PipelineResult result) {
637-
results.add(result);
638-
}
611+
execute( // Assuming you have this method
612+
new PipelineResultObserver() {
613+
final List<PipelineResult> results = new ArrayList<>();
639614

640-
@Override
641-
public void onError(Throwable t) {
642-
futureResult.setException(t);
643-
}
644-
});
615+
@Override
616+
public void onCompleted() {
617+
futureResult.set(results);
618+
}
619+
620+
@Override
621+
public void onNext(PipelineResult result) {
622+
results.add(result);
623+
}
624+
625+
@Override
626+
public void onError(Throwable t) {
627+
futureResult.setException(t);
628+
}
629+
});
645630

646-
return futureResult;
647-
} else {
648-
// Handle unsupported Firestore types
649-
throw new IllegalArgumentException("Unsupported Firestore type");
650-
}
631+
return futureResult;
651632
}
652633

653634
/**
@@ -697,60 +678,46 @@ public void onError(Throwable t) {
697678
*/
698679
@BetaApi
699680
public void execute(ApiStreamObserver<PipelineResult> observer) {
700-
if (db instanceof FirestoreImpl) {
701-
FirestoreImpl firestoreImpl = (FirestoreImpl) db;
702-
Value pipelineValue = toProto();
703-
ExecutePipelineRequest request =
704-
ExecutePipelineRequest.newBuilder()
705-
.setDatabase(firestoreImpl.getResourcePath().getDatabaseName().toString())
706-
.setStructuredPipeline(
707-
StructuredPipeline.newBuilder()
708-
.setPipeline(pipelineValue.getPipelineValue())
709-
.build())
710-
.build();
711-
712-
pipelineInternalStream(
713-
firestoreImpl,
714-
request,
715-
new PipelineResultObserver() {
716-
@Override
717-
public void onCompleted() {
718-
observer.onCompleted();
719-
}
681+
ExecutePipelineRequest request =
682+
ExecutePipelineRequest.newBuilder()
683+
.setDatabase(rpcContext.getDatabaseName())
684+
.setStructuredPipeline(StructuredPipeline.newBuilder().setPipeline(toProto()).build())
685+
.build();
686+
687+
pipelineInternalStream(
688+
request,
689+
new PipelineResultObserver() {
690+
@Override
691+
public void onCompleted() {
692+
observer.onCompleted();
693+
}
720694

721-
@Override
722-
public void onNext(PipelineResult result) {
723-
observer.onNext(result);
724-
}
695+
@Override
696+
public void onNext(PipelineResult result) {
697+
observer.onNext(result);
698+
}
725699

726-
@Override
727-
public void onError(Throwable t) {
728-
observer.onError(t);
729-
}
730-
});
731-
} else {
732-
// Handle unsupported Firestore types
733-
throw new IllegalArgumentException("Unsupported Firestore type");
734-
}
700+
@Override
701+
public void onError(Throwable t) {
702+
observer.onError(t);
703+
}
704+
});
735705
}
736706

737-
private Value toProto() {
738-
return Value.newBuilder()
739-
.setPipelineValue(
740-
com.google.firestore.v1.Pipeline.newBuilder()
741-
.addAllStages(stages.transform(StageUtils::toStageProto)))
707+
private com.google.firestore.v1.Pipeline toProto() {
708+
return com.google.firestore.v1.Pipeline.newBuilder()
709+
.addAllStages(stages.transform(StageUtils::toStageProto))
742710
.build();
743711
}
744712

745713
private void pipelineInternalStream(
746-
FirestoreImpl rpcContext,
747-
ExecutePipelineRequest request,
748-
PipelineResultObserver resultObserver) {
714+
ExecutePipelineRequest request, PipelineResultObserver resultObserver) {
749715
ResponseObserver<ExecutePipelineResponse> observer =
750716
new ResponseObserver<ExecutePipelineResponse>() {
751717
Timestamp executionTime = null;
752718
boolean firstResponse = false;
753719
int numDocuments = 0;
720+
int totalNumDocuments = 0;
754721
boolean hasCompleted = false;
755722

756723
@Override
@@ -769,15 +736,16 @@ public void onResponse(ExecutePipelineResponse response) {
769736
}
770737
if (response.getResultsCount() > 0) {
771738
numDocuments += response.getResultsCount();
772-
if (numDocuments % 100 == 0) {
739+
totalNumDocuments += response.getResultsCount();
740+
if (numDocuments > 100) {
773741
Tracing.getTracer()
774742
.getCurrentSpan()
775-
.addAnnotation("Firestore.Query: Received 100 documents");
743+
.addAnnotation("Firestore.Query: Received " + numDocuments + " documents");
744+
numDocuments = 0;
776745
}
746+
Timestamp executionTime = Timestamp.fromProto(response.getExecutionTime());
777747
for (Document doc : response.getResultsList()) {
778-
resultObserver.onNext(
779-
PipelineResult.fromDocument(
780-
rpcContext, Timestamp.fromProto(response.getExecutionTime()), doc));
748+
resultObserver.onNext(PipelineResult.fromDocument(rpcContext, executionTime, doc));
781749
}
782750
}
783751

@@ -804,7 +772,7 @@ public void onComplete() {
804772
.addAnnotation(
805773
"Firestore.ExecutePipeline: Completed",
806774
ImmutableMap.of(
807-
"numDocuments", AttributeValue.longAttributeValue((long) numDocuments)));
775+
"numDocuments", AttributeValue.longAttributeValue((long) totalNumDocuments)));
808776
resultObserver.onCompleted(executionTime);
809777
}
810778
};

‎google-cloud-firestore/src/main/java/com/google/cloud/firestore/PipelineResult.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -79,14 +79,14 @@ public String getId() {
7979
}
8080

8181
static PipelineResult fromDocument(
82-
FirestoreRpcContext<?> rpcContext, Timestamp readTime, Document document) {
82+
FirestoreRpcContext<?> rpcContext, Timestamp executionTime, Document document) {
8383
return new PipelineResult(
8484
rpcContext,
8585
document.getName().isEmpty()
8686
? null
8787
: new DocumentReference(rpcContext, ResourcePath.create(document.getName())),
8888
document.getFieldsMap(),
89-
readTime,
89+
executionTime,
9090
Timestamp.fromProto(document.getUpdateTime()),
9191
Timestamp.fromProto(document.getCreateTime()));
9292
}
@@ -442,7 +442,7 @@ public int hashCode() {
442442
@Override
443443
public String toString() {
444444
return String.format(
445-
"%s{doc=%s, fields=%s, readTime=%s, updateTime=%s, createTime=%s}",
445+
"%s{doc=%s, fields=%s, executionTime=%s, updateTime=%s, createTime=%s}",
446446
getClass().getSimpleName(), docRef, fields, executionTime, updateTime, createTime);
447447
}
448448
}

‎google-cloud-firestore/src/main/java/com/google/cloud/firestore/PipelineSource.java

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -30,11 +30,11 @@
3030
*/
3131
@BetaApi
3232
public class PipelineSource {
33-
private final Firestore db;
33+
private final FirestoreRpcContext<?> rpcContext;
3434

3535
@InternalApi
36-
PipelineSource(Firestore db) {
37-
this.db = db;
36+
PipelineSource(FirestoreRpcContext<?> rpcContext) {
37+
this.rpcContext = rpcContext;
3838
}
3939

4040
/**
@@ -46,7 +46,7 @@ public class PipelineSource {
4646
@Nonnull
4747
@BetaApi
4848
public Pipeline collection(@Nonnull String path) {
49-
return new Pipeline(this.db, new Collection(path));
49+
return new Pipeline(this.rpcContext, new Collection(path));
5050
}
5151

5252
/**
@@ -66,7 +66,7 @@ public Pipeline collectionGroup(@Nonnull String collectionId) {
6666
!collectionId.contains("/"),
6767
"Invalid collectionId '%s'. Collection IDs must not contain '/'.",
6868
collectionId);
69-
return new Pipeline(this.db, new CollectionGroup(collectionId));
69+
return new Pipeline(this.rpcContext, new CollectionGroup(collectionId));
7070
}
7171

7272
/**
@@ -80,7 +80,7 @@ public Pipeline collectionGroup(@Nonnull String collectionId) {
8080
@Nonnull
8181
@BetaApi
8282
public Pipeline database() {
83-
return new Pipeline(this.db, new Database());
83+
return new Pipeline(this.rpcContext, new Database());
8484
}
8585

8686
/**
@@ -93,6 +93,6 @@ public Pipeline database() {
9393
@Nonnull
9494
@BetaApi
9595
public Pipeline documents(DocumentReference... docs) {
96-
return new Pipeline(this.db, Documents.of(docs));
96+
return new Pipeline(this.rpcContext, Documents.of(docs));
9797
}
9898
}

‎google-cloud-firestore/src/main/java/com/google/cloud/firestore/Query.java

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2128,9 +2128,8 @@ public Pipeline pipeline() {
21282128
// From
21292129
Pipeline ppl =
21302130
this.options.getAllDescendants()
2131-
? new PipelineSource(this.getFirestore())
2132-
.collectionGroup(this.options.getCollectionId())
2133-
: new PipelineSource(this.getFirestore())
2131+
? new PipelineSource(this.rpcContext).collectionGroup(this.options.getCollectionId())
2132+
: new PipelineSource(this.rpcContext)
21342133
.collection(
21352134
this.options.getParentPath().append(this.options.getCollectionId()).getPath());
21362135

0 commit comments

Comments
 (0)