Added metrics and tracing for ReadMany operations

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ClientMetricsTest.java

@@ -31,6 +31,7 @@
 import com.azure.cosmos.models.CosmosBulkExecutionOptions;
 import com.azure.cosmos.models.CosmosBulkOperations;
 import com.azure.cosmos.models.CosmosClientTelemetryConfig;
+import com.azure.cosmos.models.CosmosItemIdentity;
 import com.azure.cosmos.models.CosmosItemOperation;
 import com.azure.cosmos.models.CosmosItemRequestOptions;
 import com.azure.cosmos.models.CosmosItemResponse;
@@ -41,6 +42,8 @@
 import com.azure.cosmos.models.CosmosMicrometerMetricsOptions;
 import com.azure.cosmos.models.CosmosPatchItemRequestOptions;
 import com.azure.cosmos.models.CosmosQueryRequestOptions;
+import com.azure.cosmos.models.CosmosReadManyRequestOptions;
+import com.azure.cosmos.models.FeedResponse;
 import com.azure.cosmos.models.ModelBridgeInternal;
 import com.azure.cosmos.models.PartitionKey;
 import com.azure.cosmos.models.SqlQuerySpec;
@@ -60,6 +63,7 @@
 import java.net.URI;
 import java.time.Duration;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.EnumSet;
 import java.util.List;
@@ -408,6 +412,95 @@ public void readItem() throws Exception {
         }
     }
 
