diff --git a/src/Common/CosmosClient.ts b/src/Common/CosmosClient.ts index a32713421..84d80b0e2 100644 --- a/src/Common/CosmosClient.ts +++ b/src/Common/CosmosClient.ts @@ -7,6 +7,7 @@ import { getErrorMessage } from "./ErrorHandlingUtils"; import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility"; import { PriorityLevel } from "../Common/Constants"; import * as PriorityBasedExecutionUtils from "../Utils/PriorityBasedExecutionUtils"; +import { AuthType } from "../AuthType"; const _global = typeof self === "undefined" ? window : self; @@ -94,6 +95,18 @@ 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, @@ -109,7 +122,7 @@ export function client(): Cosmos.CosmosClient { (options as any).plugins = [{ on: "request", plugin: requestPlugin }]; } - if (userContext.databaseAccount?.properties?.enablePriorityBasedExecution && userContext.apiType === "SQL") { + if (PriorityBasedExecutionUtils.isFeatureEnabled()) { const plugins = (options as any).plugins || []; plugins.push({ on: "request", plugin: PriorityBasedExecutionUtils.requestPlugin }); (options as any).plugins = plugins; diff --git a/src/Explorer/Panes/SettingsPane/SettingsPane.tsx b/src/Explorer/Panes/SettingsPane/SettingsPane.tsx index c3f7c2f81..fc31c7f7d 100644 --- a/src/Explorer/Panes/SettingsPane/SettingsPane.tsx +++ b/src/Explorer/Panes/SettingsPane/SettingsPane.tsx @@ -9,6 +9,7 @@ import { logConsoleInfo } from "Utils/NotificationConsoleUtils"; 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); @@ -51,8 +52,7 @@ export const SettingsPane: FunctionComponent = () => { const shouldShowGraphAutoVizOption = userContext.apiType === "Gremlin"; const shouldShowCrossPartitionOption = userContext.apiType !== "Gremlin"; const shouldShowParallelismOption = userContext.apiType !== "Gremlin"; - const shouldShowPriorityLevelOption = - userContext.databaseAccount?.properties?.enablePriorityBasedExecution && userContext.apiType === "SQL"; + const shouldShowPriorityLevelOption = PriorityBasedExecutionUtils.isFeatureEnabled(); const handlerOnSubmit = (e: MouseEvent) => { setIsExecuting(true); diff --git a/src/Platform/Hosted/extractFeatures.ts b/src/Platform/Hosted/extractFeatures.ts index 4d641c4d5..3e34ced0b 100644 --- a/src/Platform/Hosted/extractFeatures.ts +++ b/src/Platform/Hosted/extractFeatures.ts @@ -40,6 +40,7 @@ export type Features = { readonly disableCopilotPhoenixGateaway: boolean; readonly enableCopilotFullSchema: boolean; readonly copilotChatFixedMonacoEditorHeight: boolean; + readonly enablePriorityBasedExecution: boolean; // can be set via both flight and feature flag autoscaleDefault: boolean; @@ -112,6 +113,7 @@ export function extractFeatures(given = new URLSearchParams(window.location.sear disableCopilotPhoenixGateaway: "true" === get("disablecopilotphoenixgateaway"), enableCopilotFullSchema: "true" === get("enablecopilotfullschema", "true"), copilotChatFixedMonacoEditorHeight: "true" === get("copilotchatfixedmonacoeditorheight"), + enablePriorityBasedExecution: "true" === get("enableprioritybasedexecution"), }; } diff --git a/src/Utils/PriorityBasedExecutionUtils.ts b/src/Utils/PriorityBasedExecutionUtils.ts index 08c272574..1295f4e61 100644 --- a/src/Utils/PriorityBasedExecutionUtils.ts +++ b/src/Utils/PriorityBasedExecutionUtils.ts @@ -1,6 +1,15 @@ import * as Cosmos from "@azure/cosmos"; import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility"; import { PriorityLevel } from "../Common/Constants"; +import { userContext } from "../UserContext"; + +export function isFeatureEnabled(): boolean { + return ( + userContext.apiType === "SQL" && + (userContext.databaseAccount?.properties?.enablePriorityBasedExecution || + userContext.features.enablePriorityBasedExecution) + ); +} export function isRelevantRequest(requestContext: Cosmos.RequestContext): boolean { return ( diff --git a/src/hooks/useDatabaseAccounts.tsx b/src/hooks/useDatabaseAccounts.tsx index e0f62c3e3..b3bf396ae 100644 --- a/src/hooks/useDatabaseAccounts.tsx +++ b/src/hooks/useDatabaseAccounts.tsx @@ -1,7 +1,6 @@ import useSWR from "swr"; import { configContext } from "../ConfigContext"; import { DatabaseAccount } from "../Contracts/DataModels"; -import { userContext } from "../UserContext"; interface AccountListResult { nextLink: string; @@ -15,7 +14,7 @@ export async function fetchDatabaseAccounts(subscriptionId: string, accessToken: headers.append("Authorization", bearer); let accounts: Array = []; - const apiVersion = userContext.features.enableThroughputCap ? "2021-10-15-preview" : "2021-06-15"; + const apiVersion = "2023-09-15-preview"; let nextLink = `${configContext.ARM_ENDPOINT}/subscriptions/${subscriptionId}/providers/Microsoft.DocumentDB/databaseAccounts?api-version=${apiVersion}`; while (nextLink) {