RU Threshold (#1728)

* ru threshold beta

* use new ru threshold package

* fix typo

* fix merge issue

* fix package-lock.json

* fix test

* fixed settings pane test

* fixed merge issue

* sync with main

* fixed settings pane check

* fix checks

* fixed aria-label error

* fixed aria-label error

* fixed aria-label error

* fixed aria-label error

* remove learn more

---------

Co-authored-by: Asier Isayas <aisayas@microsoft.com>

src/Explorer/Panes/SettingsPane/SettingsPane.tsx

@@ -11,7 +11,13 @@ import {
 import * as Constants from "Common/Constants";
 import { InfoTooltip } from "Common/Tooltip/InfoTooltip";
 import { configContext } from "ConfigContext";
-import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility";
+import {
+  DefaultRUThreshold,
+  LocalStorageUtility,
+  StorageKey,
+  getRUThreshold,
+  ruThresholdEnabled as isRUThresholdEnabled,
+} from "Shared/StorageUtility";
 import * as StringUtility from "Shared/StringUtility";
 import { userContext } from "UserContext";
 import { logConsoleInfo } from "Utils/NotificationConsoleUtils";
@@ -35,6 +41,8 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
       ? Constants.Queries.UnlimitedPageOption
       : Constants.Queries.CustomPageOption,
   );
+  const [ruThresholdEnabled, setRUThresholdEnabled] = useState<boolean>(isRUThresholdEnabled());
+  const [ruThreshold, setRUThreshold] = useState<number>(getRUThreshold());
   const [queryTimeoutEnabled, setQueryTimeoutEnabled] = useState<boolean>(
     LocalStorageUtility.getEntryBoolean(StorageKey.QueryTimeoutEnabled),
   );
@@ -103,6 +111,7 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
       isCustomPageOptionSelected() ? customItemPerPage : Constants.Queries.unlimitedItemsPerPage,
     );
     LocalStorageUtility.setEntryNumber(StorageKey.CustomItemPerPage, customItemPerPage);
+    LocalStorageUtility.setEntryBoolean(StorageKey.RUThresholdEnabled, ruThresholdEnabled);
     LocalStorageUtility.setEntryBoolean(StorageKey.QueryTimeoutEnabled, queryTimeoutEnabled);
     LocalStorageUtility.setEntryNumber(StorageKey.RetryAttempts, retryAttempts);
     LocalStorageUtility.setEntryNumber(StorageKey.RetryInterval, retryInterval);
@@ -120,6 +129,10 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
       );
     }
 
+    if (ruThresholdEnabled) {
+      LocalStorageUtility.setEntryNumber(StorageKey.RUThreshold, ruThreshold);
+    }
+
     if (queryTimeoutEnabled) {
       LocalStorageUtility.setEntryNumber(StorageKey.QueryTimeout, queryTimeout);
       LocalStorageUtility.setEntryBoolean(
@@ -195,6 +208,17 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
     setPageOption(option.key);
   };
 
+  const handleOnRUThresholdToggleChange = (ev: React.MouseEvent<HTMLElement>, checked?: boolean): void => {
+    setRUThresholdEnabled(checked);
+  };
+
+  const handleOnRUThresholdSpinButtonChange = (ev: React.MouseEvent<HTMLElement>, newValue?: string): void => {
+    const ruThreshold = Number(newValue);
+    if (!isNaN(ruThreshold)) {
+      setRUThreshold(ruThreshold);
+    }
+  };
+
   const handleOnQueryTimeoutToggleChange = (ev: React.MouseEvent<HTMLElement>, checked?: boolean): void => {
     setQueryTimeoutEnabled(checked);
   };
@@ -259,7 +283,7 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
     ],
   };
 
