Adding Hot Partition Key Rate Limiting

.vscode/settings.json

@@ -17,12 +17,17 @@
     "test/out/**": true,
     "workers/libs/**": true
   },
+  "js/ts.tsdk.path": "node_modules/typescript/lib",
   "typescript.tsdk": "node_modules/typescript/lib",
   "editor.formatOnSave": true,
   "editor.codeActionsOnSave": {
     "source.fixAll.eslint": "explicit",
     "source.organizeImports": "explicit"
   },
+  "js/ts.preferences.importModuleSpecifier": "non-relative",
   "typescript.preferences.importModuleSpecifier": "non-relative",
-  "editor.defaultFormatter": "esbenp.prettier-vscode"
+  "editor.defaultFormatter": "esbenp.prettier-vscode",
+  "[typescriptreact]": {
+    "editor.defaultFormatter": "esbenp.prettier-vscode"
+  }
 }

src/Common/Constants.ts

@@ -88,6 +88,7 @@ export class CapabilityNames {
   public static readonly EnableDynamicDataMasking: string = "EnableDynamicDataMasking";
   public static readonly EnableNoSQLFullTextSearchPreviewFeatures: string = "EnableNoSQLFullTextSearchPreviewFeatures";
   public static readonly EnableOnlineCopyFeature: string = "EnableOnlineContainerCopy";
+  public static readonly EnableHotPartitionKeyThrottling: string = "EnableHotPartitionKeyThrottling";
 }
 
 export enum CapacityMode {

src/Common/dataAccess/readCollectionOffer.ts

@@ -112,6 +112,7 @@ const readCollectionOfferWithARM = async (databaseId: string, collectionId: stri
         : resource.softAllowedMaximumThroughput;
 
     const throughputBuckets = resource?.throughputBuckets;
+    const hotPartitionKeyRateLimitingPolicy = resource?.hotPartitionKeyRateLimitingPolicy;
 
     if (autoscaleSettings) {
       return {
@@ -123,6 +124,7 @@ const readCollectionOfferWithARM = async (databaseId: string, collectionId: stri
         instantMaximumThroughput,
         softAllowedMaximumThroughput,
         throughputBuckets,
+        hotPartitionKeyRateLimitingPolicy,
       };
     }
 
@@ -135,6 +137,7 @@ const readCollectionOfferWithARM = async (databaseId: string, collectionId: stri
       instantMaximumThroughput,
       softAllowedMaximumThroughput,
       throughputBuckets,
+      hotPartitionKeyRateLimitingPolicy,
     };
   }
 

src/Common/dataAccess/updateOffer.ts

@@ -367,6 +367,10 @@ const createUpdateOfferBody = (params: UpdateOfferParams): ThroughputSettingsUpd
     body.properties.resource.throughputBuckets = throughputBuckets;
   }
 
+  if (params.hotPartitionKeyRateLimitingPolicy !== undefined) {
+    body.properties.resource.hotPartitionKeyRateLimitingPolicy = params.hotPartitionKeyRateLimitingPolicy;
+  }
+
   return body;
 };
 
@@ -409,6 +413,10 @@ const updateOfferWithSDK = async (params: UpdateOfferParams): Promise<Offer> =>
     newOffer.content.offerAutopilotSettings = { maxThroughput: 0 };
   }
 
+  if (params.hotPartitionKeyRateLimitingPolicy !== undefined) {
+    newOffer.content.hotPartitionKeyRateLimitingPolicy = params.hotPartitionKeyRateLimitingPolicy;
+  }
+
   const sdkResponse = await client()
     .offer(params.currentOffer.id)
     // TODO Remove casting when SDK types are fixed (https://github.com/Azure/azure-sdk-for-js/issues/10660)

src/Contracts/DataModels.ts

