From 039136794d7b79669ddf7367848528165e9b01d7 Mon Sep 17 00:00:00 2001
From: Senthamil Sindhu <sindhuba@microsoft.com>
Date: Thu, 4 Jul 2024 16:01:35 -0700
Subject: [PATCH] Add fix for MPAC Tables account issue

---
 src/Common/CosmosClient.ts                    |  4 +-
 .../Panes/SettingsPane/SettingsPane.tsx       | 12 ++---
 src/hooks/useKnockoutExplorer.ts              | 52 +++++++++----------
 3 files changed, 33 insertions(+), 35 deletions(-)

diff --git a/src/Common/CosmosClient.ts b/src/Common/CosmosClient.ts
index b5e50d5cb..32f9f6e24 100644
--- a/src/Common/CosmosClient.ts
+++ b/src/Common/CosmosClient.ts
@@ -89,7 +89,7 @@ export const tokenProvider = async (requestInfo: Cosmos.RequestInfo) => {
       userContext.masterKey,
     );
     return decodeURIComponent(headers.authorization);
-  } 
+  }
 
   if (userContext.resourceToken) {
     return userContext.resourceToken;
@@ -169,7 +169,7 @@ export function client(): Cosmos.CosmosClient {
     // The header will be ignored if priority based execution is disabled on the account.
     _defaultHeaders["x-ms-cosmos-priority-level"] = PriorityLevel.Default;
   }
-    
+
   const options: Cosmos.CosmosClientOptions = {
     endpoint: endpoint() || "https://cosmos.azure.com", // CosmosClient gets upset if we pass a bad URL. This should never actually get called
     key: userContext.dataPlaneRbacEnabled ? "" : userContext.masterKey,
diff --git a/src/Explorer/Panes/SettingsPane/SettingsPane.tsx b/src/Explorer/Panes/SettingsPane/SettingsPane.tsx
index 9e6d2afd6..ae45fceaf 100644
--- a/src/Explorer/Panes/SettingsPane/SettingsPane.tsx
+++ b/src/Explorer/Panes/SettingsPane/SettingsPane.tsx
@@ -169,14 +169,14 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
         dataPlaneRbacEnabled: false,
       });
       const { databaseAccount: account, subscriptionId, resourceGroup } = userContext;
-    const keys: DatabaseAccountListKeysResult = await listKeys(subscriptionId, resourceGroup, account.name);
+      const keys: DatabaseAccountListKeysResult = await listKeys(subscriptionId, resourceGroup, account.name);
 
-    if (keys.primaryMasterKey) {
-      updateUserContext({ masterKey: keys.primaryMasterKey });
-      
-      useDataPlaneRbac.setState({ dataPlaneRbacEnabled: false });
+      if (keys.primaryMasterKey) {
+        updateUserContext({ masterKey: keys.primaryMasterKey });
+
+        useDataPlaneRbac.setState({ dataPlaneRbacEnabled: false });
+      }
     }
-  }
 
     LocalStorageUtility.setEntryBoolean(StorageKey.RUThresholdEnabled, ruThresholdEnabled);
     LocalStorageUtility.setEntryBoolean(StorageKey.QueryTimeoutEnabled, queryTimeoutEnabled);
diff --git a/src/hooks/useKnockoutExplorer.ts b/src/hooks/useKnockoutExplorer.ts
index cea0f1428..958128069 100644
--- a/src/hooks/useKnockoutExplorer.ts
+++ b/src/hooks/useKnockoutExplorer.ts
@@ -417,15 +417,20 @@ function configureEmulator(): Explorer {
   return explorer;
 }
 
- async function fetchAndUpdateKeys(subscriptionId: string, resourceGroup: string, account: string) {
+async function fetchAndUpdateKeys(subscriptionId: string, resourceGroup: string, account: string) {
   try {
-    updateUserContext({ listKeysFetchInProgress: true });
+    console.log("Starting to fetch keys...");
     const keys = await listKeys(subscriptionId, resourceGroup, account);
+    console.log("Keys fetched:", keys);
 
-    updateUserContext({ masterKey: keys.primaryMasterKey, listKeysFetchInProgress: false });
+    updateUserContext({
+      masterKey: keys.primaryMasterKey,
+    });
+
+    console.log("User context updated with master key.");
   } catch (error) {
-    updateUserContext({ listKeysFetchInProgress: false });
     console.error("Error during fetching keys or updating user context:", error);
+    throw error;
   }
 }
 
@@ -483,14 +488,8 @@ async function configurePortal(): Promise<Explorer> {
           }
 
           updateContextsFromPortalMessage(inputs);
-          explorer = new Explorer();
-          resolve(explorer);
 
-          if (userContext.apiType === "Postgres" || userContext.apiType === "SQL" || userContext.apiType === "Mongo") {
-            setTimeout(() => explorer.openNPSSurveyDialog(), 3000);
-          }
-
-          const { databaseAccount: account , subscriptionId, resourceGroup} = userContext;
+          const { databaseAccount: account, subscriptionId, resourceGroup } = userContext;
 
           if (userContext.apiType === "SQL") {
             if (LocalStorageUtility.hasItem(StorageKey.DataPlaneRbacEnabled)) {
@@ -502,10 +501,8 @@ async function configurePortal(): Promise<Explorer> {
               } else {
                 dataPlaneRbacEnabled = isDataPlaneRbacSetting === Constants.RBACOptions.setTrueRBACOption;
               }
-              if(!dataPlaneRbacEnabled) {
-                 (async () => {
-              await fetchAndUpdateKeys(subscriptionId, resourceGroup, account.name);
-            })();
+              if (!dataPlaneRbacEnabled) {
+                await fetchAndUpdateKeys(subscriptionId, resourceGroup, account.name);
               }
 
               updateUserContext({ dataPlaneRbacEnabled });
@@ -513,20 +510,21 @@ async function configurePortal(): Promise<Explorer> {
             } else {
               const dataPlaneRbacEnabled = account.properties.disableLocalAuth;
 
-              if(!dataPlaneRbacEnabled) {
-                (async () => {
-             await fetchAndUpdateKeys(subscriptionId, resourceGroup, account.name);
-           })();
-             }
+              if (!dataPlaneRbacEnabled) {
+                await fetchAndUpdateKeys(subscriptionId, resourceGroup, account.name);
+              }
 
               updateUserContext({ dataPlaneRbacEnabled });
               useDataPlaneRbac.setState({ dataPlaneRbacEnabled: dataPlaneRbacEnabled });
             }
-          } 
-          else {
-            (async () => {
-              await fetchAndUpdateKeys(subscriptionId, resourceGroup, account.name);
-            })();
+          } else {
+            await fetchAndUpdateKeys(subscriptionId, resourceGroup, account.name);
+          }
+
+          explorer = new Explorer();
+          resolve(explorer);
+          if (userContext.apiType === "Postgres" || userContext.apiType === "SQL" || userContext.apiType === "Mongo") {
+            setTimeout(() => explorer.openNPSSurveyDialog(), 3000);
           }
 
           if (openAction) {
@@ -571,9 +569,9 @@ function updateContextsFromPortalMessage(inputs: DataExplorerInputsFrame) {
   const authorizationToken = inputs.authorizationToken || "";
   const databaseAccount = inputs.databaseAccount;
 
-  if(userContext.apiType !== "SQL") {
+  if (userContext.apiType !== "SQL") {
     const masterKey = inputs.masterKey || "";
-    updateUserContext({masterKey});
+    updateUserContext({ masterKey });
   }
 
   updateConfigContext({