DE-111852 Add support for custom SQS message attributes (#102)

BrandEmbassy · Aug 29, 2024 · 761730b · 761730b
1 parent faeab36
commit 761730b
Show file tree

Hide file tree

Showing 7 changed files with 137 additions and 15 deletions.
diff --git a/src/Jobs/JobInterface.php b/src/Jobs/JobInterface.php
@@ -61,4 +61,16 @@ public function getExecutionPlannedAt(): ?DateTimeImmutable;
 
 
     public function setExecutionPlannedAt(DateTimeImmutable $executionPlannedAt): void;
+
+
+    /**
+     * @return array<string,array{DataType: string, StringValue?: string, BinaryValue?: string}>
+     */
+    public function getMessageAttributes(): array;
+
+
+    /**
+     * @param array<string,array{DataType: string, StringValue?: string, BinaryValue?: string}> $messageAttributes
+     */
+    public function setMessageAttributes(array $messageAttributes): void;
 }
diff --git a/src/Jobs/SimpleJob.php b/src/Jobs/SimpleJob.php
@@ -4,6 +4,8 @@
 
 use BE\QueueManagement\Jobs\Execution\MaximumAttemptsExceededException;
 use BE\QueueManagement\Jobs\JobDefinitions\JobDefinitionInterface;
+use BE\QueueManagement\Queue\AWSSQS\SqsMessageAttributeDataType;
+use BE\QueueManagement\Queue\AWSSQS\SqsMessageAttributeFields;
 use BrandEmbassy\DateTime\DateTimeFormatter;
 use DateTimeImmutable;
 use Doctrine\Common\Collections\Collection;
@@ -29,24 +31,32 @@ class SimpleJob implements JobInterface
 
     private ?DateTimeImmutable $executionPlannedAt;
 
+    /**
+     * @var array<string,array{DataType: string, StringValue?: string, BinaryValue?: string}>
+     */
+    private array $messageAttributes;
+
 
     /**
      * @param Collection<string, mixed> $parameters
+     * @param array<string, mixed> $messageAttributes
      */
     public function __construct(
         string $uuid,
         DateTimeImmutable $createdAt,
         int $attempts,
         JobDefinitionInterface $jobDefinition,
         Collection $parameters,
-        ?DateTimeImmutable $executionPlannedAt
+        ?DateTimeImmutable $executionPlannedAt,
+        array $messageAttributes = [],
     ) {
         $this->uuid = $uuid;
         $this->createdAt = $createdAt;
         $this->attempts = $attempts;
         $this->parameters = $parameters;
         $this->jobDefinition = $jobDefinition;
         $this->executionPlannedAt = $executionPlannedAt;
+        $this->messageAttributes = $messageAttributes;
     }
 
 
@@ -176,4 +186,41 @@ public function setExecutionPlannedAt(DateTimeImmutable $executionPlannedAt): vo
     {
         $this->executionPlannedAt = $executionPlannedAt;
     }
+
+
+    /**
+     * @return array<string,array{DataType: string, StringValue?: string, BinaryValue?: string}>
+     */
+    public function getMessageAttributes(): array
+    {
+        return $this->messageAttributes;
+    }
+
+
+    public function setMessageAttribute(
+        string $messageAttributeName,
+        string $messageAttributeValue,
+        SqsMessageAttributeDataType $messageAttributeDataType = SqsMessageAttributeDataType::STRING
+    ): void {
+        $valueKey = $messageAttributeDataType === SqsMessageAttributeDataType::BINARY
+            ? SqsMessageAttributeFields::BINARY_VALUE->value
+            : SqsMessageAttributeFields::STRING_VALUE->value;
+
+        /** @var array{DataType: string, StringValue?: string, BinaryValue?: string} $messageAttribute */
+        $messageAttribute = [
+            SqsMessageAttributeFields::DATA_TYPE->value => $messageAttributeDataType->value,
+            $valueKey => $messageAttributeValue,
+        ];
+
+        $this->messageAttributes[$messageAttributeName] = $messageAttribute;
+    }
+
+
+    /**
+     * @param array<string,array{DataType: string, StringValue?: string, BinaryValue?: string}> $messageAttributes
+     */
+    public function setMessageAttributes(array $messageAttributes): void
+    {
+        $this->messageAttributes = $messageAttributes;
+    }
 }
