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>

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);
+};