Wmt priority execution feature (#1546)

* Execute the queries with high/low priority

* improvement made to the Cosmos client and created separate plugin class for setting priority header

* added test cases for the priority execution utility

* removed unwanted code

* fix compile time issues

* fix compile time issues

* fixed lint and stylinkg issues

* fixed lint and styling issues

* skip the lint check for src/Utils/PriorityBasedExecutionUtils.ts

* incorporating review comments, added the default priority level changes

* changed the priority to default instead of low

* removed the unwanted if condition

---------

Co-authored-by: Faiz Chachiya <faizchachiya@microsoft.com>

.eslintignore

@@ -144,4 +144,5 @@ src/Explorer/Notebook/NotebookRenderer/decorators/kbd-shortcuts/index.tsx
 src/Explorer/Notebook/temp/inputs/connected-editors/codemirror.tsx
 src/Explorer/Tree/ResourceTreeAdapter.tsx
 __mocks__/monaco-editor.ts
-src/Explorer/Tree/ResourceTree.tsx
+src/Explorer/Tree/ResourceTree.tsx
+src/Utils/PriorityBasedExecutionUtils.ts

src/Common/Constants.ts

@@ -600,3 +600,9 @@ export const QueryCopilotSampleContainerSchema = {
     },
   },
 };
+
+export class PriorityLevel {
+  public static readonly High = "high";
+  public static readonly Low = "low";
+  public static readonly Default = "low";
+}

src/Common/CosmosClient.ts

@@ -4,6 +4,9 @@ import { userContext } from "../UserContext";
 import { logConsoleError } from "../Utils/NotificationConsoleUtils";
 import { EmulatorMasterKey, HttpHeaders } from "./Constants";
 import { getErrorMessage } from "./ErrorHandlingUtils";
+import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility";
+import { PriorityLevel } from "../Common/Constants";
+import * as PriorityBasedExecutionUtils from "../Utils/PriorityBasedExecutionUtils";
 
 const _global = typeof self === "undefined" ? window : self;
 
@@ -105,6 +108,13 @@ export function client(): Cosmos.CosmosClient {
   if (configContext.PROXY_PATH !== undefined) {
     (options as any).plugins = [{ on: "request", plugin: requestPlugin }];
   }
+
+  if (userContext.features.enablePriorityBasedThrottling && userContext.apiType === "SQL") {
+    const plugins = (options as any).plugins || [];
+    plugins.push({ on: "request", plugin: PriorityBasedExecutionUtils.requestPlugin });
+    (options as any).plugins = plugins;
+  }
+
   _client = new Cosmos.CosmosClient(options);
   return _client;
 }

src/Explorer/Explorer.tsx

