vercel-dramatiq 0.7.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,17 @@
1
+ """Dramatiq integration for Vercel Queue Service."""
2
+
3
+ from ._broker import (
4
+ VercelQueueBroker,
5
+ install_vercel_dramatiq_integration,
6
+ register_dramatiq_queues,
7
+ )
8
+ from ._result_backend import VercelRuntimeCacheBackend
9
+ from .version import __version__
10
+
11
+ __all__ = [
12
+ "VercelQueueBroker",
13
+ "VercelRuntimeCacheBackend",
14
+ "__version__",
15
+ "install_vercel_dramatiq_integration",
16
+ "register_dramatiq_queues",
17
+ ]
@@ -0,0 +1,587 @@
1
+ from __future__ import annotations
2
+
3
+ from typing import Any
4
+ from typing_extensions import TypedDict, Unpack
5
+
6
+ import logging
7
+ import os
8
+ import queue
9
+ import threading
10
+ import time
11
+ from collections.abc import Iterable, Mapping
12
+ from dataclasses import dataclass
13
+ from weakref import WeakKeyDictionary
14
+
15
+ import dramatiq
16
+ import dramatiq.broker as dramatiq_broker
17
+ import vercel.cache as vcache
18
+ import vercel.queue as vqs
19
+ import vercel.queue.sync as vqs_sync
20
+ from dramatiq.broker import Broker, Consumer, MessageProxy
21
+ from dramatiq.common import current_millis, dq_name
22
+ from dramatiq.errors import QueueNotFound
23
+ from dramatiq.message import Message as DramatiqMessage
24
+ from dramatiq.results import Results
25
+ from dramatiq.results.backend import ResultBackend
26
+ from dramatiq.worker import Worker
27
+
28
+ from ._result_backend import VercelRuntimeCacheBackend
29
+ from .version import __version__
30
+
31
+ DEFAULT_CONSUMER_GROUP = "dramatiq"
32
+ DEFAULT_REQUEUE_DELAY_SECONDS = 0
33
+ DEFAULT_PUSH_RETRY_DELAY_SECONDS = 1
34
+ DEFAULT_PUSH_HANDOFF_WAIT_SECONDS = 30.0
35
+ _PUSH_WAIT_POLL_INTERVAL_SECONDS = 0.05
36
+ _PUSH_SETTLE_POLL_INTERVAL_SECONDS = 0.01
37
+ _DEBUG_ENV = "VERCEL_DRAMATIQ_DEBUG"
38
+ _DEBUG_LOGGER_NAMES = (
39
+ "vercel.integrations.dramatiq",
40
+ "dramatiq",
41
+ "dramatiq.broker",
42
+ "dramatiq.worker",
43
+ )
44
+
45
+
46
+ def _is_vercel_runtime() -> bool:
47
+ try:
48
+ value = os.environ["VERCEL"]
49
+ except KeyError:
50
+ return False
51
+ return value.strip().casefold() in {"1", "yes", "on", "true"}
52
+
53
+
54
+ def _dramatiq_debug_enabled() -> bool:
55
+ return os.environ.get(_DEBUG_ENV, "").strip().casefold() in {"1", "yes", "on", "true"}
56
+
57
+
58
+ def _configure_dramatiq_debug_logging() -> None:
59
+ if not _dramatiq_debug_enabled():
60
+ return
61
+ for logger_name in _DEBUG_LOGGER_NAMES:
62
+ logging.getLogger(logger_name).setLevel(logging.DEBUG)
63
+
64
+
65
+ class VercelQueueBrokerOptions(TypedDict, total=False):
66
+ token: str | None
67
+ """Vercel API token used by the underlying Vercel Queue client."""
68
+
69
+ region: str | None
70
+ """Vercel region used by the underlying Vercel Queue client."""
71
+
72
+ base_url: str | None
73
+ """Override base URL for the underlying Vercel Queue client."""
74
+
75
+ deployment: vqs.DeploymentOption
76
+ """Deployment partition used when sending and receiving Vercel Queue messages."""
77
+
78
+ timeout: vqs.Duration | None
79
+ """Request timeout used by the underlying Vercel Queue client."""
80
+
81
+ headers: Mapping[str, str] | None
82
+ """Additional headers sent by the underlying Vercel Queue client."""
83
+
84
+ consumer_group: str
85
+ """Vercel Queue consumer group used for subscriptions and polling."""
86
+
87
+ retention: vqs.Duration | None
88
+ """Optional retention duration applied to enqueued Vercel Queue messages."""
89
+
90
+ lease_duration: vqs.Duration | None
91
+ """Optional lease duration for polled or pushed Vercel Queue messages."""
92
+
93
+ requeue_delay_seconds: int
94
+ """Visibility delay used when Dramatiq requeues or rejects a message."""
95
+
96
+ push_retry_delay_seconds: int
97
+ """Visibility delay used when push delivery cannot be handed to a worker."""
98
+
99
+ push_handoff_wait_seconds: float
100
+ """Maximum request-time wait for push worker readiness and settlement."""
101
+
102
+ queue_name_prefix: str | None
103
+ """Prefix applied to Dramatiq queue names before VQS topic sanitization."""
104
+
105
+ use_message_id_as_idempotency_key: bool
106
+ """Use Dramatiq message IDs as VQS idempotency keys when publishing."""
107
+
108
+ poll: bool
109
+ """Force polling mode when true or push mode when false."""
110
+
111
+ middleware: list[Any] | None
112
+ """Dramatiq middleware list passed to the base broker."""
113
+
114
+
115
+ @dataclass
116
+ class _TrackedMessage:
117
+ message: vqs.Message[bytes]
118
+ lease_renewal: vqs.LeaseRenewal
119
+ settlement: threading.Event
120
+
121
+
122
+ class _VercelQueueMessageProxy(MessageProxy):
123
+ def __init__(
124
+ self,
125
+ message: DramatiqMessage[Any],
126
+ *,
127
+ queue_name: str,
128
+ delivery: vqs.Message[bytes],
129
+ lease_renewal: vqs.LeaseRenewal,
130
+ ) -> None:
131
+ super().__init__(message)
132
+ self.settlement = threading.Event()
133
+ self._queue_name = queue_name
134
+ self._tracked = _TrackedMessage(
135
+ message=delivery,
136
+ lease_renewal=lease_renewal,
137
+ settlement=self.settlement,
138
+ )
139
+
140
+ def stop_lease_renewal(self) -> None:
141
+ self._tracked.lease_renewal.stop()
142
+
143
+ @property
144
+ def vqs_message(self) -> vqs.Message[bytes]:
145
+ return self._tracked.message
146
+
147
+
148
+ class _VercelQueueConsumer(Consumer):
149
+ def __init__(
150
+ self,
151
+ broker: VercelQueueBroker,
152
+ queue_name: str,
153
+ prefetch: int,
154
+ timeout: int,
155
+ ) -> None:
156
+ self.broker = broker
157
+ self.queue_name = queue_name
158
+ self.prefetch = prefetch
159
+ self.timeout = timeout
160
+ self._slots = threading.BoundedSemaphore(prefetch)
161
+ self._push_queue: queue.Queue[_VercelQueueMessageProxy] = queue.Queue()
162
+ self._closed = False
163
+ self._lock = threading.Lock()
164
+
165
+ def ack(self, message: MessageProxy) -> None:
166
+ proxy = self._as_proxy(message)
167
+ try:
168
+ self.broker.acknowledge_message(proxy.vqs_message)
169
+ finally:
170
+ self._finish_message(proxy)
171
+
172
+ def nack(self, message: MessageProxy) -> None:
173
+ proxy = self._as_proxy(message)
174
+ try:
175
+ self.broker.acknowledge_message(proxy.vqs_message)
176
+ self.broker.dead_letters_by_queue[self.queue_name].append(proxy)
177
+ finally:
178
+ self._finish_message(proxy)
179
+
180
+ def requeue(self, messages: Iterable[MessageProxy]) -> None:
181
+ for message in messages:
182
+ proxy = self._as_proxy(message)
183
+ try:
184
+ self.broker.extend_message_lease(proxy.vqs_message)
185
+ finally:
186
+ self._finish_message(proxy)
187
+
188
+ def __next__(self) -> MessageProxy | None:
189
+ if self._closed:
190
+ return None
191
+ if not self._slots.acquire(timeout=self.timeout / 1000):
192
+ return None
193
+
194
+ try:
195
+ delivery = self._next_delivery()
196
+ if delivery is None:
197
+ self._slots.release()
198
+ return None
199
+ if isinstance(delivery, _VercelQueueMessageProxy):
200
+ return delivery
201
+ return self._wrap_delivery(delivery)
202
+ except Exception:
203
+ self._slots.release()
204
+ raise
205
+
206
+ def close(self) -> None:
207
+ with self._lock:
208
+ if self._closed:
209
+ return
210
+ self._closed = True
211
+ self.broker.remove_consumer(self)
212
+
213
+ @property
214
+ def closed(self) -> bool:
215
+ return self._closed
216
+
217
+ def can_accept_push(self) -> bool:
218
+ return not self._closed and self._push_queue.qsize() < self.prefetch
219
+
220
+ def put_push_delivery(self, delivery: vqs.Message[bytes]) -> _VercelQueueMessageProxy | None:
221
+ if not self.can_accept_push():
222
+ return None
223
+ proxy = self._wrap_delivery(delivery)
224
+ self._push_queue.put(proxy)
225
+ return proxy
226
+
227
+ def _next_delivery(self) -> _VercelQueueMessageProxy | vqs.Message[bytes] | None:
228
+ try:
229
+ return self._push_queue.get_nowait()
230
+ except queue.Empty:
231
+ pass
232
+
233
+ if not self.broker.poll:
234
+ try:
235
+ return self._push_queue.get(timeout=self.timeout / 1000)
236
+ except queue.Empty:
237
+ return None
238
+
239
+ for delivery in self.broker.poll_messages(self.queue_name, self.prefetch):
240
+ return delivery
241
+ return None
242
+
243
+ def _wrap_delivery(self, delivery: vqs.Message[bytes]) -> _VercelQueueMessageProxy:
244
+ message = DramatiqMessage.decode(delivery.payload)
245
+ lease_renewal = self.broker.start_lease_renewal(delivery)
246
+ return _VercelQueueMessageProxy(
247
+ message,
248
+ queue_name=self.queue_name,
249
+ delivery=delivery,
250
+ lease_renewal=lease_renewal,
251
+ )
252
+
253
+ def _finish_message(self, message: _VercelQueueMessageProxy) -> None:
254
+ try:
255
+ message.stop_lease_renewal()
256
+ finally:
257
+ message.settlement.set()
258
+ self._slots.release()
259
+
260
+ @staticmethod
261
+ def _as_proxy(message: MessageProxy) -> _VercelQueueMessageProxy:
262
+ if not isinstance(message, _VercelQueueMessageProxy):
263
+ raise TypeError("message was not produced by this consumer")
264
+ return message
265
+
266
+
267
+ class VercelQueueBroker(Broker):
268
+ def __init__(
269
+ self,
270
+ **options: Unpack[VercelQueueBrokerOptions],
271
+ ) -> None:
272
+ """Create a Dramatiq broker backed by Vercel Queue Service.
273
+
274
+ Queue client options such as ``token``, ``region``, ``base_url``,
275
+ ``deployment``, ``timeout``, and ``headers`` are forwarded to the
276
+ underlying Vercel Queue client. Broker-specific options control the
277
+ consumer group, message retention, lease duration, requeue delay, and
278
+ Dramatiq middleware.
279
+
280
+ By default, the broker uses push delivery when ``VERCEL`` is truthy in
281
+ the environment and poll delivery otherwise. Pass ``poll=False`` or
282
+ ``poll=True`` to force a delivery mode.
283
+ """
284
+ token = options.get("token")
285
+ region = options.get("region")
286
+ base_url = options.get("base_url")
287
+ deployment = options.get("deployment", vqs.CURRENT_DEPLOYMENT)
288
+ timeout = options.get("timeout", 10.0)
289
+ headers = options.get("headers")
290
+ middleware = options.get("middleware")
291
+ super().__init__(middleware=middleware)
292
+ _configure_dramatiq_debug_logging()
293
+ self.queues: dict[str, object] = {}
294
+ self.consumer_group = options.get("consumer_group", DEFAULT_CONSUMER_GROUP)
295
+ self.retention = options.get("retention")
296
+ self.lease_duration = options.get("lease_duration")
297
+ self.requeue_delay_seconds = options.get(
298
+ "requeue_delay_seconds",
299
+ DEFAULT_REQUEUE_DELAY_SECONDS,
300
+ )
301
+ self.push_retry_delay_seconds = options.get(
302
+ "push_retry_delay_seconds",
303
+ DEFAULT_PUSH_RETRY_DELAY_SECONDS,
304
+ )
305
+ self.push_handoff_wait_seconds = options.get(
306
+ "push_handoff_wait_seconds",
307
+ DEFAULT_PUSH_HANDOFF_WAIT_SECONDS,
308
+ )
309
+ self.queue_name_prefix = options.get("queue_name_prefix") or ""
310
+ self.use_message_id_as_idempotency_key = options.get(
311
+ "use_message_id_as_idempotency_key",
312
+ False,
313
+ )
314
+ self.poll = options.get("poll", not _is_vercel_runtime())
315
+ self.dead_letters_by_queue: dict[str, list[MessageProxy]] = {}
316
+ self._queue_client = vqs_sync.QueueClient(
317
+ token=token,
318
+ region=region,
319
+ base_url=base_url,
320
+ deployment=deployment,
321
+ headers=headers,
322
+ timeout=timeout,
323
+ )
324
+ self._consumers: list[_VercelQueueConsumer] = []
325
+ self._consumers_lock = threading.RLock()
326
+ self._registered_callbacks: dict[str, Any] = {}
327
+ self._push_handoff_lock = threading.RLock()
328
+
329
+ @property
330
+ def dead_letters(self) -> list[MessageProxy]:
331
+ return [message for messages in self.dead_letters_by_queue.values() for message in messages]
332
+
333
+ def consume(self, queue_name: str, prefetch: int = 1, timeout: int = 30000) -> Consumer:
334
+ if queue_name not in self.queues:
335
+ raise QueueNotFound(queue_name)
336
+ consumer = _VercelQueueConsumer(self, queue_name, prefetch, timeout)
337
+ with self._consumers_lock:
338
+ self._consumers.append(consumer)
339
+ return consumer
340
+
341
+ def declare_queue(self, queue_name: str) -> None:
342
+ if queue_name in self.queues:
343
+ return
344
+
345
+ self.emit_before("declare_queue", queue_name)
346
+ self.queues[queue_name] = object()
347
+ self.dead_letters_by_queue.setdefault(queue_name, [])
348
+ self.emit_after("declare_queue", queue_name)
349
+ self._register_queue_callback(queue_name)
350
+
351
+ delayed_name = dq_name(queue_name)
352
+ self.queues[delayed_name] = object()
353
+ self.dead_letters_by_queue.setdefault(delayed_name, [])
354
+ self.delay_queues.add(delayed_name)
355
+ self.emit_after("declare_delay_queue", delayed_name)
356
+ self._register_queue_callback(delayed_name)
357
+
358
+ def enqueue(
359
+ self,
360
+ message: DramatiqMessage[Any],
361
+ *,
362
+ delay: int | None = None,
363
+ ) -> DramatiqMessage[Any]:
364
+ queue_name = message.queue_name
365
+ vqs_delay: vqs.Duration | None = None
366
+ if delay is not None:
367
+ queue_name = dq_name(queue_name)
368
+ message = message.copy(
369
+ queue_name=queue_name,
370
+ options={"eta": current_millis() + delay},
371
+ )
372
+ vqs_delay = delay / 1000
373
+
374
+ if queue_name not in self.queues:
375
+ raise QueueNotFound(queue_name)
376
+
377
+ self.emit_before("enqueue", message, delay)
378
+ self._queue_client.send(
379
+ self.topic_for_queue(queue_name),
380
+ message.encode(),
381
+ idempotency_key=(
382
+ message.message_id if self.use_message_id_as_idempotency_key else None
383
+ ),
384
+ retention=self.retention,
385
+ delay=vqs_delay,
386
+ )
387
+ self.emit_after("enqueue", message, delay)
388
+ return message
389
+
390
+ def close(self) -> None:
391
+ for consumer in list(self._consumers):
392
+ consumer.close()
393
+
394
+ def flush(self, queue_name: str) -> None:
395
+ del queue_name
396
+ raise NotImplementedError("Vercel Queue Service does not support queue purge")
397
+
398
+ def flush_all(self) -> None:
399
+ raise NotImplementedError("Vercel Queue Service does not support queue purge")
400
+
401
+ def join(self, queue_name: str, *, timeout: int | None = None) -> None:
402
+ del queue_name, timeout
403
+ raise NotImplementedError("Vercel Queue Service does not support queue join")
404
+
405
+ def acknowledge_message(self, message: vqs.Message[bytes]) -> None:
406
+ self._queue_client.acknowledge(message)
407
+
408
+ def extend_message_lease(self, message: vqs.Message[bytes]) -> None:
409
+ self._queue_client.retry_after(message, self.requeue_delay_seconds)
410
+
411
+ def poll_messages(self, queue_name: str, prefetch: int) -> Iterable[vqs.Message[bytes]]:
412
+ del prefetch
413
+ deliveries = self._queue_client.poll(
414
+ self.topic_for_queue(queue_name),
415
+ self.consumer_group,
416
+ limit=1,
417
+ lease_duration=self.lease_duration,
418
+ )
419
+ for delivery in deliveries:
420
+ yield delivery.accept()
421
+
422
+ def start_lease_renewal(self, message: vqs.Message[bytes]) -> vqs.LeaseRenewal:
423
+ lease_renewal = self._queue_client.run_lease_renewal(
424
+ message,
425
+ lease_duration=self.lease_duration,
426
+ )
427
+ lease_renewal.start()
428
+ return lease_renewal
429
+
430
+ def topic_for_queue(self, queue_name: str) -> vqs.Topic[bytes]:
431
+ return vqs.Topic[bytes](
432
+ vqs.sanitize_name(f"{self.queue_name_prefix}{queue_name}"),
433
+ )
434
+
435
+ def remove_consumer(self, consumer: _VercelQueueConsumer) -> None:
436
+ with self._consumers_lock:
437
+ if consumer in self._consumers:
438
+ self._consumers.remove(consumer)
439
+
440
+ def handle_push_message(self, delivery: vqs.Message[bytes]) -> None:
441
+ queue_name = self._queue_for_topic(delivery.metadata.topic)
442
+ if queue_name is None:
443
+ raise vqs.RetryAfter(self.push_retry_delay_seconds)
444
+
445
+ deadline = time.monotonic() + max(self.push_handoff_wait_seconds, 0.0)
446
+ if not self._acquire_push_handoff_lock(deadline):
447
+ raise vqs.RetryAfter(self.push_retry_delay_seconds)
448
+ try:
449
+ proxy = self._handoff_push_delivery(queue_name, delivery, deadline)
450
+ if proxy is None:
451
+ raise vqs.RetryAfter(self.push_retry_delay_seconds)
452
+ finally:
453
+ self._push_handoff_lock.release()
454
+ self._wait_for_push_settlement(proxy)
455
+ raise vqs.Handoff
456
+
457
+ def _acquire_push_handoff_lock(self, deadline: float) -> bool:
458
+ timeout = deadline - time.monotonic()
459
+ if timeout <= 0:
460
+ return self._push_handoff_lock.acquire(blocking=False)
461
+ return self._push_handoff_lock.acquire(timeout=timeout)
462
+
463
+ def _handoff_push_delivery(
464
+ self,
465
+ queue_name: str,
466
+ delivery: vqs.Message[bytes],
467
+ deadline: float,
468
+ ) -> _VercelQueueMessageProxy | None:
469
+ while True:
470
+ with self._consumers_lock:
471
+ consumers = tuple(self._consumers)
472
+ for consumer in reversed(consumers):
473
+ if consumer.queue_name != queue_name:
474
+ continue
475
+ proxy = consumer.put_push_delivery(delivery)
476
+ if proxy is not None:
477
+ return proxy
478
+ if time.monotonic() >= deadline:
479
+ return None
480
+ time.sleep(_PUSH_WAIT_POLL_INTERVAL_SECONDS)
481
+
482
+ def _wait_for_push_settlement(
483
+ self,
484
+ proxy: _VercelQueueMessageProxy,
485
+ ) -> None:
486
+ proxy.settlement.wait()
487
+
488
+ def _queue_for_topic(self, topic: str) -> str | None:
489
+ for queue_name in self.queues:
490
+ if self.topic_for_queue(queue_name).name == topic:
491
+ return queue_name
492
+ return None
493
+
494
+ def register_queue_callbacks(self) -> None:
495
+ queue_names = sorted(
496
+ self.get_declared_queues(),
497
+ key=lambda queue_name: (queue_name.endswith(".DQ"), queue_name),
498
+ )
499
+ for queue_name in queue_names:
500
+ self._register_queue_callback(queue_name)
501
+
502
+ def _register_queue_callback(self, queue_name: str) -> None:
503
+ if queue_name in self._registered_callbacks:
504
+ return
505
+ topic = self.topic_for_queue(queue_name)
506
+
507
+ def handle_queue_delivery(
508
+ message: vqs.Message[bytes],
509
+ *,
510
+ broker: VercelQueueBroker = self,
511
+ ) -> None:
512
+ vcache.prime_runtime_cache()
513
+ broker.handle_push_message(message)
514
+
515
+ handle_queue_delivery.__name__ = f"vercel_dramatiq_{topic.name}_subscriber"
516
+ vqs.subscribe(
517
+ topic=topic,
518
+ consumer_group=self.consumer_group,
519
+ retry_after=self.push_retry_delay_seconds,
520
+ )(handle_queue_delivery)
521
+ self._registered_callbacks[queue_name] = handle_queue_delivery
522
+
523
+
524
+ @dataclass
525
+ class _EmbeddedWorker:
526
+ broker: VercelQueueBroker
527
+ worker: Worker
528
+
529
+
530
+ _embedded_workers: WeakKeyDictionary[VercelQueueBroker, _EmbeddedWorker] = WeakKeyDictionary()
531
+ _embedded_workers_lock = threading.RLock()
532
+
533
+
534
+ def register_dramatiq_queues(
535
+ *,
536
+ broker: VercelQueueBroker | None = None,
537
+ start_worker: bool = True,
538
+ ) -> None:
539
+ """Register declared Dramatiq queues as VQS subscribers for push delivery."""
540
+ resolved_broker = broker or dramatiq.get_broker()
541
+ if not isinstance(resolved_broker, VercelQueueBroker):
542
+ raise TypeError("Dramatiq queue registration requires VercelQueueBroker")
543
+ broker = resolved_broker
544
+ broker.register_queue_callbacks()
545
+ if start_worker:
546
+ _start_embedded_worker(broker)
547
+
548
+
549
+ def _start_embedded_worker(broker: VercelQueueBroker) -> None:
550
+ _configure_dramatiq_debug_logging()
551
+ with _embedded_workers_lock:
552
+ if broker in _embedded_workers:
553
+ return
554
+ queues = broker.get_declared_queues() - broker.get_declared_delay_queues()
555
+ worker = Worker(broker, queues=queues, worker_threads=1, worker_timeout=100)
556
+ worker.start()
557
+ _embedded_workers[broker] = _EmbeddedWorker(broker=broker, worker=worker)
558
+
559
+
560
+ def install_vercel_dramatiq_integration(
561
+ *,
562
+ set_default_broker: bool = True,
563
+ install_results_backend: bool = True,
564
+ results_backend: ResultBackend | None = None,
565
+ **broker_options: Unpack[VercelQueueBrokerOptions],
566
+ ) -> None:
567
+ """Install Vercel Queue Service as Dramatiq's default broker.
568
+
569
+ When ``set_default_broker`` is true, this sets Dramatiq's global broker to
570
+ a new ``VercelQueueBroker`` only if no broker has been configured yet. Any
571
+ additional keyword arguments are passed to ``VercelQueueBroker``; use the
572
+ ``poll`` option there to override automatic push/poll selection.
573
+ """
574
+ if set_default_broker and dramatiq_broker.global_broker is None:
575
+ broker = VercelQueueBroker(**broker_options)
576
+ if install_results_backend:
577
+ broker.add_middleware(Results(backend=results_backend or VercelRuntimeCacheBackend()))
578
+ dramatiq.set_broker(broker)
579
+
580
+
581
+ __all__ = [
582
+ "VercelQueueBroker",
583
+ "VercelRuntimeCacheBackend",
584
+ "__version__",
585
+ "install_vercel_dramatiq_integration",
586
+ "register_dramatiq_queues",
587
+ ]
@@ -0,0 +1,95 @@
1
+ from __future__ import annotations
2
+
3
+ from typing import Any, cast
4
+
5
+ import base64
6
+ from collections.abc import Callable
7
+
8
+ from dramatiq.results.backend import Missing, ResultBackend
9
+ from vercel.cache import RuntimeCache
10
+ from vercel.queue import sanitize_name
11
+
12
+ DEFAULT_RESULT_NAMESPACE = "dramatiq-results"
13
+ _WRAPPER_MARKER = "__vercel_dramatiq_result__"
14
+ _WRAPPER_VERSION = 1
15
+
16
+
17
+ class VercelRuntimeCacheBackend(ResultBackend):
18
+ """Dramatiq result backend backed by Vercel Runtime Cache."""
19
+
20
+ def __init__(
21
+ self,
22
+ *,
23
+ namespace: str = DEFAULT_RESULT_NAMESPACE,
24
+ runtime_cache_namespace: str | None = None,
25
+ namespace_separator: str | None = None,
26
+ key_hash_function: Callable[[str], str] | None = None,
27
+ name: str | None = None,
28
+ tags: list[str] | None = None,
29
+ **kwargs: Any,
30
+ ) -> None:
31
+ super().__init__(namespace=str(sanitize_name(namespace)), **kwargs)
32
+ self.runtime_cache_namespace = runtime_cache_namespace or self.namespace
33
+ self.name = name
34
+ self.tags = tags
35
+ self._runtime_cache = RuntimeCache(
36
+ key_hash_function=key_hash_function,
37
+ namespace=self.runtime_cache_namespace,
38
+ namespace_separator=namespace_separator,
39
+ strict=True,
40
+ )
41
+
42
+ def _get(self, message_key: str) -> object:
43
+ value = self._runtime_cache.get(message_key)
44
+ if value is None:
45
+ return Missing
46
+ return self.encoder.decode(_unwrap_value(value))
47
+
48
+ def _store(self, message_key: str, result: object, ttl: int) -> None:
49
+ self._runtime_cache.set(
50
+ message_key,
51
+ _wrap_value(self.encoder.encode(cast("dict[str, Any]", result))),
52
+ self._set_options(message_key, ttl),
53
+ )
54
+
55
+ def _set_options(self, message_key: str, ttl: int) -> dict[str, object]:
56
+ options: dict[str, object] = {"name": self.name or message_key}
57
+ ttl_seconds = _ttl_milliseconds_to_seconds(ttl)
58
+ if ttl_seconds is not None:
59
+ options["ttl"] = ttl_seconds
60
+ if self.tags:
61
+ options["tags"] = self.tags
62
+ return options
63
+
64
+
65
+ def _ttl_milliseconds_to_seconds(ttl: int) -> int | None:
66
+ if ttl <= 0:
67
+ return None
68
+ return max(1, int(ttl / 1000))
69
+
70
+
71
+ def _wrap_value(value: bytes) -> dict[str, object]:
72
+ return {
73
+ _WRAPPER_MARKER: _WRAPPER_VERSION,
74
+ "encoding": "base64",
75
+ "payload": base64.b64encode(value).decode("ascii"),
76
+ }
77
+
78
+
79
+ def _unwrap_value(value: object) -> bytes:
80
+ if not isinstance(value, dict):
81
+ raise TypeError("Runtime Cache result payload is not a wrapper object")
82
+ if value.get(_WRAPPER_MARKER) != _WRAPPER_VERSION:
83
+ raise ValueError("Runtime Cache result payload has an unknown wrapper version")
84
+ if value.get("encoding") != "base64":
85
+ raise ValueError("Runtime Cache result payload has an unknown encoding")
86
+ payload = value.get("payload")
87
+ if not isinstance(payload, str):
88
+ raise TypeError("Runtime Cache result payload is not text")
89
+ return base64.b64decode(payload.encode("ascii"), validate=True)
90
+
91
+
92
+ __all__ = [
93
+ "DEFAULT_RESULT_NAMESPACE",
94
+ "VercelRuntimeCacheBackend",
95
+ ]
File without changes
@@ -0,0 +1,3 @@
1
+ """Package version metadata."""
2
+
3
+ __version__ = "0.7.0"
@@ -0,0 +1,47 @@
1
+ Metadata-Version: 2.4
2
+ Name: vercel-dramatiq
3
+ Version: 0.7.0
4
+ Summary: Dramatiq broker backed by Vercel Queue Service
5
+ License-Expression: MIT
6
+ License-File: LICENSE
7
+ Requires-Python: >=3.10
8
+ Requires-Dist: dramatiq<3,>=2.2
9
+ Requires-Dist: vercel-cache>=0.6.0
10
+ Requires-Dist: vercel-queue>=0.6.0
11
+ Description-Content-Type: text/markdown
12
+
13
+ # vercel-dramatiq
14
+
15
+ Dramatiq broker backed by Vercel Queue Service.
16
+
17
+ ```python
18
+ import dramatiq
19
+
20
+ from vercel.integrations.dramatiq import install_vercel_dramatiq_integration
21
+
22
+ install_vercel_dramatiq_integration()
23
+
24
+
25
+ @dramatiq.actor
26
+ def send_email(user_id: str) -> None: ...
27
+ ```
28
+
29
+ The installer sets Dramatiq's global broker to `VercelQueueBroker` when no
30
+ broker has been configured yet. Pass `set_default_broker=False` to opt out, or
31
+ call `dramatiq.set_broker(VercelQueueBroker(...))` directly when you need to
32
+ construct the broker yourself.
33
+
34
+ The broker uses push delivery when `VERCEL` is truthy in the environment and
35
+ poll delivery otherwise. Pass `poll=False` or `poll=True` to
36
+ `install_vercel_dramatiq_integration(...)` to force a mode, or pass the same
37
+ option to `VercelQueueBroker(...)` when constructing the broker yourself.
38
+
39
+ For Vercel push delivery, configure queue triggers for each declared queue and
40
+ its Dramatiq delay queue. For example, the `default` Dramatiq queue maps to both
41
+ the `default` topic and the sanitized `default_DDQ` delay topic. Retries and
42
+ delayed messages are delivered through the delay queue topic.
43
+
44
+ Set `VERCEL_DRAMATIQ_DEBUG=1` to enable debug logging for the integration and
45
+ Dramatiq worker loggers.
46
+
47
+ The package is standalone and depends on `vercel-queue` and Dramatiq.
@@ -0,0 +1,9 @@
1
+ vercel/integrations/dramatiq/__init__.py,sha256=0aHAhN8PTBnSbPSDIYCOBJqu-CNzilNhfnySHXvX_nE,428
2
+ vercel/integrations/dramatiq/_broker.py,sha256=pKyP8dP4kDhQyQAdXu1oV-wwyAoRaI6eqS_kWCwtYQs,20575
3
+ vercel/integrations/dramatiq/_result_backend.py,sha256=nwhRAfUoZ33sM_XW4Y_EDUtIviP4XC_LnbnSqVyjF5M,3198
4
+ vercel/integrations/dramatiq/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
5
+ vercel/integrations/dramatiq/version.py,sha256=-utTJDl5QfoT7s0Pz6Z9PS818EvYbA3bzu0eeXPCwvg,55
6
+ vercel_dramatiq-0.7.0.dist-info/METADATA,sha256=DiQdt0FDHsUaeVFnt_8UN0QTfhlPElyf0ssagWc2Wzk,1631
7
+ vercel_dramatiq-0.7.0.dist-info/WHEEL,sha256=lCkmxWfQsSc9CfIClYeavTdQeEX2toPqufh9gI35EQA,87
8
+ vercel_dramatiq-0.7.0.dist-info/licenses/LICENSE,sha256=G9LJlKrxX_SueoCDXBJxLkTGol_9OOGHk3wrxG0vfUA,1058
9
+ vercel_dramatiq-0.7.0.dist-info/RECORD,,
@@ -0,0 +1,4 @@
1
+ Wheel-Version: 1.0
2
+ Generator: hatchling 1.31.0
3
+ Root-Is-Purelib: true
4
+ Tag: py3-none-any
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) Vercel
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.