forked from seung-lab/python-task-queue
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathaws_queue_api.py
More file actions
175 lines (135 loc) · 4.09 KB
/
Copy pathaws_queue_api.py
File metadata and controls
175 lines (135 loc) · 4.09 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
import json
import re
import types
import boto3
import botocore.errorfactory
from .lib import toiter, sip
from .secrets import aws_credentials
AWS_BATCH_SIZE = 10 # send_message_batch's max batch size is 10
class AWSTaskQueueAPI(object):
def __init__(self, qurl, region_name=None):
"""
qurl: either a queue name (e.g. 'pull_queue') or a url
like https://sqs.us-east-1.amazonaws.com/DIGITS/wms-pull-queue
"""
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',
region_name=region_name,
aws_secret_access_key=credentials['AWS_SECRET_ACCESS_KEY'],
aws_access_key_id=credentials['AWS_ACCESS_KEY_ID'],
)
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):
raise NotImplementedError()
@property
def completed(self):
raise NotImplementedError()
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
resp = self.sqs.send_message_batch(
QueueUrl=self.qurl,
Entries=[ {
"Id": str(j),
"MessageBody": json.dumps(task),
"DelaySeconds": delay_seconds,
} for j, task in enumerate(batch) ],
)
total += 1
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 _request(self, num_tasks, visibility_timeout):
resp = self.sqs.receive_message(
QueueUrl=self.qurl,
AttributeNames=[
'SentTimestamp'
],
MaxNumberOfMessages=num_tasks,
MessageAttributeNames=[
'All'
],
VisibilityTimeout=visibility_timeout,
WaitTimeSeconds=20,
)
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 lease(self, seconds, num_tasks=1):
if num_tasks > 1:
raise ValueError("This library (not boto/SQS) only supports fetching one task at a time. Requested: {}.".format(num_tasks))
return self._request(num_tasks, seconds)
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
def tally(self):
raise NotImplementedError("Not supported for SQS.")
def purge(self):
# This is more efficient, but it kept freezing
# try:
# self.sqs.purge_queue(QueueUrl=self.qurl)
# except botocore.errorfactory.PurgeQueueInProgress:
while self.enqueued:
# visibility_timeout must be > 0 for delete to work
tasks = self._request(num_tasks=10, visibility_timeout=10)
for task in tasks:
self.delete(task)
return self
def list(self):
return list(self)
def __iter__(self):
return iter(self._request(num_tasks=10, visibility_timeout=0))