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>
2024-01-30 16:21:29 -05:00 · 2024-01-30 16:21:29 -05:00 · dbb0324a64
parent e207f3702b
commit dbb0324a64
10 changed files with 250 additions and 62 deletions
--- a/package-lock.json
+++ b/package-lock.json
@ -10,7 +10,7 @@
      "hasInstallScript": true,
      "dependencies": {
        "@azure/arm-cosmosdb": "9.1.0",
-        "@azure/cosmos": "4.0.0",
+        "@azure/cosmos": "4.0.1-beta.2",
        "@azure/cosmos-language-service": "0.0.5",
        "@azure/identity": "1.2.1",
        "@azure/ms-rest-nodeauth": "3.0.7",
@ -396,9 +396,9 @@
      "integrity": "sha512-AEYxH93jGFPn/a2iVAwW87VuUIkR1FVUKB77NwMF7nBTDkDrrT/Hpt/IrCJ0QXhW27jTBDcf5ZY7w6RiqTMw2Q=="
    },
    "node_modules/@azure/cosmos": {
-      "version": "4.0.0",
-      "resolved": "https://registry.npmjs.org/@azure/cosmos/-/cosmos-4.0.0.tgz",
-      "integrity": "sha512-/Z27p1+FTkmjmm8jk90zi/HrczPHw2t8WecFnsnTe4xGocWl0Z4clP0YlLUTJPhRLWYa5upwD9rMvKJkS1f1kg==",
+      "version": "4.0.1-beta.2",
+      "resolved": "https://registry.npmjs.org/@azure/cosmos/-/cosmos-4.0.1-beta.2.tgz",
+      "integrity": "sha512-iuqg/QwLQlxgRi4pnXU8JUYv+f24wkRvJ9ZZI4/sYk+DxSgkuQ194Cc2IpckpeO8z7ZpcBkVQFa82wcZVVZ8Zg==",
      "dependencies": {
        "@azure/abort-controller": "^1.0.0",
        "@azure/core-auth": "^1.3.0",
@ -408,14 +408,14 @@
        "fast-json-stable-stringify": "^2.1.0",
        "jsbi": "^3.1.3",
        "node-abort-controller": "^3.0.0",
-        "priorityqueuejs": "^1.0.0",
+        "priorityqueuejs": "^2.0.0",
        "semaphore": "^1.0.5",
        "tslib": "^2.2.0",
        "universal-user-agent": "^6.0.0",
        "uuid": "^8.3.0"
      },
      "engines": {
-        "node": ">=14.0.0"
+        "node": ">=18.0.0"
      }
    },
    "node_modules/@azure/cosmos-language-service": {
@ -33707,9 +33707,9 @@
      }
    },
    "node_modules/priorityqueuejs": {
-      "version": "1.0.0",
-      "resolved": "https://registry.npmjs.org/priorityqueuejs/-/priorityqueuejs-1.0.0.tgz",
-      "integrity": "sha512-lg++21mreCEOuGWTbO5DnJKAdxfjrdN0S9ysoW9SzdSJvbkWpkaDdpG/cdsPCsEnoLUwmd9m3WcZhngW7yKA2g=="
+      "version": "2.0.0",
+      "resolved": "https://registry.npmjs.org/priorityqueuejs/-/priorityqueuejs-2.0.0.tgz",
+      "integrity": "sha512-19BMarhgpq3x4ccvVi8k2QpJZcymo/iFUcrhPd4V96kYGovOdTsWwy7fxChYi4QY+m2EnGBWSX9Buakz+tWNQQ=="
    },
    "node_modules/prismjs": {
      "version": "1.29.0",
--- a/package.json
+++ b/package.json
@ -5,7 +5,7 @@
  "main": "index.js",
  "dependencies": {
    "@azure/arm-cosmosdb": "9.1.0",
-    "@azure/cosmos": "4.0.0",
+    "@azure/cosmos": "4.0.1-beta.2",
    "@azure/cosmos-language-service": "0.0.5",
    "@azure/identity": "1.2.1",
    "@azure/ms-rest-nodeauth": "3.0.7",
--- a/src/Common/IteratorUtilities.ts
+++ b/src/Common/IteratorUtilities.ts
@ -1,3 +1,4 @@
+import { QueryOperationOptions } from "@azure/cosmos";
 import { QueryResults } from "../Contracts/ViewModels";

 interface QueryResponse {
@ -10,13 +11,17 @@ interface QueryResponse {
 }

 export interface MinimalQueryIterator {
-  fetchNext: () => Promise<QueryResponse>;
+  fetchNext: (queryOperationOptions?: QueryOperationOptions) => Promise<QueryResponse>;
 }

 // Pick<QueryIterator<any>, "fetchNext">;

-export function nextPage(documentsIterator: MinimalQueryIterator, firstItemIndex: number): Promise<QueryResults> {
-  return documentsIterator.fetchNext().then((response) => {
+export function nextPage(
+  documentsIterator: MinimalQueryIterator,
+  firstItemIndex: number,
+  queryOperationOptions?: QueryOperationOptions,
+): Promise<QueryResults> {
+  return documentsIterator.fetchNext(queryOperationOptions).then((response) => {
    const documents = response.resources;
    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    const headers = (response as any).headers || {}; // TODO this is a private key. Remove any
--- a/src/Common/dataAccess/queryDocumentsPage.ts
+++ b/src/Common/dataAccess/queryDocumentsPage.ts
@ -1,3 +1,4 @@
+import { QueryOperationOptions } from "@azure/cosmos";
 import { QueryResults } from "../../Contracts/ViewModels";
 import { logConsoleInfo, logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
 import { getEntityName } from "../DocumentUtility";
@ -8,12 +9,13 @@ export const queryDocumentsPage = async (
  resourceName: string,
  documentsIterator: MinimalQueryIterator,
  firstItemIndex: number,
+  queryOperationOptions?: QueryOperationOptions,
 ): Promise<QueryResults> => {
  const entityName = getEntityName();
  const clearMessage = logConsoleProgress(`Querying ${entityName} for container ${resourceName}`);

  try {
-    const result: QueryResults = await nextPage(documentsIterator, firstItemIndex);
+    const result: QueryResults = await nextPage(documentsIterator, firstItemIndex, queryOperationOptions);
    const itemCount = (result.documents && result.documents.length) || 0;
    logConsoleInfo(`Successfully fetched ${itemCount} ${entityName} for container ${resourceName}`);
    return result;
--- a/src/Explorer/Panes/SettingsPane/SettingsPane.tsx
+++ b/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,6 +362,40 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
          </div>
        )}
        {userContext.apiType === "SQL" && (
+          <>
+            <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>
+                  <Toggle
+                    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 className="settingsSection">
              <div className="settingsSectionPart">
                <div>
@ -351,7 +409,7 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
                </div>
                <div>
                  <Toggle
-                  styles={queryTimeoutToggleStyles}
+                    styles={toggleStyles}
                    label="Enable query timeout"
                    onChange={handleOnQueryTimeoutToggleChange}
                    defaultChecked={queryTimeoutEnabled}
@ -368,11 +426,11 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
                      onChange={handleOnQueryTimeoutSpinButtonChange}
                      incrementButtonAriaLabel="Increase value by 1000"
                      decrementButtonAriaLabel="Decrease value by 1000"
-                    styles={queryTimeoutSpinButtonStyles}
+                      styles={spinButtonStyles}
                    />
                    <Toggle
                      label="Automatically cancel query after timeout"
-                    styles={queryTimeoutToggleStyles}
+                      styles={toggleStyles}
                      onChange={handleOnAutomaticallyCancelQueryToggleChange}
                      defaultChecked={automaticallyCancelQueryAfterTimeout}
                    />
@ -380,6 +438,7 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
                )}
              </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>
--- a/src/Explorer/Panes/SettingsPane/snapshots/SettingsPane.test.tsx.snap
+++ b/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"
    >
--- a/src/Explorer/Tabs/QueryTab/QueryTabComponent.tsx
+++ b/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,
--- a/src/Explorer/Tabs/Tabs.tsx
+++ b/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">
--- a/src/Explorer/Tabs/TriggerTabContent.tsx
+++ b/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";
--- a/src/Shared/StorageUtility.ts
+++ b/src/Shared/StorageUtility.ts
@ -1,9 +1,12 @@
 import * as LocalStorageUtility from "./LocalStorageUtility";
 import * as SessionStorageUtility from "./SessionStorageUtility";
+import * as StringUtility from "./StringUtility";

 export { LocalStorageUtility, SessionStorageUtility };
 export enum StorageKey {
  ActualItemPerPage,
+  RUThresholdEnabled,
+  RUThreshold,
  QueryTimeoutEnabled,
  QueryTimeout,
  RetryAttempts,
@ -25,3 +28,27 @@ export enum StorageKey {
  VisitedAccounts,
  PriorityLevel,
 }
+
+export const hasRUThresholdBeenConfigured = (): boolean => {
+  const ruThresholdEnabledLocalStorageRaw: string | null = LocalStorageUtility.getEntryString(
+    StorageKey.RUThresholdEnabled,
+  );
+  return ruThresholdEnabledLocalStorageRaw === "true" || ruThresholdEnabledLocalStorageRaw === "false";
+};
+
+export const ruThresholdEnabled = (): boolean => {
+  const ruThresholdEnabledLocalStorageRaw: string | null = LocalStorageUtility.getEntryString(
+    StorageKey.RUThresholdEnabled,
+  );
+  return ruThresholdEnabledLocalStorageRaw === null || StringUtility.toBoolean(ruThresholdEnabledLocalStorageRaw);
+};
+
+export const getRUThreshold = (): number => {
+  const ruThresholdRaw = LocalStorageUtility.getEntryNumber(StorageKey.RUThreshold);
+  if (ruThresholdRaw !== 0) {
+    return ruThresholdRaw;
+  }
+  return DefaultRUThreshold;
+};
+
+export const DefaultRUThreshold = 100;