For connection-string based authentication, the request would be made… (#1644)

* For connection-string based authentication, the request would be made with low priority request header and for AAD based authentication, users can select the priority-request option from the settings panel

* removed unused imports

* prettier checks

---------

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

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;

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 =
+  const shouldShowPriorityLevelOption = PriorityBasedExecutionUtils.isFeatureEnabled();
-    userContext.databaseAccount?.properties?.enablePriorityBasedExecution && userContext.apiType === "SQL";
   const handlerOnSubmit = (e: MouseEvent<HTMLButtonElement>) => {
     setIsExecuting(true);
 

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"),
   };
 }
 

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 (

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<DatabaseAccount> = [];
-  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) {