+    @Test(groups = { "fast" }, timeOut = TIMEOUT)
+    public void readManySingleItem() throws Exception {
+        this.beforeTest(CosmosMetricCategory.DEFAULT);
+        try {
+            InternalObjectNode properties = getDocumentDefinition(UUID.randomUUID().toString());
+            container.createItem(properties);
+
+            List<CosmosItemIdentity> tuplesToBeRead = new ArrayList<>();
+            tuplesToBeRead.add(new CosmosItemIdentity(
+                new PartitionKey(properties.get("mypk")),
+                properties.getId()
+            ));
+
+            FeedResponse<InternalObjectNode> readManyResponse = container.readMany(
+                tuplesToBeRead,
+                new CosmosReadManyRequestOptions(),
+                InternalObjectNode.class);
+            validateReadManyFeedResponse(Arrays.asList(properties), readManyResponse);
+
+            this.validateMetrics(
+                Tag.of(TagName.OperationStatusCode.toString(), "200"),
+                Tag.of(TagName.RequestStatusCode.toString(), "200/0"),
+                0,
+                500
+            );
+
+            this.validateMetrics(
+                Tag.of(
+                    TagName.Operation.toString(), "Document/Query/readMany"),
+                Tag.of(TagName.RequestOperationType.toString(), "Document/Read"),
+                0,
+                500
+            );
+
+            Tag queryPlanTag = Tag.of(TagName.RequestOperationType.toString(), "DocumentCollection_QueryPlan");
+            this.assertMetrics("cosmos.client.req.gw", false, queryPlanTag);
+            this.assertMetrics("cosmos.client.req.rntbd", false, queryPlanTag);
+        } finally {
+            this.afterTest();
+        }
+    }
+
+    @Test(groups = { "fast" }, timeOut = TIMEOUT)
+    public void readManyMultipleItems() throws Exception {
+        this.beforeTest(CosmosMetricCategory.DEFAULT);
+
+        List<InternalObjectNode> createdDocs = new ArrayList<>();
+        List<CosmosItemIdentity> tuplesToBeRead = new ArrayList<>();
+        try {
+            for (int i = 0; i < 20; i++) {
+                InternalObjectNode properties = getDocumentDefinition(UUID.randomUUID().toString());
+                container.createItem(properties);
+                createdDocs.add(properties);
+                tuplesToBeRead.add(new CosmosItemIdentity(
+                    new PartitionKey(properties.get("mypk")),
+                    properties.getId()
+                ));
+            }
+
+
+            FeedResponse<InternalObjectNode> readManyResponse = container.readMany(
+                tuplesToBeRead,
+                new CosmosReadManyRequestOptions(),
+                InternalObjectNode.class);
+            validateReadManyFeedResponse(createdDocs, readManyResponse);
+
+            this.validateMetrics(
+                Tag.of(TagName.OperationStatusCode.toString(), "200"),
+                Tag.of(TagName.RequestStatusCode.toString(), "200/0"),
+                0,
+                500
+            );
+
+            this.validateMetrics(
+                Tag.of(
+                    TagName.Operation.toString(), "Document/Query/readMany"),
+                Tag.of(TagName.RequestOperationType.toString(), "Document/Query"),
+                0,
+                500
+            );
+
+            Tag queryPlanTag = Tag.of(TagName.RequestOperationType.toString(), "DocumentCollection_QueryPlan");
+            this.assertMetrics("cosmos.client.req.gw", false, queryPlanTag);
+            this.assertMetrics("cosmos.client.req.rntbd", false, queryPlanTag);
+        } finally {
+            this.afterTest();
+        }
+    }
+
     private void runReadItemTestWithThresholds(
         CosmosDiagnosticsThresholds thresholds,
         boolean expectRequestMetrics
@@ -1351,6 +1444,19 @@ private void validateItemResponse(InternalObjectNode containerProperties,
             .isEqualTo(containerProperties.getId());
     }
 
+    private void validateReadManyFeedResponse(
+        List<InternalObjectNode> createdDocs,
+        FeedResponse<InternalObjectNode> readManyResponse) {
+        // Basic validation
+        assertThat(readManyResponse).isNotNull();
+        assertThat(readManyResponse.getResults()).isNotNull();
+        List<InternalObjectNode> docsFromResponse = readManyResponse.getResults();
+        assertThat(docsFromResponse).hasSize(createdDocs.size());
+        for (InternalObjectNode doc: createdDocs) {
+            assertThat(docsFromResponse.stream().anyMatch(r -> r.getId() != null && r.getId().equals(doc.getId())));
+        }
+    }
+
     private void validateItemCountMetrics(Tag expectedOperationTag) {
         if (this.getEffectiveMetricCategories().contains(MetricCategory.OperationDetails)) {
             this.assertMetrics("cosmos.client.op.maxItemCount", true, expectedOperationTag);

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosTracerTest.java

@@ -15,6 +15,7 @@
 import com.azure.cosmos.implementation.SerializationDiagnosticsContext;
 import com.azure.cosmos.implementation.TestConfigurations;
 import com.azure.cosmos.implementation.Utils;
+import com.azure.cosmos.models.CosmosItemIdentity;
 import com.azure.cosmos.models.ShowQueryMode;
 import com.azure.cosmos.implementation.directconnectivity.ReflectionUtils;
 import com.azure.cosmos.implementation.directconnectivity.StoreResponseDiagnostics;
@@ -74,6 +75,7 @@
 import java.time.Instant;
 import java.time.OffsetDateTime;
 import java.time.ZoneOffset;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.HashMap;
@@ -616,10 +618,13 @@ public void cosmosAsyncContainer(
             samplingRate);
         mockTracer.reset();
 
+        List<CosmosItemIdentity> createdDocs = new ArrayList<>();
         for (int i = 0; i < 30; i++) {
             // inserting high enough number of documents to make sure we have at least 1 doc on each
             // of the two partitions
-            item = getDocumentDefinition(ITEM_ID + "_" + i);
+            String id = ITEM_ID + "_" + i;
+            createdDocs.add(new CosmosItemIdentity(new PartitionKey(id), id));
+            item = getDocumentDefinition(id);
             requestOptions = new CosmosItemRequestOptions();
             cosmosItemResponse = cosmosAsyncContainer
                 .createItem(item, requestOptions)
@@ -883,7 +888,27 @@ public void cosmosAsyncContainer(
                 forceThresholdViolations,
                 null,
                 samplingRate);
-        mockTracer.reset();   
+        mockTracer.reset();
+
+        // read many single item
+        feedItemResponse = cosmosAsyncContainer.readMany(Arrays.asList(createdDocs.get(0)), ObjectNode.class)
+                                               .block();
+        assertThat(feedItemResponse).isNotNull();
+        verifyTracerAttributes(
+            mockTracer,
+            "readManyItems." + cosmosAsyncContainer.getId(),
+            cosmosAsyncDatabase.getId(),
+            cosmosAsyncContainer.getId(),
+            feedItemResponse.getCosmosDiagnostics(),
+            null,
+            useLegacyTracing,
+            enableRequestLevelTracing,
+            showQueryMode,
+            query,
+            forceThresholdViolations,
+            "readMany",
+            samplingRate);
+        mockTracer.reset();
     }
 
     @Test(groups = { "fast", "simple" }, dataProvider = "traceTestCaseProvider", timeOut = TIMEOUT)
@@ -1445,10 +1470,11 @@ private void verifyOTelTracerAttributes(
 
             String dbStatement = (String) attributes.get("db.statement");
 
-            if (ShowQueryMode.ALL.equals(showQueryMode)) {
+            boolean isReadMany = "readMany".equals(cosmosDiagnostics.getDiagnosticsContext().getOperationId());
+            if (!isReadMany && ShowQueryMode.ALL.equals(showQueryMode)) {
                 assertThat(attributes.get("db.statement")).isNotNull();
                 assertThat(attributes.get("db.statement")).isEqualTo(queryStatement);
-            } else if (ShowQueryMode.PARAMETERIZED_ONLY.equals(showQueryMode)
+            } else if (!isReadMany && ShowQueryMode.PARAMETERIZED_ONLY.equals(showQueryMode)
                     && null != dbStatement && dbStatement.contains("@")) {
                 assertThat(attributes.get("db.statement")).isNotNull();
                 assertThat(attributes.get("db.statement")).isEqualTo(queryStatement);

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/NonStreamingOrderByQueryVectorSearchTest.java

@@ -123,14 +123,14 @@ public void flatIndexVectorSearch() {
             .collectList().block();
         validateOrdering(6, resultDocs, false);
 
-        String euclideanSpecsQuery = String.format("SELECT DISTINCT TOP 6 c.text, VectorDistance(c.embedding, [%s], {'distanceFunction': 'euclidean'}) AS " +
+        String euclideanSpecsQuery = String.format("SELECT DISTINCT TOP 6 c.text, VectorDistance(c.embedding, [%s], false, {'distanceFunction': 'euclidean'}) AS " +
             "score FROM c ORDER BY VectorDistance(c.embedding, [%s], false, {'distanceFunction': 'euclidean'})", queryVector, queryVector);
         resultDocs = flatIndexContainer.queryItems(euclideanSpecsQuery, new CosmosQueryRequestOptions(), Document.class).byPage()
             .flatMap(feedResponse -> Flux.fromIterable(feedResponse.getResults()))
             .collectList().block();
-        validateOrdering(6, resultDocs, false);
+        validateOrdering(6, resultDocs, true);
 
-        String dotproductSpecsQuery = String.format("SELECT DISTINCT TOP 6 c.text, VectorDistance(c.embedding, [%s], {'distanceFunction': 'dotproduct'}) AS " +
+        String dotproductSpecsQuery = String.format("SELECT DISTINCT TOP 6 c.text, VectorDistance(c.embedding, [%s], false, {'distanceFunction': 'dotproduct'}) AS " +
             "score FROM c ORDER BY VectorDistance(c.embedding, [%s], false, {'distanceFunction': 'dotproduct'})", queryVector, queryVector);
         resultDocs = flatIndexContainer.queryItems(dotproductSpecsQuery, new CosmosQueryRequestOptions(), Document.class).byPage()
             .flatMap(feedResponse -> Flux.fromIterable(feedResponse.getResults()))
@@ -156,7 +156,7 @@ public void quantizedIndexVectorSearch() {
             .collectList().block();
         validateOrdering(6, resultDocs, true);
 
-        String dotproduct_specs_query = String.format("SELECT DISTINCT TOP 6 c.text, VectorDistance(c.embedding, [%s], {'distanceFunction': 'dotproduct'}) AS " +
+        String dotproduct_specs_query = String.format("SELECT DISTINCT TOP 6 c.text, VectorDistance(c.embedding, [%s], false, {'distanceFunction': 'dotproduct'}) AS " +
             "score FROM c ORDER BY VectorDistance(c.embedding, [%s], false, {'distanceFunction': 'dotproduct'})", queryVector, queryVector);
         resultDocs = quantizedIndexContainer.queryItems(dotproduct_specs_query, new CosmosQueryRequestOptions(), Document.class).byPage()
             .flatMap(feedResponse -> Flux.fromIterable(feedResponse.getResults()))

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -9,6 +9,7 @@
 #### Bugs Fixed
 
 #### Other Changes
+* Added metrics and tracing for ReadMany operations. - See [PR 41042](https://github.com/Azure/azure-sdk-for-java/pull/41042)
 
 ### 4.62.0 (2024-07-02)
 

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java

@@ -1526,6 +1526,15 @@ public <T> Mono<FeedResponse<T>> readMany(
         CosmosReadManyRequestOptions requestOptions,
         Class<T> classType) {
 
+        return withContext(context -> this.readManyInternal(itemIdentityList, requestOptions, classType, context));
+    }
+
+    private <T> Mono<FeedResponse<T>> readManyInternal(
+        List<CosmosItemIdentity> itemIdentityList,
+        CosmosReadManyRequestOptions requestOptions,
+        Class<T> classType,
+        Context context) {
+
         CosmosQueryRequestOptions queryRequestOptions = requestOptions == null
             ? new CosmosQueryRequestOptions()
             : queryOptionsAccessor.clone(readManyOptionsAccessor.getImpl(requestOptions));
@@ -1539,7 +1548,7 @@ public <T> Mono<FeedResponse<T>> readMany(
         fluxOptions.setMaxItemCount(itemIdentityList != null ? itemIdentityList.size() : 0);
         QueryFeedOperationState state = new QueryFeedOperationState(
             client,
-            this.readAllItemsSpanName,
+            this.readManyItemsSpanName,
             database.getId(),
             this.getId(),
             ResourceType.Document,
@@ -1549,9 +1558,21 @@ public <T> Mono<FeedResponse<T>> readMany(
             fluxOptions
         );
 
-        return CosmosBridgeInternal
+        Mono<FeedResponse<T>> responseMono = CosmosBridgeInternal
             .getAsyncDocumentClient(this.getDatabase())
             .readMany(itemIdentityList, BridgeInternal.getLink(this), state, classType);
+
+        RequestOptions options = queryOptionsAccessor.toRequestOptions(queryRequestOptions);
+
+        return client
+            .getDiagnosticsProvider()
+            .traceEnabledReadManyResponsePublisher(
+                itemIdentityList,
+                state,
+                responseMono,
+                client,
+                options,
+                context);
     }
 
     /**