mirror of
https://github.com/django-q2/django-q2.git
synced 2026-09-28 01:28:11 +08:00
Ilan Steemers
15155c7a99
* Black linting * Adding Black to dev dependencies and upping minimal python to 3.6.2 for compatibility * Updating packages * Removing pip-tools input and exporting requirements with poetry * Deleting old setup files and test runner * Trying 1.3.7 * Looser extras requirements to prevent conflicts * Sorted imports with isort * Added iSort to dev dependencies * Fixes localtime for naive setups
102 lines
3.2 KiB
Python
102 lines
3.2 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)
|