Address feedback comments

src/Explorer/Explorer.tsx

@@ -68,6 +68,7 @@ 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();
 
@@ -277,22 +278,17 @@ export default class Explorer {
           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) {
-          useDialog
-            .getState()
-            .showOkModalDialog(
-              "Pop up blocked",
-              "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",
-            );
+          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);
-          useDialog
-            .getState()
-            .showOkModalDialog(
-              "Failed to perform authorization",
-              `We were unable to establish authorization for this account, due to the following error: \n${errorJson}`,
-            );
+          logConsoleError(
+            `Failed to perform authorization for this account, due to the following error: \n${errorJson}`,
+          );
         }
       }
     }

src/Explorer/Menus/CommandBar/CommandBarComponentButtonFactory.tsx

@@ -20,7 +20,7 @@ import { AuthType } from "../../../AuthType";
 import * as Constants from "../../../Common/Constants";
 import { Platform, configContext } from "../../../ConfigContext";
 import * as ViewModels from "../../../Contracts/ViewModels";
-import { userContext } from "../../../UserContext";
+import { updateUserContext, userContext } from "../../../UserContext";
 import { getCollectionName, getDatabaseName } from "../../../Utils/APITypeUtils";
 import { isRunningOnNationalCloud } from "../../../Utils/CloudUtils";
 import { useSidePanel } from "../../../hooks/useSidePanel";
@@ -31,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;
 
@@ -71,11 +72,18 @@ export function createStaticCommandBarButtons(
     }
 
     if (userContext.apiType === "SQL") {
-      const addLoginForEntraIDBtn = createLoginForEntraIDButton(container);
+      const [loginButtonProps, setLoginButtonProps] = useState<CommandButtonComponentProps | undefined>(undefined);
+      const dataPlaneRbacEnabled = useDataPlaneRbac((state) => state.dataPlaneRbacEnabled);
+      const aadTokenUpdated = useDataPlaneRbac((state) => state.aadTokenUpdated);
 
-      if (addLoginForEntraIDBtn) {
+      useEffect(() => {
+        const buttonProps = createLoginForEntraIDButton(container);
+        setLoginButtonProps(buttonProps);
+      }, [dataPlaneRbacEnabled, aadTokenUpdated, container]);
+
+      if (loginButtonProps) {
         addDivider();
-        buttons.push(addLoginForEntraIDBtn);
+        buttons.push(loginButtonProps);
       }
     }
 
@@ -290,11 +298,20 @@ function createLoginForEntraIDButton(container: Explorer): CommandButtonComponen
     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: () => container.openLoginForEntraIDPopUp(),
+    onCommandClick: handleCommandClick,
     commandButtonLabel: label,
     hasPopup: true,
     ariaLabel: label,

src/Explorer/Panes/SettingsPane/SettingsPane.tsx

@@ -35,6 +35,23 @@ 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((set) => ({
+  dataPlaneRbacEnabled: false,
+}));
 
 export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
   explorer,
@@ -144,10 +161,12 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
       updateUserContext({
         dataPlaneRbacEnabled: true,
       });
+      useDataPlaneRbac.setState({ dataPlaneRbacEnabled: true });
     } else {
       updateUserContext({
         dataPlaneRbacEnabled: false,
       });
+      useDataPlaneRbac.setState({ dataPlaneRbacEnabled: false });
     }
 
     LocalStorageUtility.setEntryBoolean(StorageKey.RUThresholdEnabled, ruThresholdEnabled);
@@ -266,9 +285,10 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({
     setEnableDataPlaneRBACOption(option.key);
 
     const shouldShowWarning =
-      option.key === Constants.RBACOptions.setTrueRBACOption ||
-      (option.key === Constants.RBACOptions.setAutomaticRBACOption &&
-        userContext.databaseAccount.properties.disableLocalAuth === true);
+      (option.key === Constants.RBACOptions.setTrueRBACOption ||
+        (option.key === Constants.RBACOptions.setAutomaticRBACOption &&
+          userContext.databaseAccount.properties.disableLocalAuth === true)) &&
+      !useDataPlaneRbac.getState().aadTokenUpdated;
     setShowDataPlaneRBACWarning(shouldShowWarning);
   };
 

src/hooks/useKnockoutExplorer.ts

@@ -43,6 +43,7 @@ import { isInvalidParentFrameOrigin, shouldProcessMessage } from "../Utils/Messa
 import { listKeys } from "../Utils/arm/generatedClients/cosmos/databaseAccounts";
 import { DatabaseAccountListKeysResult } from "../Utils/arm/generatedClients/cosmos/types";
 import { applyExplorerBindings } from "../applyExplorerBindings";
+import { useDataPlaneRbac } from "Explorer/Panes/SettingsPane/SettingsPane";
 
 // This hook will create a new instance of Explorer.ts and bind it to the DOM
 // This hook has a LOT of magic, but ideally we can delete it once we have removed KO and switched entirely to React
@@ -485,6 +486,7 @@ async function configurePortal(): Promise<Explorer> {
               }
 
               updateUserContext({ dataPlaneRbacEnabled });
+              useDataPlaneRbac.setState({ dataPlaneRbacEnabled: dataPlaneRbacEnabled });
             }
           } else {
             const keys: DatabaseAccountListKeysResult = await listKeys(subscriptionId, resourceGroup, account.name);