From 9bb1d0baceeb134725b08b695a4cc10d2113d541 Mon Sep 17 00:00:00 2001 From: bogercraig <124094535+bogercraig@users.noreply.github.com> Date: Mon, 7 Apr 2025 09:29:11 -0700 Subject: [PATCH] Manual Region Selection (#2037) * Add standin region selection to settings menu. * Retrieve read and write regions from user context and populate dropdown menu. Update local storage value. Need to now connect with updating read region of primary cosmos client. * Change to only selecting region for cosmos client. Not setting up separate read and write clients. * Add read and write endpoint logging to cosmos client. * Pass changing endpoint from settings menu to client. Encountered token issues using new endpoint in client. * Rough implementation of region selection of endpoint for cosmos client. Still need to: 1 - Use separate context var to track selected region. Directly updating database account context throws off token generation by acquireMSALTokenForAccount 2 - Remove href overrides in acquireMSALTokenForAccount. * Update region selection to include global endpoint and generate a unique list of read and write endpoints. Need to continue with clearing out selected endpoint when global is selected again. Write operations stall when read region is selected even though 403 returned when region rejects operation. Need to limit feature availablility to nosql, table, gremlin (maybe). * Update cosmos client to fix bug. Clients continuously generate after changing RBAC setting. * Swapping back to default endpoint value. * Rebase on client refresh bug fix. * Enable region selection for NoSql, Table, Gremlin * Add logic to reset regional endpoint when global is selected. * Fix state changing when selecting region or resetting to global. * Rough implementation of configuring regional endpoint when DE is loaded in portal or hosted with AAD/Entra auth. * Ininitial attempt at adding error handling, but still having issues with errors caught at proxy plugin. * Added rough error handling in local requestPlugin used in local environments. Passes new error to calling code. Might need to add specific error handling for request plugin to the handleError class. * Change how request plugin returns error so existing error handling utility can process and present error. * Only enable region selection for nosql accounts. * Limit region selection to portal and hosted AAD auth. SQL accounts only. Could possibly enable on table and gremlin later. * Update error handling to account for generic error code. * Refactor error code extraction. * Update test snapshots and remove unneeded logging. * Change error handling to use only the message rather than casting to any. * Clean up debug logging in cosmos client. * Remove unused storage keys. * Use endpoint instead of region name to track selected region. Prevents having to do endpoint lookups. * Add initial button state update depending on region selection. Need to update with the API and react to user context changes. * Disable CRUD buttons when read region selected. * Default to write enabled in react. * Disable query saving when read region is selected. * Patch clientWidth error on conflicts tab. * Resolve merge conflicts from rebase. * Make sure proxy endpoints return in all cases. * Remove excess client logging and match main for ConflictsTab. * Cleaning up logging and fixing endpoint discovery bug. * Fix formatting. * Reformatting if statements with preferred formatting. * Migrate region selection to local persistence. Fixes account swapping bug. TODO: Inspect better way to reset interface elements when deleteAllStates is called. Need to react to regional endpoint being reset. * Relocate resetting interface context to helper function. * Remove legacy state storage for regional endpoint selection. * Laurent suggestion updates. --- src/Common/CosmosClient.ts | 7 +- .../Panes/SettingsPane/SettingsPane.tsx | 168 ++++++++++++++++-- .../__snapshots__/SettingsPane.test.tsx.snap | 20 +-- .../Tabs/DocumentsTabV2/DocumentsTabV2.tsx | 13 +- .../Tabs/QueryTab/QueryTabComponent.tsx | 14 +- src/Shared/AppStatePersistenceUtility.ts | 1 + src/UserContext.ts | 2 + src/hooks/useClientWriteEnabled.ts | 10 ++ src/hooks/useKnockoutExplorer.ts | 46 +++++ 9 files changed, 256 insertions(+), 25 deletions(-) create mode 100644 src/hooks/useClientWriteEnabled.ts diff --git a/src/Common/CosmosClient.ts b/src/Common/CosmosClient.ts index cf34b2279..1ecd94944 100644 --- a/src/Common/CosmosClient.ts +++ b/src/Common/CosmosClient.ts @@ -125,7 +125,11 @@ export const endpoint = () => { const location = _global.parent ? _global.parent.location : _global.location; return configContext.EMULATOR_ENDPOINT || location.origin; } - return userContext.endpoint || userContext?.databaseAccount?.properties?.documentEndpoint; + return ( + userContext.selectedRegionalEndpoint || + userContext.endpoint || + userContext?.databaseAccount?.properties?.documentEndpoint + ); }; export async function getTokenFromAuthService( @@ -203,6 +207,7 @@ export function client(): Cosmos.CosmosClient { userAgentSuffix: "Azure Portal", defaultHeaders: _defaultHeaders, connectionPolicy: { + enableEndpointDiscovery: !userContext.selectedRegionalEndpoint, retryOptions: { maxRetryAttemptCount: LocalStorageUtility.getEntryNumber(StorageKey.RetryAttempts), fixedRetryIntervalInMilliseconds: LocalStorageUtility.getEntryNumber(StorageKey.RetryInterval), diff --git a/src/Explorer/Panes/SettingsPane/SettingsPane.tsx b/src/Explorer/Panes/SettingsPane/SettingsPane.tsx index 1626f09d1..a40f4da99 100644 --- a/src/Explorer/Panes/SettingsPane/SettingsPane.tsx +++ b/src/Explorer/Panes/SettingsPane/SettingsPane.tsx @@ -6,7 +6,9 @@ import { Checkbox, ChoiceGroup, DefaultButton, + Dropdown, IChoiceGroupOption, + IDropdownOption, ISpinButtonStyles, IToggleStyles, Position, @@ -21,7 +23,15 @@ import { InfoTooltip } from "Common/Tooltip/InfoTooltip"; import { Platform, configContext } from "ConfigContext"; import { useDialog } from "Explorer/Controls/Dialog"; import { useDatabases } from "Explorer/useDatabases"; -import { deleteAllStates } from "Shared/AppStatePersistenceUtility"; +import { isFabric } from "Platform/Fabric/FabricUtil"; +import { + AppStateComponentNames, + deleteAllStates, + deleteState, + hasState, + loadState, + saveState, +} from "Shared/AppStatePersistenceUtility"; import { DefaultRUThreshold, LocalStorageUtility, @@ -37,6 +47,7 @@ import { acquireMsalTokenForAccount } from "Utils/AuthorizationUtils"; import { logConsoleError, logConsoleInfo } from "Utils/NotificationConsoleUtils"; import * as PriorityBasedExecutionUtils from "Utils/PriorityBasedExecutionUtils"; import { getReadOnlyKeys, listKeys } from "Utils/arm/generatedClients/cosmos/databaseAccounts"; +import { useClientWriteEnabled } from "hooks/useClientWriteEnabled"; import { useQueryCopilot } from "hooks/useQueryCopilot"; import { useSidePanel } from "hooks/useSidePanel"; import React, { FunctionComponent, useState } from "react"; @@ -143,6 +154,17 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({ ? LocalStorageUtility.getEntryString(StorageKey.IsGraphAutoVizDisabled) : "false", ); + const [selectedRegionalEndpoint, setSelectedRegionalEndpoint] = useState( + hasState({ + componentName: AppStateComponentNames.SelectedRegionalEndpoint, + globalAccountName: userContext.databaseAccount?.name, + }) + ? (loadState({ + componentName: AppStateComponentNames.SelectedRegionalEndpoint, + globalAccountName: userContext.databaseAccount?.name, + }) as string) + : undefined, + ); const [retryAttempts, setRetryAttempts] = useState( LocalStorageUtility.hasItem(StorageKey.RetryAttempts) ? LocalStorageUtility.getEntryNumber(StorageKey.RetryAttempts) @@ -189,6 +211,44 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({ configContext.platform !== Platform.Fabric && !isEmulator; const shouldShowPriorityLevelOption = PriorityBasedExecutionUtils.isFeatureEnabled() && !isEmulator; + + const uniqueAccountRegions = new Set(); + const regionOptions: IDropdownOption[] = []; + regionOptions.push({ + key: userContext?.databaseAccount?.properties?.documentEndpoint, + text: `Global (Default)`, + data: { + isGlobal: true, + writeEnabled: true, + }, + }); + userContext?.databaseAccount?.properties?.writeLocations?.forEach((loc) => { + if (!uniqueAccountRegions.has(loc.locationName)) { + uniqueAccountRegions.add(loc.locationName); + regionOptions.push({ + key: loc.documentEndpoint, + text: `${loc.locationName} (Read/Write)`, + data: { + isGlobal: false, + writeEnabled: true, + }, + }); + } + }); + userContext?.databaseAccount?.properties?.readLocations?.forEach((loc) => { + if (!uniqueAccountRegions.has(loc.locationName)) { + uniqueAccountRegions.add(loc.locationName); + regionOptions.push({ + key: loc.documentEndpoint, + text: `${loc.locationName} (Read)`, + data: { + isGlobal: false, + writeEnabled: false, + }, + }); + } + }); + const shouldShowCopilotSampleDBOption = userContext.apiType === "SQL" && useQueryCopilot.getState().copilotEnabled && @@ -274,6 +334,46 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({ } } + const storedRegionalEndpoint = loadState({ + componentName: AppStateComponentNames.SelectedRegionalEndpoint, + globalAccountName: userContext.databaseAccount?.name, + }) as string; + const selectedRegionIsGlobal = + selectedRegionalEndpoint === userContext?.databaseAccount?.properties?.documentEndpoint; + if (selectedRegionIsGlobal && storedRegionalEndpoint) { + deleteState({ + componentName: AppStateComponentNames.SelectedRegionalEndpoint, + globalAccountName: userContext.databaseAccount?.name, + }); + updateUserContext({ + selectedRegionalEndpoint: undefined, + writeEnabledInSelectedRegion: true, + refreshCosmosClient: true, + }); + useClientWriteEnabled.setState({ clientWriteEnabled: true }); + } else if ( + selectedRegionalEndpoint && + !selectedRegionIsGlobal && + selectedRegionalEndpoint !== storedRegionalEndpoint + ) { + saveState( + { + componentName: AppStateComponentNames.SelectedRegionalEndpoint, + globalAccountName: userContext.databaseAccount?.name, + }, + selectedRegionalEndpoint, + ); + const validWriteEndpoint = userContext.databaseAccount?.properties?.writeLocations?.find( + (loc) => loc.documentEndpoint === selectedRegionalEndpoint, + ); + updateUserContext({ + selectedRegionalEndpoint: selectedRegionalEndpoint, + writeEnabledInSelectedRegion: !!validWriteEndpoint, + refreshCosmosClient: true, + }); + useClientWriteEnabled.setState({ clientWriteEnabled: !!validWriteEndpoint }); + } + LocalStorageUtility.setEntryBoolean(StorageKey.RUThresholdEnabled, ruThresholdEnabled); LocalStorageUtility.setEntryBoolean(StorageKey.QueryTimeoutEnabled, queryTimeoutEnabled); LocalStorageUtility.setEntryNumber(StorageKey.RetryAttempts, retryAttempts); @@ -423,6 +523,10 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({ setDefaultQueryResultsView(option.key as SplitterDirection); }; + const handleOnSelectedRegionOptionChange = (ev: React.FormEvent, option: IDropdownOption): void => { + setSelectedRegionalEndpoint(option.key as string); + }; + const handleOnQueryRetryAttemptsSpinButtonChange = (ev: React.MouseEvent, newValue?: string): void => { const retryAttempts = Number(newValue); if (!isNaN(retryAttempts)) { @@ -583,9 +687,39 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({ )} + {userContext.apiType === "SQL" && userContext.authType === AuthType.AAD && !isFabric() && ( + + +
Region Selection
+ + +
+
+ Changes region the Cosmos Client uses to access account. +
+
+ Select Region + + Changes the account endpoint used to perform client operations. + +
+ option.key === selectedRegionalEndpoint)?.text + : regionOptions[0]?.text + } + onChange={handleOnSelectedRegionOptionChange} + options={regionOptions} + styles={{ root: { marginBottom: "10px" } }} + /> +
+ + + )} {userContext.apiType === "SQL" && !isEmulator && ( <> - +
Query Timeout
 @@ -626,7 +760,7 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({ - +
RU Limit
 @@ -660,7 +794,7 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({ - +
Default Query Results View
 @@ -681,8 +815,9 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({ )} + {showRetrySettings && ( - +
Retry Settings
 @@ -755,7 +890,7 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({ )} {!isEmulator && ( - +
Enable container pagination
 @@ -779,7 +914,7 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({ )} {shouldShowCrossPartitionOption && ( - +
Enable cross-partition query
 @@ -804,7 +939,7 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({ )} {shouldShowParallelismOption && ( - +
Max degree of parallelism
 @@ -837,7 +972,7 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({ )} {shouldShowPriorityLevelOption && ( - +
Priority Level
 @@ -860,7 +995,7 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({ )} {shouldShowGraphAutoVizOption && ( - +
Display Gremlin query results as: 
 @@ -881,7 +1016,7 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({ )} {shouldShowCopilotSampleDBOption && ( - +
Enable sample database
 @@ -916,7 +1051,15 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({ "Clear History", undefined, "Are you sure you want to proceed?", - () => deleteAllStates(), + () => { + deleteAllStates(); + updateUserContext({ + selectedRegionalEndpoint: undefined, + writeEnabledInSelectedRegion: true, + refreshCosmosClient: true, + }); + useClientWriteEnabled.setState({ clientWriteEnabled: true }); + }, "Cancel", undefined, <> @@ -927,6 +1070,7 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
  • Reset your customized tab layout, including the splitter positions
  • Erase your table column preferences, including any custom columns
  • Clear your filter history
    • +
  • Reset region selection to global
    • , ); diff --git a/src/Explorer/Panes/SettingsPane/__snapshots__/SettingsPane.test.tsx.snap b/src/Explorer/Panes/SettingsPane/__snapshots__/SettingsPane.test.tsx.snap index e89ee345b..577b6de5b 100644 --- a/src/Explorer/Panes/SettingsPane/__snapshots__/SettingsPane.test.tsx.snap +++ b/src/Explorer/Panes/SettingsPane/__snapshots__/SettingsPane.test.tsx.snap @@ -107,7 +107,7 @@ exports[`Settings Pane should render Default properly 1`] = `
    ; editorState: ViewModels.DocumentExplorerState; isPreferredApiMongoDB: boolean; + clientWriteEnabled: boolean; onNewDocumentClick: UiKeyboardEvent; onSaveNewDocumentClick: UiKeyboardEvent; onRevertNewDocumentClick: UiKeyboardEvent; @@ -328,6 +330,7 @@ const createUploadButton = (container: Explorer): CommandButtonComponentProps => hasPopup: true, disabled: useSelectedNode.getState().isDatabaseNodeOrNoneSelected() || + !useClientWriteEnabled.getState().clientWriteEnabled || useSelectedNode.getState().isQueryCopilotCollectionSelected(), }; }; @@ -346,6 +349,7 @@ export const getTabsButtons = ({ selectedRows, editorState, isPreferredApiMongoDB, + clientWriteEnabled, onNewDocumentClick, onSaveNewDocumentClick, onRevertNewDocumentClick, @@ -371,6 +375,7 @@ export const getTabsButtons = ({ hasPopup: false, disabled: !getNewDocumentButtonState(editorState).enabled || + !clientWriteEnabled || useSelectedNode.getState().isQueryCopilotCollectionSelected(), id: NEW_DOCUMENT_BUTTON_ID, }); @@ -388,6 +393,7 @@ export const getTabsButtons = ({ hasPopup: false, disabled: !getSaveNewDocumentButtonState(editorState).enabled || + !clientWriteEnabled || useSelectedNode.getState().isQueryCopilotCollectionSelected(), id: SAVE_BUTTON_ID, }); @@ -422,6 +428,7 @@ export const getTabsButtons = ({ hasPopup: false, disabled: !getSaveExistingDocumentButtonState(editorState).enabled || + !clientWriteEnabled || useSelectedNode.getState().isQueryCopilotCollectionSelected(), id: UPDATE_BUTTON_ID, }); @@ -454,7 +461,7 @@ export const getTabsButtons = ({ commandButtonLabel: label, ariaLabel: label, hasPopup: false, - disabled: useSelectedNode.getState().isQueryCopilotCollectionSelected(), + disabled: useSelectedNode.getState().isQueryCopilotCollectionSelected() || !clientWriteEnabled, id: DELETE_BUTTON_ID, }); } @@ -628,6 +635,7 @@ export const DocumentsTabComponent: React.FunctionComponent state.clientWriteEnabled); const [tabStateData, setTabStateData] = useState(() => readDocumentsTabSubComponentState(SubComponentName.MainTabDivider, _collection, { leftPaneWidthPercent: 35, @@ -865,6 +873,7 @@ export const DocumentsTabComponent: React.FunctionComponent void; + private unsubscribeClientWriteEnabled: () => void; componentDidMount(): void { useTabs.subscribe((state: TabsState) => { @@ -712,10 +717,17 @@ class QueryTabComponentImpl extends React.Component { + useCommandBar.getState().setContextButtons(this.getTabsButtons()); + }); } componentWillUnmount(): void { document.removeEventListener("keydown", this.handleCopilotKeyDown); + if (this.unsubscribeClientWriteEnabled) { + this.unsubscribeClientWriteEnabled(); + } } private getEditorAndQueryResult(): JSX.Element { diff --git a/src/Shared/AppStatePersistenceUtility.ts b/src/Shared/AppStatePersistenceUtility.ts index 58a02c3b3..ed354ef06 100644 --- a/src/Shared/AppStatePersistenceUtility.ts +++ b/src/Shared/AppStatePersistenceUtility.ts @@ -10,6 +10,7 @@ export enum AppStateComponentNames { MostRecentActivity = "MostRecentActivity", QueryCopilot = "QueryCopilot", DataExplorerAction = "DataExplorerAction", + SelectedRegionalEndpoint = "SelectedRegionalEndpoint", } // Subcomponent for DataExplorerAction diff --git a/src/UserContext.ts b/src/UserContext.ts index 6569d5e18..8a880f498 100644 --- a/src/UserContext.ts +++ b/src/UserContext.ts @@ -111,6 +111,8 @@ export interface UserContext { readonly isReplica?: boolean; collectionCreationDefaults: CollectionCreationDefaults; sampleDataConnectionInfo?: ParsedResourceTokenConnectionString; + readonly selectedRegionalEndpoint?: string; + readonly writeEnabledInSelectedRegion?: boolean; readonly vcoreMongoConnectionParams?: VCoreMongoConnectionParams; readonly feedbackPolicies?: AdminFeedbackPolicySettings; readonly dataPlaneRbacEnabled?: boolean; diff --git a/src/hooks/useClientWriteEnabled.ts b/src/hooks/useClientWriteEnabled.ts new file mode 100644 index 000000000..7d9d29c2e --- /dev/null +++ b/src/hooks/useClientWriteEnabled.ts @@ -0,0 +1,10 @@ +import create, { UseStore } from "zustand"; +interface ClientWriteEnabledState { + clientWriteEnabled: boolean; + setClientWriteEnabled: (writeEnabled: boolean) => void; +} + +export const useClientWriteEnabled: UseStore = create((set) => ({ + clientWriteEnabled: true, + setClientWriteEnabled: (clientWriteEnabled: boolean) => set({ clientWriteEnabled }), +})); diff --git a/src/hooks/useKnockoutExplorer.ts b/src/hooks/useKnockoutExplorer.ts index 2e29d1363..71818f572 100644 --- a/src/hooks/useKnockoutExplorer.ts +++ b/src/hooks/useKnockoutExplorer.ts @@ -17,12 +17,16 @@ import { useSelectedNode } from "Explorer/useSelectedNode"; import { isFabricMirroredKey, scheduleRefreshFabricToken } from "Platform/Fabric/FabricUtil"; import { AppStateComponentNames, + deleteState, + hasState, + loadState, OPEN_TABS_SUBCOMPONENT_NAME, readSubComponentState, } from "Shared/AppStatePersistenceUtility"; import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility"; import { isDataplaneRbacSupported } from "Utils/APITypeUtils"; import { logConsoleError } from "Utils/NotificationConsoleUtils"; +import { useClientWriteEnabled } from "hooks/useClientWriteEnabled"; import { useQueryCopilot } from "hooks/useQueryCopilot"; import { ReactTabKind, useTabs } from "hooks/useTabs"; import { useEffect, useState } from "react"; @@ -345,6 +349,9 @@ async function configureHostedWithAAD(config: AAD): Promise { `Configuring Data Explorer for ${userContext.apiType} account ${account.name}`, "Explorer/configureHostedWithAAD", ); + if (userContext.apiType === "SQL") { + checkAndUpdateSelectedRegionalEndpoint(); + } if (!userContext.features.enableAadDataPlane) { Logger.logInfo(`AAD Feature flag is not enabled for account ${account.name}`, "Explorer/configureHostedWithAAD"); if (isDataplaneRbacSupported(userContext.apiType)) { @@ -706,6 +713,10 @@ async function configurePortal(): Promise { const { databaseAccount: account, subscriptionId, resourceGroup } = userContext; + if (userContext.apiType === "SQL") { + checkAndUpdateSelectedRegionalEndpoint(); + } + let dataPlaneRbacEnabled; if (isDataplaneRbacSupported(userContext.apiType)) { if (LocalStorageUtility.hasItem(StorageKey.DataPlaneRbacEnabled)) { @@ -824,6 +835,41 @@ function updateAADEndpoints(portalEnv: PortalEnv) { } } +function checkAndUpdateSelectedRegionalEndpoint() { + const accountName = userContext.databaseAccount?.name; + if (hasState({ componentName: AppStateComponentNames.SelectedRegionalEndpoint, globalAccountName: accountName })) { + const storedRegionalEndpoint = loadState({ + componentName: AppStateComponentNames.SelectedRegionalEndpoint, + globalAccountName: accountName, + }) as string; + const validEndpoint = userContext.databaseAccount?.properties?.readLocations?.find( + (loc) => loc.documentEndpoint === storedRegionalEndpoint, + ); + const validWriteEndpoint = userContext.databaseAccount?.properties?.writeLocations?.find( + (loc) => loc.documentEndpoint === storedRegionalEndpoint, + ); + if (validEndpoint) { + updateUserContext({ + selectedRegionalEndpoint: storedRegionalEndpoint, + writeEnabledInSelectedRegion: !!validWriteEndpoint, + refreshCosmosClient: true, + }); + useClientWriteEnabled.setState({ clientWriteEnabled: !!validWriteEndpoint }); + } else { + deleteState({ componentName: AppStateComponentNames.SelectedRegionalEndpoint, globalAccountName: accountName }); + updateUserContext({ + writeEnabledInSelectedRegion: true, + }); + useClientWriteEnabled.setState({ clientWriteEnabled: true }); + } + } else { + updateUserContext({ + writeEnabledInSelectedRegion: true, + }); + useClientWriteEnabled.setState({ clientWriteEnabled: true }); + } +} + function updateContextsFromPortalMessage(inputs: DataExplorerInputsFrame) { if ( configContext.PORTAL_BACKEND_ENDPOINT &&