kmsvc-cli/internal/cli/dlq.go

package cli

import (
	"fmt"

	kmsvc "forgejo.riotpiao.com/rock/kmsvc-sdk"
	"github.com/spf13/cobra"
)

func newDLQCmd(flags *globalFlags) *cobra.Command {
	cmd := &cobra.Command{
		Use:   "dlq",
		Short: "Inspect and redrive dead-letter queues",
	}
	cmd.AddCommand(newDLQPeekCmd(flags), newDLQRedriveCmd(flags))
	return cmd
}

func newDLQPeekCmd(flags *globalFlags) *cobra.Command {
	var queue string
	var maxMessages, visibilityTimeout int32

	cmd := &cobra.Command{
		Use:   "peek",
		Short: "Receive messages from a DLQ without deleting them",
		RunE: func(cmd *cobra.Command, args []string) error {
			client, err := buildClient(cmd.Context(), flags)
			if err != nil {
				return err
			}
			defer client.Close()

			msgs, err := client.ReceiveMessage(cmd.Context(), queue, kmsvc.ReceiveOptions{
				MaxNumberOfMessages:      maxMessages,
				VisibilityTimeoutSeconds: visibilityTimeout,
			})
			if err != nil {
				return err
			}

			return renderMessages(cmd.OutOrStdout(), flags.output, msgs)
		},
	}

	cmd.Flags().StringVar(&queue, "queue", "", "DLQ name (required)")
	cmd.Flags().Int32Var(&maxMessages, "max-messages", 10, "maximum number of messages to peek (1-10)")
	cmd.Flags().Int32Var(&visibilityTimeout, "visibility-timeout", 5, "visibility timeout for the peek, in seconds — keep short so messages reappear quickly")
	cmd.MarkFlagRequired("queue")

	return cmd
}

func newDLQRedriveCmd(flags *globalFlags) *cobra.Command {
	var queue, to string
	var maxMessages int32

	cmd := &cobra.Command{
		Use:   "redrive",
		Short: "Move messages from a DLQ back to their source queue",
		Long: "Receives messages from --queue (the DLQ), sends each to --to (the source\n" +
			"queue), then deletes it from the DLQ. This is 3+ separate SDK calls, not an\n" +
			"atomic operation: if send succeeds but delete fails, the message is reported\n" +
			"as sent-but-not-removed (it may be redelivered from both queues), and the\n" +
			"command exits non-zero rather than silently continuing.",
		RunE: func(cmd *cobra.Command, args []string) error {
			client, err := buildClient(cmd.Context(), flags)
			if err != nil {
				return err
			}
			defer client.Close()

			msgs, err := client.ReceiveMessage(cmd.Context(), queue, kmsvc.ReceiveOptions{
				MaxNumberOfMessages: maxMessages,
			})
			if err != nil {
				return fmt.Errorf("receive from %s: %w", queue, err)
			}

			out := cmd.OutOrStdout()
			var failures int
			for _, m := range msgs {
				sendOut, err := client.SendMessage(cmd.Context(), kmsvc.SendMessageInput{
					QueueName: to,
					Body:      m.Body,
				})
				if err != nil {
					fmt.Fprintf(out, "redrive %s: send to %s failed, message left in DLQ: %v\n", m.MessageID, to, err)
					failures++
					continue
				}

				if err := client.DeleteMessage(cmd.Context(), queue, m.ReceiptHandle); err != nil {
					fmt.Fprintf(out, "redrive %s: sent to %s as %s, but delete from %s failed — message may be duplicated: %v\n", m.MessageID, to, sendOut.MessageID, queue, err)
					failures++
					continue
				}

				fmt.Fprintf(out, "redrive %s: sent to %s as %s, removed from %s\n", m.MessageID, to, sendOut.MessageID, queue)
			}

			if failures > 0 {
				return fmt.Errorf("%d/%d message(s) failed to fully redrive", failures, len(msgs))
			}
			return nil
		},
	}

	cmd.Flags().StringVar(&queue, "queue", "", "DLQ name (required)")
	cmd.Flags().StringVar(&to, "to", "", "source queue to redrive messages back to (required)")
	cmd.Flags().Int32Var(&maxMessages, "max-messages", 10, "maximum number of messages to redrive in this run (1-10)")
	cmd.MarkFlagRequired("queue")
	cmd.MarkFlagRequired("to")

	return cmd
}