From 9274f50500b4a4354c48169b6f843d95c647d286 Mon Sep 17 00:00:00 2001
From: Senthamil Sindhu <sindhuba@microsoft.com>
Date: Wed, 3 Jul 2024 00:02:29 -0700
Subject: [PATCH] Address bug in fetching data for Tables Account

---
 src/Common/CosmosClient.ts       | 10 ++++++++++
 src/UserContext.ts               |  1 +
 src/hooks/useKnockoutExplorer.ts | 20 ++++++++++++++++----
 3 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/src/Common/CosmosClient.ts b/src/Common/CosmosClient.ts
index 4256c31a3..682ab4e39 100644
--- a/src/Common/CosmosClient.ts
+++ b/src/Common/CosmosClient.ts
@@ -81,6 +81,12 @@ export const tokenProvider = async (requestInfo: Cosmos.RequestInfo) => {
     }
   }
 
+  let retryAttempts: number = 50;
+  while (retryAttempts > 0 && userContext.listKeysFetchInProgress) {
+    retryAttempts--;
+    await sleep(100);
+  }
+
   if (userContext.masterKey) {
     // TODO This SDK method mutates the headers object. Find a better one or fix the SDK.
     await Cosmos.setAuthorizationTokenHeaderUsingMasterKey(
@@ -118,6 +124,10 @@ export const tokenProvider = async (requestInfo: Cosmos.RequestInfo) => {
   return decodeURIComponent(result.PrimaryReadWriteToken);
 };
 
+function sleep(ms: number) {
+  return new Promise((resolve) => setTimeout(resolve, ms));
+}
+
 export const requestPlugin: Cosmos.Plugin<any> = async (requestContext, diagnosticNode, next) => {
   requestContext.endpoint = new URL(configContext.PROXY_PATH, window.location.href).href;
   requestContext.headers["x-ms-proxy-target"] = endpoint();
diff --git a/src/UserContext.ts b/src/UserContext.ts
index 4718e3915..fcf690560 100644
--- a/src/UserContext.ts
+++ b/src/UserContext.ts
@@ -73,6 +73,7 @@ export interface UserContext {
   readonly fabricContext?: FabricContext;
   readonly authType?: AuthType;
   readonly masterKey?: string;
+  readonly listKeysFetchInProgress?: boolean;
   readonly subscriptionId?: string;
   readonly resourceGroup?: string;
   readonly databaseAccount?: DatabaseAccount;
diff --git a/src/hooks/useKnockoutExplorer.ts b/src/hooks/useKnockoutExplorer.ts
index 21461d363..b98341b63 100644
--- a/src/hooks/useKnockoutExplorer.ts
+++ b/src/hooks/useKnockoutExplorer.ts
@@ -499,10 +499,9 @@ async function configurePortal(): Promise<Explorer> {
               useDataPlaneRbac.setState({ dataPlaneRbacEnabled: dataPlaneRbacEnabled });
             }
           } else {
-            const keys: DatabaseAccountListKeysResult = await listKeys(subscriptionId, resourceGroup, account.name);
-            updateUserContext({
-              masterKey: keys.primaryMasterKey,
-            });
+            (async () => {
+              await fetchAndUpdateKeys(subscriptionId, resourceGroup, account.name);
+            })();
           }
 
           if (openAction) {
@@ -530,6 +529,18 @@ async function configurePortal(): Promise<Explorer> {
   });
 }
 
+async function fetchAndUpdateKeys(subscriptionId: string, resourceGroup: string, account: string) {
+  try {
+    updateUserContext({ listKeysFetchInProgress: true });
+    const keys = await listKeys(subscriptionId, resourceGroup, account);
+
+    updateUserContext({ masterKey: keys.primaryMasterKey, listKeysFetchInProgress: false });
+  } catch (error) {
+    updateUserContext({ listKeysFetchInProgress: false });
+    console.error("Error during fetching keys or updating user context:", error);
+  }
+}
+
 function shouldForwardMessage(message: PortalMessage, messageOrigin: string) {
   // Only allow forwarding messages from the same origin
   return messageOrigin === window.document.location.origin && message.type === MessageTypes.TelemetryInfo;
@@ -567,6 +578,7 @@ function updateContextsFromPortalMessage(inputs: DataExplorerInputsFrame) {
     collectionCreationDefaults: inputs.defaultCollectionThroughput,
     isTryCosmosDBSubscription: inputs.isTryCosmosDBSubscription,
     feedbackPolicies: inputs.feedbackPolicies,
+    listKeysFetchInProgress: false,
   });
 
   if (inputs.isPostgresAccount) {