minor changes

src/Explorer/Tabs/CloudShellTab/CloudShellTerminalCore.tsx

@@ -14,7 +14,7 @@ import {
   registerCloudShellProvider,
   verifyCloudShellProviderRegistration
 } from "./Data/CloudShellClient";
-import { END_MARKER, START_MARKER } from "./ShellTypes/AbstractShellHandler";
+import { START_MARKER, AbstractShellHandler } from "./ShellTypes/AbstractShellHandler";
 import { ShellTypeHandlerFactory } from "./ShellTypes/ShellTypeFactory";
 import { AttachAddon } from "./Utils/AttachAddOn";
 import { askConfirmation, wait } from "./Utils/CommonUtils";
@@ -62,7 +62,7 @@ export const startCloudShellTerminal =
   }
 
   // Get the shell handler for this type
-  const shellHandler = ShellTypeHandlerFactory.getHandler(shellType);
+  const shellHandler = await ShellTypeHandlerFactory.getHandler(shellType);
   // Configure WebSocket connection with shell-specific commands
   const socket = await establishTerminalConnection(
     terminal,
@@ -162,7 +162,7 @@ export const provisionCloudShellSession = async (
  */
 export const establishTerminalConnection = async (
   terminal: Terminal,
-  shellHandler: any,
+  shellHandler: AbstractShellHandler,
   socketUri: string,
   provisionConsoleResponse: any,
   targetUri: string
@@ -171,15 +171,14 @@ export const establishTerminalConnection = async (
   let socket = new WebSocket(socketUri);
 
   // Get shell-specific initial commands 
-  const initCommands = await shellHandler.getInitialCommands();
+  const initCommands = shellHandler.getInitialCommands();
 
   // Configure the socket
   socket = await configureSocketConnection(socket, socketUri, terminal, initCommands, 0);
 
   const options = {
     startMarker: START_MARKER,
-    endMarker: END_MARKER,
-    terminalSuppressedData: shellHandler.getTerminalSuppressedData()
+    shellHandler: shellHandler
   };
 
   // Attach the terminal addon

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

@@ -1,43 +1,21 @@
-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 getConnectionCommand(): 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 || "";
-    }
-
+  public getInitialCommands(): string {
     const setupCommands = this.getSetUpCommands();
-
-    const config = {
-      host: getHostFromUrl(this.getEndpoint()),
-      name: dbName,
-      password: key,
-      endpoint: this.getEndpoint(),
-    };
-    const connectionCommands = this.getConnectionCommands(config);
+    const connectionCommand = this.getConnectionCommand();
 
     const allCommands = [
       START_MARKER,
       ...setupCommands,
-      END_MARKER,
-      ...connectionCommands
+      ...connectionCommand
     ];
   
     return allCommands.join("\n").concat("\n");

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

@@ -4,11 +4,18 @@
  */
 
 import { userContext } from "../../../../UserContext";
+import { getHostFromUrl } from "../Utils/CommonUtils";
 import { AbstractShellHandler } from "./AbstractShellHandler";
 
 const PACKAGE_VERSION: string = "5.0.3";
 
 export class CassandraShellHandler extends AbstractShellHandler {
+  
+  private _key: string;
+  constructor(private key: string) {
+    super(); 
+    this.key = key;
+  }
 
   public getShellName(): string {
     return "Cassandra";
@@ -31,10 +38,8 @@ export class CassandraShellHandler extends AbstractShellHandler {
     ];
   }
 
-  public getConnectionCommands(config: any): string[] {
-    return [
-      `cqlsh ${config.host} 10350 -u ${config.name} -p ${config.password} --ssl --protocol-version=4`
-    ];
+  public getConnectionCommand(): string {
+    return `cqlsh ${getHostFromUrl(this.getEndpoint())} 10350 -u ${userContext.databaseAccount?.name} -p ${this._key} --ssl --protocol-version=4`;
   }
 
   public getTerminalSuppressedData(): string {

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

@@ -5,11 +5,18 @@
 
 import { userContext } from "../../../../UserContext";
 import { AbstractShellHandler } from "./AbstractShellHandler";
+import { getHostFromUrl } from "../Utils/CommonUtils";
 
 const PACKAGE_VERSION: string = "2.3.8";
 
 export class MongoShellHandler extends AbstractShellHandler {
 
+  private _key: string;
+  constructor(private key: string) {
+    super(); 
+    this.key = key;
+  }
+
   public getShellName(): string {
     return "MongoDB";
   }
@@ -29,10 +36,8 @@ export class MongoShellHandler extends AbstractShellHandler {
     ];
   }
 
-  public getConnectionCommands(config: any): string[] {
-    return [
-      `mongosh --host ${config.host} --port 10255 --username ${config.name} --password ${config.password} --tls --tlsAllowInvalidCertificates`
-    ];
+  public getConnectionCommand(): string {
+    return `mongosh --host ${getHostFromUrl(this.getEndpoint())} --port 10255 --username ${userContext.databaseAccount?.name} --password ${this._key} --tls --tlsAllowInvalidCertificates`;
   }
 
   public getTerminalSuppressedData(): string {

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

@@ -32,10 +32,8 @@ export class PostgresShellHandler extends AbstractShellHandler {
     ];
   }
 
-  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 getConnectionCommand(): string {
+    return `psql 'read -p "Enter Database Name: " dbname && read -p "Enter Username: " username && host=${this.getEndpoint()} port=5432 dbname=$dbname user=$username sslmode=require'`;
   }
 
   public getTerminalSuppressedData(): string {

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

@@ -9,23 +9,37 @@ import { MongoShellHandler } from "./MongoShellHandler";
 import { PostgresShellHandler } from "./PostgresShellHandler";
 import { VCoreMongoShellHandler } from "./VCoreMongoShellHandler";
 import { AbstractShellHandler } from "./AbstractShellHandler";
+import { userContext } from "../../../../UserContext";
+import { listKeys } from "../../../../Utils/arm/generatedClients/cosmos/databaseAccounts";
 
 export class ShellTypeHandlerFactory {
   /**
    * Gets the appropriate handler for the given shell type
    */
-  public static getHandler(shellType: TerminalKind): AbstractShellHandler {
+  public static async getHandler(shellType: TerminalKind): Promise<AbstractShellHandler> {
     switch (shellType) {
       case TerminalKind.Postgres:
         return new PostgresShellHandler();
       case TerminalKind.Mongo:
-        return new MongoShellHandler();
+        return new MongoShellHandler(await ShellTypeHandlerFactory.getKey());
       case TerminalKind.VCoreMongo:
         return new VCoreMongoShellHandler();
       case TerminalKind.Cassandra:
-        return new CassandraShellHandler();
+        return new CassandraShellHandler(await ShellTypeHandlerFactory.getKey());
       default:
         throw new Error(`Unsupported shell type: ${shellType}`);
     }
   }
+
+  public static async getKey(): Promise<string> {
+    const dbName = userContext.databaseAccount.name;
+    let key = "";
+    if (dbName) {
+      const keys = await listKeys(userContext.subscriptionId, userContext.resourceGroup, dbName);
+      key = keys?.primaryMasterKey || "";
+    }
+
+    return key;
+  }
+  
 }

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

@@ -29,10 +29,8 @@ export class VCoreMongoShellHandler extends AbstractShellHandler {
     ];
   }
 
-  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 getConnectionCommand(): string {
+    return `read -p "Enter username: " username && mongosh "mongodb+srv://$username:@${this.getEndpoint()}/?authMechanism=SCRAM-SHA-256&retrywrites=false&maxIdleTimeMS=120000" --tls --tlsAllowInvalidCertificates`;
   }
 
   public getTerminalSuppressedData(): string {

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

@@ -1,11 +1,11 @@
 
+import { AbstractShellHandler } from 'Explorer/Tabs/CloudShellTab/ShellTypes/AbstractShellHandler';
 import { IDisposable, ITerminalAddon, Terminal } from 'xterm';
 
 interface IAttachOptions {
     bidirectional?: boolean;
     startMarker?: string;
-    endMarker?: string;
-    terminalSuppressedData?: string;
+    shellHandler?: AbstractShellHandler;
 }
 
 export class AttachAddon implements ITerminalAddon {
@@ -14,11 +14,10 @@ export class AttachAddon implements ITerminalAddon {
     private _disposables: IDisposable[] = [];
     private _socketData: string;
 
-    private _flag: boolean = true;
+    private _allowTerminalWrite: boolean = true;
 
     private _startMarker: string;
-    private _endMarker: string;
-    private _terminalSuppressedData: string;
+    private _shellHandler: AbstractShellHandler;
 
     constructor(socket: WebSocket, options?: IAttachOptions) {
         this._socket = socket;
@@ -26,10 +25,9 @@ export class AttachAddon implements ITerminalAddon {
         this._socket.binaryType = 'arraybuffer';
         this._bidirectional = !(options && options.bidirectional === false);
         this._startMarker = options?.startMarker;
-        this._endMarker = options?.endMarker;
-        this._terminalSuppressedData = options?.terminalSuppressedData;
+        this._shellHandler = options?.shellHandler;
         this._socketData = '';
-        this._flag = true;
+        this._allowTerminalWrite = true;
     }
 
     public activate(terminal: Terminal): void {
@@ -82,15 +80,21 @@ export class AttachAddon implements ITerminalAddon {
                 }
 
                 if (data.includes(this._startMarker)) {
-                    this._flag = false;
-                    terminal.write("Preparing environment...\r\n");
+                    this._allowTerminalWrite = false;
+                    terminal.write(`Preparing ${this._shellHandler.getShellName()} environment...\r\n`);
                 }
-                if (this._flag && this._terminalSuppressedData && this._terminalSuppressedData.length > 0 && !data.includes(this._terminalSuppressedData)) {
+                if (this._allowTerminalWrite && 
+                    this._shellHandler?.getTerminalSuppressedData() && 
+                    this._shellHandler?.getTerminalSuppressedData().length > 0 && 
+                    !data.includes(this._shellHandler?.getTerminalSuppressedData())) {
                     terminal.write(data);
                 }
 
-                if (data.includes(this._endMarker)) {
-                    this._flag = true;
+                console.log('data', data);
+                console.log("connection command", this._shellHandler.getConnectionCommand());
+                console.log("----")
+                if (data.includes(this._shellHandler.getConnectionCommand())) {
+                    this._allowTerminalWrite = true;
                 }
             })
         );