From 2982f057a59f52383f3b6a55b957e79249d4575b Mon Sep 17 00:00:00 2001 From: Ehsan Nasiri Date: Mon, 26 Feb 2024 17:50:28 -0800 Subject: [PATCH] feat: trace instrumentation for DocumentReference methods. --- .../cloud/firestore/DocumentReference.java | 269 ++++++++++++++---- .../google/cloud/firestore/UpdateBuilder.java | 52 ++-- .../cloud/firestore/it/ITTracingTest.java | 206 +++++++++++++- 3 files changed, 446 insertions(+), 81 deletions(-) diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/DocumentReference.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/DocumentReference.java index 4918b3974..57254bb2b 100644 --- a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/DocumentReference.java +++ b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/DocumentReference.java @@ -21,6 +21,8 @@ import com.google.api.core.InternalExtensionOnly; import com.google.api.gax.rpc.ApiException; import com.google.api.gax.rpc.ApiExceptions; +import com.google.cloud.firestore.telemetry.TraceUtil; +import com.google.cloud.firestore.telemetry.TraceUtil.Scope; import com.google.cloud.firestore.v1.FirestoreClient.ListCollectionIdsPagedResponse; import com.google.common.util.concurrent.MoreExecutors; import com.google.firestore.v1.ListCollectionIdsRequest; @@ -131,6 +133,12 @@ private ApiFuture extractFirst(ApiFuture> results) { MoreExecutors.directExecutor()); } + /** Gets the TraceUtil object associated with this DocumentReference's Firestore instance. */ + @Nonnull + private TraceUtil getTraceUtil() { + return getFirestore().getOptions().getTraceUtil(); + } + /** * Creates a new Document at the DocumentReference's Location. It fails the write if the document * exists. @@ -140,8 +148,16 @@ private ApiFuture extractFirst(ApiFuture> results) { */ @Nonnull public ApiFuture create(@Nonnull Map fields) { - WriteBatch writeBatch = rpcContext.getFirestore().batch(); - return extractFirst(writeBatch.create(this, fields).commit()); + TraceUtil.Span span = getTraceUtil().startSpan(TraceUtil.SPAN_NAME_DOC_REF_CREATE); + try (Scope ignored = span.makeCurrent()) { + WriteBatch writeBatch = rpcContext.getFirestore().batch(); + ApiFuture result = extractFirst(writeBatch.create(this, fields).commit()); + span.endAtFuture(result); + return result; + } catch (Exception error) { + span.end(error); + throw error; + } } /** @@ -153,8 +169,16 @@ public ApiFuture create(@Nonnull Map fields) { */ @Nonnull public ApiFuture create(@Nonnull Object pojo) { - WriteBatch writeBatch = rpcContext.getFirestore().batch(); - return extractFirst(writeBatch.create(this, pojo).commit()); + TraceUtil.Span span = getTraceUtil().startSpan(TraceUtil.SPAN_NAME_DOC_REF_CREATE); + try (Scope ignored = span.makeCurrent()) { + WriteBatch writeBatch = rpcContext.getFirestore().batch(); + ApiFuture result = extractFirst(writeBatch.create(this, pojo).commit()); + span.endAtFuture(result); + return result; + } catch (Exception error) { + span.end(error); + throw error; + } } /** @@ -166,8 +190,16 @@ public ApiFuture create(@Nonnull Object pojo) { */ @Nonnull public ApiFuture set(@Nonnull Map fields) { - WriteBatch writeBatch = rpcContext.getFirestore().batch(); - return extractFirst(writeBatch.set(this, fields).commit()); + TraceUtil.Span span = getTraceUtil().startSpan(TraceUtil.SPAN_NAME_DOC_REF_SET); + try (Scope ignored = span.makeCurrent()) { + WriteBatch writeBatch = rpcContext.getFirestore().batch(); + ApiFuture result = extractFirst(writeBatch.set(this, fields).commit()); + span.endAtFuture(result); + return result; + } catch (Exception error) { + span.end(error); + throw error; + } } /** @@ -182,8 +214,16 @@ public ApiFuture set(@Nonnull Map fields) { @Nonnull public ApiFuture set( @Nonnull Map fields, @Nonnull SetOptions options) { - WriteBatch writeBatch = rpcContext.getFirestore().batch(); - return extractFirst(writeBatch.set(this, fields, options).commit()); + TraceUtil.Span span = getTraceUtil().startSpan(TraceUtil.SPAN_NAME_DOC_REF_SET); + try (Scope ignored = span.makeCurrent()) { + WriteBatch writeBatch = rpcContext.getFirestore().batch(); + ApiFuture result = extractFirst(writeBatch.set(this, fields, options).commit()); + span.endAtFuture(result); + return result; + } catch (Exception error) { + span.end(error); + throw error; + } } /** @@ -195,8 +235,16 @@ public ApiFuture set( */ @Nonnull public ApiFuture set(@Nonnull Object pojo) { - WriteBatch writeBatch = rpcContext.getFirestore().batch(); - return extractFirst(writeBatch.set(this, pojo).commit()); + TraceUtil.Span span = getTraceUtil().startSpan(TraceUtil.SPAN_NAME_DOC_REF_SET); + try (Scope ignored = span.makeCurrent()) { + WriteBatch writeBatch = rpcContext.getFirestore().batch(); + ApiFuture result = extractFirst(writeBatch.set(this, pojo).commit()); + span.endAtFuture(result); + return result; + } catch (Exception error) { + span.end(error); + throw error; + } } /** @@ -210,8 +258,16 @@ public ApiFuture set(@Nonnull Object pojo) { */ @Nonnull public ApiFuture set(@Nonnull Object pojo, @Nonnull SetOptions options) { - WriteBatch writeBatch = rpcContext.getFirestore().batch(); - return extractFirst(writeBatch.set(this, pojo, options).commit()); + TraceUtil.Span span = getTraceUtil().startSpan(TraceUtil.SPAN_NAME_DOC_REF_SET); + try (Scope ignored = span.makeCurrent()) { + WriteBatch writeBatch = rpcContext.getFirestore().batch(); + ApiFuture result = extractFirst(writeBatch.set(this, pojo, options).commit()); + span.endAtFuture(result); + return result; + } catch (Exception error) { + span.end(error); + throw error; + } } /** @@ -223,8 +279,16 @@ public ApiFuture set(@Nonnull Object pojo, @Nonnull SetOptions opti */ @Nonnull public ApiFuture update(@Nonnull Map fields) { - WriteBatch writeBatch = rpcContext.getFirestore().batch(); - return extractFirst(writeBatch.update(this, fields).commit()); + TraceUtil.Span span = getTraceUtil().startSpan(TraceUtil.SPAN_NAME_DOC_REF_UPDATE); + try (Scope ignored = span.makeCurrent()) { + WriteBatch writeBatch = rpcContext.getFirestore().batch(); + ApiFuture result = extractFirst(writeBatch.update(this, fields).commit()); + span.endAtFuture(result); + return result; + } catch (Exception error) { + span.end(error); + throw error; + } } /** @@ -237,8 +301,17 @@ public ApiFuture update(@Nonnull Map fields) { */ @Nonnull public ApiFuture update(@Nonnull Map fields, Precondition options) { - WriteBatch writeBatch = rpcContext.getFirestore().batch(); - return extractFirst(writeBatch.update(this, fields, options).commit()); + TraceUtil.Span span = getTraceUtil().startSpan(TraceUtil.SPAN_NAME_DOC_REF_UPDATE); + try (Scope ignored = span.makeCurrent()) { + WriteBatch writeBatch = rpcContext.getFirestore().batch(); + ApiFuture result = + extractFirst(writeBatch.update(this, fields, options).commit()); + span.endAtFuture(result); + return result; + } catch (Exception error) { + span.end(error); + throw error; + } } /** @@ -253,8 +326,17 @@ public ApiFuture update(@Nonnull Map fields, Precon @Nonnull public ApiFuture update( @Nonnull String field, @Nullable Object value, Object... moreFieldsAndValues) { - WriteBatch writeBatch = rpcContext.getFirestore().batch(); - return extractFirst(writeBatch.update(this, field, value, moreFieldsAndValues).commit()); + TraceUtil.Span span = getTraceUtil().startSpan(TraceUtil.SPAN_NAME_DOC_REF_UPDATE); + try (Scope ignored = span.makeCurrent()) { + WriteBatch writeBatch = rpcContext.getFirestore().batch(); + ApiFuture result = + extractFirst(writeBatch.update(this, field, value, moreFieldsAndValues).commit()); + span.endAtFuture(result); + return result; + } catch (Exception error) { + span.end(error); + throw error; + } } /** @@ -269,8 +351,17 @@ public ApiFuture update( @Nonnull public ApiFuture update( @Nonnull FieldPath fieldPath, @Nullable Object value, Object... moreFieldsAndValues) { - WriteBatch writeBatch = rpcContext.getFirestore().batch(); - return extractFirst(writeBatch.update(this, fieldPath, value, moreFieldsAndValues).commit()); + TraceUtil.Span span = getTraceUtil().startSpan(TraceUtil.SPAN_NAME_DOC_REF_UPDATE); + try (Scope ignored = span.makeCurrent()) { + WriteBatch writeBatch = rpcContext.getFirestore().batch(); + ApiFuture result = + extractFirst(writeBatch.update(this, fieldPath, value, moreFieldsAndValues).commit()); + span.endAtFuture(result); + return result; + } catch (Exception error) { + span.end(error); + throw error; + } } /** @@ -289,9 +380,18 @@ public ApiFuture update( @Nonnull String field, @Nullable Object value, Object... moreFieldsAndValues) { - WriteBatch writeBatch = rpcContext.getFirestore().batch(); - return extractFirst( - writeBatch.update(this, options, field, value, moreFieldsAndValues).commit()); + TraceUtil.Span span = getTraceUtil().startSpan(TraceUtil.SPAN_NAME_DOC_REF_UPDATE); + try (Scope ignored = span.makeCurrent()) { + WriteBatch writeBatch = rpcContext.getFirestore().batch(); + ApiFuture result = + extractFirst( + writeBatch.update(this, options, field, value, moreFieldsAndValues).commit()); + span.endAtFuture(result); + return result; + } catch (Exception error) { + span.end(error); + throw error; + } } /** @@ -310,9 +410,18 @@ public ApiFuture update( @Nonnull FieldPath fieldPath, @Nullable Object value, Object... moreFieldsAndValues) { - WriteBatch writeBatch = rpcContext.getFirestore().batch(); - return extractFirst( - writeBatch.update(this, options, fieldPath, value, moreFieldsAndValues).commit()); + TraceUtil.Span span = getTraceUtil().startSpan(TraceUtil.SPAN_NAME_DOC_REF_UPDATE); + try (Scope ignored = span.makeCurrent()) { + WriteBatch writeBatch = rpcContext.getFirestore().batch(); + ApiFuture result = + extractFirst( + writeBatch.update(this, options, fieldPath, value, moreFieldsAndValues).commit()); + span.endAtFuture(result); + return result; + } catch (Exception error) { + span.end(error); + throw error; + } } /** @@ -323,8 +432,16 @@ public ApiFuture update( */ @Nonnull public ApiFuture delete(@Nonnull Precondition options) { - WriteBatch writeBatch = rpcContext.getFirestore().batch(); - return extractFirst(writeBatch.delete(this, options).commit()); + TraceUtil.Span span = getTraceUtil().startSpan(TraceUtil.SPAN_NAME_DOC_REF_DELETE); + try (Scope ignored = span.makeCurrent()) { + WriteBatch writeBatch = rpcContext.getFirestore().batch(); + ApiFuture result = extractFirst(writeBatch.delete(this, options).commit()); + span.endAtFuture(result); + return result; + } catch (Exception error) { + span.end(error); + throw error; + } } /** @@ -334,20 +451,36 @@ public ApiFuture delete(@Nonnull Precondition options) { */ @Nonnull public ApiFuture delete() { - WriteBatch writeBatch = rpcContext.getFirestore().batch(); - return extractFirst(writeBatch.delete(this).commit()); + TraceUtil.Span span = getTraceUtil().startSpan(TraceUtil.SPAN_NAME_DOC_REF_DELETE); + try (Scope ignored = span.makeCurrent()) { + WriteBatch writeBatch = rpcContext.getFirestore().batch(); + ApiFuture result = extractFirst(writeBatch.delete(this).commit()); + span.endAtFuture(result); + return result; + } catch (Exception error) { + span.end(error); + throw error; + } } /** * Reads the document referenced by this DocumentReference. If the document doesn't exist, the - * get() will return an an empty DocumentSnapshot. + * get() will return an empty DocumentSnapshot. * * @return An ApiFuture that will be resolved with the contents of the Document at this * DocumentReference, or a failure if the document does not exist. */ @Nonnull public ApiFuture get() { - return extractFirst(rpcContext.getFirestore().getAll(this)); + TraceUtil.Span span = getTraceUtil().startSpan(TraceUtil.SPAN_NAME_DOC_REF_GET); + try (Scope ignored = span.makeCurrent()) { + ApiFuture result = extractFirst(rpcContext.getFirestore().getAll(this)); + span.endAtFuture(result); + return result; + } catch (Exception error) { + span.end(error); + throw error; + } } /** @@ -360,8 +493,16 @@ public ApiFuture get() { */ @Nonnull public ApiFuture get(FieldMask fieldMask) { - return extractFirst( - rpcContext.getFirestore().getAll(new DocumentReference[] {this}, fieldMask)); + TraceUtil.Span span = getTraceUtil().startSpan(TraceUtil.SPAN_NAME_DOC_REF_GET); + try (Scope ignored = span.makeCurrent()) { + ApiFuture result = + extractFirst(rpcContext.getFirestore().getAll(new DocumentReference[] {this}, fieldMask)); + span.endAtFuture(result); + return result; + } catch (Exception error) { + span.end(error); + throw error; + } } /** @@ -372,41 +513,45 @@ public ApiFuture get(FieldMask fieldMask) { */ @Nonnull public Iterable listCollections() { - ListCollectionIdsRequest.Builder request = ListCollectionIdsRequest.newBuilder(); - request.setParent(path.toString()); - final ListCollectionIdsPagedResponse response; - try { + TraceUtil.Span span = getTraceUtil().startSpan(TraceUtil.SPAN_NAME_DOC_REF_LIST_COLLECTIONS); + try (Scope ignored = span.makeCurrent()) { + ListCollectionIdsRequest.Builder request = ListCollectionIdsRequest.newBuilder(); + request.setParent(path.toString()); + final ListCollectionIdsPagedResponse response; response = ApiExceptions.callAndTranslateApiException( rpcContext.sendRequest( request.build(), rpcContext.getClient().listCollectionIdsPagedCallable())); + Iterable result = + new Iterable() { + @Override + @Nonnull + public Iterator iterator() { + final Iterator iterator = response.iterateAll().iterator(); + return new Iterator() { + @Override + public boolean hasNext() { + return iterator.hasNext(); + } + + @Override + public CollectionReference next() { + return DocumentReference.this.collection(iterator.next()); + } + + @Override + public void remove() { + throw new UnsupportedOperationException("remove"); + } + }; + } + }; + span.end(); + return result; } catch (ApiException exception) { + span.end(exception); throw FirestoreException.forApiException(exception); } - - return new Iterable() { - @Override - @Nonnull - public Iterator iterator() { - final Iterator iterator = response.iterateAll().iterator(); - return new Iterator() { - @Override - public boolean hasNext() { - return iterator.hasNext(); - } - - @Override - public CollectionReference next() { - return DocumentReference.this.collection(iterator.next()); - } - - @Override - public void remove() { - throw new UnsupportedOperationException("remove"); - } - }; - } - }; } /** diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/UpdateBuilder.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/UpdateBuilder.java index 01fdcb68e..e5fffb35c 100644 --- a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/UpdateBuilder.java +++ b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/UpdateBuilder.java @@ -20,6 +20,8 @@ import com.google.api.core.ApiFutures; import com.google.api.core.InternalExtensionOnly; import com.google.cloud.firestore.UserDataConverter.EncodingOptions; +import com.google.cloud.firestore.telemetry.TraceUtil; +import com.google.cloud.firestore.telemetry.TraceUtil.Scope; import com.google.common.base.Preconditions; import com.google.common.util.concurrent.MoreExecutors; import com.google.firestore.v1.CommitRequest; @@ -609,24 +611,38 @@ ApiFuture> commit(@Nullable ByteString transactionId) { committed = true; - ApiFuture response = - firestore.sendRequest(request.build(), firestore.getClient().commitCallable()); - - return ApiFutures.transform( - response, - commitResponse -> { - List writeResults = - commitResponse.getWriteResultsList(); - - List result = new ArrayList<>(); - - for (com.google.firestore.v1.WriteResult writeResult : writeResults) { - result.add(WriteResult.fromProto(writeResult, commitResponse.getCommitTime())); - } - - return result; - }, - MoreExecutors.directExecutor()); + TraceUtil.Span span = + firestore + .getOptions() + .getTraceUtil() + .startSpan( + transactionId == null + ? TraceUtil.SPAN_NAME_BATCH_COMMIT + : TraceUtil.SPAN_NAME_TRANSACTION_COMMIT); + span.setAttribute("numDocuments", writes.size()); + try (Scope ignored = span.makeCurrent()) { + ApiFuture response = + firestore.sendRequest(request.build(), firestore.getClient().commitCallable()); + + ApiFuture> returnValue = + ApiFutures.transform( + response, + commitResponse -> { + List writeResults = + commitResponse.getWriteResultsList(); + List result = new ArrayList<>(); + for (com.google.firestore.v1.WriteResult writeResult : writeResults) { + result.add(WriteResult.fromProto(writeResult, commitResponse.getCommitTime())); + } + return result; + }, + MoreExecutors.directExecutor()); + span.endAtFuture(returnValue); + return returnValue; + } catch (Exception error) { + span.end(error); + throw error; + } } /** Checks whether any updates have been queued. */ diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITTracingTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITTracingTest.java index 104b6c374..3d6a8c225 100644 --- a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITTracingTest.java +++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITTracingTest.java @@ -1,5 +1,5 @@ /* - * Copyright 2023 Google LLC + * Copyright 2024 Google LLC * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -59,7 +59,10 @@ public class ITTracingTest { Logger.getLogger(com.google.cloud.firestore.it.ITBaseTest.class.getName()); private static final String SERVICE = "google.firestore.v1.Firestore/"; + private static final String BATCH_GET_DOCUMENTS_RPC_NAME = "BatchGetDocuments"; + private static final String COMMIT_RPC_NAME = "Commit"; private static final String LIST_DOCUMENTS_RPC_NAME = "ListDocuments"; + private static final String LIST_COLLECTIONS_RPC_NAME = "ListCollectionIds"; private static final String BATCH_WRITE_RPC_NAME = "BatchWrite"; // We use an InMemorySpanExporter for testing which keeps all generated trace spans @@ -244,6 +247,15 @@ void assertHasExpectedAttributes(SpanData spanData, String... additionalExpected } } + // This is a POJO used for testing APIs that take a POJO. + static class Pojo { + public int bar; + + Pojo(int bar) { + this.bar = bar; + } + } + @Test public void aggregateQueryGet() throws Exception { firestore.collection("col").count().get().get(); @@ -301,4 +313,196 @@ public void collectionListDocuments() throws Exception { assertEquals(2, spans.size()); assertSpanHierarchy(SPAN_NAME_COL_REF_LIST_DOCUMENTS, grpcSpanName(LIST_DOCUMENTS_RPC_NAME)); } + + @Test + public void docRefCreate() throws Exception { + firestore.collection("col").document().create(Collections.singletonMap("foo", "bar")).get(); + + List spans = prepareSpans(); + assertEquals(3, spans.size()); + assertSpanHierarchy( + SPAN_NAME_DOC_REF_CREATE, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME)); + } + + @Test + public void docRefCreate2() throws Exception { + firestore.collection("col").document().create(new Pojo(1)).get(); + + List spans = prepareSpans(); + assertEquals(3, spans.size()); + assertSpanHierarchy( + SPAN_NAME_DOC_REF_CREATE, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME)); + } + + @Test + public void docRefSet() throws Exception { + firestore.collection("col").document("foo").set(Collections.singletonMap("foo", "bar")).get(); + + List spans = prepareSpans(); + assertEquals(3, spans.size()); + assertSpanHierarchy( + SPAN_NAME_DOC_REF_SET, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME)); + } + + @Test + public void docRefSet2() throws Exception { + firestore + .collection("col") + .document("foo") + .set(Collections.singletonMap("foo", "bar"), SetOptions.merge()) + .get(); + + List spans = prepareSpans(); + assertEquals(3, spans.size()); + assertSpanHierarchy( + SPAN_NAME_DOC_REF_SET, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME)); + } + + @Test + public void docRefSet3() throws Exception { + firestore.collection("col").document("foo").set(new Pojo(1)).get(); + + List spans = prepareSpans(); + assertEquals(3, spans.size()); + assertSpanHierarchy( + SPAN_NAME_DOC_REF_SET, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME)); + } + + @Test + public void docRefSet4() throws Exception { + firestore.collection("col").document("foo").set(new Pojo(1), SetOptions.merge()).get(); + + List spans = prepareSpans(); + assertEquals(3, spans.size()); + assertSpanHierarchy( + SPAN_NAME_DOC_REF_SET, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME)); + } + + @Test + public void docRefUpdate() throws Exception { + firestore + .collection("col") + .document("foo") + .update(Collections.singletonMap("foo", "bar")) + .get(); + + List spans = prepareSpans(); + assertEquals(3, spans.size()); + assertSpanHierarchy( + SPAN_NAME_DOC_REF_UPDATE, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME)); + } + + @Test + public void docRefUpdate2() throws Exception { + firestore + .collection("col") + .document("foo") + .update(Collections.singletonMap("foo", "bar"), Precondition.NONE) + .get(); + + List spans = prepareSpans(); + assertEquals(3, spans.size()); + assertSpanHierarchy( + SPAN_NAME_DOC_REF_UPDATE, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME)); + } + + @Test + public void docRefUpdate3() throws Exception { + firestore.collection("col").document("foo").update("key", "value", "key2", "value2").get(); + + List spans = prepareSpans(); + assertEquals(3, spans.size()); + assertSpanHierarchy( + SPAN_NAME_DOC_REF_UPDATE, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME)); + } + + @Test + public void docRefUpdate4() throws Exception { + firestore + .collection("col") + .document("foo") + .update(FieldPath.of("key"), "value", FieldPath.of("key2"), "value2") + .get(); + + List spans = prepareSpans(); + assertEquals(3, spans.size()); + assertSpanHierarchy( + SPAN_NAME_DOC_REF_UPDATE, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME)); + } + + @Test + public void docRefUpdate5() throws Exception { + firestore + .collection("col") + .document("foo") + .update(Precondition.NONE, "key", "value", "key2", "value2") + .get(); + + List spans = prepareSpans(); + assertEquals(3, spans.size()); + assertSpanHierarchy( + SPAN_NAME_DOC_REF_UPDATE, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME)); + } + + @Test + public void docRefUpdate6() throws Exception { + firestore + .collection("col") + .document("foo") + .update(Precondition.NONE, FieldPath.of("key"), "value", FieldPath.of("key2"), "value2") + .get(); + + List spans = prepareSpans(); + assertEquals(3, spans.size()); + assertSpanHierarchy( + SPAN_NAME_DOC_REF_UPDATE, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME)); + } + + @Test + public void docRefDelete() throws Exception { + firestore.collection("col").document("doc0").delete().get(); + + List spans = prepareSpans(); + assertEquals(3, spans.size()); + assertSpanHierarchy( + SPAN_NAME_DOC_REF_DELETE, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME)); + } + + @Test + public void docRefDelete2() throws Exception { + firestore.collection("col").document("doc0").delete(Precondition.NONE).get(); + + List spans = prepareSpans(); + assertEquals(3, spans.size()); + assertSpanHierarchy( + SPAN_NAME_DOC_REF_DELETE, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME)); + } + + @Test + public void docRefGet() throws Exception { + firestore.collection("col").document("doc0").get().get(); + + List spans = prepareSpans(); + assertEquals(2, spans.size()); + assertSpanHierarchy(SPAN_NAME_DOC_REF_GET, grpcSpanName(BATCH_GET_DOCUMENTS_RPC_NAME)); + } + + @Test + public void docRefGet2() throws Exception { + firestore.collection("col").document("doc0").get(FieldMask.of("foo")).get(); + + List spans = prepareSpans(); + assertEquals(2, spans.size()); + assertSpanHierarchy(SPAN_NAME_DOC_REF_GET, grpcSpanName(BATCH_GET_DOCUMENTS_RPC_NAME)); + } + + @Test + public void docListCollections() throws Exception { + firestore.collection("col").document("doc0").listCollections(); + + List spans = prepareSpans(); + assertEquals(2, spans.size()); + assertSpanHierarchy( + SPAN_NAME_DOC_REF_LIST_COLLECTIONS, grpcSpanName(LIST_COLLECTIONS_RPC_NAME)); + } }