Pass sourceQueueName via System Message Attribute for DeadLetterQueueSourceArn

Sergii Kozlov · Sergii Kozlov · commit c2aa609ce618 · 2025-01-03T10:25:34.000-08:00
diff --git a/core/src/main/scala/org/elasticmq/MessageData.scala b/core/src/main/scala/org/elasticmq/MessageData.scala
@@ -7,6 +7,7 @@ case class MessageData(
     deliveryReceipt: Option[DeliveryReceipt],
     content: String,
     messageAttributes: Map[String, MessageAttribute],
+    messageSystemAttributes: Map[String, MessageAttribute],
     nextDelivery: MillisNextDelivery,
     created: OffsetDateTime,
     statistics: MessageStatistics,
diff --git a/core/src/main/scala/org/elasticmq/NewMessageData.scala b/core/src/main/scala/org/elasticmq/NewMessageData.scala
@@ -1,5 +1,7 @@
 package org.elasticmq
 
+import scala.collection.mutable
+
 case class NewMessageData(
     id: Option[MessageId],
     content: String,
diff --git a/core/src/main/scala/org/elasticmq/actor/queue/InternalMessage.scala b/core/src/main/scala/org/elasticmq/actor/queue/InternalMessage.scala
@@ -13,7 +13,7 @@ case class InternalMessage(
     var nextDelivery: Long,
     content: String,
     messageAttributes: Map[String, MessageAttribute],
-    messageSystemAttributes: Map[String, MessageAttribute],
+    messageSystemAttributes: mutable.HashMap[String, MessageAttribute],
     created: OffsetDateTime,
     orderIndex: Int,
     var firstReceive: Received,
@@ -61,6 +61,7 @@ case class InternalMessage(
       deliveryReceipts.lastOption.map(DeliveryReceipt(_)),
       content,
       messageAttributes,
+      messageSystemAttributes.to(Map),
       MillisNextDelivery(nextDelivery),
       created,
       MessageStatistics(firstReceive, receiveCount),
@@ -75,7 +76,7 @@ case class InternalMessage(
       Some(MessageId(id)),
       content,
       messageAttributes,
-      messageSystemAttributes,
+      messageSystemAttributes.to(Map),
       MillisNextDelivery(nextDelivery),
       messageGroupId,
       messageDeduplicationId,
@@ -97,7 +98,7 @@ object InternalMessage {
       newMessageData.nextDelivery.toMillis(now, queueData.delay.toMillis).millis,
       newMessageData.content,
       newMessageData.messageAttributes,
-      newMessageData.messageSystemAttributes,
+      newMessageData.messageSystemAttributes.to(mutable.HashMap),
       OffsetDateTime.now(),
       newMessageData.orderIndex,
       NeverReceived,
diff --git a/core/src/main/scala/org/elasticmq/actor/queue/QueueActorMessageOps.scala b/core/src/main/scala/org/elasticmq/actor/queue/QueueActorMessageOps.scala
@@ -34,7 +34,7 @@ trait QueueActorMessageOps
       case DeleteMessage(deliveryReceipt) =>
         deleteMessage(deliveryReceipt).send()
       case LookupMessage(messageId)          => messageQueue.getById(messageId.id).map(_.toMessageData)
-      case MoveMessage(message, destination) => moveMessage(message, destination).send()
+      case MoveMessage(message, destination, sourceQueueName) => moveMessage(message, destination, sourceQueueName).send()
       case DeduplicationIdsCleanup =>
         fifoMessagesHistory = fifoMessagesHistory.cleanOutdatedMessages(nowProvider)
         DoNotReply()
diff --git a/core/src/main/scala/org/elasticmq/actor/queue/operations/MoveMessageOps.scala b/core/src/main/scala/org/elasticmq/actor/queue/operations/MoveMessageOps.scala
@@ -3,13 +3,14 @@ package org.elasticmq.actor.queue.operations
 import org.elasticmq.actor.queue.{InternalMessage, QueueActorStorage, QueueEvent}
 import org.elasticmq.msg.SendMessage
 import org.elasticmq.util.Logging
-import org.elasticmq.{DeduplicationId, MoveDestination, MoveToDLQ}
+import org.elasticmq.{DeduplicationId, MoveDestination, MoveToDLQ, StringMessageAttribute}
 
 trait MoveMessageOps extends Logging {
   this: QueueActorStorage =>
 
-  def moveMessage(message: InternalMessage, destination: MoveDestination): ResultWithEvents[Unit] = {
+  def moveMessage(message: InternalMessage, destination: MoveDestination, sourceQueueName: String): ResultWithEvents[Unit] = {
 
+    message.messageSystemAttributes.put("sourceQueueName", StringMessageAttribute(sourceQueueName))
     copyMessagesToActorRef.foreach { _ ! SendMessage(message.toNewMessageData) }
 
     destination match {
diff --git a/core/src/main/scala/org/elasticmq/actor/queue/operations/ReceiveMessageOps.scala b/core/src/main/scala/org/elasticmq/actor/queue/operations/ReceiveMessageOps.scala
@@ -85,7 +85,7 @@ trait ReceiveMessageOps extends Logging {
     messageQueue.dequeue(count, deliveryTime).map { internalMessage =>
       if (queueData.deadLettersQueue.map(_.maxReceiveCount).exists(_ <= internalMessage.receiveCount)) {
         logger.debug(s"${queueData.name}: send message $internalMessage to dead letters actor $deadLettersActorRef")
-        deadLettersActorRef.foreach(_ ! MoveMessage(internalMessage, MoveToDLQ))
+        deadLettersActorRef.foreach(_ ! MoveMessage(internalMessage, MoveToDLQ, queueData.name))
         MessageToDelete(internalMessage)
       } else {
         MessageToReturn(internalMessage)
diff --git a/core/src/main/scala/org/elasticmq/msg/QueueMsg.scala b/core/src/main/scala/org/elasticmq/msg/QueueMsg.scala
@@ -42,7 +42,7 @@ case class GetQueueStatistics(deliveryTime: Long) extends QueueQueueMsg[QueueSta
 case class ClearQueue() extends QueueQueueMsg[Unit]
 
 case class SendMessage(message: NewMessageData) extends QueueMessageMsg[MessageData]
-case class MoveMessage(message: InternalMessage, moveDestination: MoveDestination) extends QueueMessageMsg[Unit]
+case class MoveMessage(message: InternalMessage, moveDestination: MoveDestination, sourceQueueName: String) extends QueueMessageMsg[Unit]
 case class UpdateVisibilityTimeout(deliveryReceipt: DeliveryReceipt, visibilityTimeout: VisibilityTimeout)
     extends QueueMessageMsg[Either[InvalidReceiptHandle, Unit]]
 case class ReceiveMessages(
diff --git a/core/src/test/scala/org/elasticmq/FifoDeduplicationIdsHistoryTest.scala b/core/src/test/scala/org/elasticmq/FifoDeduplicationIdsHistoryTest.scala
@@ -107,7 +107,7 @@ class FifoDeduplicationIdsHistoryTest extends AnyFunSuite with Matchers {
       nextDelivery = 100L,
       content = "",
       messageAttributes = Map.empty,
-      messageSystemAttributes = Map.empty,
+      messageSystemAttributes = mutable.HashMap.empty,
       created = created,
       orderIndex = 0,
       firstReceive = NeverReceived,
diff --git a/core/src/test/scala/org/elasticmq/actor/queue/InternalMessageSpec.scala b/core/src/test/scala/org/elasticmq/actor/queue/InternalMessageSpec.scala
@@ -18,7 +18,7 @@ class InternalMessageSpec extends AnyFunSuite with Matchers {
       nextDelivery = 0L,
       content = "content",
       messageAttributes = Map.empty,
-      messageSystemAttributes = Map.empty,
+      messageSystemAttributes = mutable.HashMap.empty,
       created = freezedDateTime,
       orderIndex = 100,
       firstReceive = NeverReceived,
@@ -46,7 +46,7 @@ class InternalMessageSpec extends AnyFunSuite with Matchers {
       nextDelivery = 0L,
       content = "content",
       messageAttributes = Map.empty,
-      messageSystemAttributes = Map.empty,
+      messageSystemAttributes = mutable.HashMap.empty,
       created = freezedDateTime,
       orderIndex = 100,
       firstReceive = NeverReceived,
diff --git a/core/src/test/scala/org/elasticmq/actor/queue/ReceiveRequestAttemptCacheTest.scala b/core/src/test/scala/org/elasticmq/actor/queue/ReceiveRequestAttemptCacheTest.scala
@@ -23,7 +23,7 @@ class ReceiveRequestAttemptCacheTest extends AnyFunSuite with Matchers {
       1L,
       "content",
       Map.empty,
-      Map.empty,
+      mutable.HashMap.empty,
       nowProvider.now,
       orderIndex = 0,
       NeverReceived,
@@ -74,7 +74,7 @@ class ReceiveRequestAttemptCacheTest extends AnyFunSuite with Matchers {
       1L,
       "content",
       Map.empty,
-      Map.empty,
+      mutable.HashMap.empty,
       nowProvider.now,
       orderIndex = 0,
       NeverReceived,
diff --git a/core/src/test/scala/org/elasticmq/actor/test/DataCreationHelpers.scala b/core/src/test/scala/org/elasticmq/actor/test/DataCreationHelpers.scala
@@ -35,6 +35,7 @@ trait DataCreationHelpers {
       id: String,
       content: String,
       messageAttributes: Map[String, MessageAttribute],
+      messageSystemAttributes: Map[String, MessageAttribute],
       nextDelivery: MillisNextDelivery,
       deliveryReceipt: Option[DeliveryReceipt] = None,
       messageGroupId: Option[String] = None,
@@ -46,6 +47,7 @@ trait DataCreationHelpers {
       deliveryReceipt,
       content,
       messageAttributes,
+      messageSystemAttributes,
       nextDelivery,
       OffsetDateTimeUtil.ofEpochMilli(0),
       MessageStatistics(NeverReceived, 0),
diff --git a/persistence/persistence-sql/src/main/scala/org/elasticmq/persistence/sql/DBMessage.scala b/persistence/persistence-sql/src/main/scala/org/elasticmq/persistence/sql/DBMessage.scala
@@ -7,6 +7,8 @@ import org.elasticmq.util.OffsetDateTimeUtil
 import scalikejdbc.WrappedResultSet
 import spray.json._
 
+import scala.collection.mutable
+
 case class DBMessage(
     messageId: String,
     deliveryReceipts: Array[Byte],
@@ -48,7 +50,7 @@ case class DBMessage(
       nextDelivery,
       new String(content),
       serializedAttrs,
-      Map.empty,
+      mutable.HashMap.empty,
       OffsetDateTimeUtil.ofEpochMilli(created),
       orderIndex = 0,
       firstReceive,
diff --git a/rest/rest-sqs-testing-amazon-java-sdk/src/test/scala/org/elasticmq/rest/sqs/AmazonJavaSdkTestSuite.scala b/rest/rest-sqs-testing-amazon-java-sdk/src/test/scala/org/elasticmq/rest/sqs/AmazonJavaSdkTestSuite.scala
@@ -2096,17 +2096,17 @@ class AmazonJavaSdkTestSuite extends SqsClientServerCommunication with Matchers
     secondReceiveResult.getMessages shouldBe empty
   }
 
-  test("should return DeadLetterQueueSourceArn in receive message attributes") {
+  test("should return DeadLetterQueueSourceArn in receive DLQ message attributes") {
     // given
     val messageBody = "Message 1"
-    client.createQueue(new CreateQueueRequest("testDlq")).getQueueUrl
+    val createDlqQueueResult = client.createQueue(new CreateQueueRequest("testDlq")).getQueueUrl
     val redrivePolicy = RedrivePolicy("testDlq", awsRegion, awsAccountId, 1).toJson.toString
 
     val createQueueResult = client
       .createQueue(
         new CreateQueueRequest("main")
           .withAttributes(
-            Map(redrivePolicyAttribute -> redrivePolicy).asJava
+            Map(defaultVisibilityTimeoutAttribute -> "0", redrivePolicyAttribute -> redrivePolicy).asJava
           )
       )
       .getQueueUrl
@@ -2118,10 +2118,23 @@ class AmazonJavaSdkTestSuite extends SqsClientServerCommunication with Matchers
         .withQueueUrl(createQueueResult)
         .withAttributeNames("All")
     )
+    client.receiveMessage(
+      new ReceiveMessageRequest()
+        .withQueueUrl(createQueueResult)
+        .withAttributeNames("All")
+    )
+    val receiveDlqResult = client.receiveMessage(
+      new ReceiveMessageRequest()
+        .withQueueUrl(createDlqQueueResult)
+        .withAttributeNames("All")
+    )
 
     // then
-    receiveResult.getMessages.asScala.toList.flatMap(_.getAttributes.asScala.toList) should contain(
-      ("DeadLetterQueueSourceArn", s"arn:aws:sqs:$awsRegion:$awsAccountId:testDlq")
+    receiveResult.getMessages.asScala.toList.flatMap(_.getAttributes.asScala.toList) should not contain(
+      "DeadLetterQueueSourceArn"
+    )
+    receiveDlqResult.getMessages.asScala.toList.flatMap(_.getAttributes.asScala.toList) should contain(
+      ("DeadLetterQueueSourceArn", s"arn:aws:sqs:$awsRegion:$awsAccountId:main")
     )
   }
 
diff --git a/rest/rest-sqs/src/main/scala/org/elasticmq/rest/sqs/ReceiveMessageDirectives.scala b/rest/rest-sqs/src/main/scala/org/elasticmq/rest/sqs/ReceiveMessageDirectives.scala
@@ -115,7 +115,7 @@ trait ReceiveMessageDirectives {
             ),
             Rule(AWSTraceHeaderAttribute, () => msg.tracingId.map(_.id)),
             Rule(SequenceNumberAttribute, () => msg.sequenceNumber),
-            Rule(DeadLetterQueueSourceArn, () => queueData.deadLettersQueue.map(qd => getArn(qd.name)))
+            Rule(DeadLetterQueueSourceArn, () => msg.messageSystemAttributes.get("sourceQueueName").map(sourceArn => getArn(sourceArn.asInstanceOf[StringMessageAttribute].stringValue)))
           )
         }
 
diff --git a/rest/rest-sqs/src/main/scala/org/elasticmq/rest/sqs/SendMessageDirectives.scala b/rest/rest-sqs/src/main/scala/org/elasticmq/rest/sqs/SendMessageDirectives.scala
@@ -188,7 +188,7 @@ trait SendMessageDirectives {
     val systemMessageAttributeDigest = if (message.messageSystemAttributes.isEmpty) {
       None
     } else {
-      Some(md5AttributeDigest(message.messageSystemAttributes))
+      Some(md5AttributeDigest(message.messageSystemAttributes.to(Map)))
     }
 
     for {
