Move notebook flags to zustand (#912)

src/Explorer/Notebook/useNotebook.ts

@@ -0,0 +1,106 @@
+import create, { UseStore } from "zustand";
+import { AuthType } from "../../AuthType";
+import * as Constants from "../../Common/Constants";
+import { getErrorMessage } from "../../Common/ErrorHandlingUtils";
+import * as Logger from "../../Common/Logger";
+import { configContext } from "../../ConfigContext";
+import * as DataModels from "../../Contracts/DataModels";
+import { userContext } from "../../UserContext";
+import { getAuthorizationHeader } from "../../Utils/AuthorizationUtils";
+
+interface NotebookState {
+  isNotebookEnabled: boolean;
+  isNotebooksEnabledForAccount: boolean;
+  notebookServerInfo: DataModels.NotebookWorkspaceConnectionInfo;
+  sparkClusterConnectionInfo: DataModels.SparkClusterConnectionInfo;
+  isSynapseLinkUpdating: boolean;
+  memoryUsageInfo: DataModels.MemoryUsageInfo;
+  isShellEnabled: boolean;
+  notebookBasePath: string;
+  isInitializingNotebooks: boolean;
+  setIsNotebookEnabled: (isNotebookEnabled: boolean) => void;
+  setIsNotebooksEnabledForAccount: (isNotebooksEnabledForAccount: boolean) => void;
+  setNotebookServerInfo: (notebookServerInfo: DataModels.NotebookWorkspaceConnectionInfo) => void;
+  setSparkClusterConnectionInfo: (sparkClusterConnectionInfo: DataModels.SparkClusterConnectionInfo) => void;
+  setIsSynapseLinkUpdating: (isSynapseLinkUpdating: boolean) => void;
+  setMemoryUsageInfo: (memoryUsageInfo: DataModels.MemoryUsageInfo) => void;
+  setIsShellEnabled: (isShellEnabled: boolean) => void;
+  setNotebookBasePath: (notebookBasePath: string) => void;
+  refreshNotebooksEnabledStateForAccount: () => Promise<void>;
+}
+
+export const useNotebook: UseStore<NotebookState> = create((set) => ({
+  isNotebookEnabled: false,
+  isNotebooksEnabledForAccount: false,
+  notebookServerInfo: {
+    notebookServerEndpoint: undefined,
+    authToken: undefined,
+  },
+  sparkClusterConnectionInfo: {
+    userName: undefined,
+    password: undefined,
+    endpoints: [],
+  },
+  isSynapseLinkUpdating: false,
+  memoryUsageInfo: undefined,
+  isShellEnabled: false,
+  notebookBasePath: Constants.Notebook.defaultBasePath,
+  isInitializingNotebooks: false,
+  setIsNotebookEnabled: (isNotebookEnabled: boolean) => set({ isNotebookEnabled }),
+  setIsNotebooksEnabledForAccount: (isNotebooksEnabledForAccount: boolean) => set({ isNotebooksEnabledForAccount }),
+  setNotebookServerInfo: (notebookServerInfo: DataModels.NotebookWorkspaceConnectionInfo) =>
+    set({ notebookServerInfo }),
+  setSparkClusterConnectionInfo: (sparkClusterConnectionInfo: DataModels.SparkClusterConnectionInfo) =>
+    set({ sparkClusterConnectionInfo }),
+  setIsSynapseLinkUpdating: (isSynapseLinkUpdating: boolean) => set({ isSynapseLinkUpdating }),
+  setMemoryUsageInfo: (memoryUsageInfo: DataModels.MemoryUsageInfo) => set({ memoryUsageInfo }),
+  setIsShellEnabled: (isShellEnabled: boolean) => set({ isShellEnabled }),
+  setNotebookBasePath: (notebookBasePath: string) => set({ notebookBasePath }),
+  refreshNotebooksEnabledStateForAccount: async (): Promise<void> => {
+    const { databaseAccount, authType } = userContext;
+    if (
+      authType === AuthType.EncryptedToken ||
+      authType === AuthType.ResourceToken ||
+      authType === AuthType.MasterKey
+    ) {
+      set({ isNotebooksEnabledForAccount: false });
+      return;
+    }
+
+    const firstWriteLocation =
+      databaseAccount?.properties?.writeLocations &&
+      databaseAccount?.properties?.writeLocations[0]?.locationName.toLowerCase();
+    const disallowedLocationsUri = `${configContext.BACKEND_ENDPOINT}/api/disallowedLocations`;
+    const authorizationHeader = getAuthorizationHeader();
+    try {
+      const response = await fetch(disallowedLocationsUri, {
+        method: "POST",
+        body: JSON.stringify({
+          resourceTypes: [Constants.ArmResourceTypes.notebookWorkspaces],
+        }),
+        headers: {
+          [authorizationHeader.header]: authorizationHeader.token,
+          [Constants.HttpHeaders.contentType]: "application/json",
+        },
+      });
+
+      if (!response.ok) {
+        throw new Error("Failed to fetch disallowed locations");
+      }
+
+      const disallowedLocations: string[] = await response.json();
+      if (!disallowedLocations) {
+        Logger.logInfo("No disallowed locations found", "Explorer/isNotebooksEnabledForAccount");
+        set({ isNotebooksEnabledForAccount: true });
+        return;
+      }
+
+      // firstWriteLocation should not be disallowed
+      const isAccountInAllowedLocation = firstWriteLocation && disallowedLocations.indexOf(firstWriteLocation) === -1;
+      set({ isNotebooksEnabledForAccount: isAccountInAllowedLocation });
+    } catch (error) {
+      Logger.logError(getErrorMessage(error), "Explorer/isNotebooksEnabledForAccount");
+      set({ isNotebooksEnabledForAccount: false });
+    }
+  },
+}));