-  const queryTimeoutToggleStyles: IToggleStyles = {
+  const toggleStyles: IToggleStyles = {
     label: {
       fontSize: 12,
       fontWeight: 400,
@@ -272,7 +296,7 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
     text: {},
   };
 
-  const queryTimeoutSpinButtonStyles: ISpinButtonStyles = {
+  const spinButtonStyles: ISpinButtonStyles = {
     label: {
       fontSize: 12,
       fontWeight: 400,
@@ -338,48 +362,83 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
           </div>
         )}
         {userContext.apiType === "SQL" && (
-          <div className="settingsSection">
-            <div className="settingsSectionPart">
-              <div>
-                <legend id="queryTimeoutLabel" className="settingsSectionLabel legendLabel">
-                  Query Timeout
-                </legend>
-                <InfoTooltip>
-                  When a query reaches a specified time limit, a popup with an option to cancel the query will show
-                  unless automatic cancellation has been enabled
-                </InfoTooltip>
-              </div>
-              <div>
-                <Toggle
-                  styles={queryTimeoutToggleStyles}
-                  label="Enable query timeout"
-                  onChange={handleOnQueryTimeoutToggleChange}
-                  defaultChecked={queryTimeoutEnabled}
-                />
-              </div>
-              {queryTimeoutEnabled && (
+          <>
+            <div className="settingsSection">
+              <div className="settingsSectionPart">
+                <div>
+                  <legend id="ruThresholdLabel" className="settingsSectionLabel legendLabel">
+                    RU Threshold
+                  </legend>
+                  <InfoTooltip>If a query exceeds a configured RU threshold, the query will be aborted.</InfoTooltip>
+                </div>
                 <div>
-                  <SpinButton
-                    label="Query timeout (ms)"
-                    labelPosition={Position.top}
-                    defaultValue={(queryTimeout || 5000).toString()}
-                    min={100}
-                    step={1000}
-                    onChange={handleOnQueryTimeoutSpinButtonChange}
-                    incrementButtonAriaLabel="Increase value by 1000"
-                    decrementButtonAriaLabel="Decrease value by 1000"
-                    styles={queryTimeoutSpinButtonStyles}
-                  />
                   <Toggle
-                    label="Automatically cancel query after timeout"
-                    styles={queryTimeoutToggleStyles}
-                    onChange={handleOnAutomaticallyCancelQueryToggleChange}
-                    defaultChecked={automaticallyCancelQueryAfterTimeout}
+                    styles={toggleStyles}
+                    label="Enable RU threshold"
+                    onChange={handleOnRUThresholdToggleChange}
+                    defaultChecked={ruThresholdEnabled}
                   />
                 </div>
-              )}
+                {ruThresholdEnabled && (
+                  <div>
+                    <SpinButton
+                      label="RU Threshold (RU)"
+                      labelPosition={Position.top}
+                      defaultValue={(ruThreshold || DefaultRUThreshold).toString()}
+                      min={1}
+                      step={10}
+                      onChange={handleOnRUThresholdSpinButtonChange}
+                      incrementButtonAriaLabel="Increase value by 10"
+                      decrementButtonAriaLabel="Decrease value by 10"
+                      styles={spinButtonStyles}
+                    />
+                  </div>
+                )}
+              </div>
             </div>
-          </div>
+            <div className="settingsSection">
+              <div className="settingsSectionPart">
+                <div>
+                  <legend id="queryTimeoutLabel" className="settingsSectionLabel legendLabel">
+                    Query Timeout
+                  </legend>
+                  <InfoTooltip>
+                    When a query reaches a specified time limit, a popup with an option to cancel the query will show
+                    unless automatic cancellation has been enabled
+                  </InfoTooltip>
+                </div>
+                <div>
+                  <Toggle
+                    styles={toggleStyles}
+                    label="Enable query timeout"
+                    onChange={handleOnQueryTimeoutToggleChange}
+                    defaultChecked={queryTimeoutEnabled}
+                  />
+                </div>
+                {queryTimeoutEnabled && (
+                  <div>
+                    <SpinButton
+                      label="Query timeout (ms)"
+                      labelPosition={Position.top}
+                      defaultValue={(queryTimeout || 5000).toString()}
+                      min={100}
+                      step={1000}
+                      onChange={handleOnQueryTimeoutSpinButtonChange}
+                      incrementButtonAriaLabel="Increase value by 1000"
+                      decrementButtonAriaLabel="Decrease value by 1000"
+                      styles={spinButtonStyles}
+                    />
+                    <Toggle
+                      label="Automatically cancel query after timeout"
+                      styles={toggleStyles}
+                      onChange={handleOnAutomaticallyCancelQueryToggleChange}
+                      defaultChecked={automaticallyCancelQueryAfterTimeout}
+                    />
+                  </div>
+                )}
+              </div>
+            </div>
+          </>
         )}
         <div className="settingsSection">
           <div className="settingsSectionPart">
@@ -404,7 +463,7 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
               onIncrement={(newValue) => setRetryAttempts(parseInt(newValue) + 1 || retryAttempts)}
               onDecrement={(newValue) => setRetryAttempts(parseInt(newValue) - 1 || retryAttempts)}
               onValidate={(newValue) => setRetryAttempts(parseInt(newValue) || retryAttempts)}
-              styles={queryTimeoutSpinButtonStyles}
+              styles={spinButtonStyles}
             />
             <div>
               <legend id="queryRetryAttemptsLabel" className="settingsSectionLabel legendLabel">
@@ -426,7 +485,7 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
               onIncrement={(newValue) => setRetryInterval(parseInt(newValue) + 1000 || retryInterval)}
               onDecrement={(newValue) => setRetryInterval(parseInt(newValue) - 1000 || retryInterval)}
               onValidate={(newValue) => setRetryInterval(parseInt(newValue) || retryInterval)}
-              styles={queryTimeoutSpinButtonStyles}
+              styles={spinButtonStyles}
             />
             <div>
               <legend id="queryRetryAttemptsLabel" className="settingsSectionLabel legendLabel">
@@ -448,7 +507,7 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
               onIncrement={(newValue) => setMaxWaitTimeInSeconds(parseInt(newValue) + 1 || MaxWaitTimeInSeconds)}
               onDecrement={(newValue) => setMaxWaitTimeInSeconds(parseInt(newValue) - 1 || MaxWaitTimeInSeconds)}
               onValidate={(newValue) => setMaxWaitTimeInSeconds(parseInt(newValue) || MaxWaitTimeInSeconds)}
-              styles={queryTimeoutSpinButtonStyles}
+              styles={spinButtonStyles}
             />
           </div>
         </div>

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

@@ -97,6 +97,74 @@ exports[`Settings Pane should render Default properly 1`] = `
         </div>
       </div>
     </div>
+    <div
+      className="settingsSection"
+    >
+      <div
+        className="settingsSectionPart"
+      >
+        <div>
+          <legend
+            className="settingsSectionLabel legendLabel"
+            id="ruThresholdLabel"
+          >
+            RU Threshold
+          </legend>
+          <InfoTooltip>
+            If a query exceeds a configured RU threshold, the query will be aborted.
+          </InfoTooltip>
+        </div>
+        <div>
+          <StyledToggleBase
+            defaultChecked={true}
+            label="Enable RU threshold"
+            onChange={[Function]}
+            styles={
+              Object {
+                "container": Object {},
+                "label": Object {
+                  "display": "block",
+                  "fontSize": 12,
+                  "fontWeight": 400,
+                },
+                "pill": Object {},
+                "root": Object {},
+                "text": Object {},
+                "thumb": Object {},
+              }
+            }
+          />
+        </div>
+        <div>
+          <StyledSpinButton
+            decrementButtonAriaLabel="Decrease value by 10"
+            defaultValue="100"
+            incrementButtonAriaLabel="Increase value by 10"
+            label="RU Threshold (RU)"
+            labelPosition={0}
+            min={1}
+            onChange={[Function]}
+            step={10}
+            styles={
+              Object {
+                "arrowButtonsContainer": Object {},
+                "icon": Object {},
+                "input": Object {},
+                "label": Object {
+                  "fontSize": 12,
+                  "fontWeight": 400,
+                },
+                "labelWrapper": Object {},
+                "root": Object {
+                  "paddingBottom": 10,
+                },
+                "spinButtonWrapper": Object {},
+              }
+            }
+          />
+        </div>
+      </div>
+    </div>
     <div
       className="settingsSection"
     >

src/Explorer/Tabs/QueryTab/QueryTabComponent.tsx

@@ -1,6 +1,6 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 /* eslint-disable no-console */
-import { FeedOptions } from "@azure/cosmos";
+import { FeedOptions, QueryOperationOptions } from "@azure/cosmos";
 import { useDialog } from "Explorer/Controls/Dialog";
 import { QueryCopilotFeedbackModal } from "Explorer/QueryCopilot/Modal/QueryCopilotFeedbackModal";
 import { useCopilotStore } from "Explorer/QueryCopilot/QueryCopilotContext";
@@ -10,7 +10,7 @@ import { QueryCopilotSidebar } from "Explorer/QueryCopilot/V2/Sidebar/QueryCopil
 import { QueryResultSection } from "Explorer/Tabs/QueryTab/QueryResultSection";
 import { useSelectedNode } from "Explorer/useSelectedNode";
 import { QueryConstants } from "Shared/Constants";
-import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility";
+import { LocalStorageUtility, StorageKey, getRUThreshold, ruThresholdEnabled } from "Shared/StorageUtility";
 import { Action } from "Shared/Telemetry/TelemetryConstants";
 import { QueryCopilotState, useQueryCopilot } from "hooks/useQueryCopilot";
 import { TabsState, useTabs } from "hooks/useTabs";
@@ -303,8 +303,20 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
       isExecutionError: false,
     });
 
+    let queryOperationOptions: QueryOperationOptions;
+    if (!this.isPreferredApiMongoDB && ruThresholdEnabled()) {
+      const ruThreshold: number = getRUThreshold();
+      queryOperationOptions = {
+        ruCapPerOperation: ruThreshold,
+      } as QueryOperationOptions;
+    }
     const queryDocuments = async (firstItemIndex: number) =>
-      await queryDocumentsPage(this.props.collection && this.props.collection.id(), this._iterator, firstItemIndex);
+      await queryDocumentsPage(
+        this.props.collection && this.props.collection.id(),
+        this._iterator,
+        firstItemIndex,
+        queryOperationOptions,
+      );
     this.props.tabsBaseInstance.isExecuting(true);
     this.setState({
       isExecuting: true,

src/Explorer/Tabs/Tabs.tsx

@@ -10,6 +10,7 @@ import { PostgresConnectTab } from "Explorer/Tabs/PostgresConnectTab";
 import { QuickstartTab } from "Explorer/Tabs/QuickstartTab";
 import { VcoreMongoConnectTab } from "Explorer/Tabs/VCoreMongoConnectTab";
 import { VcoreMongoQuickstartTab } from "Explorer/Tabs/VCoreMongoQuickstartTab";
+import { hasRUThresholdBeenConfigured } from "Shared/StorageUtility";
 import { userContext } from "UserContext";
 import { useTeachingBubble } from "hooks/useTeachingBubble";
 import ko from "knockout";
@@ -29,7 +30,9 @@ interface TabsProps {
 
 export const Tabs = ({ explorer }: TabsProps): JSX.Element => {
   const { openedTabs, openedReactTabs, activeTab, activeReactTab, networkSettingsWarning } = useTabs();
-
+  const [showRUThresholdMessageBar, setShowRUThresholdMessageBar] = useState<boolean>(
+    userContext.apiType === "SQL" && !hasRUThresholdBeenConfigured(),
+  );
   return (
     <div className="tabsManagerContainer">
       {networkSettingsWarning && (
@@ -54,6 +57,18 @@ export const Tabs = ({ explorer }: TabsProps): JSX.Element => {
           {networkSettingsWarning}
         </MessageBar>
       )}
+      {showRUThresholdMessageBar && (
+        <MessageBar
+          messageBarType={MessageBarType.info}
+          onDismiss={() => {
+            setShowRUThresholdMessageBar(false);
+          }}
+        >
+          {
+            "Avoid high cost queries! We automatically abort them if they exceed the set RU limit.  To adjust your limit go to Settings > RU threshold."
+          }
+        </MessageBar>
+      )}
       <div id="content" className="flexContainer hideOverflows">
         <div className="nav-tabs-margin">
           <ul className="nav nav-tabs level navTabHeight" id="navTabs" role="tablist">

src/Explorer/Tabs/TriggerTabContent.tsx

@@ -4,9 +4,9 @@ import React, { Component } from "react";
 import DiscardIcon from "../../../images/discard.svg";
 import SaveIcon from "../../../images/save-cosmos.svg";
 import * as Constants from "../../Common/Constants";
+import { getErrorMessage, getErrorStack } from "../../Common/ErrorHandlingUtils";
 import { createTrigger } from "../../Common/dataAccess/createTrigger";
 import { updateTrigger } from "../../Common/dataAccess/updateTrigger";
-import { getErrorMessage, getErrorStack } from "../../Common/ErrorHandlingUtils";
 import * as ViewModels from "../../Contracts/ViewModels";
 import { Action } from "../../Shared/Telemetry/TelemetryConstants";
 import * as TelemetryProcessor from "../../Shared/Telemetry/TelemetryProcessor";