Merge branch 'master' of https://github.com/Azure/cosmos-explorer into users/aisayas/pbe

src/Common/Constants.ts

@@ -427,12 +427,6 @@ export class JunoEndpoints {
   public static readonly Stage = "https://tools-staging.cosmos.azure.com";
 }
 
-export class PriorityLevel {
-  public static readonly High = "high";
-  public static readonly Low = "low";
-  public static readonly Default = "low";
-}
-
 export const QueryCopilotSampleDatabaseId = "CopilotSampleDb";
 export const QueryCopilotSampleContainerId = "SampleContainer";
 

src/Common/CosmosClient.ts

@@ -1,8 +1,6 @@
 import * as Cosmos from "@azure/cosmos";
 import { sendCachedDataMessage } from "Common/MessageHandler";
 import { AuthorizationToken, MessageTypes } from "Contracts/MessageTypes";
-import { AuthType } from "../AuthType";
-import { PriorityLevel } from "../Common/Constants";
 import { Platform, configContext } from "../ConfigContext";
 import { userContext } from "../UserContext";
 import { logConsoleError } from "../Utils/NotificationConsoleUtils";
@@ -109,18 +107,6 @@ export function client(): Cosmos.CosmosClient {
   _defaultHeaders["x-ms-cosmos-sdk-supportedcapabilities"] =
     SDKSupportedCapabilities.None | SDKSupportedCapabilities.PartitionMerge;
 
-  if (
-    userContext.authType === AuthType.ConnectionString ||
-    userContext.authType === AuthType.EncryptedToken ||
-    userContext.authType === AuthType.ResourceToken
-  ) {
-    // Default to low priority. Needed for non-AAD-auth scenarios
-    // where we cannot use RP API, and thus, cannot detect whether priority
-    // based execution is enabled.
-    // The header will be ignored if priority based execution is disabled on the account.
-    _defaultHeaders["x-ms-cosmos-priority-level"] = PriorityLevel.Default;
-  }
-
   const options: Cosmos.CosmosClientOptions = {
     endpoint: endpoint() || "https://cosmos.azure.com", // CosmosClient gets upset if we pass a bad URL. This should never actually get called
     key: userContext.masterKey,

src/Explorer/Explorer.tsx

@@ -689,9 +689,9 @@ export default class Explorer {
   private _initSettings() {
     if (!ExplorerSettings.hasSettingsDefined()) {
       ExplorerSettings.createDefaultSettings();
-    } else {
-      ExplorerSettings.ensurePriorityLevel();
     }
+
+    ExplorerSettings.ensurePriorityLevel();
   }
 
   public uploadFile(

src/Explorer/Panes/SettingsPane/SettingsPane.tsx

@@ -1,3 +1,4 @@
+import { PriorityLevel } from "@azure/cosmos";
 import {
   Checkbox,
   ChoiceGroup,
@@ -58,10 +59,10 @@ export const SettingsPane: FunctionComponent = () => {
       ? LocalStorageUtility.getEntryNumber(StorageKey.MaxDegreeOfParellism)
       : Constants.Queries.DefaultMaxDegreeOfParallelism,
   );
-  const [priorityLevel, setPriorityLevel] = useState<string>(
+  const [priorityLevel, setPriorityLevel] = useState<PriorityLevel>(
     LocalStorageUtility.hasItem(StorageKey.PriorityLevel)
-      ? LocalStorageUtility.getEntryString(StorageKey.PriorityLevel)
-      : Constants.PriorityLevel.Default,
+      ? (LocalStorageUtility.getEntryString(StorageKey.PriorityLevel) as PriorityLevel)
+      : PriorityLevel.Low,
   );
   const explorerVersion = configContext.gitSha;
   const shouldShowQueryPageOptions = userContext.apiType === "SQL";
@@ -81,7 +82,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());
+    LocalStorageUtility.setEntryString(StorageKey.PriorityLevel, priorityLevel);
 
     if (shouldShowGraphAutoVizOption) {
       LocalStorageUtility.setEntryBoolean(
@@ -149,15 +150,15 @@ export const SettingsPane: FunctionComponent = () => {
   ];
 
   const priorityLevelOptionList: IChoiceGroupOption[] = [
-    { key: Constants.PriorityLevel.Low, text: "Low" },
-    { key: Constants.PriorityLevel.High, text: "High" },
+    { key: PriorityLevel.Low, text: "Low" },
+    { key: PriorityLevel.High, text: "High" },
   ];
 
   const handleOnPriorityLevelOptionChange = (
     ev: React.FormEvent<HTMLInputElement>,
     option: IChoiceGroupOption,
   ): void => {
-    setPriorityLevel(option.key);
+    setPriorityLevel(option.key as PriorityLevel);
   };
 
   const handleOnPageOptionChange = (ev: React.FormEvent<HTMLInputElement>, option: IChoiceGroupOption): void => {
@@ -401,8 +402,7 @@ export const SettingsPane: FunctionComponent = () => {
                 </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.
+                  Execution.
                 </InfoTooltip>
                 <ChoiceGroup
                   ariaLabelledBy="priorityLevel"

src/Shared/ExplorerSettings.ts

@@ -1,3 +1,4 @@
+import { PriorityLevel } from "@azure/cosmos";
 import * as Constants from "../Common/Constants";
 import { LocalStorageUtility, StorageKey } from "./StorageUtility";
 
@@ -6,7 +7,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);
+  LocalStorageUtility.setEntryString(StorageKey.PriorityLevel, PriorityLevel.Low);
 };
 
 export const hasSettingsDefined = (): boolean => {
@@ -19,6 +20,6 @@ export const hasSettingsDefined = (): boolean => {
 
 export const ensurePriorityLevel = () => {
   if (!LocalStorageUtility.hasItem(StorageKey.PriorityLevel)) {
-    LocalStorageUtility.setEntryString(StorageKey.PriorityLevel, Constants.PriorityLevel.Default);
+    LocalStorageUtility.setEntryString(StorageKey.PriorityLevel, PriorityLevel.Low);
   }
 };

src/Utils/PriorityBasedExecutionUtils.test.ts

@@ -1,7 +1,7 @@
-import * as PriorityBasedExecutionUtils from "./PriorityBasedExecutionUtils";
-import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility";
-import { PriorityLevel } from "../Common/Constants";
 import * as Cosmos from "@azure/cosmos";
+import { PriorityLevel } from "@azure/cosmos";
+import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility";
+import * as PriorityBasedExecutionUtils from "./PriorityBasedExecutionUtils";
 
 describe("Priority execution utility", () => {
   it("check default priority level is Low", () => {

src/Utils/PriorityBasedExecutionUtils.ts

@@ -1,6 +1,6 @@
 import * as Cosmos from "@azure/cosmos";
+import { Constants, PriorityLevel } from "@azure/cosmos";
 import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility";
-import { PriorityLevel } from "../Common/Constants";
 import { userContext } from "../UserContext";
 
 export function isFeatureEnabled(): boolean {
@@ -21,18 +21,18 @@ export function isRelevantRequest(requestContext: Cosmos.RequestContext): boolea
 }
 
 export function getPriorityLevel(): PriorityLevel {
-  const priorityLevel = LocalStorageUtility.getEntryString(StorageKey.PriorityLevel);
-  if (priorityLevel && Object.values(PriorityLevel).includes(priorityLevel)) {
+  const priorityLevel: string = LocalStorageUtility.getEntryString(StorageKey.PriorityLevel);
+  if (priorityLevel) {
     return priorityLevel as PriorityLevel;
   } else {
-    return PriorityLevel.Default;
+    return PriorityLevel.Low;
   }
 }
 
 export const requestPlugin: Cosmos.Plugin<any> = async (requestContext, undefined, next) => {
   if (isRelevantRequest(requestContext)) {
     const priorityLevel: PriorityLevel = getPriorityLevel();
-    requestContext.headers["x-ms-cosmos-priority-level"] = priorityLevel as string;
+    requestContext.headers[Constants.HttpHeaders.PriorityLevel] = priorityLevel;
   }
   return next(requestContext);
 };