import json

import re

import types

import boto3

import botocore.errorfactory

from .lib import toiter, sip, jsonify

from .secrets import aws_credentials

from .secrets import (

AWS_DEFAULT_REGION

)

AWS_BATCH_SIZE = 10 # send_message_batch's max batch size is 10

class AWSTaskQueueAPI(object):

def __init__(self, qurl, region_name=AWS_DEFAULT_REGION, **kwargs):

"""

qurl: either a queue name (e.g. 'pull_queue') or a url

like https://sqs.us-east-1.amazonaws.com/DIGITS/wms-pull-queue

kwargs: Keywords for the underlying boto3.client constructor, other than `service_name`,

`region_name`, `aws_secret_access_key`, or `aws_access_key_id`.

"""

if 'region' in kwargs:

region_name = kwargs.pop('region')

matches = re.search(r'sqs.([\w\d-]+).amazonaws', qurl)

if matches is not None:

region_name, = matches.groups()

self.qurl = qurl

else:

self.qurl = None

credentials = aws_credentials()

self.sqs = boto3.client('sqs',

aws_secret_access_key=credentials.get('AWS_SECRET_ACCESS_KEY'),

aws_access_key_id=credentials.get('AWS_ACCESS_KEY_ID'),

region_name=region_name,

**kwargs,

)

if self.qurl is None:

try:

self.qurl = self.sqs.get_queue_url(QueueName=qurl)["QueueUrl"]

except Exception:

print(qurl)

raise

self.batch_size = AWS_BATCH_SIZE

@property

def enqueued(self):

status = self.status()

return int(status['ApproximateNumberOfMessages']) + int(status['ApproximateNumberOfMessagesNotVisible'])

@property

def inserted(self):

return float('NaN')

@property

def completed(self):

return float('NaN')

@property

def leased(self):

status = self.status()

return int(status['ApproximateNumberOfMessagesNotVisible'])

def is_empty():

return self.enqueued == 0

def status(self):

resp = self.sqs.get_queue_attributes(

QueueUrl=self.qurl,

AttributeNames=['ApproximateNumberOfMessages', 'ApproximateNumberOfMessagesNotVisible']

)

return resp['Attributes']

def insert(self, tasks, delay_seconds=0):

tasks = toiter(tasks)

total = 0

# send_message_batch's max batch size is 10

for batch in sip(tasks, self.batch_size):

if len(batch) == 0:

break

entries = [ {

"Id": str(j),

"MessageBody": jsonify(task),

"DelaySeconds": delay_seconds,

} for j, task in enumerate(batch)

]

resp = self.sqs.send_message_batch(

QueueUrl=self.qurl,

Entries=entries,

)

total += len(entries)

return total

def add_insert_count(self, ct):

pass

def rezero(self):

pass

def renew_lease(self, seconds):

raise NotImplementedError()

def cancel_lease(self, rhandle):

raise NotImplementedError()

def release_all(self):

raise NotImplementedError()

def lease(self, seconds, num_tasks=1, wait_sec=20):

if wait_sec is None:

wait_sec = 20

resp = self.sqs.receive_message(

QueueUrl=self.qurl,

AttributeNames=[

'SentTimestamp'

],

MaxNumberOfMessages=num_tasks,

MessageAttributeNames=[

'All'

],

VisibilityTimeout=seconds,

WaitTimeSeconds=wait_sec,

)

if 'Messages' not in resp:

return []

tasks = []

for msg in resp['Messages']:

task = json.loads(msg['Body'])

task['id'] = msg['ReceiptHandle']

tasks.append(task)

return tasks

def delete(self, task):

if type(task) == str:

rhandle = task

else:

try:

rhandle = task._id

except AttributeError:

rhandle = task['id']

try:

self.sqs.delete_message(

QueueUrl=self.qurl,

ReceiptHandle=rhandle,

)

except botocore.exceptions.ClientError as err:

pass

return 1

def tally(self):

pass

def purge(self, native=False):

# can throw:

# except botocore.errorfactory.PurgeQueueInProgress:

if native:

self.sqs.purge_queue(QueueUrl=self.qurl)

return

while self.enqueued:

# visibility_timeout must be > 0 for delete to work

tasks = self.lease(num_tasks=10, seconds=10)

for task in tasks:

self.delete(task)

def __iter__(self):

return iter(self.lease(num_tasks=10, seconds=0))