Add poolId as parameter to allocateContainer (#1588)

src/Explorer/Explorer.tsx

@@ -385,7 +385,7 @@ export default class Explorer {
     this._isInitializingNotebooks = false;
   }
 
-  public async allocateContainer(): Promise<void> {
+  public async allocateContainer(poolId: PoolIdType): Promise<void> {
     const notebookServerInfo = useNotebook.getState().notebookServerInfo;
     const isAllocating = useNotebook.getState().isAllocating;
     if (
@@ -395,7 +395,7 @@ export default class Explorer {
     ) {
       const provisionData: IProvisionData = {
         cosmosEndpoint: userContext?.databaseAccount?.properties?.documentEndpoint,
-        poolId: PoolIdType.QueryCopilot,
+        poolId: poolId,
       };
       const connectionStatus: ContainerConnectionInfo = {
         status: ConnectionStatusType.Connecting,
@@ -755,7 +755,7 @@ export default class Explorer {
       throw new Error(`Invalid notebookContentItem: ${notebookContentItem}`);
     }
     if (notebookContentItem.type === NotebookContentItemType.Notebook && useNotebook.getState().isPhoenixNotebooks) {
-      await this.allocateContainer();
+      await this.allocateContainer(PoolIdType.DefaultPoolId);
     }
 
     const notebookTabs = useTabs
@@ -980,7 +980,7 @@ export default class Explorer {
     }
     if (useNotebook.getState().isPhoenixNotebooks) {
       if (isGithubTree) {
-        await this.allocateContainer();
+        await this.allocateContainer(PoolIdType.DefaultPoolId);
         parent = parent || this.resourceTree.myNotebooksContentRoot;
         this.createNewNoteBook(parent, isGithubTree);
       } else {
@@ -989,7 +989,7 @@ export default class Explorer {
           undefined,
           "Create",
           async () => {
-            await this.allocateContainer();
+            await this.allocateContainer(PoolIdType.DefaultPoolId);
             parent = parent || this.resourceTree.myNotebooksContentRoot;
             this.createNewNoteBook(parent, isGithubTree);
           },
@@ -1068,7 +1068,7 @@ export default class Explorer {
 
   public async openNotebookTerminal(kind: ViewModels.TerminalKind): Promise<void> {
     if (useNotebook.getState().isPhoenixFeatures) {
-      await this.allocateContainer();
+      await this.allocateContainer(PoolIdType.DefaultPoolId);
       const notebookServerInfo = useNotebook.getState().notebookServerInfo;
       if (notebookServerInfo && notebookServerInfo.notebookServerEndpoint !== undefined) {
         this.connectToNotebookTerminal(kind);
@@ -1212,7 +1212,7 @@ export default class Explorer {
       await useNotebook.getState().getPhoenixStatus();
     }
     if (useNotebook.getState().isPhoenixNotebooks) {
-      await this.allocateContainer();
+      await this.allocateContainer(PoolIdType.DefaultPoolId);
     }
 
     // We still use github urls like https://github.com/Azure-Samples/cosmos-notebooks/blob/master/CSharp_quickstarts/GettingStarted_CSharp.ipynb
@@ -1249,7 +1249,7 @@ export default class Explorer {
         undefined,
         "Upload",
         async () => {
-          await this.allocateContainer();
+          await this.allocateContainer(PoolIdType.DefaultPoolId);
           parent = parent || this.resourceTree.myNotebooksContentRoot;
           this.uploadFilePanel(parent);
         },

src/Explorer/Menus/CommandBar/CommandBarComponentAdapter.tsx

@@ -5,10 +5,10 @@
  */
 import { CommandBar as FluentCommandBar, ICommandBarItemProps } from "@fluentui/react";
 import { useNotebook } from "Explorer/Notebook/useNotebook";
-import * as React from "react";
 import { userContext } from "UserContext";
+import * as React from "react";
 import create, { UseStore } from "zustand";
-import { ConnectionStatusType, StyleConstants } from "../../../Common/Constants";
+import { ConnectionStatusType, PoolIdType, StyleConstants } from "../../../Common/Constants";
 import { CommandButtonComponentProps } from "../../Controls/CommandButton/CommandButtonComponent";
 import Explorer from "../../Explorer";
 import { useSelectedNode } from "../../useSelectedNode";
@@ -76,7 +76,9 @@ export const CommandBar: React.FC<Props> = ({ container }: Props) => {
     (useNotebook.getState().isPhoenixNotebooks || useNotebook.getState().isPhoenixFeatures) &&
     connectionInfo?.status !== ConnectionStatusType.Connect
   ) {
-    uiFabricControlButtons.unshift(CommandBarUtil.createConnectionStatus(container, "connectionStatus"));
+    uiFabricControlButtons.unshift(
+      CommandBarUtil.createConnectionStatus(container, PoolIdType.DefaultPoolId, "connectionStatus")
+    );
   }
 
   return (

src/Explorer/Menus/CommandBar/CommandBarUtil.tsx

@@ -9,7 +9,7 @@ import {
 import * as React from "react";
 import _ from "underscore";
 import ChevronDownIcon from "../../../../images/Chevron_down.svg";
-import { StyleConstants } from "../../../Common/Constants";
+import { PoolIdType, StyleConstants } from "../../../Common/Constants";
 import { Action, ActionModifiers } from "../../../Shared/Telemetry/TelemetryConstants";
 import * as TelemetryProcessor from "../../../Shared/Telemetry/TelemetryProcessor";
 import { CommandButtonComponentProps } from "../../Controls/CommandButton/CommandButtonComponent";
@@ -204,9 +204,9 @@ export const createMemoryTracker = (key: string): ICommandBarItemProps => {
   };
 };
 
-export const createConnectionStatus = (container: Explorer, key: string): ICommandBarItemProps => {
+export const createConnectionStatus = (container: Explorer, poolId: PoolIdType, key: string): ICommandBarItemProps => {
   return {
     key,
-    onRender: () => <ConnectionStatus container={container} />,
+    onRender: () => <ConnectionStatus container={container} poolId={poolId} />,
   };
 };

src/Explorer/Menus/CommandBar/ConnectionStatusComponent.tsx

@@ -14,14 +14,15 @@ import { useId } from "@fluentui/react-hooks";
 import { ActionButton, DefaultButton } from "@fluentui/react/lib/Button";
 import * as React from "react";
 import "../../../../less/hostedexplorer.less";
-import { ConnectionStatusType, ContainerStatusType, Notebook } from "../../../Common/Constants";
+import { ConnectionStatusType, ContainerStatusType, Notebook, PoolIdType } from "../../../Common/Constants";
 import Explorer from "../../Explorer";
 import { useNotebook } from "../../Notebook/useNotebook";
 import "../CommandBar/ConnectionStatusComponent.less";
 interface Props {
   container: Explorer;
+  poolId: PoolIdType;
 }
-export const ConnectionStatus: React.FC<Props> = ({ container }: Props): JSX.Element => {
+export const ConnectionStatus: React.FC<Props> = ({ container, poolId }: Props): JSX.Element => {
   const connectionInfo = useNotebook((state) => state.connectionInfo);
   const [second, setSecond] = React.useState("00");
   const [minute, setMinute] = React.useState("00");
@@ -93,7 +94,7 @@ export const ConnectionStatus: React.FC<Props> = ({ container }: Props): JSX.Ele
     (connectionInfo.status === ConnectionStatusType.Connect || connectionInfo.status === ConnectionStatusType.Reconnect)
   ) {
     return (
-      <ActionButton className="commandReactBtn" onClick={() => container.allocateContainer()}>
+      <ActionButton className="commandReactBtn" onClick={() => container.allocateContainer(poolId)}>
         <TooltipHost content={toolTipContent}>
           <Stack className="connectionStatusContainer" horizontal>
             <Icon iconName="ConnectVirtualMachine" className="connectIcon" />
@@ -133,7 +134,9 @@ export const ConnectionStatus: React.FC<Props> = ({ container }: Props): JSX.Ele
           id={buttonId}
           className={connectionInfo.status === ConnectionStatusType.Failed ? "commandReactBtn" : "connectedReactBtn"}
           onClick={(e: React.MouseEvent<HTMLSpanElement>) =>
-            connectionInfo.status === ConnectionStatusType.Failed ? container.allocateContainer() : e.preventDefault()
+            connectionInfo.status === ConnectionStatusType.Failed
+              ? container.allocateContainer(poolId)
+              : e.preventDefault()
           }
         >
           <Stack className="connectionStatusContainer" horizontal>

src/Explorer/Panes/CopyNotebookPane/CopyNotebookPane.tsx

@@ -1,12 +1,12 @@
 import { IDropdownOption } from "@fluentui/react";
 import React, { FormEvent, FunctionComponent, useEffect, useState } from "react";
-import { HttpStatusCodes } from "../../../Common/Constants";
+import { HttpStatusCodes, PoolIdType } from "../../../Common/Constants";
 import { getErrorMessage, handleError } from "../../../Common/ErrorHandlingUtils";
 import { GitHubOAuthService } from "../../../GitHub/GitHubOAuthService";
-import { useSidePanel } from "../../../hooks/useSidePanel";
 import { IPinnedRepo, JunoClient } from "../../../Juno/JunoClient";
 import * as GitHubUtils from "../../../Utils/GitHubUtils";
 import * as NotificationConsoleUtils from "../../../Utils/NotificationConsoleUtils";
+import { useSidePanel } from "../../../hooks/useSidePanel";
 import Explorer from "../../Explorer";
 import { NotebookContentItem, NotebookContentItemType } from "../../Notebook/NotebookContentItem";
 import { useNotebook } from "../../Notebook/useNotebook";
@@ -109,7 +109,7 @@ export const CopyNotebookPane: FunctionComponent<CopyNotebookPanelProps> = ({
         };
         isGithubTree = false;
         if (useNotebook.getState().isPhoenixNotebooks) {
-          await container.allocateContainer();
+          await container.allocateContainer(PoolIdType.DefaultPoolId);
         }
         break;
 

src/Explorer/QueryCopilot/QueryCopilotTab.tsx

@@ -18,6 +18,7 @@ import {
 } from "@fluentui/react";
 import { useBoolean } from "@fluentui/react-hooks";
 import {
+  PoolIdType,
   QueryCopilotSampleContainerId,
   QueryCopilotSampleContainerSchema,
   ShortenedQueryCopilotSampleContainerSchema,
@@ -180,7 +181,7 @@ export const QueryCopilotTab: React.FC<QueryCopilotProps> = ({ explorer }: Query
   const generateSQLQuery = async (): Promise<void> => {
     try {
       if (shouldAllocateContainer && userContext.features.enableCopilotPhoenixGateaway) {
-        await explorer.allocateContainer();
+        await explorer.allocateContainer(PoolIdType.QueryCopilot);
         setShouldAllocateContainer(false);
       }
 

src/Explorer/QueryCopilot/Shared/QueryCopilotClient.ts

@@ -1,4 +1,8 @@
-import { QueryCopilotSampleContainerSchema, ShortenedQueryCopilotSampleContainerSchema } from "Common/Constants";
+import {
+  PoolIdType,
+  QueryCopilotSampleContainerSchema,
+  ShortenedQueryCopilotSampleContainerSchema,
+} from "Common/Constants";
 import { handleError } from "Common/ErrorHandlingUtils";
 import { createUri } from "Common/UrlUtility";
 import Explorer from "Explorer/Explorer";
@@ -24,7 +28,7 @@ export const SendQueryRequest = async ({
       .setChatMessages([...useQueryCopilot.getState().chatMessages, { source: 0, message: userPrompt }]);
     try {
       if (useQueryCopilot.getState().shouldAllocateContainer) {
-        await explorer.allocateContainer();
+        await explorer.allocateContainer(PoolIdType.DefaultPoolId);
         useQueryCopilot.getState().setShouldAllocateContainer(false);
       }
 
@@ -98,7 +102,7 @@ export const SubmitFeedback = async ({
       contact: contact || "",
     };
     if (shouldAllocateContainer && userContext.features.enableCopilotPhoenixGateaway) {
-      await explorer.allocateContainer();
+      await explorer.allocateContainer(PoolIdType.DefaultPoolId);
       setShouldAllocateContainer(false);
     }
     const serverInfo = useNotebook.getState().notebookServerInfo;

src/Explorer/Tabs/QuickstartTab.tsx

@@ -1,4 +1,5 @@
 import { Spinner, SpinnerSize, Stack, Text } from "@fluentui/react";
+import { PoolIdType } from "Common/Constants";
 import { configContext } from "ConfigContext";
 import { NotebookWorkspaceConnectionInfo, PostgresFirewallRule } from "Contracts/DataModels";
 import { NotebookTerminalComponent } from "Explorer/Controls/Notebook/NotebookTerminalComponent";
@@ -51,7 +52,7 @@ export const QuickstartTab: React.FC<QuickstartTabProps> = ({ explorer }: Quicks
   });
 
   useEffect(() => {
-    explorer.allocateContainer();
+    explorer.allocateContainer(PoolIdType.DefaultPoolId);
   }, []);
 
   return (

src/Explorer/Tree/Collection.ts

@@ -3,6 +3,9 @@ import { useNotebook } from "Explorer/Notebook/useNotebook";
 import * as ko from "knockout";
 import * as _ from "underscore";
 import * as Constants from "../../Common/Constants";
+import { getErrorMessage, getErrorStack } from "../../Common/ErrorHandlingUtils";
+import * as Logger from "../../Common/Logger";
+import { fetchPortalNotifications } from "../../Common/PortalNotifications";
 import { bulkCreateDocument } from "../../Common/dataAccess/bulkCreateDocument";
 import { createDocument } from "../../Common/dataAccess/createDocument";
 import { getCollectionUsageSizeInKB } from "../../Common/dataAccess/getCollectionDataUsageSize";
@@ -10,19 +13,16 @@ import { readCollectionOffer } from "../../Common/dataAccess/readCollectionOffer
 import { readStoredProcedures } from "../../Common/dataAccess/readStoredProcedures";
 import { readTriggers } from "../../Common/dataAccess/readTriggers";
 import { readUserDefinedFunctions } from "../../Common/dataAccess/readUserDefinedFunctions";
-import { getErrorMessage, getErrorStack } from "../../Common/ErrorHandlingUtils";
-import * as Logger from "../../Common/Logger";
-import { fetchPortalNotifications } from "../../Common/PortalNotifications";
 import * as DataModels from "../../Contracts/DataModels";
 import * as ViewModels from "../../Contracts/ViewModels";
 import { UploadDetailsRecord } from "../../Contracts/ViewModels";
-import { useTabs } from "../../hooks/useTabs";
 import { Action, ActionModifiers } from "../../Shared/Telemetry/TelemetryConstants";
 import * as TelemetryProcessor from "../../Shared/Telemetry/TelemetryProcessor";
 import { userContext } from "../../UserContext";
-import { SqlTriggerResource } from "../../Utils/arm/generatedClients/cosmos/types";
 import { isServerlessAccount } from "../../Utils/CapabilityUtils";
 import { logConsoleInfo } from "../../Utils/NotificationConsoleUtils";
+import { SqlTriggerResource } from "../../Utils/arm/generatedClients/cosmos/types";
+import { useTabs } from "../../hooks/useTabs";
 import Explorer from "../Explorer";
 import { useCommandBar } from "../Menus/CommandBar/CommandBarComponentAdapter";
 import { CassandraAPIDataClient, CassandraTableKey, CassandraTableKeys } from "../Tables/TableDataClient";
@@ -526,7 +526,7 @@ export default class Collection implements ViewModels.Collection {
 
   public onSchemaAnalyzerClick = async () => {
     if (useNotebook.getState().isPhoenixFeatures) {
-      await this.container.allocateContainer();
+      await this.container.allocateContainer(Constants.PoolIdType.DefaultPoolId);
     }
     useSelectedNode.getState().setSelectedNode(this);
     this.selectedSubnodeKind(ViewModels.CollectionTabKind.SchemaAnalyzer);

src/Utils/GalleryUtils.ts

@@ -1,7 +1,7 @@
 import { IChoiceGroupOption, IChoiceGroupProps, IProgressIndicatorProps } from "@fluentui/react";
 import { Notebook } from "@nteract/commutable";
 import { NotebookV4 } from "@nteract/commutable/lib/v4";
-import { HttpStatusCodes } from "../Common/Constants";
+import { HttpStatusCodes, PoolIdType } from "../Common/Constants";
 import { getErrorMessage, getErrorStack, handleError } from "../Common/ErrorHandlingUtils";
 import { TextFieldProps, useDialog } from "../Explorer/Controls/Dialog";
 import {
@@ -229,7 +229,7 @@ export function downloadItem(
     "Download",
     async () => {
       if (useNotebook.getState().isPhoenixNotebooks) {
-        await container.allocateContainer();
+        await container.allocateContainer(PoolIdType.DefaultPoolId);
       }
       const notebookServerInfo = useNotebook.getState().notebookServerInfo;
       if (notebookServerInfo && notebookServerInfo.notebookServerEndpoint !== undefined) {