Add recent changes

src/Common/CosmosClient.ts

@@ -8,11 +8,13 @@ import { AuthType } from "../AuthType";
 import { BackendApi, PriorityLevel } from "../Common/Constants";
 import * as Logger from "../Common/Logger";
 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";
 import { getErrorMessage } from "./ErrorHandlingUtils";
+import { runCommand } from "hooks/useDatabaseAccounts";
+import { acquireTokenWithMsal, getMsalInstance } from "Utils/AuthorizationUtils";
 
 const _global = typeof self === "undefined" ? window : self;
 
@@ -32,7 +34,42 @@ export const tokenProvider = async (requestInfo: Cosmos.RequestInfo) => {
       return null;
     }
     const AUTH_PREFIX = `type=aad&ver=1.0&sig=`;
-    const authorizationToken = `${AUTH_PREFIX}${userContext.aadToken}`;
+    let authorizationToken;
+
+    try {
+       authorizationToken = `${AUTH_PREFIX}${userContext.aadToken}`;
+    } catch (error) {
+      if (error.code === "ExpiredAuthenticationToken") {
+        // Renew the AAD token using runCommand
+        const newToken = await runCommand(async () => {
+          // Implement the logic to acquire a new AAD token
+          const msalInstance = await getMsalInstance();
+          const cachedAccount = msalInstance.getAllAccounts()?.[0];
+          const cachedTenantId = localStorage.getItem("cachedTenantId");
+  
+          msalInstance.setActiveAccount(cachedAccount);
+          const newAccessToken = await acquireTokenWithMsal(msalInstance, {
+            authority: `${configContext.AAD_ENDPOINT}${cachedTenantId}`,
+            scopes: [`${configContext.ARM_ENDPOINT}/.default`],
+          });
+  
+          // Update user context with the new token
+          updateUserContext({ aadToken: newAccessToken });
+          authorizationToken = `${AUTH_PREFIX}${userContext.aadToken}`;
+          return newAccessToken;
+        });
+  
+        // Retry getting the token after renewing
+        const retryResult = await getTokenFromAuthService(verb, resourceType, resourceId);
+        headers[HttpHeaders.msDate] = retryResult.XDate;
+        return decodeURIComponent(retryResult.PrimaryReadWriteToken);
+      } else {
+        console.error('An error occurred:', error.message);
+        throw error; 
+      }
+    }
+
+    
     return authorizationToken;
   }
 

src/ConfigContext.ts

@@ -5,6 +5,7 @@ import {
   MongoProxyEndpoints,
   PortalBackendEndpoints,
 } from "Common/Constants";
