Move notebook flags to zustand (#912)

src/Explorer/Notebook/NotebookContainerClient.ts

@@ -8,25 +8,26 @@ import * as DataModels from "../../Contracts/DataModels";
 import { userContext } from "../../UserContext";
 import { createOrUpdate, destroy } from "../../Utils/arm/generatedClients/cosmosNotebooks/notebookWorkspaces";
 import { logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
+import { useNotebook } from "./useNotebook";
 
 export class NotebookContainerClient {
   private clearReconnectionAttemptMessage? = () => {};
   private isResettingWorkspace: boolean;
 
-  constructor(
-    private notebookServerInfo: ko.Observable<DataModels.NotebookWorkspaceConnectionInfo>,
-    private onConnectionLost: () => void,
-    private onMemoryUsageInfoUpdate: (update: DataModels.MemoryUsageInfo) => void
-  ) {
-    if (notebookServerInfo() && notebookServerInfo().notebookServerEndpoint) {
+  constructor(private onConnectionLost: () => void) {
+    const notebookServerInfo = useNotebook.getState().notebookServerInfo;
+    if (notebookServerInfo?.notebookServerEndpoint) {
       this.scheduleHeartbeat(Constants.Notebook.heartbeatDelayMs);
     } else {
-      const subscription = notebookServerInfo.subscribe((newServerInfo: DataModels.NotebookWorkspaceConnectionInfo) => {
-        if (newServerInfo && newServerInfo.notebookServerEndpoint) {
-          this.scheduleHeartbeat(Constants.Notebook.heartbeatDelayMs);
-        }
-        subscription.dispose();
-      });
+      const unsub = useNotebook.subscribe(
+        (newServerInfo: DataModels.NotebookWorkspaceConnectionInfo) => {
+          if (newServerInfo?.notebookServerEndpoint) {
+            this.scheduleHeartbeat(Constants.Notebook.heartbeatDelayMs);
+          }
+          unsub();
+        },
+        (state) => state.notebookServerInfo
+      );
     }
   }
 
@@ -36,13 +37,14 @@ export class NotebookContainerClient {
   private scheduleHeartbeat(delayMs: number): void {
     setTimeout(() => {
       this.getMemoryUsage()
-        .then((memoryUsageInfo) => this.onMemoryUsageInfoUpdate(memoryUsageInfo))
+        .then((memoryUsageInfo) => useNotebook.getState().setMemoryUsageInfo(memoryUsageInfo))
         .finally(() => this.scheduleHeartbeat(Constants.Notebook.heartbeatDelayMs));
     }, delayMs);
   }
 
   private async getMemoryUsage(): Promise<DataModels.MemoryUsageInfo> {
-    if (!this.notebookServerInfo() || !this.notebookServerInfo().notebookServerEndpoint) {
+    const notebookServerInfo = useNotebook.getState().notebookServerInfo;
+    if (!notebookServerInfo || !notebookServerInfo.notebookServerEndpoint) {
       const error = "No server endpoint detected";
       Logger.logError(error, "NotebookContainerClient/getMemoryUsage");
       return Promise.reject(error);
@@ -98,7 +100,8 @@ export class NotebookContainerClient {
   }
 
   private async _resetWorkspace(): Promise<void> {
-    if (!this.notebookServerInfo() || !this.notebookServerInfo().notebookServerEndpoint) {
+    const notebookServerInfo = useNotebook.getState().notebookServerInfo;
+    if (!notebookServerInfo || !notebookServerInfo.notebookServerEndpoint) {
       const error = "No server endpoint detected";
       Logger.logError(error, "NotebookContainerClient/resetWorkspace");
       return Promise.reject(error);
@@ -117,15 +120,11 @@ export class NotebookContainerClient {
   }
 
   private getNotebookServerConfig(): { notebookServerEndpoint: string; authToken: string } {
-    let authToken: string,
-      notebookServerEndpoint = this.notebookServerInfo().notebookServerEndpoint,
-      token = this.notebookServerInfo().authToken;
-    if (token) {
-      authToken = `Token ${token}`;
-    }
+    const notebookServerInfo = useNotebook.getState().notebookServerInfo;
+    const authToken: string = notebookServerInfo.authToken ? `Token ${notebookServerInfo.authToken}` : undefined;
 
     return {
-      notebookServerEndpoint,
+      notebookServerEndpoint: notebookServerInfo.notebookServerEndpoint,
       authToken,
     };
   }

src/Explorer/Notebook/NotebookContentClient.ts

@@ -1,18 +1,14 @@
 import { stringifyNotebook } from "@nteract/commutable";
 import { FileType, IContent, IContentProvider, IEmptyContent, ServerConfig } from "@nteract/core";
 import { AjaxResponse } from "rxjs/ajax";
-import * as DataModels from "../../Contracts/DataModels";
 import * as StringUtils from "../../Utils/StringUtils";
 import * as FileSystemUtil from "./FileSystemUtil";
 import { NotebookContentItem, NotebookContentItemType } from "./NotebookContentItem";
 import { NotebookUtil } from "./NotebookUtil";
+import { useNotebook } from "./useNotebook";
 
 export class NotebookContentClient {
-  constructor(
-    private notebookServerInfo: ko.Observable<DataModels.NotebookWorkspaceConnectionInfo>,
-    public notebookBasePath: ko.Observable<string>,
-    private contentProvider: IContentProvider
-  ) {}
+  constructor(private contentProvider: IContentProvider) {}
 
   /**
    * This updates the item and points all the children's parent to this item
@@ -271,9 +267,10 @@ export class NotebookContentClient {
   }
 
   private getServerConfig(): ServerConfig {
+    const notebookServerInfo = useNotebook.getState().notebookServerInfo;
     return {
-      endpoint: this.notebookServerInfo().notebookServerEndpoint,
-      token: this.notebookServerInfo().authToken,
+      endpoint: notebookServerInfo.notebookServerEndpoint,
+      token: notebookServerInfo.authToken,
       crossDomain: true,
     };
   }

src/Explorer/Notebook/NotebookManager.tsx

@@ -4,13 +4,11 @@
 
 import { ImmutableNotebook } from "@nteract/commutable";
 import type { IContentProvider } from "@nteract/core";
-import ko from "knockout";
 import React from "react";
 import { contents } from "rx-jupyter";
 import { Areas, HttpStatusCodes } from "../../Common/Constants";
 import { getErrorMessage } from "../../Common/ErrorHandlingUtils";
 import * as Logger from "../../Common/Logger";
-import { MemoryUsageInfo } from "../../Contracts/DataModels";
 import { GitHubClient } from "../../GitHub/GitHubClient";
 import { GitHubContentProvider } from "../../GitHub/GitHubContentProvider";
 import { GitHubOAuthService } from "../../GitHub/GitHubOAuthService";
@@ -37,7 +35,6 @@ export type { NotebookPaneContent };
 
 export interface NotebookManagerOptions {
   container: Explorer;
-  notebookBasePath: ko.Observable<string>;
   resourceTree: ResourceTreeAdapter;
   refreshCommandBarButtons: () => void;
   refreshNotebookList: () => void;
@@ -81,17 +78,11 @@ export default class NotebookManager {
       contents.JupyterContentProvider
     );
 
-    this.notebookClient = new NotebookContainerClient(
-      this.params.container.notebookServerInfo,
-      () => this.params.container.initNotebooks(userContext?.databaseAccount),
-      (update: MemoryUsageInfo) => this.params.container.memoryUsageInfo(update)
+    this.notebookClient = new NotebookContainerClient(() =>
+      this.params.container.initNotebooks(userContext?.databaseAccount)
     );
 
-    this.notebookContentClient = new NotebookContentClient(
-      this.params.container.notebookServerInfo,
-      this.params.notebookBasePath,
-      this.notebookContentProvider
-    );
+    this.notebookContentClient = new NotebookContentClient(this.notebookContentProvider);
 
     this.gitHubOAuthService.getTokenObservable().subscribe((token) => {
       this.gitHubClient.setToken(token?.access_token);

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 });
+    }
+  },
+}));