Add Data Plane RBAC functionality (#1878)

* Fix API endpoint for CassandraProxy query API * activate Mongo Proxy and Cassandra Proxy in Prod * Add CP Prod endpoint * Run npm format and tests * Revert code * fix bug that blocked local mongo proxy and cassandra proxy development * Add prod endpoint * fix pr check tests * Remove prod * Remove prod endpoint * Remove dev endpoint * Support data plane RBAC * Support data plane RBAC * Add additional changes for Portal RBAC functionality * Address errors and checks * Cleanup DP RBAC code * Run format * Fix unit tests * Remove unnecessary code * Run npm format * Fix enableAadDataPlane feature flag behavior * Fix enable AAD dataplane feature flag behavior * Address feedback comments * Minor fix * Add new fixes * Fix Tables test * Run npm format --------- Co-authored-by: Asier Isayas <aisayas@microsoft.com>
2025-12-20 01:11:25 +00:00 · 2024-07-01 09:33:07 -07:00
parent 17754cba05
commit c30a9681fe
11 changed files with 312 additions and 31 deletions
--- a/src/Explorer/Explorer.tsx
+++ b/src/Explorer/Explorer.tsx
@@ -7,11 +7,13 @@ import { getCopilotEnabled, isCopilotFeatureRegistered } from "Explorer/QueryCop
 import { IGalleryItem } from "Juno/JunoClient";
 import { scheduleRefreshDatabaseResourceToken } from "Platform/Fabric/FabricUtil";
 import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility";
+import { acquireTokenWithMsal, getMsalInstance } from "Utils/AuthorizationUtils";
 import { allowedNotebookServerUrls, validateEndpoint } from "Utils/EndpointUtils";
 import { useQueryCopilot } from "hooks/useQueryCopilot";
 import * as ko from "knockout";
 import React from "react";
 import _ from "underscore";
+import * as msal from "@azure/msal-browser";
 import shallow from "zustand/shallow";
 import { AuthType } from "../AuthType";
 import { BindingHandlersRegisterer } from "../Bindings/BindingHandlersRegisterer";
@@ -30,14 +32,13 @@ import { PhoenixClient } from "../Phoenix/PhoenixClient";
 import * as ExplorerSettings from "../Shared/ExplorerSettings";
 import { Action, ActionModifiers } from "../Shared/Telemetry/TelemetryConstants";
 import * as TelemetryProcessor from "../Shared/Telemetry/TelemetryProcessor";
-import { isAccountNewerThanThresholdInMs, userContext } from "../UserContext";
+import { isAccountNewerThanThresholdInMs, updateUserContext, userContext } from "../UserContext";
 import { getCollectionName, getUploadName } from "../Utils/APITypeUtils";
 import { stringToBlob } from "../Utils/BlobUtils";
 import { isCapabilityEnabled } from "../Utils/CapabilityUtils";
 import { fromContentUri, toRawContentUri } from "../Utils/GitHubUtils";
 import * as NotificationConsoleUtils from "../Utils/NotificationConsoleUtils";
 import { logConsoleError, logConsoleInfo, logConsoleProgress } from "../Utils/NotificationConsoleUtils";
-import { update } from "../Utils/arm/generatedClients/cosmos/databaseAccounts";
 import { useSidePanel } from "../hooks/useSidePanel";
 import { useTabs } from "../hooks/useTabs";
 import "./ComponentRegisterer";
@@ -66,6 +67,8 @@ import { ResourceTreeAdapter } from "./Tree/ResourceTreeAdapter";
 import StoredProcedure from "./Tree/StoredProcedure";
 import { useDatabases } from "./useDatabases";
 import { useSelectedNode } from "./useSelectedNode";
+import { update } from "Utils/arm/generatedClients/cosmos/databaseAccounts";
+import { useDataPlaneRbac } from "Explorer/Panes/SettingsPane/SettingsPane";

 BindingHandlersRegisterer.registerBindingHandlers();

@@ -251,8 +254,44 @@ export default class Explorer {
    };
    useDialog.getState().openDialog(addSynapseLinkDialogProps);
    TelemetryProcessor.traceStart(Action.EnableAzureSynapseLink);
+  }

