feat: Add Output command for reading missed outputs

Author: furan917

cmd/outputs.go

@@ -0,0 +1,32 @@
+package cmd
+
+import (
+	"fmt"
+	"magecomm/messages/reader"
+
+	"github.com/spf13/cobra"
+)
+
+var OutputsCmd = &cobra.Command{
+	Use:   "outputs",
+	Short: "Drain all pending command outputs from the queue",
+	RunE: func(cmd *cobra.Command, args []string) error {
+		readerInstance, err := reader.MapReaderToEngine()
+		if err != nil {
+			return err
+		}
+
+		count, err := readerInstance.DrainOutputQueue("magerun")
+		if err != nil {
+			return err
+		}
+
+		if count == 0 {
+			fmt.Println("No outputs available")
+		} else {
+			fmt.Println("\nOutputs finished")
+		}
+
+		return nil
+	},
+}

main.go

@@ -49,6 +49,7 @@ func main() {
 	RootCmd.AddCommand(cmd.ListenCmd)
 	RootCmd.AddCommand(cmd.MagerunCmd)
 	RootCmd.AddCommand(cmd.CatCmd)
+	RootCmd.AddCommand(cmd.OutputsCmd)
 
 	RootCmd.PersistentFlags().String("config", "", "Path to config file")
 	RootCmd.PersistentFlags().Bool("debug", false, "Enable debug mode")

messages/handler/magerun.go

@@ -52,13 +52,14 @@ func (handler *MagerunHandler) ProcessMessage(messageBody string, correlationID
 		}
 	}
 
-	// Publish the output to the RMQ/SQS output queue
 	publisherClass, err := publisher.MapPublisherToEngine()
 	if err != nil {
 		logger.Warnf("Error publishing message to RMQ/SQS queue: %s", err)
 	}
+
+	outputWithCommand := fmt.Sprintf("Command: %s\n\n%s", messageBody, output)
 	logger.Infof("Publishing output to queue: %s with correlation ID: %s", queues.MapQueueToOutputQueue(MageRunQueue), correlationID)
-	_, err = publisherClass.Publish(output, queues.MapQueueToOutputQueue(MageRunQueue), correlationID)
+	_, err = publisherClass.Publish(outputWithCommand, queues.MapQueueToOutputQueue(MageRunQueue), correlationID)
 	if err != nil {
 		return err
 	}

messages/reader/interface.go

@@ -0,0 +1,5 @@
+package reader
+
+type MessageReader interface {
+	DrainOutputQueue(queueName string) (int, error)
+}

messages/reader/mapper.go

@@ -0,0 +1,22 @@
+package reader
+
+import (
+	"fmt"
+	"magecomm/config_manager"
+	"magecomm/services"
+)
+
+var Reader MessageReader
+
+func MapReaderToEngine() (MessageReader, error) {
+	engine := config_manager.GetEngine()
+
+	switch engine {
+	case services.EngineSQS:
+		return &SqsReader{}, nil
+	case services.EngineRabbitMQ:
+		return &RmqReader{}, nil
+	default:
+		return nil, fmt.Errorf("unknown engine: %s", engine)
+	}
+}

messages/reader/rmq.go

@@ -0,0 +1,58 @@
+package reader
+
+import (
+    "fmt"
+    "strings"
+
+    "magecomm/logger"
+    "magecomm/messages/queues"
+    "magecomm/services"
+
+    "github.com/streadway/amqp"
+)
+
+type RmqReader struct{}
+
+func (r *RmqReader) DrainOutputQueue(queueName string) (int, error) {
+    outputQueueName := queues.MapQueueToOutputQueue(queueName)
+
+    channel, err := services.RmqChannelPool.Get()
+    if err != nil {
+        return 0, fmt.Errorf("error getting RMQ channel: %w", err)
+    }
+    defer services.RmqChannelPool.Put(channel)
+
+    queueNameWithPrefix, err := services.CreateRmqQueue(channel, outputQueueName)
+    if err != nil {
+        return 0, fmt.Errorf("error creating queue: %w", err)
+    }
+
+    count := 0
+    for {
+        msg, ok, err := channel.Get(queueNameWithPrefix, false)
+        if err != nil {
+            return count, fmt.Errorf("error receiving message: %w", err)
+        }
+
+        if !ok {
+            break
+        }
+
+        count++
+        r.displayMessage(count, &msg)
+
+        if err := msg.Ack(false); err != nil {
+            logger.Warnf("Failed to ack message: %v", err)
+        }
+    }
+
+    return count, nil
+}
+
+func (r *RmqReader) displayMessage(index int, msg *amqp.Delivery) {
+    fmt.Printf("\n%s\n", strings.Repeat("=", 60))
+    fmt.Printf("Output %d\n", index)
+    fmt.Printf("Received: %s\n", msg.Timestamp.Format("2006-01-02 15:04:05"))
+    fmt.Printf("%s\n", strings.Repeat("=", 60))
+    fmt.Println(string(msg.Body))
+}