diff --git a/src/Queue/AWSSQS/SqsMessageAttributeDataType.php b/src/Queue/AWSSQS/SqsMessageAttributeDataType.php
@@ -0,0 +1,12 @@
+<?php declare(strict_types = 1);
+
+namespace BE\QueueManagement\Queue\AWSSQS;
+
+enum SqsMessageAttributeDataType: string
+{
+    case STRING = 'String';
+
+    case NUMBER = 'Number';
+
+    case BINARY = 'Binary';
+}
diff --git a/src/Queue/AWSSQS/SqsMessageAttributeFields.php b/src/Queue/AWSSQS/SqsMessageAttributeFields.php
@@ -0,0 +1,12 @@
+<?php declare(strict_types = 1);
+
+namespace BE\QueueManagement\Queue\AWSSQS;
+
+enum SqsMessageAttributeFields: string
+{
+    case DATA_TYPE = 'DataType';
+
+    case BINARY_VALUE = 'BinaryValue';
+
+    case STRING_VALUE = 'StringValue';
+}
diff --git a/src/Queue/AWSSQS/SqsQueueManager.php b/src/Queue/AWSSQS/SqsQueueManager.php
@@ -413,14 +413,13 @@ private function publishMessage(
             throw SqsClientException::createFromInvalidDelaySeconds($delaySeconds);
         }
 
