Add fetchAndUpdate Keys

src/Common/Constants.ts

@@ -186,9 +186,6 @@ export class CassandraProxyAPIs {
 export class Queries {
   public static CustomPageOption: string = "custom";
   public static UnlimitedPageOption: string = "unlimited";
-  public static setAutomaticRBACOption: string = "Automatic";
-  public static setTrueRBACOption: string = "True";
-  public static setFalseRBACOption: string = "False";
   public static itemsPerPage: number = 100;
   public static unlimitedItemsPerPage: number = 100; // TODO: Figure out appropriate value so it works for accounts with a large number of partitions
   public static containersPerPage: number = 50;

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";
@@ -31,15 +29,6 @@ export const tokenProvider = async (requestInfo: Cosmos.RequestInfo) => {
     }
     const AUTH_PREFIX = `type=aad&ver=1.0&sig=`;
     const authorizationToken = `${AUTH_PREFIX}${userContext.aadToken}`;
-    console.log(`Returning Auth token`);
-    return authorizationToken;
-  }
-
-  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;
   }
 
@@ -100,22 +89,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) {
@@ -199,6 +172,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/Common/MongoProxyClient.ts

@@ -710,4 +710,3 @@ 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,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;
@@ -166,8 +168,15 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
       updateUserContext({
         dataPlaneRbacEnabled: false,
       });
+      const { databaseAccount: account, subscriptionId, resourceGroup } = userContext;
+    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);
     LocalStorageUtility.setEntryBoolean(StorageKey.QueryTimeoutEnabled, queryTimeoutEnabled);

src/Explorer/Tables/TableDataClient.ts

@@ -755,10 +755,6 @@ export class CassandraAPIDataClient extends TableDataClient {
       CassandraProxyEndpoints.Prod,
     ];
     let canAccessCassandraProxy: boolean = userContext.databaseAccount.properties.publicNetworkAccess === "Enabled";
-    if (
-      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
@@ -773,4 +769,3 @@ export class CassandraAPIDataClient extends TableDataClient {
     );
   }
 }
-}

src/Platform/Hosted/extractFeatures.ts

@@ -14,7 +14,6 @@ export type Features = {
   readonly enableTtl: boolean;
   readonly executeSproc: boolean;
   readonly enableAadDataPlane: boolean;
-  readonly enableDataPlaneRbac: boolean;
   readonly enableResourceGraph: boolean;
   readonly enableKoResourceTree: boolean;
   readonly hostedDataExplorer: boolean;
@@ -70,7 +69,6 @@ export function extractFeatures(given = new URLSearchParams(window.location.sear
     canExceedMaximumValue: "true" === get("canexceedmaximumvalue"),
     cosmosdb: "true" === get("cosmosdb"),
     enableAadDataPlane: "true" === get("enableaaddataplane"),
-    enableDataPlaneRbac: "true" === get("enabledataplanerbac"),
     enableResourceGraph: "true" === get("enableresourcegraph"),
     enableChangeFeedPolicy: "true" === get("enablechangefeedpolicy"),
     enableFixedCollectionWithSharedThroughput: "true" === get("enablefixedcollectionwithsharedthroughput"),

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;

src/hooks/useKnockoutExplorer.ts

@@ -417,14 +417,25 @@ function configureEmulator(): Explorer {
   return 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);
+  }
+}
+
 async function configurePortal(): Promise<Explorer> {
   updateUserContext({
     authType: AuthType.AAD,
   });
-  
-  
   let explorer: Explorer;
-  return new Promise(async (resolve) => {
+
+  return new Promise((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")) {
@@ -437,8 +448,6 @@ async function configurePortal(): Promise<Explorer> {
         console.dir(message);
         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") {
@@ -481,7 +490,7 @@ async function configurePortal(): Promise<Explorer> {
             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)) {
@@ -493,20 +502,31 @@ async function configurePortal(): Promise<Explorer> {
               } else {
                 dataPlaneRbacEnabled = isDataPlaneRbacSetting === Constants.RBACOptions.setTrueRBACOption;
               }
+              if(!dataPlaneRbacEnabled) {
+                 (async () => {
+              await fetchAndUpdateKeys(subscriptionId, resourceGroup, account.name);
+            })();
+              }
 
               updateUserContext({ dataPlaneRbacEnabled });
               useDataPlaneRbac.setState({ dataPlaneRbacEnabled: dataPlaneRbacEnabled });
             } else {
               const dataPlaneRbacEnabled = account.properties.disableLocalAuth;
 
+              if(!dataPlaneRbacEnabled) {
+                (async () => {
+             await fetchAndUpdateKeys(subscriptionId, resourceGroup, account.name);
+           })();
+             }
+
               updateUserContext({ dataPlaneRbacEnabled });
               useDataPlaneRbac.setState({ dataPlaneRbacEnabled: dataPlaneRbacEnabled });
             }
-          } else {
+          } 
-            const keys: DatabaseAccountListKeysResult = await listKeys(subscriptionId, resourceGroup, account.name);
+          else {
-            updateUserContext({
+            (async () => {
-              masterKey: keys.primaryMasterKey,
+              await fetchAndUpdateKeys(subscriptionId, resourceGroup, account.name);
-            });
+            })();
           }
 
           if (openAction) {
@@ -531,9 +551,7 @@ async function configurePortal(): Promise<Explorer> {
     );
 
     sendReadyMessage();
-
   });
-
 }
 
 function shouldForwardMessage(message: PortalMessage, messageOrigin: string) {
@@ -553,6 +571,11 @@ function updateContextsFromPortalMessage(inputs: DataExplorerInputsFrame) {
   const authorizationToken = inputs.authorizationToken || "";
   const databaseAccount = inputs.databaseAccount;
 
+  if(userContext.apiType !== "SQL") {
+    const masterKey = inputs.masterKey || "";
+    updateUserContext({masterKey});
+  }
+
   updateConfigContext({
     BACKEND_ENDPOINT: inputs.extensionEndpoint || configContext.BACKEND_ENDPOINT,
     ARM_ENDPOINT: normalizeArmEndpoint(inputs.csmEndpoint || configContext.ARM_ENDPOINT),
@@ -573,6 +596,7 @@ function updateContextsFromPortalMessage(inputs: DataExplorerInputsFrame) {
     collectionCreationDefaults: inputs.defaultCollectionThroughput,
     isTryCosmosDBSubscription: inputs.isTryCosmosDBSubscription,
     feedbackPolicies: inputs.feedbackPolicies,
+    listKeysFetchInProgress: false,
   });
 
   if (inputs.isPostgresAccount) {