Merge branch 'master' of https://github.com/Azure/cosmos-explorer into users/artrejo/MSRC69675

src/Phoenix/PhoenixClient.ts

@@ -1,9 +1,15 @@
-import { validateEndpoint } from "Utils/EndpointValidation";
-import { ContainerStatusType, HttpHeaders, HttpStatusCodes, Notebook } from "../Common/Constants";
+import promiseRetry, { AbortError } from "p-retry";
+import { Action } from "Shared/Telemetry/TelemetryConstants";
+import { allowedJunoOrigins, validateEndpoint } from "Utils/EndpointValidation";
+import {
+  Areas,
+  ConnectionStatusType, ContainerStatusType, HttpHeaders, HttpStatusCodes, Notebook
+} from "../Common/Constants";
 import { getErrorMessage } from "../Common/ErrorHandlingUtils";
 import * as Logger from "../Common/Logger";
-import { allowedJunoOrigins, configContext } from "../ConfigContext";
+import { configContext } from "../ConfigContext";
 import {
+  ContainerConnectionInfo,
   ContainerInfo,
   IContainerData,
   IPhoenixConnectionInfoResult,
@@ -11,11 +17,17 @@ import {
   IResponse
 } from "../Contracts/DataModels";
 import { useNotebook } from "../Explorer/Notebook/useNotebook";
+import * as TelemetryProcessor from "../Shared/Telemetry/TelemetryProcessor";
 import { userContext } from "../UserContext";
 import { getAuthorizationHeader } from "../Utils/AuthorizationUtils";
 
 export class PhoenixClient {
   private containerHealthHandler: NodeJS.Timeout;
+  private retryOptions: promiseRetry.Options = {
+    retries: Notebook.retryAttempts,
+    maxTimeout: Notebook.retryAttemptDelayMs,
+    minTimeout: Notebook.retryAttemptDelayMs,
+  };
 
   public async allocateContainer(provisionData: IProvisionData): Promise<IResponse<IPhoenixConnectionInfoResult>> {
     return this.executeContainerAssignmentOperation(provisionData, "allocate");
@@ -30,8 +42,8 @@ export class PhoenixClient {
     operation: string
   ): Promise<IResponse<IPhoenixConnectionInfoResult>> {
     try {
-      const response = await fetch(`${this.getPhoenixContainerPoolingEndPoint()}/${operation}`, {
-        method: "POST",
+      const response = await fetch(`${this.getPhoenixControlPlanePathPrefix()}/containerconnections`, {
+        method: operation === "allocate" ? "POST" : "PATCH",
         headers: PhoenixClient.getHeaders(),
         body: JSON.stringify(provisionData),
       });
@@ -50,7 +62,7 @@ export class PhoenixClient {
     }
   }
 
-  public async initiateContainerHeartBeat(containerData: { forwardingId: string; dbAccountName: string }) {
+  public async initiateContainerHeartBeat(containerData: IContainerData) {
     if (this.containerHealthHandler) {
       clearTimeout(this.containerHealthHandler);
     }
@@ -65,28 +77,48 @@ export class PhoenixClient {
 
   private async getContainerStatusAsync(containerData: IContainerData): Promise<ContainerInfo> {
     try {
-      const response = await window.fetch(
-        `${this.getPhoenixContainerPoolingEndPoint()}/${containerData.dbAccountName}/${containerData.forwardingId}`,
-        {
-          method: "GET",
-          headers: PhoenixClient.getHeaders(),
+      const runContainerStatusAsync = async () => {
+        const response = await window.fetch(
+          `${this.getPhoenixControlPlanePathPrefix()}/${containerData.forwardingId}`,
+          {
+            method: "GET",
+            headers: PhoenixClient.getHeaders(),
+          }
+        );
+        if (response.status === HttpStatusCodes.OK) {
+          const containerStatus = await response.json();
+          return {
+            durationLeftInMinutes: containerStatus?.durationLeftInMinutes,
+            notebookServerInfo: containerStatus?.notebookServerInfo,
+            status: ContainerStatusType.Active,
+          };
+        } else if (response.status === HttpStatusCodes.NotFound) {
+          const error = "Disconnected from compute workspace";
+          Logger.logError(error, "");
+          const connectionStatus: ContainerConnectionInfo = {
+            status: ConnectionStatusType.Reconnect,
+          };
+          TelemetryProcessor.traceMark(Action.PhoenixHeartBeat, {
+            dataExplorerArea: Areas.Notebook,
+            message: getErrorMessage(error),
+          });
+          useNotebook.getState().resetContainerConnection(connectionStatus);
+          useNotebook.getState().setIsRefreshed(!useNotebook.getState().isRefreshed);
+          throw new AbortError(response.statusText);
         }
-      );
-      if (response.status === HttpStatusCodes.OK) {
-        const containerStatus = await response.json();
-        return {
-          durationLeftInMinutes: containerStatus?.durationLeftInMinutes,
-          notebookServerInfo: containerStatus?.notebookServerInfo,
-          status: ContainerStatusType.Active,
-        };
-      }
-      return {
-        durationLeftInMinutes: undefined,
-        notebookServerInfo: undefined,
-        status: ContainerStatusType.Disconnected,
+        throw new Error(response.statusText);
       };
+      return await promiseRetry(runContainerStatusAsync, this.retryOptions);
     } catch (error) {
-      Logger.logError(getErrorMessage(error), "PhoenixClient/getContainerStatus");
+      TelemetryProcessor.traceFailure(Action.PhoenixHeartBeat, {
+        dataExplorerArea: Areas.Notebook,
+      });
+      Logger.logError(getErrorMessage(error), "");
+      const connectionStatus: ContainerConnectionInfo = {
+        status: ConnectionStatusType.Failed,
+      };
+      useNotebook.getState().resetContainerConnection(connectionStatus);
+      useNotebook.getState().setIsRefreshed(!useNotebook.getState().isRefreshed);
       return {
         durationLeftInMinutes: undefined,
         notebookServerInfo: undefined,
@@ -95,7 +127,7 @@ export class PhoenixClient {
     }
   }
 
-  private async getContainerHealth(delayMs: number, containerData: { forwardingId: string; dbAccountName: string }) {
+  private async getContainerHealth(delayMs: number, containerData: IContainerData) {
     const containerInfo = await this.getContainerStatusAsync(containerData);
     useNotebook.getState().setContainerStatus(containerInfo);
     if (useNotebook.getState().containerStatus?.status === ContainerStatusType.Active) {
@@ -103,6 +135,19 @@ export class PhoenixClient {
     }
   }
 
+  public async isDbAcountWhitelisted(): Promise<boolean> {
+    try {
+      const response = await window.fetch(`${this.getPhoenixControlPlanePathPrefix()}`, {
+        method: "GET",
+        headers: PhoenixClient.getHeaders(),
+      });
+      return response.status === HttpStatusCodes.OK;
+    } catch (error) {
+      Logger.logError(getErrorMessage(error), "PhoenixClient/IsDbAcountWhitelisted");
+      return false;
+    }
+  }
+
   public static getPhoenixEndpoint(): string {
     const phoenixEndpoint = userContext.features.phoenixEndpoint ?? userContext.features.junoEndpoint ?? configContext.JUNO_ENDPOINT;
     if (validateEndpoint(phoenixEndpoint, allowedJunoOrigins)) {
@@ -114,8 +159,9 @@ export class PhoenixClient {
     return phoenixEndpoint;
   }
 
-  public getPhoenixContainerPoolingEndPoint(): string {
-    return `${PhoenixClient.getPhoenixEndpoint()}/api/controlplane/toolscontainer`;
+  public getPhoenixControlPlanePathPrefix(): string {
+    return `${PhoenixClient.getPhoenixEndpoint()}/api/controlplane/toolscontainer/cosmosaccounts${userContext.databaseAccount.id
+      }`;
   }
 
   private static getHeaders(): HeadersInit {