@@ -343,6 +343,11 @@ export interface Offer {
   instantMaximumThroughput?: number;
   softAllowedMaximumThroughput?: number;
   throughputBuckets?: ThroughputBucket[];
+  hotPartitionKeyRateLimitingPolicy?: HotPartitionKeyRateLimitingPolicy;
+}
+
+export interface HotPartitionKeyRateLimitingPolicy {
+  maximumPerPartitionKeyThroughputUtilizationPercent: number;
 }
 
 export interface ThroughputBucket {
@@ -359,6 +364,7 @@ export interface SDKOfferDefinition extends Resource {
     offerIsRUPerMinuteThroughputEnabled?: boolean;
     collectionThroughputInfo?: OfferThroughputInfo;
     offerAutopilotSettings?: AutoPilotOfferSettings;
+    hotPartitionKeyRateLimitingPolicy?: HotPartitionKeyRateLimitingPolicy;
   };
   resource?: string;
   offerResourceId?: string;
@@ -492,6 +498,7 @@ export interface UpdateOfferParams {
   migrateToAutoPilot?: boolean;
   migrateToManual?: boolean;
   throughputBuckets?: ThroughputBucket[];
+  hotPartitionKeyRateLimitingPolicy?: HotPartitionKeyRateLimitingPolicy | null;
 }
 
 export interface Notification {

src/Explorer/Controls/Settings/SettingsComponent.tsx

@@ -17,25 +17,25 @@ import { useIndexingPolicyStore } from "Explorer/Tabs/QueryTab/ResultsView";
 import { useDatabases } from "Explorer/useDatabases";
 import { Keys, t } from "Localization";
 import { isFabricNative } from "Platform/Fabric/FabricUtil";
-import { isVectorSearchEnabled } from "Utils/CapabilityUtils";
-import { isRunningOnPublicCloud } from "Utils/CloudUtils";
 import * as React from "react";
+import { isHotPartitionKeyThrottlingEnabled, isVectorSearchEnabled } from "Utils/CapabilityUtils";
+import { isRunningOnPublicCloud } from "Utils/CloudUtils";
 import DiscardIcon from "../../../../images/discard.svg";
 import SaveIcon from "../../../../images/save-cosmos.svg";
 import { AuthType } from "../../../AuthType";
 import * as Constants from "../../../Common/Constants";
-import { getErrorMessage, getErrorStack } from "../../../Common/ErrorHandlingUtils";
 import { getIndexTransformationProgress } from "../../../Common/dataAccess/getIndexTransformationProgress";
 import { readMongoDBCollectionThroughRP } from "../../../Common/dataAccess/readMongoDBCollection";
 import { updateCollection } from "../../../Common/dataAccess/updateCollection";
 import { updateOffer } from "../../../Common/dataAccess/updateOffer";
+import { getErrorMessage, getErrorStack } from "../../../Common/ErrorHandlingUtils";
 import * as DataModels from "../../../Contracts/DataModels";
 import * as ViewModels from "../../../Contracts/ViewModels";
 import { Action, ActionModifiers } from "../../../Shared/Telemetry/TelemetryConstants";
 import { trace, traceFailure, traceStart, traceSuccess } from "../../../Shared/Telemetry/TelemetryProcessor";
 import { userContext } from "../../../UserContext";
-import * as AutoPilotUtils from "../../../Utils/AutoPilotUtils";
 import { MongoDBCollectionResource, MongoIndex } from "../../../Utils/arm/generatedClients/cosmos/types";
+import * as AutoPilotUtils from "../../../Utils/AutoPilotUtils";
 import { CommandButtonComponentProps } from "../../Controls/CommandButton/CommandButtonComponent";
 import {
   PartitionKeyComponent,
@@ -65,16 +65,16 @@ import {
   AddMongoIndexProps,
   ChangeFeedPolicyState,
   GeospatialConfigType,
-  MongoIndexTypes,
-  SettingsV2TabTypes,
-  TtlType,
   getMongoNotification,
   getTabTitle,
   hasDatabaseSharedThroughput,
   isDataMaskingEnabled,
   isDirty,
+  MongoIndexTypes,
   parseConflictResolutionMode,
   parseConflictResolutionProcedure,
+  SettingsV2TabTypes,
+  TtlType,
 } from "./SettingsUtils";
 interface SettingsV2TabInfo {
   tab: SettingsV2TabTypes;
@@ -103,6 +103,8 @@ export interface SettingsComponentState {
   throughputBuckets: DataModels.ThroughputBucket[];
   throughputBucketsBaseline: DataModels.ThroughputBucket[];
   throughputError: string;
+  hotPartitionKeyRateLimitingPolicy: DataModels.HotPartitionKeyRateLimitingPolicy;
+  hotPartitionKeyRateLimitingPolicyBaseline: DataModels.HotPartitionKeyRateLimitingPolicy;
 
   timeToLive: TtlType;
   timeToLiveBaseline: TtlType;
@@ -225,6 +227,8 @@ export class SettingsComponent extends React.Component<SettingsComponentProps, S
       throughputBuckets: undefined,
       throughputBucketsBaseline: undefined,
       throughputError: undefined,
+      hotPartitionKeyRateLimitingPolicy: null,
+      hotPartitionKeyRateLimitingPolicyBaseline: null,
 
       timeToLive: undefined,
       timeToLiveBaseline: undefined,
@@ -495,6 +499,8 @@ export class SettingsComponent extends React.Component<SettingsComponentProps, S
       throughput: this.state.throughputBaseline,
       throughputBuckets: this.state.throughputBucketsBaseline,
       throughputBucketsBaseline: this.state.throughputBucketsBaseline,
+      hotPartitionKeyRateLimitingPolicy: this.state.hotPartitionKeyRateLimitingPolicyBaseline,
+      hotPartitionKeyRateLimitingPolicyBaseline: this.state.hotPartitionKeyRateLimitingPolicyBaseline,
       timeToLive: this.state.timeToLiveBaseline,
       timeToLiveSeconds: this.state.timeToLiveSecondsBaseline,
       displayedTtlSeconds: this.state.displayedTtlSecondsBaseline,
@@ -877,12 +883,15 @@ export class SettingsComponent extends React.Component<SettingsComponentProps, S
       ] as DataModels.ComputedProperties;
     }
     const throughputBuckets = this.offer?.throughputBuckets;
+    const hotPartitionKeyRateLimitingPolicy = this.offer?.hotPartitionKeyRateLimitingPolicy ?? null;
 
     return {
       throughput: offerThroughput,
       throughputBaseline: offerThroughput,
       throughputBuckets,
       throughputBucketsBaseline: throughputBuckets,
+      hotPartitionKeyRateLimitingPolicy,
+      hotPartitionKeyRateLimitingPolicyBaseline: hotPartitionKeyRateLimitingPolicy,
       changeFeedPolicy: changeFeedPolicy,
       changeFeedPolicyBaseline: changeFeedPolicy,
       timeToLive: timeToLive,
@@ -984,6 +993,12 @@ export class SettingsComponent extends React.Component<SettingsComponentProps, S
     this.setState({ throughputBuckets });
   };
 
+  private onHotPartitionKeyRateLimitingPolicyChange = (
+    hotPartitionKeyRateLimitingPolicy: DataModels.HotPartitionKeyRateLimitingPolicy,
+  ): void => {
+    this.setState({ hotPartitionKeyRateLimitingPolicy });
+  };
+
   private onAutoPilotSelected = (isAutoPilotSelected: boolean): void =>
     this.setState({ isAutoPilotSelected: isAutoPilotSelected });
 
@@ -999,6 +1014,9 @@ export class SettingsComponent extends React.Component<SettingsComponentProps, S
         currentOffer: this.database.offer(),
         autopilotThroughput: this.state.isAutoPilotSelected ? this.state.autoPilotThroughput : undefined,
         manualThroughput: this.state.isAutoPilotSelected ? undefined : this.state.throughput,
+        hotPartitionKeyRateLimitingPolicy: isHotPartitionKeyThrottlingEnabled()
+          ? this.state.hotPartitionKeyRateLimitingPolicy ?? null
+          : undefined,
       };
       if (this.hasProvisioningTypeChanged()) {
         if (this.state.isAutoPilotSelected) {
@@ -1232,6 +1250,9 @@ export class SettingsComponent extends React.Component<SettingsComponentProps, S
         autopilotThroughput: this.state.isAutoPilotSelected ? this.state.autoPilotThroughput : undefined,
         manualThroughput: this.state.isAutoPilotSelected ? undefined : this.state.throughput,
         throughputBuckets: this.throughputBucketsEnabled ? this.state.throughputBuckets : undefined,
+        hotPartitionKeyRateLimitingPolicy: isHotPartitionKeyThrottlingEnabled()
+          ? this.state.hotPartitionKeyRateLimitingPolicy ?? null
+          : undefined,
       };
       if (this.hasProvisioningTypeChanged()) {
         if (this.state.isAutoPilotSelected) {
@@ -1300,6 +1321,9 @@ export class SettingsComponent extends React.Component<SettingsComponentProps, S
       onScaleSaveableChange: this.onScaleSaveableChange,
       onScaleDiscardableChange: this.onScaleDiscardableChange,
       throughputError: this.state.throughputError,
+      hotPartitionKeyRateLimitingPolicy: this.state.hotPartitionKeyRateLimitingPolicy,
+      hotPartitionKeyRateLimitingPolicyBaseline: this.state.hotPartitionKeyRateLimitingPolicyBaseline,
+      onHotPartitionKeyRateLimitingPolicyChange: this.onHotPartitionKeyRateLimitingPolicyChange,
     };
     if (!this.isCollectionSettingsTab) {
       return (

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

@@ -29,6 +29,9 @@ describe("ScaleComponent", () => {
     onScaleDiscardableChange: () => {
       return;
     },
+    onHotPartitionKeyRateLimitingPolicyChange: () => {
+      return;
+    },
   };
 
   it("autoScale disabled", () => {

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

@@ -2,7 +2,7 @@ import { Link, MessageBar, MessageBarType, Stack, Text, TextField } from "@fluen
 import { Keys, t } from "Localization";
 import * as React from "react";
 import * as Constants from "../../../../Common/Constants";
-import { Platform, configContext } from "../../../../ConfigContext";
+import { configContext, Platform } from "../../../../ConfigContext";
 import * as DataModels from "../../../../Contracts/DataModels";
 import * as ViewModels from "../../../../Contracts/ViewModels";
 import * as SharedConstants from "../../../../Shared/Constants";
@@ -36,6 +36,9 @@ export interface ScaleComponentProps {
   onScaleSaveableChange: (isScaleSaveable: boolean) => void;
   onScaleDiscardableChange: (isScaleDiscardable: boolean) => void;
   throughputError?: string;
+  hotPartitionKeyRateLimitingPolicy?: DataModels.HotPartitionKeyRateLimitingPolicy;
+  hotPartitionKeyRateLimitingPolicyBaseline?: DataModels.HotPartitionKeyRateLimitingPolicy;
+  onHotPartitionKeyRateLimitingPolicyChange: (newPolicy: DataModels.HotPartitionKeyRateLimitingPolicy) => void;
 }
 
 export class ScaleComponent extends React.Component<ScaleComponentProps> {
@@ -148,6 +151,9 @@ export class ScaleComponent extends React.Component<ScaleComponentProps> {
       instantMaximumThroughput={this.offer?.instantMaximumThroughput}
       softAllowedMaximumThroughput={this.offer?.softAllowedMaximumThroughput}
       isGlobalSecondaryIndex={this.props.isGlobalSecondaryIndex}
+      hotPartitionKeyRateLimitingPolicy={this.props.hotPartitionKeyRateLimitingPolicy}
+      hotPartitionKeyRateLimitingPolicyBaseline={this.props.hotPartitionKeyRateLimitingPolicyBaseline}
+      onHotPartitionKeyRateLimitingPolicyChange={this.props.onHotPartitionKeyRateLimitingPolicyChange}
     />
   );
 

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

@@ -1,12 +1,24 @@
 import { shallow } from "enzyme";
 import React from "react";
+import * as Constants from "../../../../../Common/Constants";
 import * as DataModels from "../../../../../Contracts/DataModels";
+import { updateUserContext } from "../../../../../UserContext";
 import {
   ThroughputInputAutoPilotV3Component,
   ThroughputInputAutoPilotV3Props,
 } from "./ThroughputInputAutoPilotV3Component";
 
 describe("ThroughputInputAutoPilotV3Component", () => {
+  beforeAll(() => {
+    updateUserContext({
+      databaseAccount: {
+        properties: {
+          capabilities: [{ name: Constants.CapabilityNames.EnableHotPartitionKeyThrottling }],
+        },
+      } as DataModels.DatabaseAccount,
+    });
+  });
+
   const baseProps: ThroughputInputAutoPilotV3Props = {
     databaseAccount: {} as DataModels.DatabaseAccount,
     databaseName: "test",
@@ -45,6 +57,9 @@ describe("ThroughputInputAutoPilotV3Component", () => {
     instantMaximumThroughput: 5000,
     softAllowedMaximumThroughput: 1000000,
     isGlobalSecondaryIndex: false,
+    onHotPartitionKeyRateLimitingPolicyChange: () => {
+      return;
+    },
   };
 
   it("throughput input visible", () => {