Repository: Bogdanp/dramatiq_sqs

Issue:
Allow tagging of queues on creation
Tagging AWS resources is really useful to keep track of which resources belong to which application, and to track things like costs on a per-application basis. Currently the SQS broker doesn't allow tags to be set on the queue, but this would be useful functionality to have.

I'm happy to contribute the PR for this if you're 👍 with it.

Comments:

Patch:
diff --git a/dramatiq_sqs/broker.py b/dramatiq_sqs/broker.py
index 6ce094a..d9a2b8f 100644
--- a/dramatiq_sqs/broker.py
+++ b/dramatiq_sqs/broker.py
@@ -68,6 +68,7 @@ class SQSBroker(dramatiq.Broker):
             retention: int = MAX_MESSAGE_RETENTION,
             dead_letter: bool = False,
             max_receives: int = MAX_RECEIVES,
+            tags: Optional[Dict[str, str]] = None,
             **options,
     ) -> None:
         super().__init__(middleware=middleware)
@@ -80,6 +81,7 @@ class SQSBroker(dramatiq.Broker):
         self.queues: Dict[str, Any] = {}
         self.dead_letter: bool = dead_letter
         self.max_receives: int = max_receives
+        self.tags: Optional[Dict[str, str]] = tags
         self.sqs: Any = boto3.resource("sqs", **options)
 
     def consume(self, queue_name: str, prefetch: int = 1, timeout: int = 30000) -> dramatiq.Consumer:
@@ -104,11 +106,22 @@ class SQSBroker(dramatiq.Broker):
                     "MessageRetentionPeriod": self.retention,
                 }
             )
+            if self.tags:
+                self.sqs.meta.client.tag_queue(
+                    QueueUrl=self.queues[queue_name].url,
+                    Tags=self.tags
+                )
+
             if self.dead_letter:
                 dead_letter_queue_name = f"{prefixed_queue_name}_dlq"
                 dead_letter_queue = self.sqs.create_queue(
                     QueueName=dead_letter_queue_name
                 )
+                if self.tags:
+                    self.sqs.meta.client.tag_queue(
+                        QueueUrl=dead_letter_queue.url,
+                        Tags=self.tags
+                    )
                 redrive_policy = {
                     "deadLetterTargetArn": dead_letter_queue.attributes["QueueArn"],
                     "maxReceiveCount": str(self.max_receives)


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index c961ed5..afd5441 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -26,6 +26,9 @@ def broker():
             Pipelines(),
             Retries(min_backoff=1000, max_backoff=900000, max_retries=96),
         ],
+        tags={
+            "owner": "dramatiq_sqs_tests",
+        },
     )
     dramatiq.set_broker(broker)
     yield broker
diff --git a/tests/test_broker.py b/tests/test_broker.py
index 6183803..077a70c 100644
--- a/tests/test_broker.py
+++ b/tests/test_broker.py
@@ -6,6 +6,7 @@ import pytest
 from botocore.stub import Stubber
 
 from dramatiq_sqs import SQSBroker
+from dramatiq_sqs.broker import MAX_MESSAGE_RETENTION
 
 
 def test_can_enqueue_and_process_messages(broker, worker, queue_name):
@@ -158,3 +159,28 @@ def test_creates_dead_letter_queue():
     with stubber:
         broker.declare_queue("test")
         stubber.assert_no_pending_responses()
+
+
+def test_tags_queues_on_create():
+    # Given that I have an SQS broker with tags
+    broker = SQSBroker(
+        namespace="dramatiq_sqs_tests",
+        tags={"key1": "value1", "key2": "value2"}
+    )
+
+    # And I've stubbed out all the relevant API calls
+    stubber = Stubber(broker.sqs.meta.client)
+    stubber.add_response("create_queue", {"QueueUrl": ""})
+    stubber.add_response("tag_queue", {}, {
+        "QueueUrl": "",
+        "Tags": {
+            "key1": "value1",
+            "key2": "value2"
+        }
+    })
+
+    # When I create a queue
+    # Then the queue should have the specified tags
+    with stubber:
+        broker.declare_queue("test")
+        stubber.assert_no_pending_responses()