Added throughput bucketing

src/Explorer/Controls/Settings/SettingsComponent.tsx

@@ -86,6 +86,8 @@ export interface SettingsComponentState {
   wasAutopilotOriginallySet: boolean;
   isScaleSaveable: boolean;
   isScaleDiscardable: boolean;
+  throughputBuckets: DataModels.ThroughputBucket[];
+  throughputBucketsBaseline: DataModels.ThroughputBucket[];
   throughputError: string;
 
   timeToLive: TtlType;
@@ -158,6 +160,7 @@ export class SettingsComponent extends React.Component<SettingsComponentProps, S
   private isVectorSearchEnabled: boolean;
   private isFullTextSearchEnabled: boolean;
   private totalThroughputUsed: number;
+  private throughputBucketsEnabled: boolean;
   public mongoDBCollectionResource: MongoDBCollectionResource;
 
   constructor(props: SettingsComponentProps) {
@@ -175,6 +178,8 @@ export class SettingsComponent extends React.Component<SettingsComponentProps, S
       this.isFullTextSearchEnabled = isFullTextSearchEnabled() && !hasDatabaseSharedThroughput(this.collection);
 
       this.changeFeedPolicyVisible = userContext.features.enableChangeFeedPolicy;
+      this.throughputBucketsEnabled =
+        userContext.features.enableThroughputBuckets && userContext.authType === AuthType.AAD;
 
       // Mongo container with system partition key still treat as "Fixed"
       this.isFixedContainer =
@@ -193,6 +198,8 @@ export class SettingsComponent extends React.Component<SettingsComponentProps, S
       wasAutopilotOriginallySet: false,
       isScaleSaveable: false,
       isScaleDiscardable: false,
+      throughputBuckets: undefined,
+      throughputBucketsBaseline: undefined,
       throughputError: undefined,
 
       timeToLive: undefined,
@@ -419,6 +426,8 @@ export class SettingsComponent extends React.Component<SettingsComponentProps, S
 
     this.setState({
       throughput: this.state.throughputBaseline,
+      throughputBuckets: this.state.throughputBucketsBaseline,
+      throughputBucketsBaseline: this.state.throughputBucketsBaseline,
       timeToLive: this.state.timeToLiveBaseline,
       timeToLiveSeconds: this.state.timeToLiveSecondsBaseline,
       displayedTtlSeconds: this.state.displayedTtlSecondsBaseline,
@@ -749,9 +758,13 @@ export class SettingsComponent extends React.Component<SettingsComponentProps, S
       ] as DataModels.ComputedProperties;
     }
 
+    const throughputBuckets = this.offer?.throughputBuckets;
+
     return {
       throughput: offerThroughput,
       throughputBaseline: offerThroughput,
+      throughputBuckets,
+      throughputBucketsBaseline: throughputBuckets,
       changeFeedPolicy: changeFeedPolicy,
       changeFeedPolicyBaseline: changeFeedPolicy,
       timeToLive: timeToLive,
@@ -839,6 +852,10 @@ export class SettingsComponent extends React.Component<SettingsComponentProps, S
     this.setState({ throughput: newThroughput, throughputError });
   };
 
+  private onThroughputBucketChange = (throughputBuckets: DataModels.ThroughputBucket[]): void => {
+    this.setState({ throughputBuckets });
+  };
+
   private onAutoPilotSelected = (isAutoPilotSelected: boolean): void =>
     this.setState({ isAutoPilotSelected: isAutoPilotSelected });
 
@@ -1036,6 +1053,7 @@ export class SettingsComponent extends React.Component<SettingsComponentProps, S
         currentOffer: this.collection.offer(),
         autopilotThroughput: this.state.isAutoPilotSelected ? this.state.autoPilotThroughput : undefined,
         manualThroughput: this.state.isAutoPilotSelected ? undefined : this.state.throughput,
+        ...(this.state.throughputBuckets && { throughputBuckets: this.state.throughputBuckets }),
       };
       if (this.hasProvisioningTypeChanged()) {
         if (this.state.isAutoPilotSelected) {
@@ -1102,6 +1120,10 @@ export class SettingsComponent extends React.Component<SettingsComponentProps, S
       onMaxAutoPilotThroughputChange: this.onMaxAutoPilotThroughputChange,
       onScaleSaveableChange: this.onScaleSaveableChange,
       onScaleDiscardableChange: this.onScaleDiscardableChange,
+      throughputBucketsBaseline: this.state.throughputBucketsBaseline,
+      throughputBuckets: this.state.throughputBuckets,
+      enableThroughputBuckets: this.isCollectionSettingsTab && this.throughputBucketsEnabled,
+      onThroughputBucketChange: this.onThroughputBucketChange,
       throughputError: this.state.throughputError,
     };
 

src/Explorer/Controls/Settings/SettingsSubComponents/ScaleComponent.tsx

@@ -16,6 +16,7 @@ import {
   titleAndInputStackProps,
 } from "../SettingsRenderUtils";
 import { hasDatabaseSharedThroughput } from "../SettingsUtils";
+import { ThroughputBucketsComponent } from "./ThroughputInputComponents/ThroughputBucketsComponent";
 import { ThroughputInputAutoPilotV3Component } from "./ThroughputInputComponents/ThroughputInputAutoPilotV3Component";
 
 export interface ScaleComponentProps {
@@ -33,10 +34,21 @@ export interface ScaleComponentProps {
   onMaxAutoPilotThroughputChange: (newThroughput: number) => void;
   onScaleSaveableChange: (isScaleSaveable: boolean) => void;
   onScaleDiscardableChange: (isScaleDiscardable: boolean) => void;
+  throughputBuckets: DataModels.ThroughputBucket[];
+  throughputBucketsBaseline: DataModels.ThroughputBucket[];
+  enableThroughputBuckets: boolean;
+  onThroughputBucketChange: (throughputBuckets: DataModels.ThroughputBucket[]) => void;
   throughputError?: string;
 }
 
-export class ScaleComponent extends React.Component<ScaleComponentProps> {
+interface ScaleComponentState {
+  isThroughputSaveable: boolean;
+  isBucketsSaveable: boolean;
+  isThroughputDiscardable: boolean;
+  isBucketsDiscardable: boolean;
+}
+
+export class ScaleComponent extends React.Component<ScaleComponentProps, ScaleComponentState> {
   private isEmulator: boolean;
   private offer: DataModels.Offer;
   private databaseId: string;
@@ -48,6 +60,13 @@ export class ScaleComponent extends React.Component<ScaleComponentProps> {
     this.offer = this.props.database?.offer() || this.props.collection?.offer();
     this.databaseId = this.props.database?.id() || this.props.collection.databaseId;
     this.collectionId = this.props.collection?.id();
+    this.state = {
+      isThroughputSaveable: false,
+      isBucketsSaveable: false,
+      isThroughputDiscardable: false,
+      isBucketsDiscardable: false,
+    };
+    console.log(this.offer);
   }
 
   public isAutoScaleEnabled = (): boolean => {
@@ -137,8 +156,8 @@ export class ScaleComponent extends React.Component<ScaleComponentProps> {
       maxAutoPilotThroughputBaseline={this.props.autoPilotThroughputBaseline}
       onMaxAutoPilotThroughputChange={this.props.onMaxAutoPilotThroughputChange}
       spendAckChecked={false}
-      onScaleSaveableChange={this.props.onScaleSaveableChange}
-      onScaleDiscardableChange={this.props.onScaleDiscardableChange}
+      onScaleSaveableChange={this.handleThroughputSaveableChange}
+      onScaleDiscardableChange={this.handleThroughputDiscardableChange}
       usageSizeInKB={this.props.collection?.usageSizeInKB()}
       throughputError={this.props.throughputError}
       instantMaximumThroughput={this.offer?.instantMaximumThroughput}
@@ -167,6 +186,39 @@ export class ScaleComponent extends React.Component<ScaleComponentProps> {
     );
   }
 
+  private updateScaleSettingsState = (updates: Partial<ScaleComponentState>) => {
+    this.setState(
+      (prevState) => {
+        const hasChanges = Object.keys(updates).some(
+          (key) => prevState[key as keyof ScaleComponentState] !== updates[key as keyof ScaleComponentState],
+        );
+        return hasChanges ? { ...prevState, ...updates } : null;
+      },
+      () => {
+        const isSaveable = this.state.isThroughputSaveable || this.state.isBucketsSaveable;
+        const isDiscardable = this.state.isThroughputDiscardable || this.state.isBucketsDiscardable;
+        this.props.onScaleSaveableChange(isSaveable);
+        this.props.onScaleDiscardableChange(isDiscardable);
+      },
+    );
+  };
+
+  private handleThroughputSaveableChange = (isSaveable: boolean) => {
+    this.updateScaleSettingsState({ isThroughputSaveable: isSaveable });
+  };
+
+  private handleThroughputDiscardableChange = (isDiscardable: boolean) => {
+    this.updateScaleSettingsState({ isThroughputDiscardable: isDiscardable });
+  };
+
+  private handleBucketsSaveableChange = (isSaveable: boolean) => {
+    this.updateScaleSettingsState({ isBucketsSaveable: isSaveable });
+  };
+
+  private handleBucketsDiscardableChange = (isDiscardable: boolean) => {
+    this.updateScaleSettingsState({ isBucketsDiscardable: isDiscardable });
+  };
+
   public render(): JSX.Element {
     return (
       <Stack {...subComponentStackProps}>
@@ -182,6 +234,15 @@ export class ScaleComponent extends React.Component<ScaleComponentProps> {
           <MessageBar messageBarType={MessageBarType.warning}>{this.getInitialNotificationElement()}</MessageBar>
         )}
         {!this.isAutoScaleEnabled() && <Stack {...subComponentStackProps}>{this.getThroughputInputComponent()}</Stack>}
+        {this.props.enableThroughputBuckets && (
+          <ThroughputBucketsComponent
+            currentBuckets={this.props.throughputBuckets}
+            throughputBucketsBaseline={this.props.throughputBucketsBaseline}
+            onBucketsChange={this.props.onThroughputBucketChange}
+            onSaveableChange={this.handleBucketsSaveableChange}
+            onDiscardableChange={this.handleBucketsDiscardableChange}
+          />
+        )}
 
         {/* TODO: Replace link with call to the Azure Support blade */}
         {this.isAutoScaleEnabled() && (

src/Explorer/Controls/Settings/SettingsSubComponents/ThroughputInputComponents/ThroughputBucketsComponent.test.tsx

@@ -0,0 +1,129 @@
+import "@testing-library/jest-dom";
+import { fireEvent, render, screen } from "@testing-library/react";
+import React from "react";
+import { ThroughputBucketsComponent } from "./ThroughputBucketsComponent";
+
+describe("ThroughputBucketsComponent", () => {
+  const mockOnBucketsChange = jest.fn();
+  const mockOnSaveableChange = jest.fn();
+  const mockOnDiscardableChange = jest.fn();
+
+  const defaultProps = {
+    currentBuckets: [
+      { id: 1, maxThroughputPercentage: 50 },
+      { id: 2, maxThroughputPercentage: 60 },
+    ],
+    throughputBucketsBaseline: [
+      { id: 1, maxThroughputPercentage: 40 },
+      { id: 2, maxThroughputPercentage: 50 },
+    ],
+    onBucketsChange: mockOnBucketsChange,
+    onSaveableChange: mockOnSaveableChange,
+    onDiscardableChange: mockOnDiscardableChange,
+  };
+
+  beforeEach(() => {
+    jest.clearAllMocks();
+  });
+
+  it("renders 5 buckets with default values when input buckets are missing", () => {
+    render(<ThroughputBucketsComponent {...defaultProps} currentBuckets={[{ id: 1, maxThroughputPercentage: 50 }]} />);
+
+    expect(screen.getAllByText(/Bucket \d+/)).toHaveLength(5);
+    expect(screen.getByDisplayValue("50")).toBeInTheDocument();
+    expect(screen.getAllByDisplayValue("100").length).toBe(4);
+  });
+
+  it("renders buckets in the correct order even if input is unordered", () => {
+    const unorderedBuckets = [
+      { id: 2, maxThroughputPercentage: 60 },
+      { id: 1, maxThroughputPercentage: 50 },
+    ];
+    render(<ThroughputBucketsComponent {...defaultProps} currentBuckets={unorderedBuckets} />);
+
+    const bucketLabels = screen.getAllByText(/Bucket \d+/).map((el) => el.textContent);
+    expect(bucketLabels).toEqual(["Bucket 1", "Bucket 2", "Bucket 3", "Bucket 4", "Bucket 5"]);
+    expect(screen.getByDisplayValue("50")).toBeInTheDocument();
+    expect(screen.getByDisplayValue("60")).toBeInTheDocument();
+    expect(screen.getAllByDisplayValue("100").length).toBe(3);
+  });
+
+  it("calls onBucketsChange when a bucket value changes", () => {
+    render(<ThroughputBucketsComponent {...defaultProps} />);
+
+    const input = screen.getByDisplayValue("50");
+    fireEvent.change(input, { target: { value: "70" } });
+
+    expect(mockOnBucketsChange).toHaveBeenCalledWith([
+      { id: 1, maxThroughputPercentage: 70 },
+      { id: 2, maxThroughputPercentage: 60 },
+      { id: 3, maxThroughputPercentage: 100 },
+      { id: 4, maxThroughputPercentage: 100 },
+      { id: 5, maxThroughputPercentage: 100 },
+    ]);
+  });
+
+  it("triggers onSaveableChange and onDiscardableChange when values change", () => {
+    render(<ThroughputBucketsComponent {...defaultProps} />);
+
+    const input = screen.getByDisplayValue("50");
+    fireEvent.change(input, { target: { value: "80" } });
+
+    expect(mockOnSaveableChange).toHaveBeenCalledWith(true);
+    expect(mockOnDiscardableChange).toHaveBeenCalledWith(true);
+  });
+
+  it("ensures buckets revert to default when no buckets are provided", () => {
+    render(<ThroughputBucketsComponent {...defaultProps} currentBuckets={[]} />);
+
+    expect(screen.getAllByText(/Bucket \d+/)).toHaveLength(5);
+    expect(screen.getAllByDisplayValue("100")).toHaveLength(5);
+  });
+
+  it("updates state consistently after multiple changes to different buckets", () => {
+    render(<ThroughputBucketsComponent {...defaultProps} />);
+
+    const input1 = screen.getByDisplayValue("50");
+    fireEvent.change(input1, { target: { value: "70" } });
+
+    const input2 = screen.getByDisplayValue("60");
+    fireEvent.change(input2, { target: { value: "80" } });
+
+    expect(mockOnBucketsChange).toHaveBeenCalledWith([
+      { id: 1, maxThroughputPercentage: 70 },
+      { id: 2, maxThroughputPercentage: 80 },
+      { id: 3, maxThroughputPercentage: 100 },
+      { id: 4, maxThroughputPercentage: 100 },
+      { id: 5, maxThroughputPercentage: 100 },
+    ]);
+  });
+
+  it("resets to baseline when currentBuckets are reset", () => {
+    const { rerender } = render(<ThroughputBucketsComponent {...defaultProps} />);
+
+    const input1 = screen.getByDisplayValue("50");
+    fireEvent.change(input1, { target: { value: "70" } });
+
+    expect(mockOnBucketsChange).toHaveBeenCalledWith([
+      { id: 1, maxThroughputPercentage: 70 },
+      { id: 2, maxThroughputPercentage: 60 },
+      { id: 3, maxThroughputPercentage: 100 },
+      { id: 4, maxThroughputPercentage: 100 },
+      { id: 5, maxThroughputPercentage: 100 },
+    ]);
+
+    rerender(<ThroughputBucketsComponent {...defaultProps} currentBuckets={defaultProps.throughputBucketsBaseline} />);
+
+    expect(screen.getByDisplayValue("40")).toBeInTheDocument();
+    expect(screen.getByDisplayValue("50")).toBeInTheDocument();
+  });
+
+  it("does not call onBucketsChange when value remains unchanged", () => {
+    render(<ThroughputBucketsComponent {...defaultProps} />);
+
+    const input = screen.getByDisplayValue("50");
+    fireEvent.change(input, { target: { value: "50" } });
+
+    expect(mockOnBucketsChange).not.toHaveBeenCalled();
+  });
+});

src/Explorer/Controls/Settings/SettingsSubComponents/ThroughputInputComponents/ThroughputBucketsComponent.tsx

@@ -0,0 +1,92 @@
+import { Icon, Label, Slider, Stack, TextField } from "@fluentui/react";
+import { ThroughputBucket } from "Contracts/DataModels";
+import React, { FC, useEffect, useState } from "react";
+import { isDirty } from "../../SettingsUtils";
+
+const MAX_BUCKET_SIZES = 5;
+
+const DEFAULT_BUCKETS = Array.from({ length: MAX_BUCKET_SIZES }, (_, i) => ({
+  id: i + 1,
+  maxThroughputPercentage: 100,
+}));
+
+interface ThroughputBucketsComponentProps {
+  currentBuckets: ThroughputBucket[];
+  throughputBucketsBaseline: ThroughputBucket[];
+  onBucketsChange: (updatedBuckets: ThroughputBucket[]) => void;
+  onSaveableChange: (isSaveable: boolean) => void;
+  onDiscardableChange: (isDiscardable: boolean) => void;
+}
+
+export const ThroughputBucketsComponent: FC<ThroughputBucketsComponentProps> = ({
+  currentBuckets,
+  throughputBucketsBaseline,
+  onBucketsChange,
+  onSaveableChange,
+  onDiscardableChange,
+}) => {
+  const getThroughputBuckets = (buckets: ThroughputBucket[]): ThroughputBucket[] => {
+    return DEFAULT_BUCKETS.map(
+      (defaultBucket) => buckets.find((bucket) => bucket.id === defaultBucket.id) || defaultBucket,
+    );
+  };
+
+  const [throughputBuckets, setThroughputBuckets] = useState<ThroughputBucket[]>(getThroughputBuckets(currentBuckets));
+
+  useEffect(() => {
+    const isChanged = isDirty(currentBuckets, getThroughputBuckets(throughputBucketsBaseline));
+    isChanged && setThroughputBuckets(getThroughputBuckets(currentBuckets));
+  }, [currentBuckets]);
+
+  useEffect(() => {
+    const isChanged = isDirty(throughputBuckets, getThroughputBuckets(throughputBucketsBaseline));
+    onSaveableChange(isChanged);
+    onDiscardableChange(isChanged);
+  }, [throughputBuckets]);
+
+  const handleBucketChange = (id: number, newValue: number) => {
+    const updatedBuckets = throughputBuckets.map((bucket) =>
+      bucket.id === id ? { ...bucket, maxThroughputPercentage: newValue } : bucket,
+    );
+    setThroughputBuckets(updatedBuckets);
+    const settingsChanged = isDirty(updatedBuckets, throughputBuckets);
+    settingsChanged && onBucketsChange(updatedBuckets);
+  };
+
+  return (
+    <Stack tokens={{ childrenGap: "m" }} styles={{ root: { width: "70%", maxWidth: 700 } }}>
+      <Label>Throughput buckets</Label>
+      <Stack>
+        {throughputBuckets?.map((bucket) => (
+          <Stack key={bucket.id} horizontal tokens={{ childrenGap: 8 }} verticalAlign="center">
+            <Slider
+              min={1}
+              max={100}
+              step={1}
+              value={bucket.maxThroughputPercentage}
+              onChange={(newValue) => handleBucketChange(bucket.id, newValue)}
+              showValue={false}
+              label={`Bucket ${bucket.id}`}
+              styles={{ root: { flex: 2, maxWidth: 400 } }}
+            />
+            <TextField
+              value={bucket.maxThroughputPercentage.toString()}
+              onChange={(event, newValue) => handleBucketChange(bucket.id, parseInt(newValue || "0", 10))}
+              type="number"
+              suffix="%"
+              styles={{
+                fieldGroup: { width: 80 },
+              }}
+            />
+            {bucket.id === 1 && (
+              <Stack horizontal tokens={{ childrenGap: 4 }} verticalAlign="center">
+                <Icon iconName="TagSolid" />
+                <span>Data Explorer Query Bucket</span>
+              </Stack>
+            )}
+          </Stack>
+        ))}
+      </Stack>
+    </Stack>
+  );
+};

src/Explorer/Controls/Settings/SettingsUtils.tsx

@@ -11,7 +11,8 @@ export type isDirtyTypes =
   | DataModels.IndexingPolicy
   | DataModels.ComputedProperties
   | DataModels.VectorEmbedding[]
-  | DataModels.FullTextPolicy;
+  | DataModels.FullTextPolicy
+  | DataModels.ThroughputBucket[];
 export const TtlOff = "off";
 export const TtlOn = "on";
 export const TtlOnNoDefault = "on-nodefault";

src/Explorer/Tabs/QueryTab/QueryTabComponent.tsx

@@ -1,6 +1,6 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 /* eslint-disable no-console */
-import { FeedOptions, QueryOperationOptions } from "@azure/cosmos";
+import { FeedOptions } from "@azure/cosmos";
 import QueryError, { createMonacoErrorLocationResolver, createMonacoMarkersForQueryErrors } from "Common/QueryError";
 import { SplitterDirection } from "Common/Splitter";
 import { Platform, configContext } from "ConfigContext";
@@ -18,7 +18,7 @@ import { CosmosFluentProvider } from "Explorer/Theme/ThemeUtil";
 import { useSelectedNode } from "Explorer/useSelectedNode";
 import { KeyboardAction } from "KeyboardShortcuts";
 import { QueryConstants } from "Shared/Constants";
-import { LocalStorageUtility, StorageKey, getRUThreshold, ruThresholdEnabled } from "Shared/StorageUtility";
+import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility";
 import { Action } from "Shared/Telemetry/TelemetryConstants";
 import { Allotment } from "allotment";
 import { QueryCopilotState, useQueryCopilot } from "hooks/useQueryCopilot";
@@ -368,20 +368,8 @@ class QueryTabComponentImpl extends React.Component<QueryTabComponentImplProps,
       isExecutionError: false,
     });
 
-    let queryOperationOptions: QueryOperationOptions;
-    if (userContext.apiType === "SQL" && 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,
-        queryOperationOptions,
-      );
+      await queryDocumentsPage(this.props.collection && this.props.collection.id(), this._iterator, firstItemIndex);
     this.props.tabsBaseInstance.isExecuting(true);
     this.setState({
       isExecuting: true,