Upgrade Cosmos SDK to v4.3

package.json

@@ -5,7 +5,7 @@
   "main": "index.js",
   "dependencies": {
     "@azure/arm-cosmosdb": "9.1.0",
-    "@azure/cosmos": "4.2.0-beta.1",
+    "@azure/cosmos": "4.3.0",
     "@azure/cosmos-language-service": "0.0.5",
     "@azure/identity": "4.5.0",
     "@azure/msal-browser": "2.14.2",

src/Common/IteratorUtilities.ts

@@ -1,4 +1,3 @@
-import { QueryOperationOptions } from "@azure/cosmos";
 import { Action } from "Shared/Telemetry/TelemetryConstants";
 import * as Constants from "../Common/Constants";
 import { QueryResults } from "../Contracts/ViewModels";
@@ -14,18 +13,14 @@ interface QueryResponse {
 }
 
 export interface MinimalQueryIterator {
-  fetchNext: (queryOperationOptions?: QueryOperationOptions) => Promise<QueryResponse>;
+  fetchNext: () => Promise<QueryResponse>;
 }
 
 // Pick<QueryIterator<any>, "fetchNext">;
 
-export function nextPage(
+export function nextPage(documentsIterator: MinimalQueryIterator, firstItemIndex: number): Promise<QueryResults> {
-  documentsIterator: MinimalQueryIterator,
-  firstItemIndex: number,
-  queryOperationOptions?: QueryOperationOptions,
-): Promise<QueryResults> {
   TelemetryProcessor.traceStart(Action.ExecuteQuery);
-  return documentsIterator.fetchNext(queryOperationOptions).then((response) => {
+  return documentsIterator.fetchNext().then((response) => {
     TelemetryProcessor.traceSuccess(Action.ExecuteQuery, { dataExplorerArea: Constants.Areas.Tab });
     const documents = response.resources;
     // eslint-disable-next-line @typescript-eslint/no-explicit-any

src/Common/QueryError.ts

@@ -1,5 +1,4 @@
 import { monaco } from "Explorer/LazyMonaco";
-import { getRUThreshold, ruThresholdEnabled } from "Shared/StorageUtility";
 
 export enum QueryErrorSeverity {
   Error = "Error",
@@ -103,20 +102,9 @@ export interface ErrorEnrichment {
   learnMoreUrl?: string;
 }
 
-const REPLACEMENT_MESSAGES: Record<string, (original: string) => string> = {
+const REPLACEMENT_MESSAGES: Record<string, (original: string) => string> = {};
-  OPERATION_RU_LIMIT_EXCEEDED: (original) => {
-    if (ruThresholdEnabled()) {
-      const threshold = getRUThreshold();
-      return `Query exceeded the Request Unit (RU) limit of ${threshold} RUs. You can change this limit in Data Explorer settings.`;
-    }
-    return original;
-  },
-};
 
-const HELP_LINKS: Record<string, string> = {
+const HELP_LINKS: Record<string, string> = {};
-  OPERATION_RU_LIMIT_EXCEEDED:
-    "https://learn.microsoft.com/en-us/azure/cosmos-db/data-explorer#configure-request-unit-threshold",
-};
 
 export default class QueryError {
   message: string;

src/Common/dataAccess/queryDocuments.ts

@@ -26,6 +26,7 @@ export const getCommonQueryOptions = (options: FeedOptions): FeedOptions => {
     options.maxItemCount ||
     (storedItemPerPageSetting !== undefined && storedItemPerPageSetting) ||
     Queries.itemsPerPage;
+  options.enableQueryControl = LocalStorageUtility.getEntryBoolean(StorageKey.QueryControlEnabled);
   options.maxDegreeOfParallelism = LocalStorageUtility.getEntryNumber(StorageKey.MaxDegreeOfParellism);
   options.disableNonStreamingOrderByQuery = !isVectorSearchEnabled();
   return options;

src/Common/dataAccess/queryDocumentsPage.ts

@@ -1,4 +1,3 @@
-import { QueryOperationOptions } from "@azure/cosmos";
 import { QueryResults } from "../../Contracts/ViewModels";
 import { logConsoleInfo, logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
 import { getEntityName } from "../DocumentUtility";
@@ -9,13 +8,12 @@ export const queryDocumentsPage = async (
   resourceName: string,
   documentsIterator: MinimalQueryIterator,
   firstItemIndex: number,
-  queryOperationOptions?: QueryOperationOptions,
 ): Promise<QueryResults> => {
   const entityName = getEntityName();
   const clearMessage = logConsoleProgress(`Querying ${entityName} for container ${resourceName}`);
 
   try {
-    const result: QueryResults = await nextPage(documentsIterator, firstItemIndex, queryOperationOptions);
+    const result: QueryResults = await nextPage(documentsIterator, firstItemIndex);
     const itemCount = (result.documents && result.documents.length) || 0;
     logConsoleInfo(`Successfully fetched ${itemCount} ${entityName} for container ${resourceName}`);
     return result;

src/Explorer/Panes/SettingsPane/SettingsPane.tsx

@@ -180,6 +180,11 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
       ? LocalStorageUtility.getEntryNumber(StorageKey.MaxWaitTimeInSeconds)
       : Constants.Queries.DefaultMaxWaitTimeInSeconds,
   );
+  const [queryControlEnabled, setQueryControlEnabled] = useState<boolean>(
+    LocalStorageUtility.hasItem(StorageKey.QueryControlEnabled)
+      ? LocalStorageUtility.getEntryString(StorageKey.QueryControlEnabled) === "true"
+      : false,
+  );
   const [maxDegreeOfParallelism, setMaxDegreeOfParallelism] = useState<number>(
     LocalStorageUtility.hasItem(StorageKey.MaxDegreeOfParellism)
       ? LocalStorageUtility.getEntryNumber(StorageKey.MaxDegreeOfParellism)
@@ -204,6 +209,7 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
     !isEmulator;
   const shouldShowGraphAutoVizOption = userContext.apiType === "Gremlin" && !isEmulator;
   const shouldShowCrossPartitionOption = userContext.apiType !== "Gremlin" && !isEmulator;
+  const shouldShowEnhancedQueryControl = userContext.apiType === "SQL";
   const shouldShowParallelismOption = userContext.apiType !== "Gremlin" && !isEmulator;
   const showEnableEntraIdRbac =
     isDataplaneRbacSupported(userContext.apiType) &&
@@ -381,6 +387,7 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
     LocalStorageUtility.setEntryNumber(StorageKey.MaxWaitTimeInSeconds, MaxWaitTimeInSeconds);
     LocalStorageUtility.setEntryString(StorageKey.ContainerPaginationEnabled, containerPaginationEnabled.toString());
     LocalStorageUtility.setEntryString(StorageKey.IsCrossPartitionQueryEnabled, crossPartitionQueryEnabled.toString());
+    LocalStorageUtility.setEntryString(StorageKey.QueryControlEnabled, queryControlEnabled.toString());
     LocalStorageUtility.setEntryNumber(StorageKey.MaxDegreeOfParellism, maxDegreeOfParallelism);
     LocalStorageUtility.setEntryString(StorageKey.PriorityLevel, priorityLevel.toString());
     LocalStorageUtility.setEntryString(StorageKey.CopilotSampleDBEnabled, copilotSampleDBEnabled.toString());
@@ -410,6 +417,7 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
       `Updated items per page setting to ${LocalStorageUtility.getEntryNumber(StorageKey.ActualItemPerPage)}`,
     );
     logConsoleInfo(`${crossPartitionQueryEnabled ? "Enabled" : "Disabled"} cross-partition query feed option`);
+    logConsoleInfo(`${queryControlEnabled ? "Enabled" : "Disabled"} query control option`);
     logConsoleInfo(
       `Updated the max degree of parallelism query feed option to ${LocalStorageUtility.getEntryNumber(
         StorageKey.MaxDegreeOfParellism,
@@ -760,7 +768,6 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
                     )}
                   </AccordionPanel>
                 </AccordionItem>
-
                 <AccordionItem value="5">
                   <AccordionHeader>
                     <div className={styles.header}>RU Limit</div>
@@ -943,6 +950,38 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
                 </AccordionPanel>
               </AccordionItem>
             )}
+            {shouldShowEnhancedQueryControl && (
+              <AccordionItem value="10">
+                <AccordionHeader>
+                  <div className={styles.header}>Enhanced query control</div>
+                </AccordionHeader>
+                <AccordionPanel>
+                  <div className={styles.settingsSectionContainer}>
+                    <div className={styles.settingsSectionDescription}>
+                      Query up to the max degree of parallelism.
+                      <a
+                        href="https://learn.microsoft.com/en-us/azure/cosmos-db/nosql/performance-tips-query-sdk?tabs=v3&pivots=programming-language-nodejs#enhanced-query-control"
+                        target="_blank"
+                        rel="noopener noreferrer"
+                      >
+                        {" "}
+                        Learn more{" "}
+                      </a>
+                    </div>
+                    <Checkbox
+                      styles={{
+                        label: { padding: 0 },
+                      }}
+                      className="padding"
+                      ariaLabel="EnableQueryControl"
+                      checked={queryControlEnabled}
+                      onChange={() => setQueryControlEnabled(!queryControlEnabled)}
+                      label="Enable query control"
+                    />
+                  </div>
+                </AccordionPanel>
+              </AccordionItem>
+            )}
             {shouldShowParallelismOption && (
               <AccordionItem value="10">
                 <AccordionHeader>

src/Explorer/Tabs/QueryTab/QueryTabComponent.tsx

@@ -1,6 +1,6 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 /* eslint-disable no-console */
-import { FeedOptions, QueryOperationOptions } from "@azure/cosmos";
+import { FeedOptions } from "@azure/cosmos";
 import { AuthType } from "AuthType";
 import QueryError, { createMonacoErrorLocationResolver, createMonacoMarkersForQueryErrors } from "Common/QueryError";
 import { SplitterDirection } from "Common/Splitter";
@@ -19,7 +19,7 @@ import { CosmosFluentProvider } from "Explorer/Theme/ThemeUtil";
 import { useSelectedNode } from "Explorer/useSelectedNode";
 import { KeyboardAction } from "KeyboardShortcuts";
 import { QueryConstants } from "Shared/Constants";
-import { LocalStorageUtility, StorageKey, getRUThreshold, ruThresholdEnabled } from "Shared/StorageUtility";
+import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility";
 import { Action } from "Shared/Telemetry/TelemetryConstants";
 import { Allotment } from "allotment";
 import { useClientWriteEnabled } from "hooks/useClientWriteEnabled";
@@ -370,21 +370,8 @@ class QueryTabComponentImpl extends React.Component<QueryTabComponentImplProps,
       isExecutionError: false,
     });
 
-    let queryOperationOptions: QueryOperationOptions;
-    if (userContext.apiType === "SQL" && ruThresholdEnabled()) {
-      const ruThreshold: number = getRUThreshold();
-      queryOperationOptions = {
-        ruCapPerOperation: ruThreshold,
-      } as QueryOperationOptions;
-    }
-
     const queryDocuments = async (firstItemIndex: number) =>
-      await queryDocumentsPage(
+      await queryDocumentsPage(this.props.collection && this.props.collection.id(), this._iterator, firstItemIndex);
-        this.props.collection && this.props.collection.id(),
-        this._iterator,
-        firstItemIndex,
-        queryOperationOptions,
-      );
     this.props.tabsBaseInstance.isExecuting(true);
     this.setState({
       isExecuting: true,

src/Shared/StorageUtility.ts

@@ -21,6 +21,7 @@ export enum StorageKey {
   DatabaseAccountId,
   EncryptedKeyToken,
   IsCrossPartitionQueryEnabled,
+  QueryControlEnabled,
   MaxDegreeOfParellism,
   IsGraphAutoVizDisabled,
   TenantId,

src/Utils/QueryUtils.ts

@@ -1,4 +1,7 @@
 import { PartitionKey, PartitionKeyDefinition } from "@azure/cosmos";
+import { getRUThreshold, ruThresholdEnabled } from "Shared/StorageUtility";
+import { userContext } from "UserContext";
+import { logConsoleError } from "Utils/NotificationConsoleUtils";
 import * as DataModels from "../Contracts/DataModels";
 import * as ViewModels from "../Contracts/ViewModels";
 
@@ -86,6 +89,15 @@ export const queryPagesUntilContentPresent = async (
     results.roundTrips = roundTrips;
     results.requestCharge = Number(results.requestCharge) + netRequestCharge;
     netRequestCharge = Number(results.requestCharge);
+
+    if (results.hasMoreResults && userContext.apiType === "SQL" && ruThresholdEnabled()) {
+      const ruThreshold: number = getRUThreshold();
+      if (netRequestCharge > ruThreshold) {
+        logConsoleError(`Request discontinued after exceeding the Request Unit threshold of ${ruThreshold}.`);
+        return results;
+      }
+    }
+
     const resultsMetadata = {
       hasMoreResults: results.hasMoreResults,
       itemCount: results.itemCount,