-    // TODO: return result
+  public async openLoginForEntraIDPopUp(): Promise<void> {
+    if (userContext.databaseAccount.properties?.documentEndpoint) {
+      const hrefEndpoint = new URL(userContext.databaseAccount.properties.documentEndpoint).href.replace(
+        /\/$/,
+        "/.default",
+      );
+      const msalInstance = await getMsalInstance();
+
+      try {
+        const response = await msalInstance.loginPopup({
+          redirectUri: configContext.msalRedirectURI,
+          scopes: [],
+        });
+        localStorage.setItem("cachedTenantId", response.tenantId);
+        const cachedAccount = msalInstance.getAllAccounts()?.[0];
+        msalInstance.setActiveAccount(cachedAccount);
+        const aadToken = await acquireTokenWithMsal(msalInstance, {
+          forceRefresh: true,
+          scopes: [hrefEndpoint],
+          authority: `${configContext.AAD_ENDPOINT}${localStorage.getItem("cachedTenantId")}`,
+        });
+        updateUserContext({ aadToken: aadToken });
+        useDataPlaneRbac.setState({ aadTokenUpdated: true });
+      } catch (error) {
+        if (error instanceof msal.AuthError && error.errorCode === msal.BrowserAuthErrorMessage.popUpWindowError.code) {
+          logConsoleError(
+            "We were unable to establish authorization for this account, due to pop-ups being disabled in the browser.\nPlease enable pop-ups for this site and try again",
+          );
+        } else {
+          const errorJson = JSON.stringify(error);
+          logConsoleError(
+            `Failed to perform authorization for this account, due to the following error: \n${errorJson}`,
+          );
+        }
+      }
+    }
  }

  public openNPSSurveyDialog(): void {
--- a/src/Explorer/Menus/CommandBar/CommandBarComponentButtonFactory.tsx
+++ b/src/Explorer/Menus/CommandBar/CommandBarComponentButtonFactory.tsx
@@ -15,6 +15,7 @@ import OpenQueryFromDiskIcon from "../../../../images/OpenQueryFromDisk.svg";
 import OpenInTabIcon from "../../../../images/open-in-tab.svg";
 import SettingsIcon from "../../../../images/settings_15x15.svg";
 import SynapseIcon from "../../../../images/synapse-link.svg";
+import EntraIDIcon from "../../../../images/EntraID.svg";
 import { AuthType } from "../../../AuthType";
 import * as Constants from "../../../Common/Constants";
 import { Platform, configContext } from "../../../ConfigContext";
@@ -30,9 +31,10 @@ import { OpenFullScreen } from "../../OpenFullScreen";
 import { AddDatabasePanel } from "../../Panes/AddDatabasePanel/AddDatabasePanel";
 import { BrowseQueriesPane } from "../../Panes/BrowseQueriesPane/BrowseQueriesPane";
 import { LoadQueryPane } from "../../Panes/LoadQueryPane/LoadQueryPane";
-import { SettingsPane } from "../../Panes/SettingsPane/SettingsPane";
+import { SettingsPane, useDataPlaneRbac } from "../../Panes/SettingsPane/SettingsPane";
 import { useDatabases } from "../../useDatabases";
 import { SelectedNodeState, useSelectedNode } from "../../useSelectedNode";
+import { useEffect, useState } from "react";

 let counter = 0;

@@ -69,6 +71,22 @@ export function createStaticCommandBarButtons(
      }
    }

+    if (userContext.apiType === "SQL") {
+      const [loginButtonProps, setLoginButtonProps] = useState<CommandButtonComponentProps | undefined>(undefined);
+      const dataPlaneRbacEnabled = useDataPlaneRbac((state) => state.dataPlaneRbacEnabled);
+      const aadTokenUpdated = useDataPlaneRbac((state) => state.aadTokenUpdated);
+
+      useEffect(() => {
+        const buttonProps = createLoginForEntraIDButton(container);
+        setLoginButtonProps(buttonProps);
+      }, [dataPlaneRbacEnabled, aadTokenUpdated, container]);
+
+      if (loginButtonProps) {
+        addDivider();
+        buttons.push(loginButtonProps);
+      }
+    }
+
    if (userContext.apiType !== "Tables") {
      newCollectionBtn.children = [createNewCollectionGroup(container)];
      const newDatabaseBtn = createNewDatabase(container);
@@ -275,6 +293,31 @@ function createOpenSynapseLinkDialogButton(container: Explorer): CommandButtonCo
  };
 }

