Add code for arm token refresh

This commit is contained in:
Senthamil Sindhu 2024-07-23 11:14:19 -07:00
parent 521ff39eb0
commit 3b86a9477f

View File

@ -3,6 +3,8 @@ import { QueryRequestOptions, QueryResponse } from "Contracts/AzureResourceGraph
import useSWR from "swr"; import useSWR from "swr";
import { configContext } from "../ConfigContext"; import { configContext } from "../ConfigContext";
import { DatabaseAccount } from "../Contracts/DataModels"; import { DatabaseAccount } from "../Contracts/DataModels";
import { acquireTokenWithMsal, getMsalInstance } from "Utils/AuthorizationUtils";
import React from "react";
/* eslint-disable @typescript-eslint/no-explicit-any */ /* eslint-disable @typescript-eslint/no-explicit-any */
interface AccountListResult { interface AccountListResult {
@ -72,14 +74,24 @@ export async function fetchDatabaseAccountsFromGraph(
}); });
if (!response.ok) { if (!response.ok) {
throw new Error(await response.text()); // throw new Error(await response.text());
} // }
const queryResponse: QueryResponse = (await response.json()) as QueryResponse; const queryResponse: QueryResponse = (await response.json()) as QueryResponse;
skipToken = queryResponse.$skipToken; skipToken = queryResponse.$skipToken;
queryResponse.data?.map((databaseAccount: any) => { queryResponse.data?.map((databaseAccount: any) => {
databaseAccounts.push(databaseAccount as DatabaseAccount); databaseAccounts.push(databaseAccount as DatabaseAccount);
}); });
} else {
try{
console.log("Token expired");
await acquireNewTokenAndRetry(body);
}
catch (error) {
throw new Error(error);
}
}
} while (skipToken); } while (skipToken);
return databaseAccounts.sort((a, b) => a.name.localeCompare(b.name)); return databaseAccounts.sort((a, b) => a.name.localeCompare(b.name));
@ -92,3 +104,53 @@ export function useDatabaseAccounts(subscriptionId: string, armToken: string): D
); );
return data; return data;
} }
async function acquireNewTokenAndRetry(body: any) {
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`],
});
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);
});
} 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;
}
}