Address bug in fetching data for Tables Account

src/Common/Constants.ts

@@ -186,6 +186,9 @@ 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

@@ -31,9 +31,18 @@ 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;
+  }
+  
   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);
@@ -81,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(
@@ -118,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();

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/Tables/TableDataClient.ts

@@ -755,6 +755,10 @@ 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
@@ -769,3 +773,4 @@ export class CassandraAPIDataClient extends TableDataClient {
     );
   }
 }
+}

src/Platform/Hosted/extractFeatures.ts

@@ -14,6 +14,7 @@ export type Features = {
   readonly enableTtl: boolean;
   readonly executeSproc: boolean;
   readonly enableAadDataPlane: boolean;
+  readonly enableDataPlaneRbac: boolean;
   readonly enableResourceGraph: boolean;
   readonly enableKoResourceTree: boolean;
   readonly hostedDataExplorer: boolean;
@@ -69,6 +70,7 @@ 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

@@ -421,8 +421,10 @@ 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")) {
@@ -435,6 +437,8 @@ 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") {
@@ -458,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 (
@@ -499,10 +503,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) {
@@ -525,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) {
@@ -567,6 +584,7 @@ function updateContextsFromPortalMessage(inputs: DataExplorerInputsFrame) {
     collectionCreationDefaults: inputs.defaultCollectionThroughput,
     isTryCosmosDBSubscription: inputs.isTryCosmosDBSubscription,
     feedbackPolicies: inputs.feedbackPolicies,
+    listKeysFetchInProgress: false,
   });
 
   if (inputs.isPostgresAccount) {