Add support for refreshing access tokens

src/Contracts/FabricMessageTypes.ts

@@ -4,6 +4,7 @@
 export enum FabricMessageTypes {
   GetAuthorizationToken = "GetAuthorizationToken",
   GetAllResourceTokens = "GetAllResourceTokens",
+  GetAccessToken = "GetAccessToken",
   Ready = "Ready",
 }
 

src/Contracts/FabricMessagesContract.ts

@@ -73,6 +73,14 @@ export type FabricMessageV3 =
       message: {
         visible: boolean;
       };
+    }
+  | {
+      type: "accessToken";
+      message: {
+        id: string;
+        error: string | undefined;
+        data: { accessToken: string };
+      };
     };
 
 export enum CosmosDbArtifactType {

src/Explorer/Explorer.tsx

@@ -11,7 +11,7 @@ import { IGalleryItem } from "Juno/JunoClient";
 import {
   isFabricMirrored,
   isFabricMirroredKey,
-  scheduleRefreshDatabaseResourceToken,
+  scheduleRefreshFabricToken,
 } from "Platform/Fabric/FabricUtil";
 import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility";
 import { acquireMsalTokenForAccount } from "Utils/AuthorizationUtils";
@@ -356,7 +356,7 @@ export default class Explorer {
 
   public onRefreshResourcesClick = async (): Promise<void> => {
     if (isFabricMirroredKey()) {
-      scheduleRefreshDatabaseResourceToken(true).then(() => this.refreshAllDatabases());
+      scheduleRefreshFabricToken(true).then(() => this.refreshAllDatabases());
       return;
     }
 

src/Platform/Fabric/FabricUtil.ts

@@ -12,27 +12,66 @@ let timeoutId: NodeJS.Timeout | undefined;
 // Prevents multiple parallel requests during DEBOUNCE_DELAY_MS
 let lastRequestTimestamp: number | undefined = undefined;
 
-const requestDatabaseResourceTokens = async (): Promise<void> => {
+/**
+ * Request fabric token:
+ * - Mirrored key and AAD: Database Resource Tokens
+ * - Native: AAD token
+ * @returns
+ */
+const requestFabricToken = async (): Promise<void> => {
   if (lastRequestTimestamp !== undefined && lastRequestTimestamp + DEBOUNCE_DELAY_MS > Date.now()) {
     return;
   }
 
   if (!userContext.fabricContext || !userContext.databaseAccount) {
+    // This should not happen
+    logConsoleError("Fabric context or database account is missing: cannot request tokens");
     return;
   }
 
   lastRequestTimestamp = Date.now();
   try {
-    const resourceTokenInfo = await sendCachedDataMessage<ResourceTokenInfo>(
-      FabricMessageTypes.GetAllResourceTokens,
-      [],
-      userContext.fabricContext.artifactInfo?.connectionId,
-    );
-
-    if (!userContext.databaseAccount.properties.documentEndpoint) {
-      userContext.databaseAccount.properties.documentEndpoint = resourceTokenInfo.endpoint;
+    if (isFabricMirrored()) {
+      await requestAndStoreDatabaseResourceTokens();
+    } else if (isFabricNative()) {
+      await requestAndStoreAccessToken();
     }
 
+    scheduleRefreshFabricToken();
+  } catch (error) {
+    logConsoleError(error as string);
+    throw error;
+  } finally {
+    lastRequestTimestamp = undefined;
+  }
+};
+
+const requestAndStoreDatabaseResourceTokens = async (): Promise<void> => {
+  const resourceTokenInfo = await sendCachedDataMessage<ResourceTokenInfo>(
+    FabricMessageTypes.GetAllResourceTokens,
+    [],
+    userContext.fabricContext.artifactInfo?.connectionId,
+  );
+
+  if (!userContext.databaseAccount.properties.documentEndpoint) {
+    userContext.databaseAccount.properties.documentEndpoint = resourceTokenInfo.endpoint;
+  }
+
+  if (resourceTokenInfo.credentialType === "OAuth2") {
+    // Mirrored AAD
+    updateUserContext({
+      fabricContext: {
+        ...userContext.fabricContext,
+        databaseName: resourceTokenInfo.databaseId,
+        artifactInfo: undefined,
+        isReadOnly: resourceTokenInfo.isReadOnly ?? userContext.fabricContext.isReadOnly,
+      },
+      databaseAccount: { ...userContext.databaseAccount },
+      aadToken: resourceTokenInfo.accessToken,
+    });
+  } else {
+    // TODO: In Fabric contract V2, credentialType is undefined. For V3, it is "Key". Check for "Key" when V3 is supported for Fabric Mirroring Key
+    // Mirrored key
     updateUserContext({
       fabricContext: {
         ...userContext.fabricContext,
@@ -45,21 +84,27 @@ const requestDatabaseResourceTokens = async (): Promise<void> => {
       },
       databaseAccount: { ...userContext.databaseAccount },
     });
-    scheduleRefreshDatabaseResourceToken();
-  } catch (error) {
-    logConsoleError(error as string);
-    throw error;
-  } finally {
-    lastRequestTimestamp = undefined;
   }
 };
 
+const requestAndStoreAccessToken = async (): Promise<void> => {
+  const accessTokenInfo = await sendCachedDataMessage<{ accessToken: string }>(
+    FabricMessageTypes.GetAccessToken,
+    [],
+    userContext.fabricContext.artifactInfo?.connectionId,
+  );
+
+  updateUserContext({
+    aadToken: accessTokenInfo.accessToken,
+  });
+};
+
 /**
  * Check token validity and schedule a refresh if necessary
  * @param tokenTimestamp
  * @returns
  */
-export const scheduleRefreshDatabaseResourceToken = (refreshNow?: boolean): Promise<void> => {
+export const scheduleRefreshFabricToken = (refreshNow?: boolean): Promise<void> => {
   return new Promise((resolve) => {
     if (timeoutId !== undefined) {
       clearTimeout(timeoutId);
@@ -68,7 +113,7 @@ export const scheduleRefreshDatabaseResourceToken = (refreshNow?: boolean): Prom
 
     timeoutId = setTimeout(
       () => {
-        requestDatabaseResourceTokens().then(resolve);
+        requestFabricToken().then(resolve);
       },
       refreshNow ? 0 : TOKEN_VALIDITY_MS,
     );
@@ -77,7 +122,7 @@ export const scheduleRefreshDatabaseResourceToken = (refreshNow?: boolean): Prom
 
 export const checkDatabaseResourceTokensValidity = (tokenTimestamp: number): void => {
   if (tokenTimestamp + TOKEN_VALIDITY_MS < Date.now()) {
-    scheduleRefreshDatabaseResourceToken(true);
+    scheduleRefreshFabricToken(true);
   }
 };
 

src/hooks/useKnockoutExplorer.ts

@@ -14,7 +14,7 @@ import { useDialog } from "Explorer/Controls/Dialog";
 import Explorer from "Explorer/Explorer";
 import { useDataPlaneRbac } from "Explorer/Panes/SettingsPane/SettingsPane";
 import { useSelectedNode } from "Explorer/useSelectedNode";
-import { scheduleRefreshDatabaseResourceToken } from "Platform/Fabric/FabricUtil";
+import { isFabricMirroredKey, scheduleRefreshFabricToken } from "Platform/Fabric/FabricUtil";
 import {
   AppStateComponentNames,
   OPEN_TABS_SUBCOMPONENT_NAME,
@@ -154,7 +154,7 @@ async function configureFabric(): Promise<Explorer> {
               };
 
               explorer = createExplorerFabricLegacy(initializationMessage, data.version);
-              await scheduleRefreshDatabaseResourceToken(true);
+              await scheduleRefreshFabricToken(true);
               resolve(explorer);
               await explorer.refreshAllDatabases();
               if (userContext.fabricContext.isVisible) {
@@ -169,9 +169,12 @@ async function configureFabric(): Promise<Explorer> {
               if (initializationMessage.artifactType === CosmosDbArtifactType.MIRRORED_KEY) {
                 // Do not show Home tab for Mirrored
                 useTabs.getState().closeReactTab(ReactTabKind.Home);
-                await scheduleRefreshDatabaseResourceToken(true);
               }
 
+              // All tokens used in fabric expire
+              // For Mirrored key, we need the token right away to get the database and containers list.
+              await scheduleRefreshFabricToken(isFabricMirroredKey());
+
               resolve(explorer);
               await explorer.refreshAllDatabases();
 
@@ -188,7 +191,8 @@ async function configureFabric(): Promise<Explorer> {
             explorer.onNewCollectionClicked();
             break;
           case "authorizationToken":
-          case "allResourceTokens_v2": {
+          case "allResourceTokens_v2":
+          case "accessToken": {
             handleCachedDataMessage(data);
             break;
           }