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

@@ -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),

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<string>(
+    hasState({
+      componentName: AppStateComponentNames.SelectedRegionalEndpoint,
+      globalAccountName: userContext.databaseAccount?.name,
+    })
+      ? (loadState({
+          componentName: AppStateComponentNames.SelectedRegionalEndpoint,
+          globalAccountName: userContext.databaseAccount?.name,
+        }) as string)
+      : undefined,
+  );
   const [retryAttempts, setRetryAttempts] = useState<number>(
     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<string>();
+  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<HTMLInputElement>, option: IDropdownOption): void => {
+    setSelectedRegionalEndpoint(option.key as string);
+  };
+
   const handleOnQueryRetryAttemptsSpinButtonChange = (ev: React.MouseEvent<HTMLElement>, newValue?: string): void => {
     const retryAttempts = Number(newValue);
     if (!isNaN(retryAttempts)) {
@@ -583,9 +687,39 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
               </AccordionPanel>
             </AccordionItem>
           )}
+          {userContext.apiType === "SQL" && userContext.authType === AuthType.AAD && !isFabric() && (
+            <AccordionItem value="3">
+              <AccordionHeader>
+                <div className={styles.header}>Region Selection</div>
+              </AccordionHeader>
+              <AccordionPanel>
+                <div className={styles.settingsSectionContainer}>
+                  <div className={styles.settingsSectionDescription}>
+                    Changes region the Cosmos Client uses to access account.
+                  </div>
+                  <div>
+                    <span className={styles.subHeader}>Select Region</span>
+                    <InfoTooltip className={styles.headerIcon}>
+                      Changes the account endpoint used to perform client operations.
+                    </InfoTooltip>
+                  </div>
+                  <Dropdown
+                    placeholder={
+                      selectedRegionalEndpoint
+                        ? regionOptions.find((option) => option.key === selectedRegionalEndpoint)?.text
+                        : regionOptions[0]?.text
+                    }
+                    onChange={handleOnSelectedRegionOptionChange}
+                    options={regionOptions}
+                    styles={{ root: { marginBottom: "10px" } }}
+                  />
+                </div>
+              </AccordionPanel>
+            </AccordionItem>
+          )}
           {userContext.apiType === "SQL" && !isEmulator && (
             <>
-              <AccordionItem value="3">
+              <AccordionItem value="4">
                 <AccordionHeader>
                   <div className={styles.header}>Query Timeout</div>
                 </AccordionHeader>
@@ -626,7 +760,7 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
                 </AccordionPanel>
               </AccordionItem>
 
-              <AccordionItem value="4">
+              <AccordionItem value="5">
                 <AccordionHeader>
                   <div className={styles.header}>RU Limit</div>
                 </AccordionHeader>
@@ -660,7 +794,7 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
                 </AccordionPanel>
               </AccordionItem>
 
-              <AccordionItem value="5">
+              <AccordionItem value="6">
                 <AccordionHeader>
                   <div className={styles.header}>Default Query Results View</div>
                 </AccordionHeader>
@@ -681,8 +815,9 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
               </AccordionItem>
             </>
           )}