-        $messageAttributes = [
-            SqsSendingMessageFields::QUEUE_URL => [
-                'DataType' => 'String',
-                // queueName might be handy here if we want to consume
-                // from multiple queues in parallel via promises.
-                // Then we need queue in message directly so that we can delete it.
-                'StringValue' => $prefixedQueueName,
-            ],
+        $messageAttributes = $job->getMessageAttributes();
+        $messageAttributes[SqsSendingMessageFields::QUEUE_URL] = [
+            SqsMessageAttributeFields::DATA_TYPE->value => SqsMessageAttributeDataType::STRING->value,
+            // queueName might be handy here if we want to consume
+            // from multiple queues in parallel via promises.
+            // Then we need queue in message directly so that we can delete it.
+            SqsMessageAttributeFields::STRING_VALUE->value => $prefixedQueueName,
         ];
 
         if (SqsMessage::isTooBig($messageBody, $messageAttributes)) {

diff --git a/tests/Jobs/ExampleJob.php b/tests/Jobs/ExampleJob.php
@@ -28,8 +28,14 @@ class ExampleJob extends SimpleJob
     public const PARAMETER_FOO = 'foo';
 
 
-    public function __construct(?JobDefinitionInterface $jobDefinition = null, string $bar = 'bar')
-    {
+    /**
+     * @param array<string,array{DataType: string, StringValue?: string, BinaryValue?: string}> $messageAttributes
+     */
+    public function __construct(
+        ?JobDefinitionInterface $jobDefinition = null,
+        string $bar = 'bar',
+        array $messageAttributes = [],
+    ) {
         /**
          * Prevent phpstan error Template type T on class Doctrine\Common\Collections\Collection is not covariant
          * @var array<string,mixed> $parameters
@@ -43,6 +49,7 @@ public function __construct(?JobDefinitionInterface $jobDefinition = null, strin
             $jobDefinition ?? ExampleJobDefinition::create(),
             new ArrayCollection($parameters),
             null,
+            $messageAttributes,
         );
     }
 

diff --git a/tests/Queue/AWSSQS/SqsQueueManagerTest.php b/tests/Queue/AWSSQS/SqsQueueManagerTest.php
@@ -41,6 +41,17 @@ class SqsQueueManagerTest extends TestCase
 
     private const QUEUE_URL = 'https://sqs.eu-central-1.amazonaws.com/583027123456/MyQueue1';
 
+    private const CUSTOM_MESSAGE_ATTRIBUTE = 'customMessageAttribute';
+
+    private const CUSTOM_MESSAGE_ATTRIBUTE_VALUE = 'customMessageAttributeValue';
+
+    public const CUSTOM_MESSAGE_ATTRIBUTES = [
+        self::CUSTOM_MESSAGE_ATTRIBUTE => [
+            'DataType' => 'String',
+            'StringValue' => self::CUSTOM_MESSAGE_ATTRIBUTE_VALUE,
+        ],
+    ];
+
     private const RECEIPT_HANDLE = 'AQEBMJRLDYbo...BYSvLGdGU9t8Q==';
 
     private const S3_BUCKET_NAME = 'thisIsS3Bucket';
@@ -116,12 +127,20 @@ public function testPushWithInvalidCharacters(string $queueName, string $queueNa
                 ->withQueueName($queueName),
             'This is random text.',
         );
+        $exampleJobWithInvalidCharacter->setMessageAttribute(
+            self::CUSTOM_MESSAGE_ATTRIBUTE,
+            self::CUSTOM_MESSAGE_ATTRIBUTE_VALUE,
+        );
 
         $exampleJobWithValidCharacter = new ExampleJob(
             ExampleJobDefinition::create()
                 ->withQueueName($queueName),
             'This is random text.',
         );
+        $exampleJobWithValidCharacter->setMessageAttribute(
+            self::CUSTOM_MESSAGE_ATTRIBUTE,
+            self::CUSTOM_MESSAGE_ATTRIBUTE_VALUE,
+        );
 
         $this->sqsClientMock->expects('sendMessage')
             ->with(
@@ -172,7 +191,7 @@ public function testPushWithTooBigMessage(string $queueName, string $queueNamePr
         $this->sqsClientMock->expects('sendMessage')
             ->with(
                 Mockery::on(
-                    fn(array $message): bool => $this->messageCheckOk($message, $messageBody, 0),
+                    fn(array $message): bool => $this->messageCheckOk($message, $messageBody, 0, false),
                 ),
             )
             ->andReturn($this->createSqsSendMessageResultMock());
@@ -287,6 +306,10 @@ public function testPushDelayedWithJobDelayOverSqsMaxDelayLimitAndEventDispatche
             ->with(Mockery::on(fn($event) => $event instanceof MessageSentEvent
                 && $event->delayInSeconds === 900
                 && $event->messageAttributes === [
+                    'customMessageAttribute' => [
+                        'DataType' => 'String',
+                        'StringValue' => 'customMessageAttributeValue',
+                    ],
                     'QueueUrl' => [
                         'DataType' => 'String',
                         'StringValue' => $queueNamePrefix . $queueName,
@@ -647,12 +670,20 @@ private function getSampleSqsMessages(): array
     /**
      * @param array<string, mixed> $message
      */
-    private function messageCheckOk(array $message, string $messageBody, int $delay): bool
-    {
+    private function messageCheckOk(
+        array $message,
+        string $messageBody,
+        int $delay,
+        bool $checkCustomAttribute = true
+    ): bool {
         return $message['MessageBody'] === $messageBody
             && $message[SqsSendingMessageFields::DELAY_SECONDS] === $delay
             && $message[SqsSendingMessageFields::QUEUE_URL] === self::QUEUE_URL
-            && $message[SqsSendingMessageFields::MESSAGE_ATTRIBUTES][SqsSendingMessageFields::QUEUE_URL]['StringValue'] === self::QUEUE_URL;
+            && $message[SqsSendingMessageFields::MESSAGE_ATTRIBUTES][SqsSendingMessageFields::QUEUE_URL]['StringValue'] === self::QUEUE_URL
+            && (
+                $checkCustomAttribute === false
+                || $message[SqsSendingMessageFields::MESSAGE_ATTRIBUTES][self::CUSTOM_MESSAGE_ATTRIBUTE]['StringValue'] === self::CUSTOM_MESSAGE_ATTRIBUTE_VALUE
+            );
     }
 
 
@@ -661,6 +692,8 @@ private function createExampleJob(string $queueName): ExampleJob
         return new ExampleJob(
             ExampleJobDefinition::create()
                 ->withQueueName($queueName),
+            'bar',
+            self::CUSTOM_MESSAGE_ATTRIBUTES,
         );
     }