messages/reader/sqs.go

@@ -0,0 +1,87 @@
+package reader
+
+import (
+	"fmt"
+	"strconv"
+	"strings"
+	"time"
+
+	"magecomm/logger"
+	"magecomm/messages/queues"
+	"magecomm/services"
+
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/service/sqs"
+)
+
+type SqsReader struct{}
+
+func (r *SqsReader) DrainOutputQueue(queueName string) (int, error) {
+	outputQueueName := queues.MapQueueToOutputQueue(queueName)
+
+	sqsConnection, err := services.GetSQSConnection()
+	if err != nil {
+		return 0, fmt.Errorf("unable to get SQS connection: %w", err)
+	}
+	defer services.ReleaseSQSConnection(sqsConnection)
+
+	if err := sqsConnection.Connect(); err != nil {
+		return 0, fmt.Errorf("error connecting to SQS: %w", err)
+	}
+
+	sqsClient := sqsConnection.Client
+	queueURL, err := services.CreateSQSQueueIfNotExists(sqsClient, outputQueueName)
+	if err != nil {
+		return 0, fmt.Errorf("error getting queue URL: %w", err)
+	}
+
+	count := 0
+	for {
+		result, err := sqsClient.ReceiveMessage(&sqs.ReceiveMessageInput{
+			QueueUrl:              aws.String(queueURL),
+			MaxNumberOfMessages:   aws.Int64(10),
+			WaitTimeSeconds:       aws.Int64(1),
+			AttributeNames:        []*string{aws.String("SentTimestamp")},
+			MessageAttributeNames: []*string{aws.String("All")},
+		})
+
+		if err != nil {
+			return count, fmt.Errorf("error receiving messages: %w", err)
+		}
+
+		if len(result.Messages) == 0 {
+			break
+		}
+
+		for _, msg := range result.Messages {
+			count++
+			r.displayMessage(count, msg)
+
+			_, err := sqsClient.DeleteMessage(&sqs.DeleteMessageInput{
+				QueueUrl:      aws.String(queueURL),
+				ReceiptHandle: msg.ReceiptHandle,
+			})
+			if err != nil {
+				logger.Warnf("Failed to delete message: %v", err)
+			}
+		}
+	}
+
+	return count, nil
+}
+
+func (r *SqsReader) displayMessage(index int, msg *sqs.Message) {
+	fmt.Printf("\n%s\n", strings.Repeat("=", 60))
+	fmt.Printf("Output %d\n", index)
+
+	if msg.Attributes["SentTimestamp"] != nil {
+		timestamp := *msg.Attributes["SentTimestamp"]
+		if ms, err := strconv.ParseInt(timestamp, 10, 64); err == nil {
+			t := time.Unix(ms/1000, 0)
+			fmt.Printf("Received: %s\n", t.Format("2006-01-02 15:04:05"))
+		}
+	}
+
+	fmt.Printf("%s\n", strings.Repeat("=", 60))
+	fmt.Println(*msg.Body)
+}

services/rabbitmq_pool.go

