Address bug in fetching data for Tables Account

src/Common/CosmosClient.ts

@@ -3,10 +3,12 @@ import { getAuthorizationTokenUsingResourceTokens } from "Common/getAuthorizatio
 import { AuthorizationToken } from "Contracts/FabricMessageTypes";
 import { checkDatabaseResourceTokensValidity } from "Platform/Fabric/FabricUtil";
 import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility";
+import { listKeys } from "Utils/arm/generatedClients/cosmos/databaseAccounts";
+import { DatabaseAccountListKeysResult } from "Utils/arm/generatedClients/cosmos/types";
 import { AuthType } from "../AuthType";
 import { PriorityLevel } from "../Common/Constants";
 import { Platform, configContext } from "../ConfigContext";
-import { userContext } from "../UserContext";
+import { updateUserContext, userContext } from "../UserContext";
 import { logConsoleError } from "../Utils/NotificationConsoleUtils";
 import * as PriorityBasedExecutionUtils from "../Utils/PriorityBasedExecutionUtils";
 import { EmulatorMasterKey, HttpHeaders } from "./Constants";
@@ -33,14 +35,14 @@ export const tokenProvider = async (requestInfo: Cosmos.RequestInfo) => {
     return authorizationToken;
   }
 
-  if (userContext.dataPlaneRbacEnabled && userContext.authorizationToken) {
-    console.log(` Getting Portal Auth token `);
+  if ((userContext.dataPlaneRbacEnabled) && userContext.authorizationToken) {
+    console.log(` Getting Portal Auth token `)
     const AUTH_PREFIX = `type=aad&ver=1.0&sig=`;
     const authorizationToken = `${AUTH_PREFIX}${userContext.authorizationToken}`;
     console.log(`Returning Portal Auth token`);
     return authorizationToken;
   }
-
+  
   if (configContext.platform === Platform.Emulator) {
     // TODO This SDK method mutates the headers object. Find a better one or fix the SDK.
     await Cosmos.setAuthorizationTokenHeaderUsingMasterKey(verb, resourceId, resourceType, headers, EmulatorMasterKey);
@@ -88,6 +90,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(
@@ -98,6 +106,22 @@ export const tokenProvider = async (requestInfo: Cosmos.RequestInfo) => {
       userContext.masterKey,
     );
     return decodeURIComponent(headers.authorization);
+  } else if (userContext.dataPlaneRbacEnabled == false) {
+    const { databaseAccount: account, subscriptionId, resourceGroup } = userContext;
+    const keys: DatabaseAccountListKeysResult = await listKeys(subscriptionId, resourceGroup, account.name);
+
+    if (keys.primaryMasterKey) {
+      updateUserContext({ masterKey: keys.primaryMasterKey });
+      // TODO This SDK method mutates the headers object. Find a better one or fix the SDK.
+      await Cosmos.setAuthorizationTokenHeaderUsingMasterKey(
+        verb,
+        resourceId,
+        resourceType,
+        headers,
+        keys.primaryMasterKey,
+      );
+      return decodeURIComponent(headers.authorization);
+    }
   }
 
   if (userContext.resourceToken) {
@@ -109,6 +133,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();
@@ -161,11 +189,8 @@ enum SDKSupportedCapabilities {
 let _client: Cosmos.CosmosClient;
 
 export function client(): Cosmos.CosmosClient {
-  if (_client) {
-    if (!userContext.hasDataPlaneRbacSettingChanged) {
-      return _client;
-    }
-  }
+  if (_client) return _client;
+
   let _defaultHeaders: Cosmos.CosmosHeaders = {};
   _defaultHeaders["x-ms-cosmos-sdk-supportedcapabilities"] =
     SDKSupportedCapabilities.None | SDKSupportedCapabilities.PartitionMerge;
@@ -184,7 +209,6 @@ export function client(): Cosmos.CosmosClient {
 
   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,
     tokenProvider,
     userAgentSuffix: "Azure Portal",
     defaultHeaders: _defaultHeaders,

src/Common/MongoProxyClient.ts

@@ -710,3 +710,4 @@ async function errorHandling(response: Response, action: string, params: unknown
 export function getARMCreateCollectionEndpoint(params: DataModels.MongoParameters): string {
   return `subscriptions/${params.sid}/resourceGroups/${params.rg}/providers/Microsoft.DocumentDB/databaseAccounts/${userContext.databaseAccount.name}/mongodbDatabases/${params.db}/collections/${params.coll}`;
 }
+

src/Explorer/Panes/SettingsPane/SettingsPane.tsx

@@ -36,8 +36,6 @@ import Explorer from "../../Explorer";
 import { RightPaneForm, RightPaneFormProps } from "../RightPaneForm/RightPaneForm";
 import { AuthType } from "AuthType";
 import create, { UseStore } from "zustand";
-import { DatabaseAccountListKeysResult } from "@azure/arm-cosmosdb/esm/models";
-import { listKeys } from "Utils/arm/generatedClients/cosmos/databaseAccounts";
 
 export interface DataPlaneRbacState {
   dataPlaneRbacEnabled: boolean;
@@ -162,24 +160,13 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
     ) {
       updateUserContext({
         dataPlaneRbacEnabled: true,
-        hasDataPlaneRbacSettingChanged: true,
       });
       useDataPlaneRbac.setState({ dataPlaneRbacEnabled: true });
     } else {
       updateUserContext({
         dataPlaneRbacEnabled: false,
-        hasDataPlaneRbacSettingChanged: true,
       });
-      const { databaseAccount: account, subscriptionId, resourceGroup } = userContext;
-      if (!userContext.features.enableAadDataPlane) {
-        const keys: DatabaseAccountListKeysResult = await listKeys(subscriptionId, resourceGroup, account.name);
-
-        if (keys.primaryMasterKey) {
-          updateUserContext({ masterKey: keys.primaryMasterKey });
-
-          useDataPlaneRbac.setState({ dataPlaneRbacEnabled: false });
-        }
-      }
+      useDataPlaneRbac.setState({ dataPlaneRbacEnabled: false });
     }
 
     LocalStorageUtility.setEntryBoolean(StorageKey.RUThresholdEnabled, ruThresholdEnabled);

src/Explorer/Tables/TableDataClient.ts

@@ -759,18 +759,18 @@ export class CassandraAPIDataClient extends TableDataClient {
       configContext.CASSANDRA_PROXY_ENDPOINT !== CassandraProxyEndpoints.Development &&
       userContext.databaseAccount.properties.ipRules?.length > 0
     ) {
-      if (
-        configContext.CASSANDRA_PROXY_ENDPOINT !== CassandraProxyEndpoints.Development &&
-        userContext.databaseAccount.properties.ipRules?.length > 0
-      ) {
-        canAccessCassandraProxy = canAccessCassandraProxy && configContext.CASSANDRA_PROXY_OUTBOUND_IPS_ALLOWLISTED;
-      }
-
-      return (
-        canAccessCassandraProxy &&
-        configContext.NEW_CASSANDRA_APIS?.includes(api) &&
-        activeCassandraProxyEndpoints.includes(configContext.CASSANDRA_PROXY_ENDPOINT)
-      );
+    if (
+      configContext.CASSANDRA_PROXY_ENDPOINT !== CassandraProxyEndpoints.Development &&
+      userContext.databaseAccount.properties.ipRules?.length > 0
+    ) {
+      canAccessCassandraProxy = canAccessCassandraProxy && configContext.CASSANDRA_PROXY_OUTBOUND_IPS_ALLOWLISTED;
     }
+
+    return (
+      canAccessCassandraProxy &&
+      configContext.NEW_CASSANDRA_APIS?.includes(api) &&
+      activeCassandraProxyEndpoints.includes(configContext.CASSANDRA_PROXY_ENDPOINT)
+    );
   }
 }
+}

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;
@@ -102,7 +103,6 @@ export interface UserContext {
   readonly vcoreMongoConnectionParams?: VCoreMongoConnectionParams;
   readonly feedbackPolicies?: AdminFeedbackPolicySettings;
   readonly dataPlaneRbacEnabled?: boolean;
-  readonly hasDataPlaneRbacSettingChanged?: boolean;
 }
 
 export type ApiType = "SQL" | "Mongo" | "Gremlin" | "Tables" | "Cassandra" | "Postgres" | "VCoreMongo";

src/hooks/useKnockoutExplorer.ts

@@ -41,6 +41,7 @@ import { Node, PortalEnv, updateUserContext, userContext } from "../UserContext"
 import { acquireTokenWithMsal, getAuthorizationHeader, getMsalInstance } from "../Utils/AuthorizationUtils";
 import { isInvalidParentFrameOrigin, shouldProcessMessage } from "../Utils/MessageValidation";
 import { listKeys } from "../Utils/arm/generatedClients/cosmos/databaseAccounts";
+import { DatabaseAccountListKeysResult } from "../Utils/arm/generatedClients/cosmos/types";
 import { applyExplorerBindings } from "../applyExplorerBindings";
 import { useDataPlaneRbac } from "Explorer/Panes/SettingsPane/SettingsPane";
 
@@ -287,26 +288,19 @@ async function configureHostedWithAAD(config: AAD): Promise<Explorer> {
           } else {
             dataPlaneRbacEnabled = isDataPlaneRbacSetting === Constants.RBACOptions.setTrueRBACOption;
           }
-          if (!dataPlaneRbacEnabled) {
-            await fetchAndUpdateKeys(subscriptionId, resourceGroup, account.name);
-          }
 
           updateUserContext({ dataPlaneRbacEnabled });
         } else {
           const dataPlaneRbacEnabled = account.properties.disableLocalAuth;
-          if (!dataPlaneRbacEnabled) {
-            await fetchAndUpdateKeys(subscriptionId, resourceGroup, account.name);
-          }
 
           updateUserContext({ dataPlaneRbacEnabled });
           useDataPlaneRbac.setState({ dataPlaneRbacEnabled: dataPlaneRbacEnabled });
         }
       } else {
-        await fetchAndUpdateKeys(subscriptionId, resourceGroup, account.name);
-      }
-    } else {
-      if (!account.properties.disableLocalAuth) {
-        await fetchAndUpdateKeys(subscriptionId, resourceGroup, account.name);
+        const keys: DatabaseAccountListKeysResult = await listKeys(subscriptionId, resourceGroup, account.name);
+        updateUserContext({
+          masterKey: keys.primaryMasterKey,
+        });
       }
     }
   } catch (e) {
@@ -423,26 +417,14 @@ function configureEmulator(): Explorer {
   return explorer;
 }
 
-async function fetchAndUpdateKeys(subscriptionId: string, resourceGroup: string, account: string) {
-  try {
-    const keys = await listKeys(subscriptionId, resourceGroup, account);
-
-    updateUserContext({
-      masterKey: keys.primaryMasterKey,
-    });
-  } catch (error) {
-    console.error("Error during fetching keys or updating user context:", error);
-    throw error;
-  }
-}
-
 async function configurePortal(): Promise<Explorer> {
   updateUserContext({
     authType: AuthType.AAD,
   });
-
+  
+  
   let explorer: Explorer;
-  return new Promise((resolve) => {
+  return new Promise(async (resolve) => {
     // In development mode, try to load the iframe message from session storage.
     // This allows webpack hot reload to function properly in the portal
     if (process.env.NODE_ENV === "development" && !window.location.search.includes("disablePortalInitCache")) {
@@ -456,6 +438,7 @@ async function configurePortal(): Promise<Explorer> {
         updateContextsFromPortalMessage(message);
         explorer = new Explorer();
 
+
         // In development mode, save the iframe message from the portal in session storage.
         // This allows webpack hot reload to funciton properly
         if (process.env.NODE_ENV === "development") {
@@ -479,7 +462,7 @@ async function configurePortal(): Promise<Explorer> {
 
         // Check for init message
         const message: PortalMessage = event.data?.data;
-        const inputs = message?.inputs;
+        const inputs = message?.inputs; 
         const openAction = message?.openAction;
         if (inputs) {
           if (
@@ -491,6 +474,12 @@ 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;
 
@@ -504,33 +493,19 @@ async function configurePortal(): Promise<Explorer> {
               } else {
                 dataPlaneRbacEnabled = isDataPlaneRbacSetting === Constants.RBACOptions.setTrueRBACOption;
               }
-              if (!dataPlaneRbacEnabled) {
-                await fetchAndUpdateKeys(subscriptionId, resourceGroup, account.name);
-              }
 
               updateUserContext({ dataPlaneRbacEnabled });
               useDataPlaneRbac.setState({ dataPlaneRbacEnabled: dataPlaneRbacEnabled });
             } else {
               const dataPlaneRbacEnabled = account.properties.disableLocalAuth;
 
-              if (!dataPlaneRbacEnabled) {
-                await fetchAndUpdateKeys(subscriptionId, resourceGroup, account.name);
-              }
-
               updateUserContext({ dataPlaneRbacEnabled });
               useDataPlaneRbac.setState({ dataPlaneRbacEnabled: dataPlaneRbacEnabled });
             }
           } else {
-            if (userContext.apiType !== "Postgres" && userContext.apiType !== "VCoreMongo") {
-              await listKeys(subscriptionId, resourceGroup, account.name);
-            }
-          }
-
-          explorer = new Explorer();
-          resolve(explorer);
-
-          if (userContext.apiType === "Postgres" || userContext.apiType === "SQL" || userContext.apiType === "Mongo") {
-            setTimeout(() => explorer.openNPSSurveyDialog(), 3000);
+            (async () => {
+              await fetchAndUpdateKeys(subscriptionId, resourceGroup, account.name);
+            })();
           }
 
           if (openAction) {
@@ -553,9 +528,23 @@ async function configurePortal(): Promise<Explorer> {
       },
       false,
     );
-
+    
     sendReadyMessage();
+
   });
+
+}
+
+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) {
@@ -595,6 +584,7 @@ function updateContextsFromPortalMessage(inputs: DataExplorerInputsFrame) {
     collectionCreationDefaults: inputs.defaultCollectionThroughput,
     isTryCosmosDBSubscription: inputs.isTryCosmosDBSubscription,
     feedbackPolicies: inputs.feedbackPolicies,
+    listKeysFetchInProgress: false,
   });
 
   if (inputs.isPostgresAccount) {