diff --git a/src/Common/Constants.ts b/src/Common/Constants.ts index 46b890d29..f342dac64 100644 --- a/src/Common/Constants.ts +++ b/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"; diff --git a/src/Common/CosmosClient.ts b/src/Common/CosmosClient.ts index b8a9573a9..33504c768 100644 --- a/src/Common/CosmosClient.ts +++ b/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, diff --git a/src/Explorer/Explorer.tsx b/src/Explorer/Explorer.tsx index 23ee27e3e..6271ca9ac 100644 --- a/src/Explorer/Explorer.tsx +++ b/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( diff --git a/src/Explorer/Panes/SettingsPane/SettingsPane.tsx b/src/Explorer/Panes/SettingsPane/SettingsPane.tsx index fc31c7f7d..fcf2b1d29 100644 --- a/src/Explorer/Panes/SettingsPane/SettingsPane.tsx +++ b/src/Explorer/Panes/SettingsPane/SettingsPane.tsx @@ -1,3 +1,4 @@ +import { PriorityLevel } from "@azure/cosmos"; import { Checkbox, ChoiceGroup, IChoiceGroupOption, SpinButton } from "@fluentui/react"; import * as Constants from "Common/Constants"; import { InfoTooltip } from "Common/Tooltip/InfoTooltip"; @@ -6,10 +7,10 @@ import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility"; import * as StringUtility from "Shared/StringUtility"; import { userContext } from "UserContext"; import { logConsoleInfo } from "Utils/NotificationConsoleUtils"; +import * as PriorityBasedExecutionUtils from "Utils/PriorityBasedExecutionUtils"; import { useSidePanel } from "hooks/useSidePanel"; import React, { FunctionComponent, MouseEvent, useState } from "react"; import { RightPaneForm, RightPaneFormProps } from "../RightPaneForm/RightPaneForm"; -import * as PriorityBasedExecutionUtils from "Utils/PriorityBasedExecutionUtils"; export const SettingsPane: FunctionComponent = () => { const closeSidePanel = useSidePanel((state) => state.closeSidePanel); @@ -42,10 +43,10 @@ export const SettingsPane: FunctionComponent = () => { ? LocalStorageUtility.getEntryNumber(StorageKey.MaxDegreeOfParellism) : Constants.Queries.DefaultMaxDegreeOfParallelism, ); - const [priorityLevel, setPriorityLevel] = useState( + const [priorityLevel, setPriorityLevel] = useState( 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"; @@ -64,7 +65,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( @@ -125,15 +126,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, option: IChoiceGroupOption, ): void => { - setPriorityLevel(option.key); + setPriorityLevel(option.key as PriorityLevel); }; const handleOnPageOptionChange = (ev: React.FormEvent, option: IChoiceGroupOption): void => { @@ -289,8 +290,7 @@ export const SettingsPane: FunctionComponent = () => { Sets the priority level for data-plane requests from Data Explorer when using Priority-Based - Execution. If "None" is selected, Data Explorer will not specify priority level, and the - server-side default priority level will be used. + Execution. { 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); } }; diff --git a/src/Utils/PriorityBasedExecutionUtils.test.ts b/src/Utils/PriorityBasedExecutionUtils.test.ts index 40c9881fc..d6beaf4c0 100644 --- a/src/Utils/PriorityBasedExecutionUtils.test.ts +++ b/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", () => { diff --git a/src/Utils/PriorityBasedExecutionUtils.ts b/src/Utils/PriorityBasedExecutionUtils.ts index efb80386d..4091889e7 100644 --- a/src/Utils/PriorityBasedExecutionUtils.ts +++ b/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 = 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); };