+
           {showRetrySettings && (
-            <AccordionItem value="6">
+            <AccordionItem value="7">
               <AccordionHeader>
                 <div className={styles.header}>Retry Settings</div>
               </AccordionHeader>
@@ -755,7 +890,7 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
             </AccordionItem>
           )}
           {!isEmulator && (
-            <AccordionItem value="7">
+            <AccordionItem value="8">
               <AccordionHeader>
                 <div className={styles.header}>Enable container pagination</div>
               </AccordionHeader>
@@ -779,7 +914,7 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
             </AccordionItem>
           )}
           {shouldShowCrossPartitionOption && (
-            <AccordionItem value="8">
+            <AccordionItem value="9">
               <AccordionHeader>
                 <div className={styles.header}>Enable cross-partition query</div>
               </AccordionHeader>
@@ -804,7 +939,7 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
             </AccordionItem>
           )}
           {shouldShowParallelismOption && (
-            <AccordionItem value="9">
+            <AccordionItem value="10">
               <AccordionHeader>
                 <div className={styles.header}>Max degree of parallelism</div>
               </AccordionHeader>
@@ -837,7 +972,7 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
             </AccordionItem>
           )}
           {shouldShowPriorityLevelOption && (
-            <AccordionItem value="10">
+            <AccordionItem value="11">
               <AccordionHeader>
                 <div className={styles.header}>Priority Level</div>
               </AccordionHeader>
@@ -860,7 +995,7 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
             </AccordionItem>
           )}
           {shouldShowGraphAutoVizOption && (
-            <AccordionItem value="11">
+            <AccordionItem value="12">
               <AccordionHeader>
                 <div className={styles.header}>Display Gremlin query results as:&nbsp;</div>
               </AccordionHeader>
@@ -881,7 +1016,7 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
             </AccordionItem>
           )}
           {shouldShowCopilotSampleDBOption && (
-            <AccordionItem value="12">
+            <AccordionItem value="13">
               <AccordionHeader>
                 <div className={styles.header}>Enable sample database</div>
               </AccordionHeader>
@@ -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 }> = ({
                       <li>Reset your customized tab layout, including the splitter positions</li>
                       <li>Erase your table column preferences, including any custom columns</li>
                       <li>Clear your filter history</li>
+                      <li>Reset region selection to global</li>
                     </ul>
                   </>,
                 );

src/Explorer/Panes/SettingsPane/__snapshots__/SettingsPane.test.tsx.snap

@@ -107,7 +107,7 @@ exports[`Settings Pane should render Default properly 1`] = `
         </AccordionPanel>
       </AccordionItem>
       <AccordionItem
-        value="3"
+        value="4"
       >
         <AccordionHeader>
           <div
@@ -148,7 +148,7 @@ exports[`Settings Pane should render Default properly 1`] = `
         </AccordionPanel>
       </AccordionItem>
       <AccordionItem
-        value="4"
+        value="5"
       >
         <AccordionHeader>
           <div
@@ -219,7 +219,7 @@ exports[`Settings Pane should render Default properly 1`] = `
         </AccordionPanel>
       </AccordionItem>
       <AccordionItem
-        value="5"
+        value="6"
       >
         <AccordionHeader>
           <div
@@ -281,7 +281,7 @@ exports[`Settings Pane should render Default properly 1`] = `
         </AccordionPanel>
       </AccordionItem>
       <AccordionItem
-        value="6"
+        value="7"
       >
         <AccordionHeader>
           <div
@@ -423,7 +423,7 @@ exports[`Settings Pane should render Default properly 1`] = `
         </AccordionPanel>
       </AccordionItem>
       <AccordionItem
-        value="7"
+        value="8"
       >
         <AccordionHeader>
           <div
@@ -459,7 +459,7 @@ exports[`Settings Pane should render Default properly 1`] = `
         </AccordionPanel>
       </AccordionItem>
       <AccordionItem
-        value="8"
+        value="9"
       >
         <AccordionHeader>
           <div
@@ -495,7 +495,7 @@ exports[`Settings Pane should render Default properly 1`] = `
         </AccordionPanel>
       </AccordionItem>
       <AccordionItem
-        value="9"
+        value="10"
       >
         <AccordionHeader>
           <div
@@ -575,7 +575,7 @@ exports[`Settings Pane should render Gremlin properly 1`] = `
       className="customAccordion ___1uf6361_0000000 fz7g6wx"
     >
       <AccordionItem
-        value="6"
+        value="7"
       >
         <AccordionHeader>
           <div
@@ -717,7 +717,7 @@ exports[`Settings Pane should render Gremlin properly 1`] = `
         </AccordionPanel>
       </AccordionItem>
       <AccordionItem
-        value="7"
+        value="8"
       >
         <AccordionHeader>
           <div
@@ -753,7 +753,7 @@ exports[`Settings Pane should render Gremlin properly 1`] = `
         </AccordionPanel>
       </AccordionItem>
       <AccordionItem
-        value="11"
+        value="12"
       >
         <AccordionHeader>
           <div

src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2.tsx

@@ -49,6 +49,7 @@ import { Action } from "Shared/Telemetry/TelemetryConstants";
 import { userContext } from "UserContext";
 import { logConsoleError, logConsoleInfo } from "Utils/NotificationConsoleUtils";
 import { Allotment } from "allotment";
+import { useClientWriteEnabled } from "hooks/useClientWriteEnabled";
 import React, { KeyboardEventHandler, useCallback, useEffect, useMemo, useRef, useState } from "react";
 import { format } from "react-string-format";
 import DeleteDocumentIcon from "../../../../images/DeleteDocument.svg";
@@ -305,6 +306,7 @@ export type ButtonsDependencies = {
   selectedRows: Set<TableRowId>;
   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<IDocumentsTabCompone
   );
 
   // State
+  const clientWriteEnabled = useClientWriteEnabled((state) => state.clientWriteEnabled);
   const [tabStateData, setTabStateData] = useState<TabDivider>(() =>
     readDocumentsTabSubComponentState<TabDivider>(SubComponentName.MainTabDivider, _collection, {
       leftPaneWidthPercent: 35,
@@ -865,6 +873,7 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
       selectedRows,
       editorState,
       isPreferredApiMongoDB,
+      clientWriteEnabled,
       onNewDocumentClick,
       onSaveNewDocumentClick,
       onRevertNewDocumentClick,
@@ -1284,6 +1293,7 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
         selectedRows,
         editorState,
         isPreferredApiMongoDB,
+        clientWriteEnabled,
         onNewDocumentClick,
         onSaveNewDocumentClick,
         onRevertNewDocumentClick,
@@ -1296,6 +1306,7 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
       selectedRows,
       editorState,
       isPreferredApiMongoDB,
+      clientWriteEnabled,
       onNewDocumentClick,
       onSaveNewDocumentClick,
       onRevertNewDocumentClick,

src/Explorer/Tabs/QueryTab/QueryTabComponent.tsx

@@ -1,6 +1,7 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 /* eslint-disable no-console */
 import { FeedOptions, QueryOperationOptions } from "@azure/cosmos";
+import { AuthType } from "AuthType";
 import QueryError, { createMonacoErrorLocationResolver, createMonacoMarkersForQueryErrors } from "Common/QueryError";
 import { SplitterDirection } from "Common/Splitter";
 import { Platform, configContext } from "ConfigContext";
@@ -21,6 +22,7 @@ import { QueryConstants } from "Shared/Constants";
 import { LocalStorageUtility, StorageKey, getRUThreshold, ruThresholdEnabled } from "Shared/StorageUtility";
 import { Action } from "Shared/Telemetry/TelemetryConstants";
 import { Allotment } from "allotment";
+import { useClientWriteEnabled } from "hooks/useClientWriteEnabled";
 import { QueryCopilotState, useQueryCopilot } from "hooks/useQueryCopilot";
 import { TabsState, useTabs } from "hooks/useTabs";
 import React, { Fragment, createRef } from "react";
@@ -484,7 +486,9 @@ class QueryTabComponentImpl extends React.Component<QueryTabComponentImplProps,
           commandButtonLabel: label,
           ariaLabel: label,
           hasPopup: false,
-          disabled: !this.saveQueryButton.enabled,
+          disabled:
+            !this.saveQueryButton.enabled ||
+            (!useClientWriteEnabled.getState().clientWriteEnabled && userContext.authType === AuthType.AAD),
         });
       }
 
@@ -696,6 +700,7 @@ class QueryTabComponentImpl extends React.Component<QueryTabComponentImplProps,
   }
 
   private unsubscribeCopilotSidebar: () => void;
+  private unsubscribeClientWriteEnabled: () => void;
 
   componentDidMount(): void {
     useTabs.subscribe((state: TabsState) => {
@@ -712,10 +717,17 @@ class QueryTabComponentImpl extends React.Component<QueryTabComponentImplProps,
 
     useCommandBar.getState().setContextButtons(this.getTabsButtons());
     document.addEventListener("keydown", this.handleCopilotKeyDown);
+
+    this.unsubscribeClientWriteEnabled = useClientWriteEnabled.subscribe(() => {
+      useCommandBar.getState().setContextButtons(this.getTabsButtons());
+    });
   }
 
   componentWillUnmount(): void {
     document.removeEventListener("keydown", this.handleCopilotKeyDown);
+    if (this.unsubscribeClientWriteEnabled) {
+      this.unsubscribeClientWriteEnabled();
+    }
   }
 
   private getEditorAndQueryResult(): JSX.Element {

src/Shared/AppStatePersistenceUtility.ts

@@ -10,6 +10,7 @@ export enum AppStateComponentNames {
   MostRecentActivity = "MostRecentActivity",
   QueryCopilot = "QueryCopilot",
   DataExplorerAction = "DataExplorerAction",
+  SelectedRegionalEndpoint = "SelectedRegionalEndpoint",
 }
 
 // Subcomponent for DataExplorerAction

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;

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<ClientWriteEnabledState> = create((set) => ({
+  clientWriteEnabled: true,
+  setClientWriteEnabled: (clientWriteEnabled: boolean) => set({ clientWriteEnabled }),
+}));

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<Explorer> {
       `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<Explorer> {
 
           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 &&