-
Notifications
You must be signed in to change notification settings - Fork 9
Expand file tree
/
Copy pathgoog_pubsub_api.py
More file actions
192 lines (148 loc) · 5.63 KB
/
Copy pathgoog_pubsub_api.py
File metadata and controls
192 lines (148 loc) · 5.63 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
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
import json
import re
import types
from google.cloud import pubsub_v1
from google.api_core.exceptions import ClientError
from google.pubsub_v1.types import PullRequest
from google.pubsub_v1.types import AcknowledgeRequest
from .secrets import google_credentials
from .lib import toiter, sip, jsonify
import tenacity
PUBSUB_BATCH_SIZE = 10 # send_message_batch's max batch size is 10
class ClientSideError(Exception):
pass
retry = tenacity.retry(
reraise=True,
stop=tenacity.stop_after_attempt(4),
wait=tenacity.wait_random_exponential(0.5, 60.0),
retry=tenacity.retry_if_not_exception_type(ClientSideError),
)
class PubSubTaskQueueAPI(object):
def __init__(self, qurl, **kwargs):
"""
qurl: a topic or subscription location
conforms to this format projects/{project_id}/topics/{topic_id}/subscriptions/{subscription_id}
kwargs: Keywords for the underlying boto3.client constructor, other than `service_name`,
`region_name`, `aws_secret_access_key`, or `aws_access_key_id`.
"""
pattern = r"^projects/(?P<project_id>[\w\d-]+)/topics/(?P<topic_id>[\w\d-]+)/subscriptions/(?P<subscription_id>[\w\d-]+)$"
matches = re.match(pattern, qurl)
if matches is None:
raise ValueError(
"qurl does not conform to the required format (projects/{project_id}/topics/{topic_id}/subscriptions/{subscription_id})"
)
matches = re.search(r"projects/([\w\d-]+)/", qurl)
self.project_id = matches.group(1)
matches = re.search(r"topics/([\w\d-]+)", qurl)
self.topic_id = matches.group(1)
matches = re.search(r"subscriptions/([\w\d-]+)", qurl)
self.subscription_id = matches.group(1)
project_name, credentials = google_credentials()
self.subscriber = pubsub_v1.SubscriberClient(credentials=credentials)
self.publisher = pubsub_v1.PublisherClient(credentials=credentials)
self._topic_path = self.publisher.topic_path(self.project_id, self.topic_id)
self._subscription_path = self.subscriber.subscription_path(
self.project_id, self.subscription_id
)
self.batch_size = PUBSUB_BATCH_SIZE
@property
def enqueued(self):
raise float("Nan")
@property
def inserted(self):
return float("NaN")
@property
def completed(self):
return float("NaN")
@property
def leased(self):
return float("NaN")
def is_empty(self):
return self.enqueued == 0
@retry
def insert(self, tasks, delay_seconds=0):
tasks = toiter(tasks)
def publish_batch(batch):
if not batch:
return 0
futures = []
for task in batch:
data = jsonify(task).encode("utf-8")
future = self.publisher.publish(self._topic_path, data)
futures.append(future)
# Wait for all messages to be published
for future in futures:
try:
# Blocks until the message is published
future.result()
except Exception as e:
raise ClientError(e)
return len(futures)
total = 0
# send_message_batch's max batch size is 10
for batch in sip(tasks, self.batch_size):
if len(batch) == 0:
break
total += publish_batch(batch)
return total
def add_insert_count(self, ct):
pass
def rezero(self):
pass
@retry
def renew_lease(self, task, seconds):
self.subscriber.modify_ack_deadline(
self._subscription_path,
[task.id],
seconds,
)
def cancel_lease(self, task):
self.subscriber.acknowledge(self._subscription_path, [task.id])
def release_all(self):
raise NotImplementedError()
def lease(self, seconds, num_tasks=1, wait_sec=20):
# Pull messages from the subscription
request = PullRequest(
subscription=self._subscription_path, max_messages=num_tasks
)
response = self.subscriber.pull(request)
tasks = []
for received_message in response.received_messages:
# Load the message data as JSON
task = json.loads(received_message.message.data.decode("utf-8"))
# Store the acknowledgement ID in the task
task["id"] = received_message.ack_id
tasks.append(task)
return tasks
def delete(self, task):
if isinstance(task, str):
ack_id = task
else:
try:
ack_id = task._id
except AttributeError:
ack_id = task["id"]
request = AcknowledgeRequest(
subscription=self._subscription_path, ack_ids=[ack_id]
)
self.subscriber.acknowledge(request=request)
return 1
def tally(self):
pass
def purge(self, native=False):
while True:
# Pull messages from the subscription
response = self.subscriber.pull(
self._subscription_path, max_messages=self.batch_size
)
if not response.received_messages:
# No more messages, break the loop
break
# Acknowledge all received messages
ack_ids = [msg.ack_id for msg in response.received_messages]
request = AcknowledgeRequest(
subscription=self._subscription_path, ack_ids=ack_ids
)
self.subscriber.acknowledge(request=request)
def __iter__(self):
return iter(self.lease(num_tasks=10, seconds=0))