diff --git a/package.json b/package.json index 4a046a012..755254b2b 100644 --- a/package.json +++ b/package.json @@ -5,7 +5,7 @@ "main": "index.js", "dependencies": { "@azure/arm-cosmosdb": "9.1.0", - "@azure/cosmos": "4.2.0", + "@azure/cosmos": "4.2.0-beta.1", "@azure/cosmos-language-service": "0.0.5", "@azure/identity": "1.5.2", "@azure/ms-rest-nodeauth": "3.1.1", diff --git a/src/Common/IteratorUtilities.ts b/src/Common/IteratorUtilities.ts index f85ad7fb2..6283488b8 100644 --- a/src/Common/IteratorUtilities.ts +++ b/src/Common/IteratorUtilities.ts @@ -1,3 +1,4 @@ +import { QueryOperationOptions } from "@azure/cosmos"; import { QueryResults } from "../Contracts/ViewModels"; interface QueryResponse { @@ -10,13 +11,17 @@ interface QueryResponse { } export interface MinimalQueryIterator { - fetchNext: () => Promise; + fetchNext: (queryOperationOptions?: QueryOperationOptions) => Promise; } // Pick, "fetchNext">; -export function nextPage(documentsIterator: MinimalQueryIterator, firstItemIndex: number): Promise { - return documentsIterator.fetchNext().then((response) => { +export function nextPage( + documentsIterator: MinimalQueryIterator, + firstItemIndex: number, + queryOperationOptions?: QueryOperationOptions, +): Promise { + return documentsIterator.fetchNext(queryOperationOptions).then((response) => { const documents = response.resources; // eslint-disable-next-line @typescript-eslint/no-explicit-any const headers = (response as any).headers || {}; // TODO this is a private key. Remove any diff --git a/src/Common/dataAccess/queryDocumentsPage.ts b/src/Common/dataAccess/queryDocumentsPage.ts index 556ed290c..17e84ba28 100644 --- a/src/Common/dataAccess/queryDocumentsPage.ts +++ b/src/Common/dataAccess/queryDocumentsPage.ts @@ -1,3 +1,4 @@ +import { QueryOperationOptions } from "@azure/cosmos"; import { QueryResults } from "../../Contracts/ViewModels"; import { logConsoleInfo, logConsoleProgress } from "../../Utils/NotificationConsoleUtils"; import { getEntityName } from "../DocumentUtility"; @@ -8,12 +9,13 @@ export const queryDocumentsPage = async ( resourceName: string, documentsIterator: MinimalQueryIterator, firstItemIndex: number, + queryOperationOptions?: QueryOperationOptions, ): Promise => { const entityName = getEntityName(); const clearMessage = logConsoleProgress(`Querying ${entityName} for container ${resourceName}`); try { - const result: QueryResults = await nextPage(documentsIterator, firstItemIndex); + const result: QueryResults = await nextPage(documentsIterator, firstItemIndex, queryOperationOptions); const itemCount = (result.documents && result.documents.length) || 0; logConsoleInfo(`Successfully fetched ${itemCount} ${entityName} for container ${resourceName}`); return result; diff --git a/src/Common/dataAccess/updateOffer.ts b/src/Common/dataAccess/updateOffer.ts index fde4e6814..4d26ca68d 100644 --- a/src/Common/dataAccess/updateOffer.ts +++ b/src/Common/dataAccess/updateOffer.ts @@ -39,10 +39,7 @@ import { migrateTableToManualThroughput, updateTableThroughput, } from "../../Utils/arm/generatedClients/cosmos/tableResources"; -import { - ThroughputSettingsGetResults, - ThroughputSettingsUpdateParameters, -} from "../../Utils/arm/generatedClients/cosmos/types"; +import { ThroughputSettingsUpdateParameters } from "../../Utils/arm/generatedClients/cosmos/types"; import { logConsoleInfo, logConsoleProgress } from "../../Utils/NotificationConsoleUtils"; import { HttpHeaders } from "../Constants"; import { client } from "../CosmosClient"; @@ -149,28 +146,23 @@ const updateSqlContainerOffer = async (params: UpdateOfferParams): Promise const { subscriptionId, resourceGroup, databaseAccount } = userContext; const accountName = databaseAccount.name; - let updatedOffer: ThroughputSettingsGetResults; - if (params.migrateToAutoPilot) { - updatedOffer = (await migrateSqlContainerToAutoscale( + await migrateSqlContainerToAutoscale( subscriptionId, resourceGroup, accountName, params.databaseId, params.collectionId, - )) as ThroughputSettingsGetResults; - params.autopilotThroughput = updatedOffer.properties?.resource?.autoscaleSettings?.maxThroughput; + ); } else if (params.migrateToManual) { - updatedOffer = (await migrateSqlContainerToManualThroughput( + await migrateSqlContainerToManualThroughput( subscriptionId, resourceGroup, accountName, params.databaseId, params.collectionId, - )) as ThroughputSettingsGetResults; - params.manualThroughput = updatedOffer.properties?.resource?.throughput; - } - if (params.throughputBuckets || !(params.migrateToAutoPilot || params.migrateToManual)) { + ); + } else { const body: ThroughputSettingsUpdateParameters = createUpdateOfferBody(params); await updateSqlContainerThroughput( subscriptionId, diff --git a/src/Explorer/Controls/Settings/SettingsComponent.test.tsx b/src/Explorer/Controls/Settings/SettingsComponent.test.tsx index a97ee8f45..5a7a47def 100644 --- a/src/Explorer/Controls/Settings/SettingsComponent.test.tsx +++ b/src/Explorer/Controls/Settings/SettingsComponent.test.tsx @@ -1,5 +1,7 @@ +import { AuthType } from "AuthType"; import { shallow } from "enzyme"; import ko from "knockout"; +import { Features } from "Platform/Hosted/extractFeatures"; import React from "react"; import { updateCollection } from "../../../Common/dataAccess/updateCollection"; import { updateOffer } from "../../../Common/dataAccess/updateOffer"; @@ -247,4 +249,42 @@ describe("SettingsComponent", () => { expect(conflictResolutionPolicy.mode).toEqual(DataModels.ConflictResolutionMode.Custom); expect(conflictResolutionPolicy.conflictResolutionProcedure).toEqual(expectSprocPath); }); + + it("should save throughput bucket changes when Save button is clicked", async () => { + updateUserContext({ + apiType: "SQL", + features: { enableThroughputBuckets: true } as Features, + authType: AuthType.AAD, + }); + + const wrapper = shallow(); + + const settingsComponentInstance = wrapper.instance() as SettingsComponent; + const isEnabled = settingsComponentInstance["throughputBucketsEnabled"]; + expect(isEnabled).toBe(true); + + wrapper.setState({ + isThroughputBucketsSaveable: true, + throughputBuckets: [ + { id: 1, maxThroughputPercentage: 70 }, + { id: 2, maxThroughputPercentage: 60 }, + ], + }); + + await settingsComponentInstance.onSaveClick(); + + expect(updateOffer).toHaveBeenCalledWith({ + databaseId: collection.databaseId, + collectionId: collection.id(), + currentOffer: expect.any(Object), + autopilotThroughput: collection.offer().autoscaleMaxThroughput, + manualThroughput: collection.offer().manualThroughput, + throughputBuckets: [ + { id: 1, maxThroughputPercentage: 70 }, + { id: 2, maxThroughputPercentage: 60 }, + ], + }); + + expect(wrapper.state("isThroughputBucketsSaveable")).toBe(false); + }); }); diff --git a/src/Explorer/Controls/Settings/SettingsSubComponents/ThroughputInputComponents/ThroughputBucketsComponent.test.tsx b/src/Explorer/Controls/Settings/SettingsSubComponents/ThroughputInputComponents/ThroughputBucketsComponent.test.tsx index 998014179..460dfa252 100644 --- a/src/Explorer/Controls/Settings/SettingsSubComponents/ThroughputInputComponents/ThroughputBucketsComponent.test.tsx +++ b/src/Explorer/Controls/Settings/SettingsSubComponents/ThroughputInputComponents/ThroughputBucketsComponent.test.tsx @@ -6,7 +6,6 @@ import { ThroughputBucketsComponent } from "./ThroughputBucketsComponent"; describe("ThroughputBucketsComponent", () => { const mockOnBucketsChange = jest.fn(); const mockOnSaveableChange = jest.fn(); - const mockOnDiscardableChange = jest.fn(); const defaultProps = { currentBuckets: [ @@ -19,19 +18,15 @@ describe("ThroughputBucketsComponent", () => { ], onBucketsChange: mockOnBucketsChange, onSaveableChange: mockOnSaveableChange, - onDiscardableChange: mockOnDiscardableChange, }; beforeEach(() => { jest.clearAllMocks(); }); - it("renders 5 buckets with default values when input buckets are missing", () => { - render(); - - expect(screen.getAllByText(/Bucket \d+/)).toHaveLength(5); - expect(screen.getByDisplayValue("50")).toBeInTheDocument(); - expect(screen.getAllByDisplayValue("100").length).toBe(4); + it("renders the correct number of buckets", () => { + render(); + expect(screen.getAllByText(/Group \d+/)).toHaveLength(5); }); it("renders buckets in the correct order even if input is unordered", () => { @@ -41,16 +36,36 @@ describe("ThroughputBucketsComponent", () => { ]; render(); - const bucketLabels = screen.getAllByText(/Bucket \d+/).map((el) => el.textContent); - expect(bucketLabels).toEqual(["Bucket 1", "Bucket 2", "Bucket 3", "Bucket 4", "Bucket 5"]); + const bucketLabels = screen.getAllByText(/Group \d+/).map((el) => el.textContent); + expect(bucketLabels).toEqual(["Group 1 (Data Explorer Query Bucket)", "Group 2", "Group 3", "Group 4", "Group 5"]); + }); + + it("renders all provided buckets even if they exceed the max default bucket count", () => { + const oversizedBuckets = [ + { id: 1, maxThroughputPercentage: 50 }, + { id: 2, maxThroughputPercentage: 60 }, + { id: 3, maxThroughputPercentage: 70 }, + { id: 4, maxThroughputPercentage: 80 }, + { id: 5, maxThroughputPercentage: 90 }, + { id: 6, maxThroughputPercentage: 100 }, + { id: 7, maxThroughputPercentage: 40 }, + ]; + + render(); + + expect(screen.getAllByText(/Group \d+/)).toHaveLength(7); + expect(screen.getByDisplayValue("50")).toBeInTheDocument(); expect(screen.getByDisplayValue("60")).toBeInTheDocument(); - expect(screen.getAllByDisplayValue("100").length).toBe(3); + expect(screen.getByDisplayValue("70")).toBeInTheDocument(); + expect(screen.getByDisplayValue("80")).toBeInTheDocument(); + expect(screen.getByDisplayValue("90")).toBeInTheDocument(); + expect(screen.getByDisplayValue("100")).toBeInTheDocument(); + expect(screen.getByDisplayValue("40")).toBeInTheDocument(); }); it("calls onBucketsChange when a bucket value changes", () => { render(); - const input = screen.getByDisplayValue("50"); fireEvent.change(input, { target: { value: "70" } }); @@ -63,21 +78,12 @@ describe("ThroughputBucketsComponent", () => { ]); }); - it("triggers onSaveableChange and onDiscardableChange when values change", () => { + it("triggers onSaveableChange when values change", () => { render(); - const input = screen.getByDisplayValue("50"); fireEvent.change(input, { target: { value: "80" } }); expect(mockOnSaveableChange).toHaveBeenCalledWith(true); - expect(mockOnDiscardableChange).toHaveBeenCalledWith(true); - }); - - it("ensures buckets revert to default when no buckets are provided", () => { - render(); - - expect(screen.getAllByText(/Bucket \d+/)).toHaveLength(5); - expect(screen.getAllByDisplayValue("100")).toHaveLength(5); }); it("updates state consistently after multiple changes to different buckets", () => { @@ -100,18 +106,9 @@ describe("ThroughputBucketsComponent", () => { it("resets to baseline when currentBuckets are reset", () => { const { rerender } = render(); - const input1 = screen.getByDisplayValue("50"); fireEvent.change(input1, { target: { value: "70" } }); - expect(mockOnBucketsChange).toHaveBeenCalledWith([ - { id: 1, maxThroughputPercentage: 70 }, - { id: 2, maxThroughputPercentage: 60 }, - { id: 3, maxThroughputPercentage: 100 }, - { id: 4, maxThroughputPercentage: 100 }, - { id: 5, maxThroughputPercentage: 100 }, - ]); - rerender(); expect(screen.getByDisplayValue("40")).toBeInTheDocument(); @@ -120,10 +117,61 @@ describe("ThroughputBucketsComponent", () => { it("does not call onBucketsChange when value remains unchanged", () => { render(); - const input = screen.getByDisplayValue("50"); fireEvent.change(input, { target: { value: "50" } }); expect(mockOnBucketsChange).not.toHaveBeenCalled(); }); + + it("disables input and slider when maxThroughputPercentage is 100", () => { + render( + , + ); + + const disabledInputs = screen.getAllByDisplayValue("100"); + expect(disabledInputs.length).toBeGreaterThan(0); + expect(disabledInputs[0]).toBeDisabled(); + + const sliders = screen.getAllByRole("slider"); + expect(sliders.length).toBeGreaterThan(0); + expect(sliders[0]).toHaveAttribute("aria-disabled", "true"); + expect(sliders[1]).toHaveAttribute("aria-disabled", "false"); + }); + + it("toggles bucket value between 50 and 100 with switch", () => { + render(); + const toggles = screen.getAllByRole("switch"); + + fireEvent.click(toggles[0]); + + expect(mockOnBucketsChange).toHaveBeenCalledWith([ + { id: 1, maxThroughputPercentage: 100 }, + { id: 2, maxThroughputPercentage: 60 }, + { id: 3, maxThroughputPercentage: 100 }, + { id: 4, maxThroughputPercentage: 100 }, + { id: 5, maxThroughputPercentage: 100 }, + ]); + + fireEvent.click(toggles[0]); + + expect(mockOnBucketsChange).toHaveBeenCalledWith([ + { id: 1, maxThroughputPercentage: 50 }, + { id: 2, maxThroughputPercentage: 60 }, + { id: 3, maxThroughputPercentage: 100 }, + { id: 4, maxThroughputPercentage: 100 }, + { id: 5, maxThroughputPercentage: 100 }, + ]); + }); + + it("ensures default buckets are used when no buckets are provided", () => { + render(); + expect(screen.getAllByText(/Group \d+/)).toHaveLength(5); + expect(screen.getAllByDisplayValue("100")).toHaveLength(5); + }); }); diff --git a/src/Explorer/Controls/Settings/SettingsSubComponents/ThroughputInputComponents/ThroughputBucketsComponent.tsx b/src/Explorer/Controls/Settings/SettingsSubComponents/ThroughputInputComponents/ThroughputBucketsComponent.tsx index cc60cb8eb..a9408b1e4 100644 --- a/src/Explorer/Controls/Settings/SettingsSubComponents/ThroughputInputComponents/ThroughputBucketsComponent.tsx +++ b/src/Explorer/Controls/Settings/SettingsSubComponents/ThroughputInputComponents/ThroughputBucketsComponent.tsx @@ -43,13 +43,11 @@ export const ThroughputBucketsComponent: FC = ( useEffect(() => { setThroughputBuckets(getThroughputBuckets(currentBuckets)); onSaveableChange(false); - // onDiscardableChange(false); }, [currentBuckets]); useEffect(() => { const isChanged = isDirty(throughputBuckets, getThroughputBuckets(throughputBucketsBaseline)); onSaveableChange(isChanged); - // onDiscardableChange(isChanged); }, [throughputBuckets]); const handleBucketChange = (id: number, newValue: number) => { @@ -92,10 +90,6 @@ export const ThroughputBucketsComponent: FC = ( }} disabled={bucket.maxThroughputPercentage === 100} /> - {/* onToggle(bucket.id, bucket.maxThroughputPercentage === 100)} - > */} = ( onChange={(event, checked) => onToggle(bucket.id, checked)} styles={{ root: { marginBottom: 0 }, text: { fontSize: 12 } }} > - {/* {bucket.id === 1 && ( - - - Data Explorer Query Bucket - - )} */} ))} diff --git a/src/Explorer/Tabs/QueryTab/QueryTabComponent.tsx b/src/Explorer/Tabs/QueryTab/QueryTabComponent.tsx index 0ad3e20e0..821aebfe5 100644 --- a/src/Explorer/Tabs/QueryTab/QueryTabComponent.tsx +++ b/src/Explorer/Tabs/QueryTab/QueryTabComponent.tsx @@ -1,6 +1,6 @@ /* eslint-disable @typescript-eslint/no-explicit-any */ /* eslint-disable no-console */ -import { FeedOptions } from "@azure/cosmos"; +import { FeedOptions, QueryOperationOptions } from "@azure/cosmos"; import QueryError, { createMonacoErrorLocationResolver, createMonacoMarkersForQueryErrors } from "Common/QueryError"; import { SplitterDirection } from "Common/Splitter"; import { Platform, configContext } from "ConfigContext"; @@ -18,7 +18,7 @@ import { CosmosFluentProvider } from "Explorer/Theme/ThemeUtil"; import { useSelectedNode } from "Explorer/useSelectedNode"; import { KeyboardAction } from "KeyboardShortcuts"; import { QueryConstants } from "Shared/Constants"; -import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility"; +import { LocalStorageUtility, StorageKey, getRUThreshold, ruThresholdEnabled } from "Shared/StorageUtility"; import { Action } from "Shared/Telemetry/TelemetryConstants"; import { Allotment } from "allotment"; import { QueryCopilotState, useQueryCopilot } from "hooks/useQueryCopilot"; @@ -368,8 +368,21 @@ class QueryTabComponentImpl extends React.Component - await queryDocumentsPage(this.props.collection && this.props.collection.id(), this._iterator, firstItemIndex); + await queryDocumentsPage( + this.props.collection && this.props.collection.id(), + this._iterator, + firstItemIndex, + queryOperationOptions, + ); this.props.tabsBaseInstance.isExecuting(true); this.setState({ isExecuting: true,