Files
django-q2/django_q/brokers/aws_sqs.py
T
Gouri JavedandJaved Gouri e205674d31 added long polling support (#506)
* added long pooling support for  sqs broker

* changed the parameter name

* added long polling support ,

introduce new parameter in broker sqs block
parameter for long polling is receive_message_wait_time_seconds

eg :
Q_CLUSTER = {
      'name': 'test-queue',
      'sqs': {
         'aws_region': AWS_S3_REGION_NAME,
         'aws_access_key_id': AWS_ACCESS_KEY_ID,
         'aws_secret_access_key': AWS_SECRET_ACCESS_KEY,
         'receive_message_wait_time_seconds':20
     }
}

* added receive_message_wait_time_seconds parameter in test_broker.py

Co-authored-by: Javed Gouri <you@example.com>
2021-03-21 12:04:16 +01:00

100 lines
3.1 KiB
Python

from boto3 import Session
from botocore.client import ClientError
from django_q.brokers import Broker
from django_q.conf import Conf
QUEUE_DOES_NOT_EXIST = "AWS.SimpleQueueService.NonExistentQueue"
class Sqs(Broker):
def __init__(self, list_key: str = Conf.PREFIX):
self.sqs = None
super(Sqs, self).__init__(list_key)
self.queue = self.get_queue()
def __setstate__(self, state):
super(Sqs, self).__setstate__(state)
self.sqs = None
self.queue = self.get_queue()
def enqueue(self, task):
response = self.queue.send_message(MessageBody=task)
return response.get("MessageId")
def dequeue(self):
# sqs supports max 10 messages in bulk
if Conf.BULK > 10:
Conf.BULK = 10
params = {"MaxNumberOfMessages": Conf.BULK, "VisibilityTimeout": Conf.RETRY}
# sqs long polling
sqs_config = Conf.SQS
if "receive_message_wait_time_seconds" in sqs_config:
wait_time_second = sqs_config.get("receive_message_wait_time_seconds", 20)
# validation of parameter
if not isinstance(wait_time_second, int):
raise ValueError("receive_message_wait_time_seconds should be int")
if wait_time_second > 20:
raise ValueError("receive_message_wait_time_seconds is invalid. Reason: Must be >= 0 and <= 20")
params.update({"WaitTimeSeconds": wait_time_second})
tasks = self.queue.receive_messages(**params)
if tasks:
return [(t.receipt_handle, t.body) for t in tasks]
def acknowledge(self, task_id):
return self.delete(task_id)
def queue_size(self) -> int:
return int(self.queue.attributes["ApproximateNumberOfMessages"])
def lock_size(self) -> int:
return int(self.queue.attributes["ApproximateNumberOfMessagesNotVisible"])
def delete(self, task_id):
message = self.sqs.Message(self.queue.url, task_id)
message.delete()
def fail(self, task_id):
self.delete(task_id)
def delete_queue(self):
self.queue.delete()
def purge_queue(self):
self.queue.purge()
def ping(self) -> bool:
return "sqs" in self.connection.get_available_resources()
def info(self) -> str:
return "AWS SQS"
@staticmethod
def get_connection(list_key: str = Conf.PREFIX) -> Session:
config = Conf.SQS
if "aws_region" in config:
config["region_name"] = config["aws_region"]
del config["aws_region"]
if 'receive_message_wait_time_seconds' in config:
del config["receive_message_wait_time_seconds"]
return Session(**config)
def get_queue(self):
self.sqs = self.connection.resource("sqs")
try:
# try to return an existing queue by name. If the queue does not
# exist try to create it.
return self.sqs.get_queue_by_name(QueueName=self.list_key)
except ClientError as exp:
if exp.response["Error"]["Code"] != QUEUE_DOES_NOT_EXIST:
raise exp
return self.sqs.create_queue(QueueName=self.list_key)