refactored code

src/Explorer/Tabs/CloudShellTab/CloudShellTerminalCore.tsx

@@ -14,7 +14,8 @@ import {
   registerCloudShellProvider,
   verifyCloudShellProviderRegistration
 } from "./Data/CloudShellClient";
-import { ShellTypeHandler } from "./ShellTypes/ShellTypeFactory";
+import { END_MARKER, START_MARKER } from "./ShellTypes/AbstractShellHandler";
+import { ShellTypeHandlerFactory } from "./ShellTypes/ShellTypeFactory";
 import { AttachAddon } from "./Utils/AttachAddOn";
 import { askConfirmation, wait } from "./Utils/CommonUtils";
 import { getNormalizedRegion } from "./Utils/RegionUtils";
@@ -55,7 +56,7 @@ export const startCloudShellTerminal =
   }
 
   // Get the shell handler for this type
-  const shellHandler = ShellTypeHandler.getHandler(shellType);
+  const shellHandler = ShellTypeHandlerFactory.getHandler(shellType);
   // Configure WebSocket connection with shell-specific commands
   const socket = await establishTerminalConnection(
     terminal,
@@ -169,8 +170,14 @@ export const establishTerminalConnection = async (
   // Configure the socket
   socket = await configureSocketConnection(socket, socketUri, terminal, initCommands, 0);
 
+  const options = {
+    startMarker: START_MARKER,
+    endMarker: END_MARKER,
+    terminalSuppressedData: shellHandler.getTerminalSuppressedData()
+  };
+
   // Attach the terminal addon
-  const attachAddon = new AttachAddon(socket);
+  const attachAddon = new AttachAddon(socket, options);
   terminal.loadAddon(attachAddon);
 
   // Authorize the session
@@ -205,7 +212,7 @@ export const configureSocketConnection = async (
 
   socket.onerror = async () => {
     if (socketRetryCount < MAX_RETRY_COUNT && socket.readyState !== WebSocket.CLOSED) {
-      await configureSocketConnection(socket, uri, terminal, initCommands, socketRetryCount + 1, attachAddon);
+      await configureSocketConnection(socket, uri, terminal, initCommands, socketRetryCount + 1);
     } else {
       socket.close();
     }

src/Explorer/Tabs/CloudShellTab/ShellTypes/AbstractShellHandler.tsx

@@ -0,0 +1,45 @@
+import { userContext } from "../../../../UserContext";
+import { listKeys } from "../../../../Utils/arm/generatedClients/cosmos/databaseAccounts";
+import { getHostFromUrl } from "../Utils/CommonUtils";
+
+export const START_MARKER = `echo "START INITIALIZATION" > /dev/null`;
+export const END_MARKER = `echo "END INITIALIZATION" > /dev/null`;
+
+export abstract class AbstractShellHandler {
+  
+  abstract getShellName(): string;
+  abstract getSetUpCommands(): string[];
+  abstract getConnectionCommands(config: any): string[];
+  abstract getEndpoint(): string;
+  abstract getTerminalSuppressedData(): string;
+
+  public async getInitialCommands(): Promise<string> {
+    const dbAccount = userContext.databaseAccount;
+    const dbName = dbAccount.name;
+
+    let key = "";
+    if (dbName) {
+      const keys = await listKeys(userContext.subscriptionId, userContext.resourceGroup, dbName);
+      key = keys?.primaryMasterKey || "";
+    }
+
+    const setupCommands = this.getSetUpCommands();
+
+    const config = {
+      host: getHostFromUrl(this.getEndpoint()),
+      name: dbName,
+      password: key,
+      endpoint: this.getEndpoint(),
+    };
+    const connectionCommands = this.getConnectionCommands(config);
+
+    const allCommands = [
+      START_MARKER,
+      ...setupCommands,
+      END_MARKER,
+      ...connectionCommands
+    ];
+  
+    return allCommands.join("\n").concat("\n");
+  }
+}

src/Explorer/Tabs/CloudShellTab/ShellTypes/CassandraShellHandler.tsx

@@ -4,49 +4,40 @@
  */
 
 import { userContext } from "../../../../UserContext";
-import { listKeys } from "../../../../Utils/arm/generatedClients/cosmos/databaseAccounts";
-import { getHostFromUrl } from "../Utils/CommonUtils";
-import { ShellTypeConfig } from "./ShellTypeFactory";
+import { AbstractShellHandler } from "./AbstractShellHandler";
 
-export class CassandraShellHandler implements ShellTypeConfig {
+const PACKAGE_VERSION: string = "5.0.3";
+
+export class CassandraShellHandler extends AbstractShellHandler {
 
   public getShellName(): string {
     return "Cassandra";
   }
 
-  public async getInitialCommands(): Promise<string> {
-    const dbAccount = userContext.databaseAccount;
-    const endpoint = dbAccount.properties.cassandraEndpoint;
-    
-    // Get database key
-    const dbName = dbAccount.name;
-    let key = "";
-    if (dbName) {
-      const keys = await listKeys(userContext.subscriptionId, userContext.resourceGroup, dbName);
-      key = keys?.primaryMasterKey || "";
+  public getEndpoint(): string {
+    return userContext.databaseAccount?.properties?.cassandraEndpoint;
   }
 
-    const config = {
-      host: getHostFromUrl(endpoint),
-      name: dbAccount.name,
-      password: key,
-      endpoint: endpoint
-    };
-
-    return this.getCommands(config).join("\n").concat("\n");
-  }
-
-  private getCommands(config: any): string[] {
+  public getSetUpCommands(): string[] {
     return [
       "if ! command -v cqlsh &> /dev/null; then echo '⚠️ cqlsh not found. Installing...'; fi",
-      "if ! command -v cqlsh &> /dev/null; then curl -LO https://archive.apache.org/dist/cassandra/5.0.3/apache-cassandra-5.0.3-bin.tar.gz; fi",
-      "if ! command -v cqlsh &> /dev/null; then tar -xvzf apache-cassandra-5.0.3-bin.tar.gz; fi",
-      "if ! command -v cqlsh &> /dev/null; then mkdir -p ~/cassandra && mv apache-cassandra-5.0.3/* ~/cassandra/; fi",
+      `if ! command -v cqlsh &> /dev/null; then curl -LO https://archive.apache.org/dist/cassandra/${PACKAGE_VERSION}/apache-cassandra-${PACKAGE_VERSION}-bin.tar.gz; fi`,
+      `if ! command -v cqlsh &> /dev/null; then tar -xvzf apache-cassandra-${PACKAGE_VERSION}-bin.tar.gz; fi`,
+      `if ! command -v cqlsh &> /dev/null; then mkdir -p ~/cassandra && mv apache-cassandra-${PACKAGE_VERSION}/* ~/cassandra/; fi`,
       "if ! command -v cqlsh &> /dev/null; then echo 'export PATH=$HOME/cassandra/bin:$PATH' >> ~/.bashrc; fi",
       "if ! command -v cqlsh &> /dev/null; then echo 'export SSL_VERSION=TLSv1_2' >> ~/.bashrc; fi",
       "if ! command -v cqlsh &> /dev/null; then echo 'export SSL_VALIDATE=false' >> ~/.bashrc; fi",
-      "source ~/.bashrc",
+      "source ~/.bashrc"
+    ];
+  }
+
+  public getConnectionCommands(config: any): string[] {
+    return [
       `cqlsh ${config.host} 10350 -u ${config.name} -p ${config.password} --ssl --protocol-version=4`
     ];
   }
+
+  public getTerminalSuppressedData(): string {
+    return "Non-Generic MongoDB Shell";
+  }
 }

src/Explorer/Tabs/CloudShellTab/ShellTypes/MongoShellHandler.tsx

@@ -4,49 +4,38 @@
  */
 
 import { userContext } from "../../../../UserContext";
-import { listKeys } from "../../../../Utils/arm/generatedClients/cosmos/databaseAccounts";
-import { getHostFromUrl } from "../Utils/CommonUtils";
-import { ShellTypeConfig } from "./ShellTypeFactory";
+import { AbstractShellHandler } from "./AbstractShellHandler";
 
-export class MongoShellHandler implements ShellTypeConfig {
+const PACKAGE_VERSION: string = "2.3.8";
+
+export class MongoShellHandler extends AbstractShellHandler {
 
   public getShellName(): string {
     return "MongoDB";
   }
 
-  public async getInitialCommands(): Promise<string> {
-    const dbAccount = userContext.databaseAccount;
-    const endpoint = dbAccount.properties.mongoEndpoint;
-    
-    // Get database key
-    const dbName = dbAccount.name;
-    let key = "";
-    if (dbName) {
-      const keys = await listKeys(userContext.subscriptionId, userContext.resourceGroup, dbName);
-      key = keys?.primaryMasterKey || "";
+  public getEndpoint(): string {
+    return userContext.databaseAccount?.properties?.mongoEndpoint;
   }
 
-    const config = {
-      host: getHostFromUrl(endpoint),
-      name: dbAccount.name,
-      password: key,
-      endpoint: endpoint
-    };
-
-    return this.getCommands(config).join("\n").concat("\n");
-  }
-
-  private getCommands(config: any): string[] {
+  public getSetUpCommands(): string[] {
     return [
-      "echo 'START INITIALIZATION'",
       "if ! command -v mongosh &> /dev/null; then echo '⚠️ mongosh not found. Installing...'; fi",
-      "if ! command -v mongosh &> /dev/null; then curl -LO https://downloads.mongodb.com/compass/mongosh-2.3.8-linux-x64.tgz; fi",
-      "if ! command -v mongosh &> /dev/null; then tar -xvzf mongosh-2.3.8-linux-x64.tgz; fi",
-      "if ! command -v mongosh &> /dev/null; then mkdir -p ~/mongosh && mv mongosh-2.3.8-linux-x64/* ~/mongosh/; fi",
+      `if ! command -v mongosh &> /dev/null; then curl -LO https://downloads.mongodb.com/compass/mongosh-${PACKAGE_VERSION}-linux-x64.tgz; fi`,
+      `if ! command -v mongosh &> /dev/null; then tar -xvzf mongosh-${PACKAGE_VERSION}-linux-x64.tgz; fi`,
+      `if ! command -v mongosh &> /dev/null; then mkdir -p ~/mongosh && mv mongosh-${PACKAGE_VERSION}-linux-x64/* ~/mongosh/; fi`,
       "if ! command -v mongosh &> /dev/null; then echo 'export PATH=$HOME/mongosh/bin:$PATH' >> ~/.bashrc; fi",
-      "source ~/.bashrc",
-      "echo 'END INITIALIZATION'",
+      "source ~/.bashrc"
+    ];
+  }
+
+  public getConnectionCommands(config: any): string[] {
+    return [
       `mongosh --host ${config.host} --port 10255 --username ${config.name} --password ${config.password} --tls --tlsAllowInvalidCertificates`
     ];
   }
+
+  public getTerminalSuppressedData(): string {
+    return "Non-Generic MongoDB Shell";
+  }
 }

src/Explorer/Tabs/CloudShellTab/ShellTypes/PostgresShellHandler.tsx

@@ -4,50 +4,41 @@
  */
 
 import { userContext } from "../../../../UserContext";
-import { listKeys } from "../../../../Utils/arm/generatedClients/cosmos/databaseAccounts";
-import { getHostFromUrl } from "../Utils/CommonUtils";
-import { ShellTypeConfig } from "./ShellTypeFactory";
+import { AbstractShellHandler } from "./AbstractShellHandler";
 
-export class PostgresShellHandler implements ShellTypeConfig {
+const PACKAGE_VERSION: string = "15.2";
+
+export class PostgresShellHandler extends AbstractShellHandler {
 
   public getShellName(): string {
     return "PostgreSQL";
   }
 
-  public async getInitialCommands(): Promise<string> {
-    const dbAccount = userContext.databaseAccount;
-    const endpoint = dbAccount.properties.postgresqlEndpoint;
-    
-    // Get database key
-    const dbName = dbAccount.name;
-    let key = "";
-    if (dbName) {
-      const keys = await listKeys(userContext.subscriptionId, userContext.resourceGroup, dbName);
-      key = keys?.primaryMasterKey || "";
+  public getEndpoint(): string {
+    return userContext.databaseAccount?.properties?.postgresqlEndpoint;
   }
 
-    const config = {
-      host: getHostFromUrl(endpoint),
-      name: dbAccount.name,
-      password: key,
-      endpoint: endpoint
-    };
-
-    return this.getCommands(config).join("\n").concat("\n");
-  }
-
-  private getCommands(config: any): string[] {
+  public getSetUpCommands(): string[] {
     return [
       "if ! command -v psql &> /dev/null; then echo '⚠️ psql not found. Installing...'; fi",
-      "if ! command -v psql &> /dev/null; then curl -LO https://ftp.postgresql.org/pub/source/v15.2/postgresql-15.2.tar.bz2; fi",
-      "if ! command -v psql &> /dev/null; then tar -xvjf postgresql-15.2.tar.bz2; fi",
+      `if ! command -v psql &> /dev/null; then curl -LO https://ftp.postgresql.org/pub/source/v${PACKAGE_VERSION}/postgresql-${PACKAGE_VERSION}.tar.bz2; fi`,
+      `if ! command -v psql &> /dev/null; then tar -xvjf postgresql-${PACKAGE_VERSION}.tar.bz2; fi`,
       "if ! command -v psql &> /dev/null; then mkdir -p ~/pgsql; fi",
       "if ! command -v psql &> /dev/null; then curl -LO https://ftp.gnu.org/gnu/readline/readline-8.1.tar.gz; fi",
       "if ! command -v psql &> /dev/null; then tar -xvzf readline-8.1.tar.gz; fi",
       "if ! command -v psql &> /dev/null; then cd readline-8.1 && ./configure --prefix=$HOME/pgsql; fi",
       "if ! command -v psql &> /dev/null; then echo 'export PATH=$HOME/pgsql/bin:$PATH' >> ~/.bashrc; fi",
-      "source ~/.bashrc",
+      "source ~/.bashrc"
+    ];
+  }
+
+  public getConnectionCommands(config: any): string[] {
+    return [
       `psql 'read -p "Enter Database Name: " dbname && read -p "Enter Username: " username && host=${config.endpoint} port=5432 dbname=$dbname user=$username sslmode=require'`
     ];
   }
+
+  public getTerminalSuppressedData(): string {
+    return "";
+  }
 }

src/Explorer/Tabs/CloudShellTab/ShellTypes/ShellTypeFactory.tsx

@@ -8,17 +8,13 @@ import { CassandraShellHandler } from "./CassandraShellHandler";
 import { MongoShellHandler } from "./MongoShellHandler";
 import { PostgresShellHandler } from "./PostgresShellHandler";
 import { VCoreMongoShellHandler } from "./VCoreMongoShellHandler";
+import { AbstractShellHandler } from "./AbstractShellHandler";
 
-export interface ShellTypeConfig {
-  getShellName(): string;
-  getInitialCommands(): Promise<string>;
-}
-
-export class ShellTypeHandler {
+export class ShellTypeHandlerFactory {
   /**
    * Gets the appropriate handler for the given shell type
    */
-  public static getHandler(shellType: TerminalKind): ShellTypeConfig {
+  public static getHandler(shellType: TerminalKind): AbstractShellHandler {
     switch (shellType) {
       case TerminalKind.Postgres:
         return new PostgresShellHandler();
@@ -32,21 +28,4 @@ export class ShellTypeHandler {
         throw new Error(`Unsupported shell type: ${shellType}`);
     }
   }
-
-  /**
-   * Gets the display name for a shell type
-   */
-  public static getShellNameForDisplay(terminalKind: TerminalKind): string {
-    switch (terminalKind) {
-      case TerminalKind.Postgres:
-        return "PostgreSQL";
-      case TerminalKind.Mongo:
-      case TerminalKind.VCoreMongo:
-        return "MongoDB";
-      case TerminalKind.Cassandra:
-        return "Cassandra";
-      default:
-        return "";
-    }
-  }
 }

src/Explorer/Tabs/CloudShellTab/ShellTypes/VCoreMongoShellHandler.tsx

@@ -4,47 +4,38 @@
  */
 
 import { userContext } from "../../../../UserContext";
-import { listKeys } from "../../../../Utils/arm/generatedClients/cosmos/databaseAccounts";
-import { getHostFromUrl } from "../Utils/CommonUtils";
-import { ShellTypeConfig } from "./ShellTypeFactory";
+import { AbstractShellHandler } from "./AbstractShellHandler";
 
-export class VCoreMongoShellHandler implements ShellTypeConfig {
+const PACKAGE_VERSION: string = "2.3.8";
+
+export class VCoreMongoShellHandler extends AbstractShellHandler {
 
   public getShellName(): string {
     return "MongoDB VCore";
   }
 
-  public async getInitialCommands(): Promise<string> {
-    const dbAccount = userContext.databaseAccount;
-    const endpoint = dbAccount.properties.vcoreMongoEndpoint;
-    
-    // Get database key
-    const dbName = dbAccount.name;
-    let key = "";
-    if (dbName) {
-      const keys = await listKeys(userContext.subscriptionId, userContext.resourceGroup, dbName);
-      key = keys?.primaryMasterKey || "";
+  public getEndpoint(): string {
+    return userContext.databaseAccount?.properties?.vcoreMongoEndpoint;
   }
 
-    const config = {
-      host: getHostFromUrl(endpoint),
-      name: dbAccount.name,
-      password: key,
-      endpoint: endpoint
-    };
-
-    return this.getCommands(config).join("\n").concat("\n");
-  }
-
-  private getCommands(config: any): string[] {
+  public getSetUpCommands(): string[] {
     return [
       "if ! command -v mongosh &> /dev/null; then echo '⚠️ mongosh not found. Installing...'; fi",
-      "if ! command -v mongosh &> /dev/null; then curl -LO https://downloads.mongodb.com/compass/mongosh-2.3.8-linux-x64.tgz; fi",
-      "if ! command -v mongosh &> /dev/null; then tar -xvzf mongosh-2.3.8-linux-x64.tgz; fi",
-      "if ! command -v mongosh &> /dev/null; then mkdir -p ~/mongosh && mv mongosh-2.3.8-linux-x64/* ~/mongosh/; fi",
+      `if ! command -v mongosh &> /dev/null; then curl -LO https://downloads.mongodb.com/compass/mongosh-${PACKAGE_VERSION}-linux-x64.tgz; fi`,
+      `if ! command -v mongosh &> /dev/null; then tar -xvzf mongosh-${PACKAGE_VERSION}-linux-x64.tgz; fi`,
+      `if ! command -v mongosh &> /dev/null; then mkdir -p ~/mongosh && mv mongosh-${PACKAGE_VERSION}-linux-x64/* ~/mongosh/; fi`,
       "if ! command -v mongosh &> /dev/null; then echo 'export PATH=$HOME/mongosh/bin:$PATH' >> ~/.bashrc; fi",
-      "source ~/.bashrc",
+      "source ~/.bashrc"
+    ];
+  }
+
+  public getConnectionCommands(config: any): string[] {
+    return [
       `read -p "Enter username: " username && mongosh "mongodb+srv://$username:@${config.endpoint}/?authMechanism=SCRAM-SHA-256&retrywrites=false&maxIdleTimeMS=120000" --tls --tlsAllowInvalidCertificates`
     ];
   }
+
+  public getTerminalSuppressedData(): string {
+    return "Non-Generic MongoDB Shell";
+  }
 }

src/Explorer/Tabs/CloudShellTab/Utils/AttachAddOn.tsx

@@ -3,6 +3,9 @@ import { IDisposable, ITerminalAddon, Terminal } from 'xterm';
 
 interface IAttachOptions {
     bidirectional?: boolean;
+    startMarker?: string;
+    endMarker?: string;
+    terminalSuppressedData?: string;
 }
 
 export class AttachAddon implements ITerminalAddon {
@@ -11,11 +14,20 @@ export class AttachAddon implements ITerminalAddon {
     private _disposables: IDisposable[] = [];
     private _socketData: string;
 
+    private _flag: boolean = true;
+
+    private _startMarker: string;
+    private _endMarker: string;
+    private _terminalSuppressedData: string;
+
     constructor(socket: WebSocket, options?: IAttachOptions) {
         this._socket = socket;
         // always set binary type to arraybuffer, we do not handle blobs
         this._socket.binaryType = 'arraybuffer';
         this._bidirectional = !(options && options.bidirectional === false);
+        this._startMarker = options?.startMarker;
+        this._endMarker = options?.endMarker;
+        this._terminalSuppressedData = options?.terminalSuppressedData;
         this._socketData = '';
     }
 
@@ -67,7 +79,18 @@ export class AttachAddon implements ITerminalAddon {
                     this._socketData += data;
                     data = '';
                 }
+
+                if (data.includes(this._startMarker)) {
+                    this._flag = false;
+                    terminal.write("Preparing environment...\r\n");
+                }
+                if (this._flag && this._terminalSuppressedData && this._terminalSuppressedData.length > 0 && !data.includes(this._terminalSuppressedData)) {
                     terminal.write(data);
+                }
+
+                if (data.includes(this._endMarker)) {
+                    this._flag = true;
+                }
             })
         );