From 478467bda5463365d2d4a1a4072b292cf20ffbea Mon Sep 17 00:00:00 2001
From: Senthamil Sindhu <sindhuba@microsoft.com>
Date: Tue, 25 Jun 2024 16:27:40 -0700
Subject: [PATCH] Fix  enable AAD dataplane feature flag behavior

---
 src/Common/CosmosClient.ts       |  8 ++++----
 src/hooks/useKnockoutExplorer.ts | 32 +++++++++++++++++---------------
 2 files changed, 21 insertions(+), 19 deletions(-)

diff --git a/src/Common/CosmosClient.ts b/src/Common/CosmosClient.ts
index 0dc2737d0..9d35b15e5 100644
--- a/src/Common/CosmosClient.ts
+++ b/src/Common/CosmosClient.ts
@@ -19,10 +19,10 @@ const _global = typeof self === "undefined" ? window : self;
 export const tokenProvider = async (requestInfo: Cosmos.RequestInfo) => {
   const { verb, resourceId, resourceType, headers } = requestInfo;
 
-  if (
-    (userContext.features.enableAadDataPlane && userContext.databaseAccount.properties.disableLocalAuth) ||
-    (userContext.dataPlaneRbacEnabled && userContext.apiType === "SQL")
-  ) {
+  const aadDataPlaneFeatureEnabled =
+    userContext.features.enableAadDataPlane && userContext.databaseAccount.properties.disableLocalAuth;
+  const dataPlaneRBACOptionEnabled = userContext.dataPlaneRbacEnabled && userContext.apiType === "SQL";
+  if (aadDataPlaneFeatureEnabled || (!userContext.features.enableAadDataPlane && dataPlaneRBACOptionEnabled)) {
     if (!userContext.aadToken) {
       logConsoleError(
         `AAD token does not exist. Please use "Login for Entra ID" prior to performing Entra ID RBAC operations`,
diff --git a/src/hooks/useKnockoutExplorer.ts b/src/hooks/useKnockoutExplorer.ts
index 754daa4da..b899a4a84 100644
--- a/src/hooks/useKnockoutExplorer.ts
+++ b/src/hooks/useKnockoutExplorer.ts
@@ -276,24 +276,26 @@ async function configureHostedWithAAD(config: AAD): Promise<Explorer> {
       databaseAccount: config.databaseAccount,
     });
 
-    if (userContext.apiType === "SQL") {
-      if (LocalStorageUtility.hasItem(StorageKey.DataPlaneRbacEnabled)) {
-        const isDataPlaneRbacSetting = LocalStorageUtility.getEntryString(StorageKey.DataPlaneRbacEnabled);
+    if (!userContext.features.enableAadDataPlane) {
+      if (userContext.apiType === "SQL") {
+        if (LocalStorageUtility.hasItem(StorageKey.DataPlaneRbacEnabled)) {
+          const isDataPlaneRbacSetting = LocalStorageUtility.getEntryString(StorageKey.DataPlaneRbacEnabled);
 
-        let dataPlaneRbacEnabled;
-        if (isDataPlaneRbacSetting === Constants.RBACOptions.setAutomaticRBACOption) {
-          dataPlaneRbacEnabled = account.properties.disableLocalAuth;
-        } else {
-          dataPlaneRbacEnabled = isDataPlaneRbacSetting === Constants.RBACOptions.setTrueRBACOption;
+          let dataPlaneRbacEnabled;
+          if (isDataPlaneRbacSetting === Constants.RBACOptions.setAutomaticRBACOption) {
+            dataPlaneRbacEnabled = account.properties.disableLocalAuth;
+          } else {
+            dataPlaneRbacEnabled = isDataPlaneRbacSetting === Constants.RBACOptions.setTrueRBACOption;
+          }
+
+          updateUserContext({ dataPlaneRbacEnabled });
         }
-
-        updateUserContext({ dataPlaneRbacEnabled });
+      } else {
+        const keys: DatabaseAccountListKeysResult = await listKeys(subscriptionId, resourceGroup, account.name);
+        updateUserContext({
+          masterKey: keys.primaryMasterKey,
+        });
       }
-    } else {
-      const keys: DatabaseAccountListKeysResult = await listKeys(subscriptionId, resourceGroup, account.name);
-      updateUserContext({
-        masterKey: keys.primaryMasterKey,
-      });
     }
   } catch (e) {
     if (userContext.features.enableAadDataPlane) {