From 11959bb7b0d05d78dd97b61d80e6abed8a9ae612 Mon Sep 17 00:00:00 2001
From: Laurent Nguyen <laurent.nguyen@microsoft.com>
Date: Wed, 16 Sep 2026 14:30:19 +0200
Subject: [PATCH] Enhance logging and error handling in sdk read collections
 api (#2602)

* Enhance logging and error handling in readCollections function

* Add defensive measure to avoid breaking runtime.

---------

Co-authored-by: Laurent Nguyen <languye@microsoft.com>
---
 src/Common/dataAccess/readCollections.test.ts | 97 +++++++++++++++----
 src/Common/dataAccess/readCollections.ts      | 13 ++-
 2 files changed, 91 insertions(+), 19 deletions(-)

diff --git a/src/Common/dataAccess/readCollections.test.ts b/src/Common/dataAccess/readCollections.test.ts
index 50348b574..d3998ecfc 100644
--- a/src/Common/dataAccess/readCollections.test.ts
+++ b/src/Common/dataAccess/readCollections.test.ts
@@ -1,13 +1,39 @@
 jest.mock("../../Utils/arm/request");
 jest.mock("../CosmosClient");
+jest.mock("../Logger");
+jest.mock("../ErrorHandlingUtils", () => ({ handleError: jest.fn() }));
+jest.mock("../../Utils/NotificationConsoleUtils");
+import * as Logger from "Common/Logger";
 import { AuthType } from "../../AuthType";
 import { DatabaseAccount } from "../../Contracts/DataModels";
 import { updateUserContext } from "../../UserContext";
+import { logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
 import { armRequest } from "../../Utils/arm/request";
 import { client } from "../CosmosClient";
+import { handleError } from "../ErrorHandlingUtils";
 import { readCollections } from "./readCollections";
 
 describe("readCollections", () => {
+  const clearMessage = jest.fn();
+  const fetchAll = jest.fn();
+  const readAll = jest.fn(() => ({ fetchAll }));
+  const database = jest.fn(() => ({ containers: { readAll } }));
+  const diagnostics = {
+    clientSideRequestStatistics: {
+      requestDurationInMs: 123,
+      locationEndpointsContacted: ["https://test.documents.azure.com"],
+      retryDiagnostics: { failedAttempts: [{ statusCode: 429 }] },
+    },
+    diagnosticNode: { data: { responsePayload: "not-for-logging" } },
+  };
+
+  beforeEach(() => {
+    jest.clearAllMocks();
+    (logConsoleProgress as jest.Mock).mockReturnValue(clearMessage);
+    (client as jest.Mock).mockReturnValue({ database });
+    updateUserContext({ authType: AuthType.MasterKey });
+  });
+
   beforeAll(() => {
     updateUserContext({
       databaseAccount: {
@@ -23,26 +49,61 @@ describe("readCollections", () => {
     });
     await readCollections("database");
     expect(armRequest).toHaveBeenCalled();
+    expect(client).not.toHaveBeenCalled();
+    expect(clearMessage).toHaveBeenCalledTimes(1);
   });
 
-  it("should call SDK if not logged in with non-AAD method", async () => {
-    updateUserContext({
-      authType: AuthType.MasterKey,
+  it("should log SDK request statistics and return collections for non-AAD authentication", async () => {
+    const resources = [{ id: "container" }];
+    fetchAll.mockResolvedValue({ resources, diagnostics });
+
+    await expect(readCollections("database")).resolves.toBe(resources);
+
+    expect(database).toHaveBeenCalledWith("database");
+    expect(readAll).toHaveBeenCalledWith();
+    expect(fetchAll).toHaveBeenCalledTimes(1);
+    expect(Logger.logInfo).toHaveBeenLastCalledWith(
+      expect.stringContaining(`diagnostics=${JSON.stringify(diagnostics.clientSideRequestStatistics)}`),
+      "readCollections",
+    );
+    expect(JSON.stringify(jest.mocked(Logger.logInfo).mock.calls)).not.toContain("not-for-logging");
+    expect(clearMessage).toHaveBeenCalledTimes(1);
+  });
+
+  it("should log SDK failure statistics and rethrow the original error", async () => {
+    const error = Object.assign(new Error("fetchAll failed"), { diagnostics });
+    fetchAll.mockRejectedValue(error);
+
+    await expect(readCollections("database")).rejects.toBe(error);
+
+    expect(Logger.logError).toHaveBeenCalledWith(
+      `readCollections: fetchAll failed for database database, diagnostics=${JSON.stringify(
+        diagnostics.clientSideRequestStatistics,
+      )}`,
+      "readCollections",
+    );
+    expect(handleError).toHaveBeenCalledWith(
+      error,
+      "ReadCollections",
+      "Error while querying containers for database database",
+    );
+    expect(clearMessage).toHaveBeenCalledTimes(1);
+  });
+
+  it("should preserve error handling when diagnostics are unavailable", async () => {
+    const error = new Error("client initialization failed");
+    (client as jest.Mock).mockImplementationOnce(() => {
+      throw error;
     });
-    (client as jest.Mock).mockReturnValue({
-      database: () => {
-        return {
-          containers: {
-            readAll: () => {
-              return {
-                fetchAll: (): unknown => [],
-              };
-            },
-          },
-        };
-      },
-    });
-    await readCollections("database");
-    expect(client).toHaveBeenCalled();
+
+    await expect(readCollections("database")).rejects.toBe(error);
+
+    expect(Logger.logError).not.toHaveBeenCalled();
+    expect(handleError).toHaveBeenCalledWith(
+      error,
+      "ReadCollections",
+      "Error while querying containers for database database",
+    );
+    expect(clearMessage).toHaveBeenCalledTimes(1);
   });
 });
diff --git a/src/Common/dataAccess/readCollections.ts b/src/Common/dataAccess/readCollections.ts
index d90832219..caf317a16 100644
--- a/src/Common/dataAccess/readCollections.ts
+++ b/src/Common/dataAccess/readCollections.ts
@@ -78,7 +78,9 @@ export async function readCollections(databaseId: string): Promise<DataModels.Co
     const sdkResponse = await client().database(databaseId).containers.readAll().fetchAll();
     Logger.logInfo(
       `readCollections: fetchAll completed for database ${databaseId}, count=${sdkResponse.resources
-        ?.length}, durationMs=${Date.now() - fetchAllStart}`,
+        ?.length}, durationMs=${Date.now() - fetchAllStart}, diagnostics=${JSON.stringify(
+        sdkResponse.diagnostics?.clientSideRequestStatistics,
+      )}`,
       "readCollections",
     );
     traceSuccess(
@@ -88,6 +90,15 @@ export async function readCollections(databaseId: string): Promise<DataModels.Co
     );
     return sdkResponse.resources as DataModels.Collection[];
   } catch (error) {
+    const diagnostics = error instanceof Error && "diagnostics" in error ? error.diagnostics : undefined;
+    if (diagnostics && typeof diagnostics === "object" && "clientSideRequestStatistics" in diagnostics) {
+      Logger.logError(
+        `readCollections: fetchAll failed for database ${databaseId}, diagnostics=${JSON.stringify(
+          diagnostics.clientSideRequestStatistics,
+        )}`,
+        "readCollections",
+      );
+    }
     traceFailure(Action.ReadCollections, { databaseId, error: error?.message }, startKey);
     handleError(error, "ReadCollections", `Error while querying containers for database ${databaseId}`);
     throw error;