Merge branch 'master' of https://github.com/Azure/cosmos-explorer

src/Common/CosmosClient.ts

@@ -3,12 +3,10 @@ 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 { updateUserContext, userContext } from "../UserContext";
+import { userContext } from "../UserContext";
 import { logConsoleError } from "../Utils/NotificationConsoleUtils";
 import * as PriorityBasedExecutionUtils from "../Utils/PriorityBasedExecutionUtils";
 import { EmulatorMasterKey, HttpHeaders } from "./Constants";
@@ -100,22 +98,6 @@ 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) {
@@ -179,8 +161,11 @@ enum SDKSupportedCapabilities {
 let _client: Cosmos.CosmosClient;
 
 export function client(): Cosmos.CosmosClient {
-  if (_client) return _client;
-
+  if (_client) {
+    if (!userContext.hasDataPlaneRbacSettingChanged) {
+      return _client;
+    }
+  }
   let _defaultHeaders: Cosmos.CosmosHeaders = {};
   _defaultHeaders["x-ms-cosmos-sdk-supportedcapabilities"] =
     SDKSupportedCapabilities.None | SDKSupportedCapabilities.PartitionMerge;
@@ -199,6 +184,7 @@ 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/Explorer/Panes/SettingsPane/SettingsPane.tsx

@@ -36,6 +36,8 @@ 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;
@@ -160,13 +162,24 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
     ) {
       updateUserContext({
         dataPlaneRbacEnabled: true,
+        hasDataPlaneRbacSettingChanged: true,
       });
       useDataPlaneRbac.setState({ dataPlaneRbacEnabled: true });
     } else {
       updateUserContext({
         dataPlaneRbacEnabled: false,
+        hasDataPlaneRbacSettingChanged: true,
       });
-      useDataPlaneRbac.setState({ dataPlaneRbacEnabled: false });
+      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 });
+        }
+      }
     }
 
     LocalStorageUtility.setEntryBoolean(StorageKey.RUThresholdEnabled, ruThresholdEnabled);

src/UserContext.ts

@@ -102,6 +102,7 @@ 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,7 +41,6 @@ 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";
 
@@ -288,19 +287,26 @@ 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 {
-        const keys: DatabaseAccountListKeysResult = await listKeys(subscriptionId, resourceGroup, account.name);
-        updateUserContext({
-          masterKey: keys.primaryMasterKey,
-        });
+        await fetchAndUpdateKeys(subscriptionId, resourceGroup, account.name);
+      }
+    } else {
+      if (!account.properties.disableLocalAuth) {
+        await fetchAndUpdateKeys(subscriptionId, resourceGroup, account.name);
       }
     }
   } catch (e) {
@@ -417,6 +423,19 @@ 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,
@@ -474,12 +493,6 @@ 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;
 
@@ -493,20 +506,31 @@ 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 {
-            const keys: DatabaseAccountListKeysResult = await listKeys(subscriptionId, resourceGroup, account.name);
-            updateUserContext({
-              masterKey: keys.primaryMasterKey,
-            });
+            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) {