@@ -2,9 +2,13 @@ package services
 
 import (
 	"errors"
-	"github.com/streadway/amqp"
+	"fmt"
+
 	"magecomm/config_manager"
+	"magecomm/logger"
 	"sync"
+
+	"github.com/streadway/amqp"
 )
 
 var (
@@ -20,7 +24,8 @@ type RabbitMQConnectionPool struct {
 }
 
 type RabbitMQChannelPool struct {
-	pool *sync.Pool
+	pool     *sync.Pool
+	connPool *RabbitMQConnectionPool
 }
 
 func Close() {
@@ -32,18 +37,19 @@ func Close() {
 func NewRabbitMQConnectionPool(initialSize int) *RabbitMQConnectionPool {
 	p := &sync.Pool{
 		New: func() interface{} {
-			rmqConn := NewRabbitMQConnection()
-			if err := rmqConn.Connect(""); err != nil {
-				return err
-			}
-			return rmqConn
+			return NewRabbitMQConnection()
 		},
 	}
 
 	cp := &RabbitMQConnectionPool{pool: p}
 
 	for i := 0; i < initialSize; i++ {
-		p.Put(p.New())
+		conn := NewRabbitMQConnection()
+		if err := conn.Connect(""); err != nil {
+			logger.Warnf("Failed to create initial RMQ connection: %v", err)
+			continue
+		}
+		p.Put(conn)
 	}
 
 	return cp
@@ -52,54 +58,92 @@ func NewRabbitMQConnectionPool(initialSize int) *RabbitMQConnectionPool {
 func NewRabbitMQChannelPool(connPool *RabbitMQConnectionPool, initialSize int) *RabbitMQChannelPool {
 	p := &sync.Pool{
 		New: func() interface{} {
-			conn, err := connPool.Get()
-			if err != nil {
-				return err
-			}
-			channel, err := conn.CreateChannel()
-			if err != nil {
-				return err
-			}
-			return channel
+			return nil
 		},
 	}
 
-	cp := &RabbitMQChannelPool{pool: p}
+	cp := &RabbitMQChannelPool{
+		pool:     p,
+		connPool: connPool,
+	}
 
 	for i := 0; i < initialSize; i++ {
-		p.Put(p.New())
+		conn, err := connPool.Get()
+		if err != nil {
+			logger.Warnf("Failed to get connection for initial channel: %v", err)
+			continue
+		}
+		channel, err := conn.CreateChannel()
+		if err != nil {
+			logger.Warnf("Failed to create initial RMQ channel: %v", err)
+			connPool.Put(conn)
+			continue
+		}
+		connPool.Put(conn)
+		p.Put(channel)
 	}
 
 	return cp
 }
 
 func (cp *RabbitMQConnectionPool) Get() (*RabbitMQConnection, error) {
-	conn := cp.pool.Get()
-	if conn == nil {
+	obj := cp.pool.Get()
+	if obj == nil {
 		return nil, ErrConnectionPoolClosed
 	}
-	return conn.(*RabbitMQConnection), nil
+
+	conn, ok := obj.(*RabbitMQConnection)
+	if !ok || conn == nil {
+		conn = NewRabbitMQConnection()
+	}
+
+	if conn.Connection == nil || conn.Connection.IsClosed() {
+		if err := conn.Connect(""); err != nil {
+			return nil, fmt.Errorf("failed to connect to RabbitMQ: %w", err)
+		}
+	}
+
+	return conn, nil
 }
 
 func (cp *RabbitMQConnectionPool) Put(conn *RabbitMQConnection) {
-	cp.pool.Put(conn)
+	if conn != nil && conn.Connection != nil && !conn.Connection.IsClosed() {
+		cp.pool.Put(conn)
+	}
 }
 
 func (cp *RabbitMQConnectionPool) Close() {
 	cp.pool = nil
 }
 
 func (cp *RabbitMQChannelPool) Get() (*amqp.Channel, error) {
+	obj := cp.pool.Get()
 
-	channel := cp.pool.Get()
-	if channel == nil {
-		return nil, ErrChannelPoolClosed
+	if obj != nil {
+		if channel, ok := obj.(*amqp.Channel); ok && channel != nil {
+			return channel, nil
+		}
 	}
-	return channel.(*amqp.Channel), nil
+
+	conn, err := cp.connPool.Get()
+	if err != nil {
+		return nil, fmt.Errorf("failed to get connection: %w", err)
+	}
+
+	channel, err := conn.CreateChannel()
+	if err != nil {
+		cp.connPool.Put(conn)
+		return nil, fmt.Errorf("failed to create channel: %w", err)
+	}
+
+	cp.connPool.Put(conn)
+	return channel, nil
 }
 
 func (cp *RabbitMQChannelPool) Put(channel *amqp.Channel) {
-	cp.pool.Put(channel)
+	if channel != nil {
+		cp.pool.Put(channel)
+	}
 }
 
 func (cp *RabbitMQChannelPool) Close() {