ARTEMIS-5168 Improve remoting to brokers from Artemis shell

apache · Nov 22, 2024 · 78c816b · 78c816b
1 parent 1ba7bec
commit 78c816b
Show file tree

Hide file tree

Showing 5 changed files with 60 additions and 3 deletions.
diff --git a/artemis-cli/src/main/java/org/apache/activemq/artemis/cli/Artemis.java b/artemis-cli/src/main/java/org/apache/activemq/artemis/cli/Artemis.java
@@ -275,7 +275,7 @@ public static CommandLine buildCommand(boolean includeInstanceCommands, boolean
       commandLine.addSubcommand(new QueueGroup(commandLine));
       commandLine.addSubcommand(new AddressGroup(commandLine));
 
-      if (shellEnabled) {
+      if (Shell.inShell()) {
          commandLine.addSubcommand(new Connect());
          commandLine.addSubcommand(new Disconnect());
       }

diff --git a/artemis-cli/src/main/java/org/apache/activemq/artemis/cli/Shell.java b/artemis-cli/src/main/java/org/apache/activemq/artemis/cli/Shell.java
@@ -67,6 +67,8 @@ public void run() {
    }
 
    private static ThreadLocal<AtomicBoolean> IN_SHELL = ThreadLocal.withInitial(() -> new AtomicBoolean(false));
+   private static ThreadLocal<AtomicBoolean> CONNECTED = ThreadLocal.withInitial(() -> new AtomicBoolean(false));
+   private static ThreadLocal<String> PROMPT = new ThreadLocal<>();
 
    public static boolean inShell() {
       return IN_SHELL.get().get();
@@ -76,6 +78,14 @@ public static void setInShell() {
       IN_SHELL.get().set(true);
    }
 
+   public static boolean isConnected() {
+      return CONNECTED.get().get();
+   }
+
+   public static void setConnected(boolean connected) {
+      CONNECTED.get().set(connected);
+   }
+
    public static void runShell(boolean printBanner) {
       try {
          setInShell();
@@ -104,7 +114,6 @@ public static void runShell(boolean printBanner) {
                .build();
             factory.setTerminal(terminal);
 
-            String prompt = org.apache.activemq.artemis.cli.Terminal.YELLOW_UNICODE + Artemis.getNameFromBanner() + " > " + org.apache.activemq.artemis.cli.Terminal.CLEAR_UNICODE;
             String rightPrompt = null;
 
             if (printBanner) {
@@ -121,14 +130,19 @@ public static void runShell(boolean printBanner) {
                   // We build a new command every time, as they could have state from previous executions
                   systemRegistry.setCommandRegistries(new PicocliCommands(Artemis.buildCommand(isInstance, !isInstance, false)));
                   systemRegistry.cleanUp();
-                  line = reader.readLine(prompt, rightPrompt, (MaskingCallback) null, null);
+                  line = reader.readLine(getPrompt(), rightPrompt, (MaskingCallback) null, null);
                   systemRegistry.execute(line);
                } catch (InterruptedException e) {
                   e.printStackTrace();
                   // Ignore
                } catch (UserInterruptException userInterruptException) {
                   // ignore
                } catch (EndOfFileException e) {
+                  if (isConnected()) {
+                     //if connected, [Ctrl + D] tries to disconnect instead of close
+                     systemRegistry.execute("disconnect");
+                     continue;
+                  }
                   return;
                } catch (Exception e) {
                   systemRegistry.trace(e);
@@ -154,4 +168,27 @@ private static void printBanner() {
       System.out.print(org.apache.activemq.artemis.cli.Terminal.CLEAR_UNICODE);
    }
 
+   private static String getPrompt() {
+      if (PROMPT.get() == null) {
+         setDefaultPrompt();
+      }
+
+      return PROMPT.get();
+   }
+
+   public static void setDefaultPrompt() {
+      try {
+         setPrompt(Artemis.getNameFromBanner());
+      } catch (Exception e) {
+         System.out.println("Error when getting prompt name from banner:");
+         e.printStackTrace();
+
+         setPrompt("Artemis Shell");
+      }
+   }
+
+   public static void setPrompt(String prompt) {
+      PROMPT.set(org.apache.activemq.artemis.cli.Terminal.YELLOW_UNICODE + prompt + " > " + org.apache.activemq.artemis.cli.Terminal.CLEAR_UNICODE);
+   }
+
 }
diff --git a/artemis-cli/src/main/java/org/apache/activemq/artemis/cli/commands/Connect.java b/artemis-cli/src/main/java/org/apache/activemq/artemis/cli/commands/Connect.java
@@ -17,6 +17,7 @@
 
 package org.apache.activemq.artemis.cli.commands;
 
+import org.apache.activemq.artemis.cli.Shell;
 import org.apache.activemq.artemis.cli.commands.messages.ConnectionAbstract;
 import picocli.CommandLine;
 
@@ -30,6 +31,11 @@ public Object execute(ActionContext context) throws Exception {
          CONNECTION_INFORMATION.remove();
          createConnectionFactory();
          context.out.println("Connection Successful!");
+
+         if (Shell.inShell()) {
+            Shell.setConnected(true);
+         }
+
       } catch (Exception e) {
          context.out.println("Connection Failure!");
          e.printStackTrace();

diff --git a/artemis-cli/src/main/java/org/apache/activemq/artemis/cli/commands/Disconnect.java b/artemis-cli/src/main/java/org/apache/activemq/artemis/cli/commands/Disconnect.java
@@ -17,6 +17,7 @@
 
 package org.apache.activemq.artemis.cli.commands;
 
+import org.apache.activemq.artemis.cli.Shell;
 import org.apache.activemq.artemis.cli.commands.messages.ConnectionAbstract;
 import picocli.CommandLine;
 
@@ -28,6 +29,12 @@ public Object execute(ActionContext context) throws Exception {
       super.execute(context);
       CONNECTION_INFORMATION.remove();
       context.out.println("Connection information cleared!");
+
+      if (Shell.inShell()) {
+         Shell.setDefaultPrompt();
+         Shell.setConnected(false);
+      }
+
       return null;
    }
 }
diff --git a/...ava/org/apache/activemq/artemis/cli/commands/messages/ConnectionConfigurationAbtract.java b/...ava/org/apache/activemq/artemis/cli/commands/messages/ConnectionConfigurationAbtract.java
@@ -137,6 +137,13 @@ protected void saveConnectionInfo(String brokerURL, String user, String password
          this.brokerURL = brokerURL;
          this.user = user;
          this.password = password;
+
+         if (user != null) {
+            Shell.setPrompt(user + "@" + brokerURL);
+         } else {
+            Shell.setPrompt(brokerURL);
+         }
+
       }
    }