+function createLoginForEntraIDButton(container: Explorer): CommandButtonComponentProps {
+  if (configContext.platform !== Platform.Portal) {
+    return undefined;
+  }
+
+  const handleCommandClick = async () => {
+    await container.openLoginForEntraIDPopUp();
+    useDataPlaneRbac.setState({ dataPlaneRbacEnabled: true });
+  };
+
+  if (!userContext.dataPlaneRbacEnabled || userContext.aadToken) {
+    return undefined;
+  }
+
+  const label = "Login for Entra ID RBAC";
+  return {
+    iconSrc: EntraIDIcon,
+    iconAlt: label,
+    onCommandClick: handleCommandClick,
+    commandButtonLabel: label,
+    hasPopup: true,
+    ariaLabel: label,
+  };
+}
+
 function createNewDatabase(container: Explorer): CommandButtonComponentProps {
  const label = "New " + getDatabaseName();
  return {
--- a/src/Explorer/Panes/SettingsPane/SettingsPane.tsx
+++ b/src/Explorer/Panes/SettingsPane/SettingsPane.tsx
@@ -4,14 +4,18 @@ import {
  IChoiceGroupOption,
  ISpinButtonStyles,
  IToggleStyles,
+  Icon,
+  MessageBar,
+  MessageBarType,
  Position,
  SpinButton,
  Toggle,
+  TooltipHost,
 } from "@fluentui/react";
 import * as Constants from "Common/Constants";
 import { SplitterDirection } from "Common/Splitter";
 import { InfoTooltip } from "Common/Tooltip/InfoTooltip";
-import { configContext } from "ConfigContext";
+import { Platform, configContext } from "ConfigContext";
 import { useDatabases } from "Explorer/useDatabases";
 import {
  DefaultRUThreshold,
@@ -22,7 +26,7 @@ import {
  ruThresholdEnabled as isRUThresholdEnabled,
 } from "Shared/StorageUtility";
 import * as StringUtility from "Shared/StringUtility";
-import { userContext } from "UserContext";
+import { updateUserContext, userContext } from "UserContext";
 import { logConsoleInfo } from "Utils/NotificationConsoleUtils";
 import * as PriorityBasedExecutionUtils from "Utils/PriorityBasedExecutionUtils";
 import { useQueryCopilot } from "hooks/useQueryCopilot";
@@ -30,6 +34,24 @@ import { useSidePanel } from "hooks/useSidePanel";
 import React, { FunctionComponent, useState } from "react";
 import Explorer from "../../Explorer";
 import { RightPaneForm, RightPaneFormProps } from "../RightPaneForm/RightPaneForm";
+import { AuthType } from "AuthType";
+import create, { UseStore } from "zustand";
+
+export interface DataPlaneRbacState {
+  dataPlaneRbacEnabled: boolean;
+  aadTokenUpdated: boolean;
+
+  getState?: () => DataPlaneRbacState;
+
+  setDataPlaneRbacEnabled: (dataPlaneRbacEnabled: boolean) => void;
+  setAadDataPlaneUpdated: (aadTokenUpdated: boolean) => void;
+}
+
+type DataPlaneRbacStore = UseStore<Partial<DataPlaneRbacState>>;
+
+export const useDataPlaneRbac: DataPlaneRbacStore = create(() => ({
+  dataPlaneRbacEnabled: false,
+}));

 export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
  explorer,
@@ -44,6 +66,14 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
      ? Constants.Queries.UnlimitedPageOption
      : Constants.Queries.CustomPageOption,
  );
+
+  const [enableDataPlaneRBACOption, setEnableDataPlaneRBACOption] = useState<string>(
+    LocalStorageUtility.hasItem(StorageKey.DataPlaneRbacEnabled)
+      ? LocalStorageUtility.getEntryString(StorageKey.DataPlaneRbacEnabled)
+      : Constants.RBACOptions.setAutomaticRBACOption,
+  );
+  const [showDataPlaneRBACWarning, setShowDataPlaneRBACWarning] = useState<boolean>(false);
+
  const [ruThresholdEnabled, setRUThresholdEnabled] = useState<boolean>(isRUThresholdEnabled());
  const [ruThreshold, setRUThreshold] = useState<number>(getRUThreshold());
  const [queryTimeoutEnabled, setQueryTimeoutEnabled] = useState<boolean>(
@@ -119,7 +149,26 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
      StorageKey.ActualItemPerPage,
      isCustomPageOptionSelected() ? customItemPerPage : Constants.Queries.unlimitedItemsPerPage,
    );
+
    LocalStorageUtility.setEntryNumber(StorageKey.CustomItemPerPage, customItemPerPage);
+
+    LocalStorageUtility.setEntryString(StorageKey.DataPlaneRbacEnabled, enableDataPlaneRBACOption);
+    if (
+      enableDataPlaneRBACOption === Constants.RBACOptions.setTrueRBACOption ||
+      (enableDataPlaneRBACOption === Constants.RBACOptions.setAutomaticRBACOption &&
+        userContext.databaseAccount.properties.disableLocalAuth)
+    ) {
+      updateUserContext({
+        dataPlaneRbacEnabled: true,
+      });
+      useDataPlaneRbac.setState({ dataPlaneRbacEnabled: true });
+    } else {
+      updateUserContext({
+        dataPlaneRbacEnabled: false,
+      });
+      useDataPlaneRbac.setState({ dataPlaneRbacEnabled: false });
+    }
+
    LocalStorageUtility.setEntryBoolean(StorageKey.RUThresholdEnabled, ruThresholdEnabled);
    LocalStorageUtility.setEntryBoolean(StorageKey.QueryTimeoutEnabled, queryTimeoutEnabled);
    LocalStorageUtility.setEntryNumber(StorageKey.RetryAttempts, retryAttempts);
@@ -207,6 +256,12 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
    { key: Constants.PriorityLevel.High, text: "High" },
  ];

+  const dataPlaneRBACOptionsList: IChoiceGroupOption[] = [
+    { key: Constants.RBACOptions.setAutomaticRBACOption, text: "Automatic" },
+    { key: Constants.RBACOptions.setTrueRBACOption, text: "True" },
+    { key: Constants.RBACOptions.setFalseRBACOption, text: "False" },
+  ];
+
  const defaultQueryResultsViewOptionList: IChoiceGroupOption[] = [
    { key: SplitterDirection.Vertical, text: "Vertical" },
    { key: SplitterDirection.Horizontal, text: "Horizontal" },
@@ -223,6 +278,20 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
    setPageOption(option.key);
  };

+  const handleOnDataPlaneRBACOptionChange = (
+    ev: React.FormEvent<HTMLInputElement>,
+    option: IChoiceGroupOption,
+  ): void => {
+    setEnableDataPlaneRBACOption(option.key);
+
+    const shouldShowWarning =
+      (option.key === Constants.RBACOptions.setTrueRBACOption ||
+        (option.key === Constants.RBACOptions.setAutomaticRBACOption &&
+          userContext.databaseAccount.properties.disableLocalAuth === true)) &&
+      !useDataPlaneRbac.getState().aadTokenUpdated;
+    setShowDataPlaneRBACWarning(shouldShowWarning);
+  };
+
  const handleOnRUThresholdToggleChange = (ev: React.MouseEvent<HTMLElement>, checked?: boolean): void => {
    setRUThresholdEnabled(checked);
  };
@@ -383,6 +452,54 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
            </div>
          </div>
        )}
