/
awssqstools.py
103 lines (82 loc) · 3.27 KB
/
awssqstools.py
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
# -*- coding: utf-8 -*-
''' AwsSQSTools '''
from base64 import b64encode
from boto.sqs import regions
from boto.sqs.connection import SQSConnection
class AwsSQSTools(SQSConnection):
''' AwsSQSTools
:param str aws_access_key_id: aws_access_key_id
:param str aws_secret_access_key: aws_secret_access_key
:param str region: region.
:param str queue_name: AWS SQS queue name.
'''
def __init__(self, aws_access_key_id, aws_secret_access_key, region,
queue_name):
super(AwsSQSTools, self).__init__(aws_access_key_id,
aws_secret_access_key,
region=[i for i in regions() if i.name == region][0])
self.queue = self.create_queue(queue_name)
def get_all_queues(self, *args, **kwargs):
''' Get all queue
:rtype: list
:returns: A list of :class:`boto.sqs.queue.Queue` instances.
'''
return super(AwsSQSTools, self).get_all_queues(*args, **kwargs)
def write(self, body):
''' Write a message into queue
:param str body: message body
:rtype: :class:`boto.sqs.message.Message`
:returns: :class:`boto.sqs.message.Message`
'''
return self.queue.write(self.queue.new_message(body))
def write_batch(self, body_list):
''' Write huge messages will be split 10 messages
in each single request.
:param list body_list: a list of message raw data.
:rtype: list
:returns: :class:`boto.sqs.batchresults.BatchResults` in list
.. todo:: default args.
'''
result = []
for loops in xrange(len(body_list) / 10 + 1):
batch = []
for no_msg, body in enumerate(body_list[10*loops:10*(loops+1)]):
batch.append((no_msg, b64encode(body), 0))
result.append(self.queue.write_batch(batch))
return result
def get_messages(self, num_messages=10, *args, **kwargs):
''' Get messages from queue
:param int num_messages: get message in one time.
:rtype: list
:returns: A list of :class:`boto.sqs.message.Message` instances
.. todo::
- Need to increase `get_messages` concurrency.
- Maybe auto convert json object.
'''
for msg in self.queue.get_messages(num_messages, *args, **kwargs):
yield msg
if __name__ == '__main__':
print 'remove comment before use'
#import json
#import setting
#from datetime import datetime
#SQS = AwsSQSTools(setting.ID, setting.KEY, setting.REGION, 'test_toomore')
#print SQS.queue
#print SQS.get_all_queues('test_')
#print dir(SQS.queue)
#print [SQS.write(str(i)) for i in xrange(10)]
# ----- test write / write_batch ----- #
#print SQS.queue.clear()
#t1 = datetime.now()
#print [SQS.write(json.dumps(i)) for i in [dict(name=u'國'),]*23]
#t2 = datetime.now()
#print SQS.write_batch([json.dumps(i) for i in [dict(name=u'國'),]*23])
#t3 = datetime.now()
#print t2 - t1, t3 - t2
# ----- test get_messages ----- #
#t1 = datetime.now()
#print list(SQS.get_messages())
#t2 = datetime.now()
#print list(SQS.get_messages())
#t3 = datetime.now()
#print t2 - t1, t3 - t2