Set default throughput based on account's workload type (#2021)

* assign default throughput based on workload type

* combined common logic

* fix unit tests

* add tests

* update tests

* npm run format

* Update ci.yml

---------

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

.github/workflows/ci.yml

@@ -83,7 +83,7 @@ jobs:
       - run: npm ci
       - run: npm run build:contracts
       - name: Restore Build Cache
-        uses: actions/cache@v2
+        uses: actions/cache@v4
         with:
           path: .cache
           key: ${{ runner.os }}-build-cache

src/Common/Constants.ts

@@ -97,6 +97,12 @@ export enum CapacityMode {
   Serverless = "Serverless",
 }
 
+export enum WorkloadType {
+  Learning = "Learning",
+  DevelopmentTesting = "Development/Testing",
+  Production = "Production",
+  None = "None",
+}
 // flight names returned from the portal are always lowercase
 export class Flights {
   public static readonly SettingsV2 = "settingsv2";
@@ -119,6 +125,7 @@ export class AfecFeatures {
 
 export class TagNames {
   public static defaultExperience: string = "defaultExperience";
+  public static WorkloadType: string = "hidden-workload-type";
 }
 
 export class MongoDBAccounts {

src/Common/DatabaseAccountUtility.test.ts

@@ -0,0 +1,34 @@
+import { WorkloadType } from "Common/Constants";
+import { getWorkloadType } from "Common/DatabaseAccountUtility";
+import { DatabaseAccount, Tags } from "Contracts/DataModels";
+import { updateUserContext } from "UserContext";
+
+describe("Database Account Utility", () => {
+  describe("Workload Type", () => {
+    beforeEach(() => {
+      updateUserContext({
+        databaseAccount: {
+          tags: {} as Tags,
+        } as DatabaseAccount,
+      });
+    });
+
+    it("Workload Type should return Learning", () => {
+      updateUserContext({
+        databaseAccount: {
+          tags: {
+            "hidden-workload-type": WorkloadType.Learning,
+          } as Tags,
+        } as DatabaseAccount,
+      });
+
+      const workloadType: WorkloadType = getWorkloadType();
+      expect(workloadType).toBe(WorkloadType.Learning);
+    });
+
+    it("Workload Type should return None", () => {
+      const workloadType: WorkloadType = getWorkloadType();
+      expect(workloadType).toBe(WorkloadType.None);
+    });
+  });
+});

src/Common/DatabaseAccountUtility.ts

@@ -1,3 +1,5 @@
+import { TagNames, WorkloadType } from "Common/Constants";
+import { Tags } from "Contracts/DataModels";
 import { userContext } from "../UserContext";
 
 function isVirtualNetworkFilterEnabled() {
@@ -15,3 +17,12 @@ function isPrivateEndpointConnectionsEnabled() {
 export function isPublicInternetAccessAllowed(): boolean {
   return !isVirtualNetworkFilterEnabled() && !isIpRulesEnabled() && !isPrivateEndpointConnectionsEnabled();
 }
+
+export function getWorkloadType(): WorkloadType {
+  const tags: Tags = userContext?.databaseAccount?.tags;
+  const workloadType: WorkloadType = tags && (tags[TagNames.WorkloadType] as WorkloadType);
+  if (!workloadType) {
+    return WorkloadType.None;
+  }
+  return workloadType;
+}

src/Contracts/DataModels.ts

@@ -6,6 +6,7 @@ export interface ArmEntity {
   location: string;
   type: string;
   kind: string;
+  tags?: Tags;
 }
 
 export interface DatabaseAccount extends ArmEntity {
@@ -663,3 +664,5 @@ export interface FeatureRegistration {
     state: string;
   };
 }
+
+export type Tags = { [key: string]: string };

src/Explorer/Controls/ThroughputInput/ThroughputInput.tsx

@@ -1,4 +1,5 @@
 import { Checkbox, DirectionalHint, Link, Stack, Text, TextField, TooltipHost } from "@fluentui/react";
+import { getWorkloadType } from "Common/DatabaseAccountUtility";
 import { useDatabases } from "Explorer/useDatabases";
 import React, { FunctionComponent, useEffect, useState } from "react";
 import * as Constants from "../../../Common/Constants";
@@ -34,10 +35,15 @@ export const ThroughputInput: FunctionComponent<ThroughputInputProps> = ({
   setIsThroughputCapExceeded,
   onCostAcknowledgeChange,
 }: ThroughputInputProps) => {
+  const defaultThroughput: number =
+    isFreeTier ||
+    isQuickstart ||
+    [Constants.WorkloadType.Learning, Constants.WorkloadType.DevelopmentTesting].includes(getWorkloadType())
+      ? AutoPilotUtils.autoPilotThroughput1K
+      : AutoPilotUtils.autoPilotThroughput4K;
+
   const [isAutoscaleSelected, setIsAutoScaleSelected] = useState<boolean>(true);
-  const [throughput, setThroughput] = useState<number>(
-    isFreeTier || isQuickstart ? AutoPilotUtils.autoPilotThroughput1K : AutoPilotUtils.autoPilotThroughput4K,
-  );
+  const [throughput, setThroughput] = useState<number>(defaultThroughput);
   const [isCostAcknowledged, setIsCostAcknowledged] = useState<boolean>(false);
   const [throughputError, setThroughputError] = useState<string>("");
   const [totalThroughputUsed, setTotalThroughputUsed] = useState<number>(0);
@@ -47,7 +53,6 @@ export const ThroughputInput: FunctionComponent<ThroughputInputProps> = ({
 
   const throughputCap = userContext.databaseAccount?.properties.capacity?.totalThroughputLimit;
   const numberOfRegions = userContext.databaseAccount?.properties.locations?.length || 1;
-
   useEffect(() => {
     // throughput cap check for the initial state
     let totalThroughput = 0;
@@ -157,9 +162,6 @@ export const ThroughputInput: FunctionComponent<ThroughputInputProps> = ({
 
   const handleOnChangeMode = (event: React.ChangeEvent<HTMLInputElement>, mode: string): void => {
     if (mode === "Autoscale") {
-      const defaultThroughput = isFreeTier
-        ? AutoPilotUtils.autoPilotThroughput1K
-        : AutoPilotUtils.autoPilotThroughput4K;
       setThroughput(defaultThroughput);
       setIsAutoScaleSelected(true);
       setThroughputValue(defaultThroughput);