From 6be46bff28b6e7f7bf8c0a32cb42bb64acffc438 Mon Sep 17 00:00:00 2001 From: sunghyunkang1111 <114709653+sunghyunkang1111@users.noreply.github.com> Date: Wed, 14 May 2025 12:26:30 -0500 Subject: [PATCH] Throughput bucketing build release (#2149) * Added AFEC check in userContext (#2140) * Added AFEC check in userContext * Update unit tests and fix Group to Bucket * Hide throughput bucket settings for nonshared, nondedicated throughput container (#2146) --- .../Settings/SettingsComponent.test.tsx | 3 +- .../Controls/Settings/SettingsComponent.tsx | 15 +++--- .../ThroughputBucketsComponent.test.tsx | 16 +++++-- .../ThroughputBucketsComponent.tsx | 2 +- src/Explorer/Explorer.tsx | 6 +++ src/UserContext.ts | 1 + src/Utils/FeatureRegistrationUtils.ts | 48 +++++++++++++++++++ 7 files changed, 74 insertions(+), 17 deletions(-) create mode 100644 src/Utils/FeatureRegistrationUtils.ts diff --git a/src/Explorer/Controls/Settings/SettingsComponent.test.tsx b/src/Explorer/Controls/Settings/SettingsComponent.test.tsx index 5a7a47def..825c3fb14 100644 --- a/src/Explorer/Controls/Settings/SettingsComponent.test.tsx +++ b/src/Explorer/Controls/Settings/SettingsComponent.test.tsx @@ -1,7 +1,6 @@ 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"; @@ -253,7 +252,7 @@ describe("SettingsComponent", () => { it("should save throughput bucket changes when Save button is clicked", async () => { updateUserContext({ apiType: "SQL", - features: { enableThroughputBuckets: true } as Features, + throughputBucketsEnabled: true, authType: AuthType.AAD, }); diff --git a/src/Explorer/Controls/Settings/SettingsComponent.tsx b/src/Explorer/Controls/Settings/SettingsComponent.tsx index 65000e7ee..3f286d822 100644 --- a/src/Explorer/Controls/Settings/SettingsComponent.tsx +++ b/src/Explorer/Controls/Settings/SettingsComponent.tsx @@ -191,10 +191,7 @@ export class SettingsComponent extends React.Component, diff --git a/src/Explorer/Controls/Settings/SettingsSubComponents/ThroughputInputComponents/ThroughputBucketsComponent.test.tsx b/src/Explorer/Controls/Settings/SettingsSubComponents/ThroughputInputComponents/ThroughputBucketsComponent.test.tsx index 460dfa252..2ad5c819e 100644 --- a/src/Explorer/Controls/Settings/SettingsSubComponents/ThroughputInputComponents/ThroughputBucketsComponent.test.tsx +++ b/src/Explorer/Controls/Settings/SettingsSubComponents/ThroughputInputComponents/ThroughputBucketsComponent.test.tsx @@ -26,7 +26,7 @@ describe("ThroughputBucketsComponent", () => { it("renders the correct number of buckets", () => { render(); - expect(screen.getAllByText(/Group \d+/)).toHaveLength(5); + expect(screen.getAllByText(/Bucket \d+/)).toHaveLength(5); }); it("renders buckets in the correct order even if input is unordered", () => { @@ -36,8 +36,14 @@ describe("ThroughputBucketsComponent", () => { ]; render(); - 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"]); + const bucketLabels = screen.getAllByText(/Bucket \d+/).map((el) => el.textContent); + expect(bucketLabels).toEqual([ + "Bucket 1 (Data Explorer Query Bucket)", + "Bucket 2", + "Bucket 3", + "Bucket 4", + "Bucket 5", + ]); }); it("renders all provided buckets even if they exceed the max default bucket count", () => { @@ -53,7 +59,7 @@ describe("ThroughputBucketsComponent", () => { render(); - expect(screen.getAllByText(/Group \d+/)).toHaveLength(7); + expect(screen.getAllByText(/Bucket \d+/)).toHaveLength(7); expect(screen.getByDisplayValue("50")).toBeInTheDocument(); expect(screen.getByDisplayValue("60")).toBeInTheDocument(); @@ -171,7 +177,7 @@ describe("ThroughputBucketsComponent", () => { it("ensures default buckets are used when no buckets are provided", () => { render(); - expect(screen.getAllByText(/Group \d+/)).toHaveLength(5); + expect(screen.getAllByText(/Bucket \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 a9408b1e4..a0b85cabf 100644 --- a/src/Explorer/Controls/Settings/SettingsSubComponents/ThroughputInputComponents/ThroughputBucketsComponent.tsx +++ b/src/Explorer/Controls/Settings/SettingsSubComponents/ThroughputInputComponents/ThroughputBucketsComponent.tsx @@ -76,7 +76,7 @@ export const ThroughputBucketsComponent: FC = ( value={bucket.maxThroughputPercentage} onChange={(newValue) => handleBucketChange(bucket.id, newValue)} showValue={false} - label={`Group ${bucket.id}${bucket.id === 1 ? " (Data Explorer Query Bucket)" : ""}`} + label={`Bucket ${bucket.id}${bucket.id === 1 ? " (Data Explorer Query Bucket)" : ""}`} styles={{ root: { flex: 2, maxWidth: 400 } }} disabled={bucket.maxThroughputPercentage === 100} /> diff --git a/src/Explorer/Explorer.tsx b/src/Explorer/Explorer.tsx index 30c051662..5bf42e9aa 100644 --- a/src/Explorer/Explorer.tsx +++ b/src/Explorer/Explorer.tsx @@ -12,6 +12,7 @@ import { isFabricMirrored, isFabricMirroredKey, scheduleRefreshFabricToken } fro import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility"; import { acquireMsalTokenForAccount } from "Utils/AuthorizationUtils"; import { allowedNotebookServerUrls, validateEndpoint } from "Utils/EndpointUtils"; +import { featureRegistered } from "Utils/FeatureRegistrationUtils"; import { update } from "Utils/arm/generatedClients/cosmos/databaseAccounts"; import { useQueryCopilot } from "hooks/useQueryCopilot"; import * as ko from "knockout"; @@ -1169,6 +1170,11 @@ export default class Explorer { await this.initNotebooks(userContext.databaseAccount); } + if (userContext.authType === AuthType.AAD && userContext.apiType === "SQL") { + const throughputBucketsEnabled = await featureRegistered(userContext.subscriptionId, "ThroughputBucketing"); + updateUserContext({ throughputBucketsEnabled }); + } + this.refreshSampleData(); } diff --git a/src/UserContext.ts b/src/UserContext.ts index 8a880f498..8a723ad91 100644 --- a/src/UserContext.ts +++ b/src/UserContext.ts @@ -117,6 +117,7 @@ export interface UserContext { readonly feedbackPolicies?: AdminFeedbackPolicySettings; readonly dataPlaneRbacEnabled?: boolean; readonly refreshCosmosClient?: boolean; + throughputBucketsEnabled?: boolean; } export type ApiType = "SQL" | "Mongo" | "Gremlin" | "Tables" | "Cassandra" | "Postgres" | "VCoreMongo"; diff --git a/src/Utils/FeatureRegistrationUtils.ts b/src/Utils/FeatureRegistrationUtils.ts new file mode 100644 index 000000000..9cc5676c9 --- /dev/null +++ b/src/Utils/FeatureRegistrationUtils.ts @@ -0,0 +1,48 @@ +import { configContext } from "ConfigContext"; +import { FeatureRegistration } from "Contracts/DataModels"; +import { AuthorizationTokenHeaderMetadata } from "Contracts/ViewModels"; +import { getAuthorizationHeader } from "Utils/AuthorizationUtils"; + +export const featureRegistered = async (subscriptionId: string, feature: string) => { + const api_version = "2021-07-01"; + const url = `${configContext.ARM_ENDPOINT}/subscriptions/${subscriptionId}/providers/Microsoft.Features/featureProviders/Microsoft.DocumentDB/subscriptionFeatureRegistrations/${feature}?api-version=${api_version}`; + const authorizationHeader: AuthorizationTokenHeaderMetadata = getAuthorizationHeader(); + const headers = { [authorizationHeader.header]: authorizationHeader.token }; + + let response; + + try { + response = await _fetchWithTimeout(url, headers); + } catch (error) { + return false; + } + + if (!response?.ok) { + return false; + } + + const featureRegistration = (await response?.json()) as FeatureRegistration; + return featureRegistration?.properties?.state === "Registered"; +}; + +async function _fetchWithTimeout( + url: string, + headers: { + [x: string]: string; + }, +) { + const timeout = 10000; + const options = { timeout }; + + const controller = new AbortController(); + const id = setTimeout(() => controller.abort(), timeout); + + const response = await window.fetch(url, { + headers, + ...options, + signal: controller.signal, + }); + clearTimeout(id); + + return response; +}