@@ -650,6 +650,8 @@ export default class Explorer {
   private _initSettings() {
     if (!ExplorerSettings.hasSettingsDefined()) {
       ExplorerSettings.createDefaultSettings();
+    } else {
+      ExplorerSettings.ensurePriorityLevel();
     }
   }
 

src/Explorer/Panes/SettingsPane/SettingsPane.tsx

@@ -41,12 +41,18 @@ export const SettingsPane: FunctionComponent = () => {
       ? LocalStorageUtility.getEntryNumber(StorageKey.MaxDegreeOfParellism)
       : Constants.Queries.DefaultMaxDegreeOfParallelism
   );
+  const [priorityLevel, setPriorityLevel] = useState<string>(
+    LocalStorageUtility.hasItem(StorageKey.PriorityLevel)
+      ? LocalStorageUtility.getEntryString(StorageKey.PriorityLevel)
+      : Constants.PriorityLevel.Default
+  );
   const explorerVersion = configContext.gitSha;
   const shouldShowQueryPageOptions = userContext.apiType === "SQL";
   const shouldShowGraphAutoVizOption = userContext.apiType === "Gremlin";
   const shouldShowCrossPartitionOption = userContext.apiType !== "Gremlin";
   const shouldShowParallelismOption = userContext.apiType !== "Gremlin";
-
+  const shouldShowPriorityLevelOption =
+    userContext.features.enablePriorityBasedThrottling && userContext.apiType === "SQL";
   const handlerOnSubmit = (e: MouseEvent<HTMLButtonElement>) => {
     setIsExecuting(true);
 
@@ -58,6 +64,7 @@ export const SettingsPane: FunctionComponent = () => {
     LocalStorageUtility.setEntryString(StorageKey.ContainerPaginationEnabled, containerPaginationEnabled.toString());
     LocalStorageUtility.setEntryString(StorageKey.IsCrossPartitionQueryEnabled, crossPartitionQueryEnabled.toString());
     LocalStorageUtility.setEntryNumber(StorageKey.MaxDegreeOfParellism, maxDegreeOfParallelism);
+    LocalStorageUtility.setEntryString(StorageKey.PriorityLevel, priorityLevel.toString());
 
     if (shouldShowGraphAutoVizOption) {
       LocalStorageUtility.setEntryBoolean(
@@ -76,6 +83,7 @@ export const SettingsPane: FunctionComponent = () => {
         StorageKey.MaxDegreeOfParellism
       )}`
     );
+    logConsoleInfo(`Updated priority level setting to ${LocalStorageUtility.getEntryString(StorageKey.PriorityLevel)}`);
 
     if (shouldShowGraphAutoVizOption) {
       logConsoleInfo(
@@ -116,6 +124,18 @@ export const SettingsPane: FunctionComponent = () => {
     { key: "true", text: "JSON" },
   ];
 
+  const priorityLevelOptionList: IChoiceGroupOption[] = [
+    { key: Constants.PriorityLevel.Low, text: "Low" },
+    { key: Constants.PriorityLevel.High, text: "High" },
+  ];
+
+  const handleOnPriorityLevelOptionChange = (
+    ev: React.FormEvent<HTMLInputElement>,
+    option: IChoiceGroupOption
+  ): void => {
+    setPriorityLevel(option.key);
+  };
+
   const handleOnPageOptionChange = (ev: React.FormEvent<HTMLInputElement>, option: IChoiceGroupOption): void => {
     setPageOption(option.key);
   };
@@ -260,6 +280,29 @@ export const SettingsPane: FunctionComponent = () => {
             </div>
           </div>
         )}
+        {shouldShowPriorityLevelOption && (
+          <div className="settingsSection">
+            <div className="settingsSectionPart">
+              <fieldset>
+                <legend id="priorityLevel" className="settingsSectionLabel legendLabel">
+                  Priority Level
+                </legend>
+                <InfoTooltip>
+                  Sets the priority level for data-plane requests from Data Explorer when using Priority-Based
+                  Execution. If &quot;None&quot; is selected, Data Explorer will not specify priority level, and the
+                  server-side default priority level will be used.
+                </InfoTooltip>
+                <ChoiceGroup
+                  ariaLabelledBy="priorityLevel"
+                  selectedKey={priorityLevel}
+                  options={priorityLevelOptionList}
+                  styles={choiceButtonStyles}
+                  onChange={handleOnPriorityLevelOptionChange}
+                />
+              </fieldset>
+            </div>
+          </div>
+        )}
         {shouldShowGraphAutoVizOption && (
           <div className="settingsSection">
             <div className="settingsSectionPart">

src/Platform/Hosted/extractFeatures.ts

@@ -36,6 +36,7 @@ export type Features = {
   readonly enableLegacyMongoShellV2Debug: boolean;
   readonly loadLegacyMongoShellFromBE: boolean;
   readonly enableCopilot: boolean;
+  readonly enablePriorityBasedThrottling: boolean;
   readonly enableNPSSurvey: boolean;
   readonly copilotVersion?: string;
 
@@ -106,6 +107,7 @@ export function extractFeatures(given = new URLSearchParams(window.location.sear
     enableLegacyMongoShellV2Debug: "true" === get("enablelegacymongoshellv2debug"),
     loadLegacyMongoShellFromBE: "true" === get("loadlegacymongoshellfrombe"),
     enableCopilot: "true" === get("enablecopilot"),
+    enablePriorityBasedThrottling: "true" === get("enableprioritybasedthrottling"),
     enableNPSSurvey: "true" === get("enablenpssurvey"),
     copilotVersion: get("copilotVersion") ? get("copilotVersion") : "v1.0",
   };

src/Shared/ExplorerSettings.ts

@@ -6,6 +6,7 @@ export const createDefaultSettings = () => {
   LocalStorageUtility.setEntryNumber(StorageKey.CustomItemPerPage, Constants.Queries.itemsPerPage);
   LocalStorageUtility.setEntryString(StorageKey.IsCrossPartitionQueryEnabled, "true");
   LocalStorageUtility.setEntryNumber(StorageKey.MaxDegreeOfParellism, Constants.Queries.DefaultMaxDegreeOfParallelism);
+  LocalStorageUtility.setEntryString(StorageKey.PriorityLevel, Constants.PriorityLevel.Default);
 };
 
 export const hasSettingsDefined = (): boolean => {
@@ -15,3 +16,9 @@ export const hasSettingsDefined = (): boolean => {
     LocalStorageUtility.hasItem(StorageKey.MaxDegreeOfParellism)
   );
 };
+
+export const ensurePriorityLevel = () => {
+  if (!LocalStorageUtility.hasItem(StorageKey.PriorityLevel)) {
+    LocalStorageUtility.setEntryString(StorageKey.PriorityLevel, Constants.PriorityLevel.Default);
+  }
+};

src/Shared/StorageUtility.ts

@@ -16,4 +16,5 @@ export enum StorageKey {
   SetPartitionKeyUndefined,
   GalleryCalloutDismissed,
   VisitedAccounts,
+  PriorityLevel,
 }

src/Utils/PriorityBasedExecutionUtils.test.ts

@@ -0,0 +1,59 @@
+import * as PriorityBasedExecutionUtils from "./PriorityBasedExecutionUtils";
+import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility";
+import { PriorityLevel } from "../Common/Constants";
+import * as Cosmos from "@azure/cosmos";
+
+describe("Priority execution utility", () => {
+  it("check default priority level is Low", () => {
+    expect(PriorityBasedExecutionUtils.getPriorityLevel()).toEqual(PriorityLevel.Low);
+  });
+
+  it("check the priority level is returned as present in local storage", () => {
+    LocalStorageUtility.setEntryString(StorageKey.PriorityLevel, PriorityLevel.High);
+    expect(PriorityBasedExecutionUtils.getPriorityLevel()).toEqual(PriorityLevel.High);
+  });
+
+  it("check relevant request based on different resource types", () => {
+    const requestContext1: Cosmos.RequestContext = {
+      resourceType: Cosmos.ResourceType.item,
+      globalEndpointManager: null,
+      connectionPolicy: null,
+      requestAgent: null,
+      method: null,
+      options: null,
+      plugins: null,
+    };
+    const requestContext2: Cosmos.RequestContext = {
+      resourceType: Cosmos.ResourceType.conflicts,
+      globalEndpointManager: null,
+      connectionPolicy: null,
+      requestAgent: null,
+      method: null,
+      options: null,
+      plugins: null,
+    };
+    const requestContext3: Cosmos.RequestContext = {
+      resourceType: Cosmos.ResourceType.sproc,
+      operationType: Cosmos.OperationType.Execute,
+      globalEndpointManager: null,
+      connectionPolicy: null,
+      requestAgent: null,
+      method: null,
+      options: null,
+      plugins: null,
+    };
+    const requestContext4: Cosmos.RequestContext = {
+      resourceType: Cosmos.ResourceType.database,
+      globalEndpointManager: null,
+      connectionPolicy: null,
+      requestAgent: null,
+      method: null,
+      options: null,
+      plugins: null,
+    };
+    expect(PriorityBasedExecutionUtils.isRelevantRequest(requestContext1)).toEqual(true);
+    expect(PriorityBasedExecutionUtils.isRelevantRequest(requestContext2)).toEqual(true);
+    expect(PriorityBasedExecutionUtils.isRelevantRequest(requestContext3)).toEqual(true);
+    expect(PriorityBasedExecutionUtils.isRelevantRequest(requestContext4)).toEqual(false);
+  });
+});

src/Utils/PriorityBasedExecutionUtils.ts

@@ -0,0 +1,29 @@
+import * as Cosmos from "@azure/cosmos";
+import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility";
+import { PriorityLevel } from "../Common/Constants";
+
+export function isRelevantRequest(requestContext: Cosmos.RequestContext): boolean {
+  return (
+    requestContext.resourceType === Cosmos.ResourceType.item ||
+    requestContext.resourceType === Cosmos.ResourceType.conflicts ||
+    (requestContext.resourceType === Cosmos.ResourceType.sproc &&
+      requestContext.operationType === Cosmos.OperationType.Execute)
+  );
+}
+
+export function getPriorityLevel(): PriorityLevel {
+  const priorityLevel = LocalStorageUtility.getEntryString(StorageKey.PriorityLevel);
+  if (priorityLevel && Object.values(PriorityLevel).includes(priorityLevel)) {
+    return priorityLevel as PriorityLevel;
+  } else {
+    return PriorityLevel.Default;
+  }
+}
+
+export const requestPlugin: Cosmos.Plugin<any> = async (requestContext, next) => {
+  if (isRelevantRequest(requestContext)) {
+    const priorityLevel: PriorityLevel = getPriorityLevel();
+    requestContext.headers["x-ms-cosmos-priority-level"] = priorityLevel as string;
+  }
+  return next(requestContext);
+};