Repository navigation
Expand file tree
/
Copy path_request_queue_client.py
More file actions
331 lines (271 loc) · 12.2 KB
/
Copy path_request_queue_client.py
File metadata and controls
331 lines (271 loc) · 12.2 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
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
from __future__ import annotations
from collections import OrderedDict
from datetime import datetime, timezone
from logging import getLogger
from typing import TYPE_CHECKING
from typing_extensions import Self, override
from crawlee import Request
from crawlee._utils.crypto import crypto_random_object_id
from crawlee._utils.raise_if_too_many_kwargs import raise_if_too_many_kwargs
from crawlee.storage_clients._base import RequestQueueClient
from crawlee.storage_clients.models import AddRequestsResponse, ProcessedRequest, RequestQueueMetadata
if TYPE_CHECKING:
from collections.abc import Sequence
logger = getLogger(__name__)
class MemoryRequestQueueClient(RequestQueueClient):
"""Memory implementation of the request queue client.
No data is persisted between process runs, which means all requests are lost when the program terminates.
This implementation is primarily useful for testing, development, and short-lived crawler runs where
persistence is not required.
This client provides fast access to request data but is limited by available memory and does not support
data sharing across different processes.
"""
def __init__(
self,
*,
metadata: RequestQueueMetadata,
) -> None:
"""Initialize a new instance.
Preferably use the `MemoryRequestQueueClient.open` class method to create a new instance.
"""
self._metadata = metadata
# The three stores below are keyed by unique key and disjoint - a request known to the queue lives in
# exactly one of them, so together they also serve as the lookup by unique key.
self._pending_requests = OrderedDict[str, Request]()
"""Pending requests are those that have been added to the queue but not yet fetched for processing.
Ordered from the front of the queue to its end, which keeps both fetching and repositioning a request
to the forefront O(1).
"""
self._handled_requests = dict[str, Request]()
"""Handled requests are those that have been processed and marked as handled."""
self._in_progress_requests = dict[str, Request]()
"""In-progress requests are those that have been fetched but not yet marked as handled or reclaimed."""
@override
async def get_metadata(self) -> RequestQueueMetadata:
return self._metadata
@classmethod
async def open(
cls,
*,
id: str | None,
name: str | None,
alias: str | None,
) -> Self:
"""Open or create a new memory request queue client.
This method creates a new in-memory request queue instance. Unlike persistent storage implementations,
memory queues don't check for existing queues with the same name or ID since all data exists only
in memory and is lost when the process terminates.
Alias does not have any effect on the memory storage client implementation, because unnamed storages
are supported by default, since data are not persisted.
Args:
id: The ID of the request queue. If not provided, a random ID will be generated.
name: The name of the request queue for named (global scope) storages.
alias: The alias of the request queue for unnamed (run scope) storages.
Returns:
An instance for the opened or created storage client.
Raises:
ValueError: If both name and alias are provided.
"""
# Validate input parameters.
raise_if_too_many_kwargs(id=id, name=name, alias=alias)
# Create a new queue
queue_id = id or crypto_random_object_id()
now = datetime.now(timezone.utc)
metadata = RequestQueueMetadata(
id=queue_id,
name=name,
created_at=now,
accessed_at=now,
modified_at=now,
had_multiple_clients=False,
handled_request_count=0,
pending_request_count=0,
total_request_count=0,
)
return cls(metadata=metadata)
@override
async def drop(self) -> None:
self._pending_requests.clear()
self._handled_requests.clear()
self._in_progress_requests.clear()
await self._update_metadata(
update_modified_at=True,
update_accessed_at=True,
new_handled_request_count=0,
new_pending_request_count=0,
new_total_request_count=0,
)
@override
async def purge(self) -> None:
self._pending_requests.clear()
self._handled_requests.clear()
self._in_progress_requests.clear()
await self._update_metadata(
update_modified_at=True,
update_accessed_at=True,
new_pending_request_count=0,
new_handled_request_count=0,
new_total_request_count=0,
)
@override
async def add_batch_of_requests(
self,
requests: Sequence[Request],
*,
forefront: bool = False,
) -> AddRequestsResponse:
processed_requests = []
new_total_request_count = self._metadata.total_request_count
new_pending_request_count = self._metadata.pending_request_count
for request in requests:
# Check which of the stores, if any, the request is already in.
was_already_handled = request.unique_key in self._handled_requests
is_in_progress = request.unique_key in self._in_progress_requests
was_already_present = was_already_handled or is_in_progress or request.unique_key in self._pending_requests
# If the request is already in the queue and handled, don't add it again.
if was_already_handled:
processed_requests.append(
ProcessedRequest(
unique_key=request.unique_key,
was_already_present=True,
was_already_handled=True,
)
)
continue
# If the request is already in progress, don't add it again.
if is_in_progress:
processed_requests.append(
ProcessedRequest(
unique_key=request.unique_key,
was_already_present=True,
was_already_handled=False,
)
)
continue
# A new request is appended to the end of the queue. A re-add of a still-pending request keeps the
# originally enqueued object: the incoming duplicate is typically a freshly built one that lost the
# state accumulated so far (e.g. `retry_count`).
if not was_already_present:
self._pending_requests[request.unique_key] = request
new_total_request_count += 1
new_pending_request_count += 1
# The only effect a re-add may have is repositioning the request to the front of the queue.
if forefront:
self._pending_requests.move_to_end(request.unique_key, last=False)
processed_requests.append(
ProcessedRequest(
unique_key=request.unique_key,
was_already_present=was_already_present,
was_already_handled=False,
)
)
await self._update_metadata(
update_accessed_at=True,
update_modified_at=True,
new_total_request_count=new_total_request_count,
new_pending_request_count=new_pending_request_count,
)
return AddRequestsResponse(
processed_requests=processed_requests,
unprocessed_requests=[],
)
@override
async def fetch_next_request(self) -> Request | None:
if not self._pending_requests:
return None
_, request = self._pending_requests.popitem(last=False)
# Mark as in progress.
self._in_progress_requests[request.unique_key] = request
return request
@override
async def get_request(self, unique_key: str) -> Request | None:
await self._update_metadata(update_accessed_at=True)
return (
self._pending_requests.get(unique_key)
or self._in_progress_requests.get(unique_key)
or self._handled_requests.get(unique_key)
)
@override
async def mark_request_as_handled(self, request: Request) -> ProcessedRequest | None:
# Check if the request is in progress.
if request.unique_key not in self._in_progress_requests:
return None
# Set handled_at timestamp if not already set.
if not request.was_already_handled:
request.handled_at = datetime.now(timezone.utc)
# Move request to handled storage.
self._handled_requests[request.unique_key] = request
# Remove from in-progress.
del self._in_progress_requests[request.unique_key]
# Update metadata.
await self._update_metadata(
new_handled_request_count=self._metadata.handled_request_count + 1,
new_pending_request_count=self._metadata.pending_request_count - 1,
update_modified_at=True,
)
return ProcessedRequest(
unique_key=request.unique_key,
was_already_present=True,
was_already_handled=True,
)
@override
async def reclaim_request(
self,
request: Request,
*,
forefront: bool = False,
) -> ProcessedRequest | None:
# Check if the request is in progress.
if request.unique_key not in self._in_progress_requests:
return None
# Remove from in-progress.
del self._in_progress_requests[request.unique_key]
# Add the request back to the pending queue. Unlike a re-add, a reclaim carries the state accumulated
# while the request was in progress, so the reclaimed object supersedes the one that was fetched.
self._pending_requests[request.unique_key] = request
if forefront:
self._pending_requests.move_to_end(request.unique_key, last=False)
# Update metadata timestamps.
await self._update_metadata(update_modified_at=True)
return ProcessedRequest(
unique_key=request.unique_key,
was_already_present=True,
was_already_handled=False,
)
@override
async def is_empty(self) -> bool:
await self._update_metadata(update_accessed_at=True)
# Queue is empty if there are no pending requests.
return len(self._pending_requests) == 0
@override
async def is_finished(self) -> bool:
# Queue is finished if it is empty and there are no in-progress requests.
return await self.is_empty() and len(self._in_progress_requests) == 0
async def _update_metadata(
self,
*,
update_accessed_at: bool = False,
update_modified_at: bool = False,
new_handled_request_count: int | None = None,
new_pending_request_count: int | None = None,
new_total_request_count: int | None = None,
) -> None:
"""Update the request queue metadata with current information.
Args:
update_accessed_at: If True, update the `accessed_at` timestamp to the current time.
update_modified_at: If True, update the `modified_at` timestamp to the current time.
new_handled_request_count: If provided, set the handled request count to this value.
new_pending_request_count: If provided, set the pending request count to this value.
new_total_request_count: If provided, set the total request count to this value.
"""
now = datetime.now(timezone.utc)
if update_accessed_at:
self._metadata.accessed_at = now
if update_modified_at:
self._metadata.modified_at = now
if new_handled_request_count is not None:
self._metadata.handled_request_count = new_handled_request_count
if new_pending_request_count is not None:
self._metadata.pending_request_count = new_pending_request_count
if new_total_request_count is not None:
self._metadata.total_request_count = new_total_request_count