From 495daaf00b0c9c1b7dc06ab3728342afa28ddc9d Mon Sep 17 00:00:00 2001
From: Laurent Nguyen <laurent.nguyen@microsoft.com>
Date: Thu, 10 Apr 2025 15:23:29 +0200
Subject: [PATCH] Trace query closer to iterator operation

---
 src/Common/IteratorUtilities.ts               |  5 +++
 src/Common/dataAccess/queryDocuments.ts       |  3 --
 .../Tabs/DocumentsTabV2/DocumentsTabV2.tsx    | 32 +++++++++++--------
 3 files changed, 23 insertions(+), 17 deletions(-)

diff --git a/src/Common/IteratorUtilities.ts b/src/Common/IteratorUtilities.ts
index 6283488b8..0cbe8ea63 100644
--- a/src/Common/IteratorUtilities.ts
+++ b/src/Common/IteratorUtilities.ts
@@ -1,5 +1,8 @@
 import { QueryOperationOptions } from "@azure/cosmos";
+import { Action } from "Shared/Telemetry/TelemetryConstants";
+import * as Constants from "../Common/Constants";
 import { QueryResults } from "../Contracts/ViewModels";
+import * as TelemetryProcessor from "../Shared/Telemetry/TelemetryProcessor";
 
 interface QueryResponse {
   // [Todo] remove any
@@ -21,7 +24,9 @@ export function nextPage(
   firstItemIndex: number,
   queryOperationOptions?: QueryOperationOptions,
 ): Promise<QueryResults> {
+  TelemetryProcessor.traceStart(Action.ExecuteQuery);
   return documentsIterator.fetchNext(queryOperationOptions).then((response) => {
+    TelemetryProcessor.traceSuccess(Action.ExecuteQuery, { dataExplorerArea: Constants.Areas.Tab });
     const documents = response.resources;
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     const headers = (response as any).headers || {}; // TODO this is a private key. Remove any
diff --git a/src/Common/dataAccess/queryDocuments.ts b/src/Common/dataAccess/queryDocuments.ts
index 7538d214e..223fe987d 100644
--- a/src/Common/dataAccess/queryDocuments.ts
+++ b/src/Common/dataAccess/queryDocuments.ts
@@ -1,8 +1,6 @@
 import { FeedOptions, ItemDefinition, QueryIterator, Resource } from "@azure/cosmos";
-import { Action, ActionModifiers } from "Shared/Telemetry/TelemetryConstants";
 import { isVectorSearchEnabled } from "Utils/CapabilityUtils";
 import { LocalStorageUtility, StorageKey } from "../../Shared/StorageUtility";
-import * as TelemetryProcessor from "../../Shared/Telemetry/TelemetryProcessor";
 import { Queries } from "../Constants";
 import { client } from "../CosmosClient";
 
@@ -13,7 +11,6 @@ export const queryDocuments = (
   options: FeedOptions,
 ): QueryIterator<ItemDefinition & Resource> => {
   options = getCommonQueryOptions(options);
-  TelemetryProcessor.trace(Action.ExecuteQuery, ActionModifiers.Submit);
   return client().database(databaseId).container(containerId).items.query(query, options);
 };
 
diff --git a/src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2.tsx b/src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2.tsx
index cef323b0e..1095cf29a 100644
--- a/src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2.tsx
+++ b/src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2.tsx
@@ -495,7 +495,11 @@ const getNewDocumentButtonState = (editorState: ViewModels.DocumentExplorerState
 const _loadNextPageInternal = (
   iterator: QueryIterator<ItemDefinition & Resource>,
 ): Promise<DataModels.DocumentId[]> => {
-  return iterator.fetchNext().then((response) => response.resources);
+  TelemetryProcessor.traceStart(Action.ExecuteQuery);
+  return iterator.fetchNext().then((response) => {
+    TelemetryProcessor.traceSuccess(Action.ExecuteQuery, { dataExplorerArea: Constants.Areas.Tab });
+    return response.resources;
+  });
 };
 
 // Export to expose to unit tests
@@ -1264,8 +1268,8 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
         ? `the selected ${selectedRows.size} items`
         : "the selected item"
       : isPlural
-      ? `the selected ${selectedRows.size} documents`
-      : "the selected document";
+        ? `the selected ${selectedRows.size} documents`
+        : "the selected document";
     const msg = `Are you sure you want to delete ${documentName}?`;
 
     useDialog
@@ -2232,17 +2236,17 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
             )}
             {(bulkDeleteProcess.failedIds.length > 0 ||
               (bulkDeleteProcess.throttledIds.length > 0 && bulkDeleteMode !== "inProgress")) && (
-              <MessageBar intent="error" style={{ marginBottom: tokens.spacingVerticalL }}>
-                <MessageBarBody>
-                  <MessageBarTitle>Error</MessageBarTitle>
-                  Failed to delete{" "}
-                  {bulkDeleteMode === "inProgress"
-                    ? bulkDeleteProcess.failedIds.length
-                    : bulkDeleteProcess.failedIds.length + bulkDeleteProcess.throttledIds.length}{" "}
-                  document(s).
-                </MessageBarBody>
-              </MessageBar>
-            )}
+                <MessageBar intent="error" style={{ marginBottom: tokens.spacingVerticalL }}>
+                  <MessageBarBody>
+                    <MessageBarTitle>Error</MessageBarTitle>
+                    Failed to delete{" "}
+                    {bulkDeleteMode === "inProgress"
+                      ? bulkDeleteProcess.failedIds.length
+                      : bulkDeleteProcess.failedIds.length + bulkDeleteProcess.throttledIds.length}{" "}
+                    document(s).
+                  </MessageBarBody>
+                </MessageBar>
+              )}
             {bulkDeleteProcess.hasBeenThrottled && (
               <MessageBar intent="warning">
                 <MessageBarBody>