diff --git a/src/Common/CosmosClient.ts b/src/Common/CosmosClient.ts index 4256c31a3..682ab4e39 100644 --- a/src/Common/CosmosClient.ts +++ b/src/Common/CosmosClient.ts @@ -81,6 +81,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 +124,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 = async (requestContext, diagnosticNode, next) => { requestContext.endpoint = new URL(configContext.PROXY_PATH, window.location.href).href; requestContext.headers["x-ms-proxy-target"] = endpoint(); diff --git a/src/UserContext.ts b/src/UserContext.ts index 4718e3915..fcf690560 100644 --- a/src/UserContext.ts +++ b/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; diff --git a/src/hooks/useKnockoutExplorer.ts b/src/hooks/useKnockoutExplorer.ts index 21461d363..b98341b63 100644 --- a/src/hooks/useKnockoutExplorer.ts +++ b/src/hooks/useKnockoutExplorer.ts @@ -499,10 +499,9 @@ async function configurePortal(): Promise { 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) { @@ -530,6 +529,18 @@ async function configurePortal(): Promise { }); } +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) { // Only allow forwarding messages from the same origin return messageOrigin === window.document.location.origin && message.type === MessageTypes.TelemetryInfo; @@ -567,6 +578,7 @@ function updateContextsFromPortalMessage(inputs: DataExplorerInputsFrame) { collectionCreationDefaults: inputs.defaultCollectionThroughput, isTryCosmosDBSubscription: inputs.isTryCosmosDBSubscription, feedbackPolicies: inputs.feedbackPolicies, + listKeysFetchInProgress: false, }); if (inputs.isPostgresAccount) {