+        {userContext.apiType === "SQL" && userContext.authType === AuthType.AAD && (
+          <>
+            <div className="settingsSection">
+              <div className="settingsSectionPart">
+                <fieldset>
+                  <legend id="enableDataPlaneRBACOptions" className="settingsSectionLabel legendLabel">
+                    Enable Entra ID RBAC
+                  </legend>
+                  <TooltipHost
+                    content={
+                      <>
+                        Choose Automatic to enable Entra ID RBAC automatically. True/False to force enable/disable Entra
+                        ID RBAC.
+                        <a
+                          href="https://learn.microsoft.com/en-us/azure/cosmos-db/how-to-setup-rbac#use-data-explorer"
+                          target="_blank"
+                          rel="noopener noreferrer"
+                        >
+                          {" "}
+                          Learn more{" "}
+                        </a>
+                      </>
+                    }
+                  >
+                    <Icon iconName="Info" ariaLabel="Info tooltip" className="panelInfoIcon" tabIndex={0} />
+                  </TooltipHost>
+                  {showDataPlaneRBACWarning && configContext.platform === Platform.Portal && (
+                    <MessageBar
+                      messageBarType={MessageBarType.warning}
+                      isMultiline={true}
+                      onDismiss={() => setShowDataPlaneRBACWarning(false)}
+                      dismissButtonAriaLabel="Close"
+                    >
+                      Please click on &quot;Login for Entra ID RBAC&quot; prior to performing Entra ID RBAC operations
+                    </MessageBar>
+                  )}
+                  <ChoiceGroup
+                    ariaLabelledBy="enableDataPlaneRBACOptions"
+                    options={dataPlaneRBACOptionsList}
+                    styles={choiceButtonStyles}
+                    selectedKey={enableDataPlaneRBACOption}
+                    onChange={handleOnDataPlaneRBACOptionChange}
+                  />
+                </fieldset>
+              </div>
+            </div>
+          </>
+        )}
        {userContext.apiType === "SQL" && (
          <>
            <div className="settingsSection">