+import { userContext } from "UserContext";
 import {
   allowedAadEndpoints,
   allowedArcadiaEndpoints,
@@ -38,6 +39,7 @@ export interface ConfigContext {
   gitSha?: string;
   proxyPath?: string;
   AAD_ENDPOINT: string;
+  ENVIRONMENT: string;
   ARM_AUTH_AREA: string;
   ARM_ENDPOINT: string;
   EMULATOR_ENDPOINT?: string;
@@ -93,7 +95,7 @@ let configContext: Readonly<ConfigContext> = {
   ], // Webpack injects this at build time
   gitSha: process.env.GIT_SHA,
   hostedExplorerURL: "https://cosmos.azure.com/",
-  AAD_ENDPOINT: "https://login.microsoftonline.com/",
+  AAD_ENDPOINT: "",
   ARM_AUTH_AREA: "https://management.azure.com/",
   ARM_ENDPOINT: "https://management.azure.com/",
   ARM_API_VERSION: "2016-06-01",

src/UserContext.ts

@@ -80,6 +80,7 @@ export interface UserContext {
   readonly endpoint?: string;
   readonly aadToken?: string;
   readonly accessToken?: string;
+  readonly armToken?: string;
   readonly authorizationToken?: string;
   readonly resourceToken?: string;
   readonly subscriptionType?: SubscriptionType;

src/hooks/useAADAuth.ts

@@ -3,6 +3,7 @@ import { useBoolean } from "@fluentui/react-hooks";
 import * as React from "react";
 import { configContext } from "../ConfigContext";
 import { acquireTokenWithMsal, getMsalInstance } from "../Utils/AuthorizationUtils";
+import { updateUserContext } from "UserContext";
 
 const msalInstance = await getMsalInstance();
 
@@ -79,7 +80,7 @@ export function useAADAuth(): ReturnType {
         authority: `${configContext.AAD_ENDPOINT}${tenantId}`,
         scopes: [`${configContext.ARM_ENDPOINT}/.default`],
       });
-
+      updateUserContext({ armToken: armToken});
       setArmToken(armToken);
       setAuthFailure(null);
     } catch (error) {

src/hooks/useDatabaseAccounts.tsx

@@ -1,8 +1,8 @@
 import { HttpHeaders } from "Common/Constants";
-import { QueryRequestOptions, QueryResponse } from "Contracts/AzureResourceGraph";
 import useSWR from "swr";
-import { configContext } from "../ConfigContext";
-import { DatabaseAccount } from "../Contracts/DataModels";
+import { acquireTokenWithMsal, getMsalInstance } from "Utils/AuthorizationUtils";
+import React from "react";
+import { updateUserContext, userContext } from "UserContext";
 /* eslint-disable  @typescript-eslint/no-explicit-any */
 
 interface AccountListResult {
@@ -34,11 +34,10 @@ export async function fetchDatabaseAccounts(subscriptionId: string, accessToken:
 }
 
 export async function fetchDatabaseAccountsFromGraph(
-  subscriptionId: string,
-  accessToken: string,
+  subscriptionId: string
 ): Promise<DatabaseAccount[]> {
   const headers = new Headers();
-  const bearer = `Bearer ${accessToken}`;
+  const bearer = `Bearer ${userContext.armToken}`;
 
   headers.append("Authorization", bearer);
   headers.append(HttpHeaders.contentType, "application/json");
@@ -46,8 +45,9 @@ export async function fetchDatabaseAccountsFromGraph(
   const apiVersion = "2021-03-01";
   const managementResourceGraphAPIURL = `${configContext.ARM_ENDPOINT}providers/Microsoft.ResourceGraph/resources?api-version=${apiVersion}`;
 
-  const databaseAccounts: DatabaseAccount[] = [];
+  let databaseAccounts: DatabaseAccount[] = [];
   let skipToken: string;
+  console.log("Old ARM Token", userContext.armToken);
   do {
     const body = {
       query: databaseAccountsQuery,
@@ -75,20 +75,165 @@ export async function fetchDatabaseAccountsFromGraph(
       throw new Error(await response.text());
     }
     
+
     const queryResponse: QueryResponse = (await response.json()) as QueryResponse;
     skipToken = queryResponse.$skipToken;
     queryResponse.data?.map((databaseAccount: any) => {
       databaseAccounts.push(databaseAccount as DatabaseAccount);
     });
-  } while (skipToken);
   
+  //  else {
+  //   try{
+  //     console.log("Token expired");
+  //     databaseAccounts = await acquireNewTokenAndRetry(body);
+  //   }
+    // catch (error) {
+    //  throw new Error(error);
+    // }
+  
+  //}
+  } while (skipToken);
   return databaseAccounts.sort((a, b) => a.name.localeCompare(b.name));
 }
 
-export function useDatabaseAccounts(subscriptionId: string, armToken: string): DatabaseAccount[] | undefined {
+export function useDatabaseAccounts(subscriptionId: string): DatabaseAccount[] | undefined {
   const { data } = useSWR(
-    () => (armToken && subscriptionId ? ["databaseAccounts", subscriptionId, armToken] : undefined),
-    (_, subscriptionId, armToken) => fetchDatabaseAccountsFromGraph(subscriptionId, armToken),
+    () => ( subscriptionId ? ["databaseAccounts", subscriptionId] : undefined),
+    (_, subscriptionId) => runCommand(fetchDatabaseAccountsFromGraph, subscriptionId),
   );
   return data;
 }
+
+
+
+// Define the types for your responses
+interface DatabaseAccount {
+  name: string;
+  id: string;
+  // Add other relevant fields as per your use case
+}
+
+interface Subscription {
+  displayName: string;
+  subscriptionId: string;
+  state: string;
+}
+
+interface QueryRequestOptions {
+  $top?: number;
+  $skipToken?: string;
+  $allowPartialScopes?: boolean;
+}
+
+// Define the configuration context and headers if not already defined
+const configContext = {
+  ARM_ENDPOINT: 'https://management.azure.com/',
+  AAD_ENDPOINT: 'https://login.microsoftonline.com/'
+};
+
+interface QueryResponse {
+  data?: any[];
+  $skipToken?: string;
+}
+
+export async function runCommand<T>(
+  fn: (...args: any[]) => Promise<T>,
+  ...args: any[]
+): Promise<T> {
+  try {
+      // Attempt to execute the function passed as an argument
+      const result = await fn(...args);
+      console.log('Successfully executed function:', result);
+      return result;
+
+  } catch (error) {
+      // Handle any error that is thrown during the execution of the function
+      //(error.code === "ExpiredAuthenticationToken") 
+      if(error) {
+        console.log('Creating new token');
+        const msalInstance = await getMsalInstance();
+    
+        const cachedAccount = msalInstance.getAllAccounts()?.[0];
+        const cachedTenantId = localStorage.getItem("cachedTenantId");
+        
+          
+            msalInstance.setActiveAccount(cachedAccount);
+        
+            const newAccessToken = await acquireTokenWithMsal(msalInstance, {
+              authority: `${configContext.AAD_ENDPOINT}${cachedTenantId}`,
+              scopes: [`${configContext.ARM_ENDPOINT}/.default`],
+            });
+            
+  console.log("Latest ARM Token", userContext.armToken);
+            updateUserContext({armToken: newAccessToken});
+            const result = await fn(...args);
+            return result;
+      }
+      else {
+          console.error('An error occurred:', error.message);
+          throw new error; 
+      }
+      
+  }
+}
+
+// Running the functions using runCommand
+
+const accessToken = 'your-access-token';
+const subscriptionId = 'your-subscription-id';
+
+//runCommand(fetchDatabaseAccountsFromGraph, subscriptionId, accessToken);
+//runCommand(fetchSubscriptionsFromGraph, accessToken);
+
+async function acquireNewTokenAndRetry(body: any) : Promise<DatabaseAccount[]> {
+  try {
+    const msalInstance = await getMsalInstance();
+    
+const cachedAccount = msalInstance.getAllAccounts()?.[0];
+const cachedTenantId = localStorage.getItem("cachedTenantId");
+
+   // const [tenantId, setTenantId] = React.useState<string>(cachedTenantId);
+
+  
+    msalInstance.setActiveAccount(cachedAccount);
+
+    const newAccessToken = await acquireTokenWithMsal(msalInstance, {
+      authority: `${configContext.AAD_ENDPOINT}${cachedTenantId}`,
+      scopes: [`${configContext.ARM_ENDPOINT}/.default`],
+    });
+    console.log("New ARM Token", newAccessToken);
+    const newBearer = `Bearer ${newAccessToken}`;
+    const newHeaders = new Headers();
+    newHeaders.append("Authorization", newBearer);
+    newHeaders.append(HttpHeaders.contentType, "application/json");
+    const apiVersion = "2021-03-01";
+    const managementResourceGraphAPIURL = `${configContext.ARM_ENDPOINT}providers/Microsoft.ResourceGraph/resources?api-version=${apiVersion}`;
+    
+  const databaseAccounts: DatabaseAccount[] = [];
+  let skipToken: string;
+  
+
+    // Retry the request with the new token
+    const response = await fetch(managementResourceGraphAPIURL, {
+      method: "POST",
+      headers: newHeaders,
+      body: JSON.stringify(body),
+    });
+
+    if (response.ok) {
+      // Handle successful response with new token
+      const queryResponse: QueryResponse = await response.json();
+      skipToken = queryResponse.$skipToken;
+      queryResponse.data?.forEach((databaseAccount: any) => {
+        databaseAccounts.push(databaseAccount as DatabaseAccount);
+      });
+      return databaseAccounts;
+    } else {
+      throw new Error(`Failed to fetch data after acquiring new token. Status: ${response.status}, ${await response.text()}`);
+    }
+  } catch (error) {
+    console.error("Error acquiring new token and retrying:", error);
+    throw error;
+  }
+}
+

src/hooks/useKnockoutExplorer.ts

@@ -5,6 +5,7 @@ import { FabricMessageTypes } from "Contracts/FabricMessageTypes";
 import { FABRIC_RPC_VERSION, FabricMessageV2 } from "Contracts/FabricMessagesContract";
 import Explorer from "Explorer/Explorer";
 import { useDataPlaneRbac } from "Explorer/Panes/SettingsPane/SettingsPane";
+import { useDataPlaneRbac } from "Explorer/Panes/SettingsPane/SettingsPane";
 import { useSelectedNode } from "Explorer/useSelectedNode";
 import { scheduleRefreshDatabaseResourceToken } from "Platform/Fabric/FabricUtil";
 import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility";
@@ -18,6 +19,7 @@ import { AuthType } from "../AuthType";
 import { AccountKind, Flights } from "../Common/Constants";
 import { normalizeArmEndpoint } from "../Common/EnvironmentUtility";
 import * as Logger from "../Common/Logger";
+import * as Logger from "../Common/Logger";
 import { handleCachedDataMessage, sendMessage, sendReadyMessage } from "../Common/MessageHandler";
 import { Platform, configContext, updateConfigContext } from "../ConfigContext";
 import { ActionType, DataExplorerAction, TabKind } from "../Contracts/ActionContracts";
@@ -464,6 +466,7 @@ export async function fetchAndUpdateKeys(subscriptionId: string, resourceGroup:
   Logger.logInfo(`Fetching keys for ${userContext.apiType} account ${account}`, "Explorer/fetchAndUpdateKeys");
   let keys;
   try {
+    keys = await listKeys(subscriptionId, resourceGroup, account);
     keys = await listKeys(subscriptionId, resourceGroup, account);
     Logger.logInfo(`Keys fetched for ${userContext.apiType} account ${account}`, "Explorer/fetchAndUpdateKeys");
     updateUserContext({
@@ -487,6 +490,23 @@ export async function fetchAndUpdateKeys(subscriptionId: string, resourceGroup:
       );
       throw error;
     }
+    if (error.code === "AuthorizationFailed") {
+      keys = await getReadOnlyKeys(subscriptionId, resourceGroup, account);
+      Logger.logInfo(
+        `Read only Keys fetched for ${userContext.apiType} account ${account}`,
+        "Explorer/fetchAndUpdateKeys",
+      );
+      updateUserContext({
+        masterKey: keys.primaryReadonlyMasterKey,
+      });
+    } else {
+      logConsoleError(`Error occurred fetching keys for the account." ${error.message}`);
+      Logger.logError(
+        `Error during fetching keys or updating user context: ${error} for ${userContext.apiType} account ${account}`,
+        "Explorer/fetchAndUpdateKeys",
+      );
+      throw error;
+    }
   }
 }
 

src/hooks/useSubscriptions.tsx

@@ -3,6 +3,7 @@ import { QueryRequestOptions, QueryResponse } from "Contracts/AzureResourceGraph
 import useSWR from "swr";
 import { configContext } from "../ConfigContext";
 import { Subscription } from "../Contracts/DataModels";
+import { acquireTokenWithMsal, getMsalInstance } from "Utils/AuthorizationUtils";
 /* eslint-disable  @typescript-eslint/no-explicit-any */
 
 interface SubscriptionListResult {
@@ -92,3 +93,5 @@ export function useSubscriptions(armToken: string): Subscription[] | undefined {
   );
   return data;
 }
+
+