Add support for psql shell (#1324)

Add support for psql shell
- add new terminal type
- handle missing documentEndpoint property

src/Contracts/ViewModels.ts

@@ -372,6 +372,7 @@ export enum TerminalKind {
   Default = 0,
   Mongo = 1,
   Cassandra = 2,
+  Postgres = 3,
 }
 
 export interface DataExplorerInputsFrame {

src/Explorer/Controls/Notebook/NotebookTerminalComponent.tsx

@@ -74,6 +74,8 @@ export class NotebookTerminalComponent extends React.Component<NotebookTerminalC
         this.props.databaseAccount?.properties.mongoEndpoint || this.props.databaseAccount?.properties.documentEndpoint;
     } else if (StringUtils.endsWith(notebookServerEndpoint, "cassandra")) {
       terminalEndpoint = this.props.databaseAccount?.properties.cassandraEndpoint;
+    } else if (StringUtils.endsWith(notebookServerEndpoint, "postgresql")) {
+      return (this.props.databaseAccount?.properties as any).postgresqlEndpoint;
     }
 
     if (terminalEndpoint) {

src/Explorer/Controls/Settings/__snapshots__/SettingsComponent.test.tsx.snap

@@ -35,6 +35,7 @@ exports[`SettingsComponent renders 1`] = `
                 "onRefreshDatabasesKeyPress": [Function],
                 "onRefreshResourcesClick": [Function],
                 "phoenixClient": PhoenixClient {
+                  "armResourceId": undefined,
                   "retryOptions": Object {
                     "maxTimeout": 5000,
                     "minTimeout": 5000,
@@ -111,6 +112,7 @@ exports[`SettingsComponent renders 1`] = `
                 "onRefreshDatabasesKeyPress": [Function],
                 "onRefreshResourcesClick": [Function],
                 "phoenixClient": PhoenixClient {
+                  "armResourceId": undefined,
                   "retryOptions": Object {
                     "maxTimeout": 5000,
                     "minTimeout": 5000,

src/Explorer/Explorer.tsx

@@ -93,7 +93,8 @@ export default class Explorer {
       dataExplorerArea: Constants.Areas.ResourceTree,
     });
     this._isInitializingNotebooks = false;
-    this.phoenixClient = new PhoenixClient();
+
+    this.phoenixClient = new PhoenixClient(userContext?.databaseAccount?.id);
     useNotebook.subscribe(
       () => this.refreshCommandBarButtons(),
       (state) => state.isNotebooksEnabledForAccount
@@ -353,7 +354,7 @@ export default class Explorer {
         (notebookServerInfo && notebookServerInfo.notebookServerEndpoint === undefined))
     ) {
       const provisionData: IProvisionData = {
-        cosmosEndpoint: userContext.databaseAccount.properties.documentEndpoint,
+        cosmosEndpoint: userContext?.databaseAccount?.properties?.documentEndpoint,
         poolId: PoolIdType.DefaultPoolId,
       };
       const connectionStatus: ContainerConnectionInfo = {
@@ -1058,6 +1059,10 @@ export default class Explorer {
         title = "Cassandra Shell";
         break;
 
+      case ViewModels.TerminalKind.Postgres:
+        title = "PSQL Shell";
+        break;
+
       default:
         throw new Error("Terminal kind: ${kind} not supported");
     }

src/Explorer/Menus/CommandBar/CommandBarComponentButtonFactory.tsx

@@ -85,11 +85,14 @@ export function createStaticCommandBarButtons(
       (userContext.apiType === "Mongo" &&
         useNotebook.getState().isShellEnabled &&
         selectedNodeState.isDatabaseNodeOrNoneSelected()) ||
-      userContext.apiType === "Cassandra"
+      userContext.apiType === "Cassandra" ||
+      userContext.apiType === "Postgres"
     ) {
       notebookButtons.push(createDivider());
       if (userContext.apiType === "Cassandra") {
         notebookButtons.push(createOpenCassandraTerminalButton(container));
+      } else if (userContext.apiType === "Postgres") {
+        notebookButtons.push(createOpenPsqlTerminalButton(container));
       } else {
         notebookButtons.push(createOpenMongoTerminalButton(container));
       }
@@ -523,6 +526,28 @@ function createOpenCassandraTerminalButton(container: Explorer): CommandButtonCo
   };
 }
 
+function createOpenPsqlTerminalButton(container: Explorer): CommandButtonComponentProps {
+  const label = "Open PSQL Shell";
+  const disableButton =
+    !useNotebook.getState().isNotebooksEnabledForAccount && !useNotebook.getState().isNotebookEnabled;
+  return {
+    iconSrc: HostedTerminalIcon,
+    iconAlt: label,
+    onCommandClick: () => {
+      if (useNotebook.getState().isNotebookEnabled) {
+        container.openNotebookTerminal(ViewModels.TerminalKind.Postgres);
+      }
+    },
+    commandButtonLabel: label,
+    hasPopup: false,
+    disabled: disableButton,
+    ariaLabel: label,
+    tooltipText: !disableButton
+      ? ""
+      : "This feature is not yet available in your account's region. View supported regions here: https://aka.ms/cosmos-enable-notebooks.",
+  };
+}
+
 function createNotebookWorkspaceResetButton(container: Explorer): CommandButtonComponentProps {
   const label = "Reset Workspace";
   return {

src/Explorer/Notebook/NotebookContainerClient.ts

@@ -23,7 +23,7 @@ export class NotebookContainerClient {
   private scheduleTimerId: NodeJS.Timeout;
 
   constructor(private onConnectionLost: () => void) {
-    this.phoenixClient = new PhoenixClient();
+    this.phoenixClient = new PhoenixClient(userContext?.databaseAccount?.id);
     this.retryOptions = {
       retries: Notebook.retryAttempts,
       maxTimeout: Notebook.retryAttemptDelayMs,

src/Explorer/Notebook/useNotebook.ts

@@ -307,7 +307,7 @@ export const useNotebook: UseStore<NotebookState> = create((set, get) => ({
       let isPhoenixFeatures = false;
 
       const isPublicInternetAllowed = isPublicInternetAccessAllowed();
-      const phoenixClient = new PhoenixClient();
+      const phoenixClient = new PhoenixClient(userContext?.databaseAccount?.id);
       const dbAccountAllowedInfo = await phoenixClient.getDbAccountAllowedStatus();
 
       if (dbAccountAllowedInfo.status === HttpStatusCodes.OK) {

src/Explorer/Panes/GitHubReposPanel/__snapshots__/GitHubReposPanel.test.tsx.snap

@@ -24,6 +24,7 @@ exports[`GitHub Repos Panel should render Default properly 1`] = `
             "onRefreshDatabasesKeyPress": [Function],
             "onRefreshResourcesClick": [Function],
             "phoenixClient": PhoenixClient {
+              "armResourceId": undefined,
               "retryOptions": Object {
                 "maxTimeout": 5000,
                 "minTimeout": 5000,

src/Explorer/Panes/StringInputPane/__snapshots__/StringInputPane.test.tsx.snap

@@ -14,6 +14,7 @@ exports[`StringInput Pane should render Create new directory properly 1`] = `
       "onRefreshDatabasesKeyPress": [Function],
       "onRefreshResourcesClick": [Function],
       "phoenixClient": PhoenixClient {
+        "armResourceId": undefined,
         "retryOptions": Object {
           "maxTimeout": 5000,
           "minTimeout": 5000,

src/Explorer/Tabs/TerminalTab.tsx

@@ -95,6 +95,10 @@ export default class TerminalTab extends TabsBase {
         endpointSuffix = "cassandra";
         break;
 
+      case ViewModels.TerminalKind.Postgres:
+        endpointSuffix = "postgresql";
+        break;
+
       default:
         throw new Error(`Terminal kind: ${options.kind} not supported`);
     }

src/Phoenix/PhoenixClient.ts

@@ -32,6 +32,7 @@ import { userContext } from "../UserContext";
 import { getAuthorizationHeader } from "../Utils/AuthorizationUtils";
 
 export class PhoenixClient {
+  private armResourceId: string;
   private containerHealthHandler: NodeJS.Timeout;
   private retryOptions: promiseRetry.Options = {
     retries: Notebook.retryAttempts,
@@ -39,6 +40,10 @@ export class PhoenixClient {
     minTimeout: Notebook.retryAttemptDelayMs,
   };
 
+  constructor(armResourceId: string) {
+    this.armResourceId = armResourceId;
+  }
+
   public async allocateContainer(provisionData: IProvisionData): Promise<IResponse<IPhoenixServiceInfo>> {
     return this.executeContainerAssignmentOperation(provisionData, "allocate");
   }
@@ -214,22 +219,21 @@ export class PhoenixClient {
     }
   }
 
-  public static getPhoenixEndpoint(): string {
-    const phoenixEndpoint =
+  private getPhoenixControlPlanePathPrefix(): string {
+    if (!this.armResourceId) {
+      throw new Error("The Phoenix client was not initialized properly: missing ARM resourcce id");
+    }
+
+    const toolsEndpoint =
       userContext.features.phoenixEndpoint ?? userContext.features.junoEndpoint ?? configContext.JUNO_ENDPOINT;
-    if (!validateEndpoint(phoenixEndpoint, allowedJunoOrigins)) {
-      const error = `${phoenixEndpoint} not allowed as juno endpoint`;
+
+    if (!validateEndpoint(toolsEndpoint, allowedJunoOrigins)) {
+      const error = `${toolsEndpoint} not allowed as tools endpoint`;
       console.error(error);
       throw new Error(error);
     }
 
-    return phoenixEndpoint;
-  }
-
-  public getPhoenixControlPlanePathPrefix(): string {
-    return `${PhoenixClient.getPhoenixEndpoint()}/api/controlplane/toolscontainer/cosmosaccounts${
-      userContext.databaseAccount.id
-    }`;
+    return `${toolsEndpoint}/api/controlplane/toolscontainer/cosmosaccounts${this.armResourceId}`;
   }
 
   private static getHeaders(): HeadersInit {

src/Terminal/JupyterLabAppFactory.ts

@@ -24,6 +24,10 @@ export class JupyterLabAppFactory {
     this.isShellStarted = content?.includes("Connected to") && content?.includes("cqlsh");
   }
 
+  private isPostgresShellStarted(content: string | undefined) {
+    this.isShellStarted = content?.includes("citus=>");
+  }
+
   constructor(closeTab: () => void) {
     this.onShellExited = closeTab;
     this.isShellStarted = false;
@@ -36,6 +40,9 @@ export class JupyterLabAppFactory {
       case "Cassandra":
         this.checkShellStarted = this.isCassandraShellStarted;
         break;
+      case "Postgres":
+        this.checkShellStarted = this.isPostgresShellStarted;
+        break;
     }
   }