Completed SQS Implementation

moto/sqs/models.py

@@ -2,6 +2,7 @@ from __future__ import unicode_literals
 
 import base64
 import hashlib
+import json
 import re
 import six
 import struct
@@ -9,6 +10,7 @@ from xml.sax.saxutils import escape
 
 import boto.sqs
 
+from moto.core.exceptions import RESTError
 from moto.core import BaseBackend, BaseModel
 from moto.core.utils import camelcase_to_underscores, get_random_message_id, unix_time, unix_time_millis
 from .utils import generate_receipt_handle
@@ -166,11 +168,14 @@ class Queue(BaseModel):
                             'ReceiveMessageWaitTimeSeconds',
                             'VisibilityTimeout',
                             'WaitTimeSeconds']
+    ALLOWED_PERMISSIONS = ('*', 'ChangeMessageVisibility', 'DeleteMessage', 'GetQueueAttributes',
+                          'GetQueueUrl', 'ReceiveMessage', 'SendMessage')
 
     def __init__(self, name, region, **kwargs):
         self.name = name
         self.visibility_timeout = int(kwargs.get('VisibilityTimeout', 30))
         self.region = region
+        self.tags = {}
 
         self._messages = []
 
@@ -189,14 +194,42 @@ class Queue(BaseModel):
         self.message_retention_period = int(kwargs.get('MessageRetentionPeriod', 86400 * 4))  # four days
         self.queue_arn = 'arn:aws:sqs:{0}:123456789012:{1}'.format(self.region, self.name)
         self.receive_message_wait_time_seconds = int(kwargs.get('ReceiveMessageWaitTimeSeconds', 0))
+        self.permissions = {}
 
         # wait_time_seconds will be set to immediate return messages
         self.wait_time_seconds = int(kwargs.get('WaitTimeSeconds', 0))
 
+        self.redrive_policy = {}
+        self.dead_letter_queue = None
+
+        if 'RedrivePolicy' in kwargs:
+            self._setup_dlq(kwargs['RedrivePolicy'])
+
         # Check some conditions
         if self.fifo_queue and not self.name.endswith('.fifo'):
             raise MessageAttributesInvalid('Queue name must end in .fifo for FIFO queues')
 
+    def _setup_dlq(self, policy_json):
+        try:
+            self.redrive_policy = json.loads(policy_json)
+        except ValueError:
+            raise RESTError('InvalidParameterValue', 'Redrive policy does not contain valid json')
+
+        if 'deadLetterTargetArn' not in self.redrive_policy:
+            raise RESTError('InvalidParameterValue', 'Redrive policy does not contain deadLetterTargetArn')
+        if 'maxReceiveCount' not in self.redrive_policy:
+            raise RESTError('InvalidParameterValue', 'Redrive policy does not contain maxReceiveCount')
+
+        for queue in sqs_backends[self.region].queues.values():
+            if queue.queue_arn == self.redrive_policy['deadLetterTargetArn']:
+                self.dead_letter_queue = queue
+
+                if self.fifo_queue and not queue.fifo_queue:
+                    raise RESTError('InvalidParameterCombination', 'Fifo queues cannot use non fifo dead letter queues')
+                break
+        else:
+            raise RESTError('AWS.SimpleQueueService.NonExistentQueue', 'Could not find DLQ for {0}'.format(self.redrive_policy['deadLetterTargetArn']))
+
     @classmethod
     def create_from_cloudformation_json(cls, resource_name, cloudformation_json, region_name):
         properties = cloudformation_json['Properties']
@@ -382,9 +415,14 @@ class SQSBackend(BaseBackend):
                 time.sleep(0.001)
                 continue
 
+            messages_to_dlq = []
             for message in queue.messages:
                 if not message.visible:
                     continue
+                if queue.dead_letter_queue is not None and message.approximate_receive_count >= queue.redrive_policy['maxReceiveCount']:
+                    messages_to_dlq.append(message)
+                    continue
+
                 message.mark_received(
                     visibility_timeout=visibility_timeout
                 )
@@ -392,6 +430,10 @@ class SQSBackend(BaseBackend):
                 if len(result) >= count:
                     break
 
+            for message in messages_to_dlq:
+                queue._messages.remove(message)
+                queue.dead_letter_queue.add_message(message)
+
         return result
 
     def delete_message(self, queue_name, receipt_handle):
@@ -419,6 +461,49 @@ class SQSBackend(BaseBackend):
         queue = self.get_queue(queue_name)
         queue._messages = []
 
+    def list_dead_letter_source_queues(self, queue_name):
+        dlq = self.get_queue(queue_name)
+
+        queues = []
+        for queue in self.queues.values():
+            if queue.dead_letter_queue is dlq:
+                queues.append(queue)
+
+        return queues
+
+    def add_permission(self, queue_name, actions, account_ids, label):
+        queue = self.get_queue(queue_name)
+
+        if actions is None or len(actions) == 0:
+            raise RESTError('InvalidParameterValue', 'Need at least one Action')
+        if account_ids is None or len(account_ids) == 0:
+            raise RESTError('InvalidParameterValue', 'Need at least one Account ID')
+
+        if not all([item in Queue.ALLOWED_PERMISSIONS for item in actions]):
+            raise RESTError('InvalidParameterValue', 'Invalid permissions')
+
+        queue.permissions[label] = (account_ids, actions)
+
+    def remove_permission(self, queue_name, label):
+        queue = self.get_queue(queue_name)
+
+        if label not in queue.permissions:
+            raise RESTError('InvalidParameterValue', 'Permission doesnt exist for the given label')
+
+        del queue.permissions[label]
+
+    def tag_queue(self, queue_name, tags):
+        queue = self.get_queue(queue_name)
+        queue.tags.update(tags)
+
+    def untag_queue(self, queue_name, tag_keys):
+        queue = self.get_queue(queue_name)
+        for key in tag_keys:
+            try:
+                del queue.tags[key]
+            except KeyError:
+                pass
+
 
 sqs_backends = {}
 for region in boto.sqs.regions():