Skip to content

Commit f0d49a3

Browse files
danfimovdanfimov
committed
feat: multiqueue support
1 parent 86160c0 commit f0d49a3

15 files changed

Lines changed: 665 additions & 199 deletions

‎.github/dependabot.yml‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
version: 2
2+
updates:
3+
- package-ecosystem: "uv"
4+
directory: "/"
5+
schedule:
6+
interval: "monthly"
7+
cooldown:
8+
default-days: 30
9+
- package-ecosystem: "github-actions"
10+
directory: "/"
11+
schedule:
12+
interval: "monthly"
13+
cooldown:
14+
default-days: 30

‎README.md‎

Lines changed: 190 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,191 @@
11
# taskiq-mongodb
2-
Taskiq integration for MongoDB
2+
3+
MongoDB broker and result backend for [TaskIQ](https://github.com/taskiq-python/taskiq).
4+
5+
## Installation
6+
7+
```bash
8+
pip install taskiq-mongodb
9+
```
10+
11+
Requires Python 3.11+ and a MongoDB instance (a standalone `mongod` is enough — no replica set needed).
12+
13+
## Basic example
14+
15+
```python
16+
# broker.py
17+
import asyncio
18+
import os
19+
20+
from taskiq_mongodb import MongoBroker, MongoResultBackend
21+
22+
MONGO_URI = os.environ.get("TASKIQ_MONGODB_URI", "mongodb://root:password@localhost:27017")
23+
24+
broker = MongoBroker(MONGO_URI, "my_app").with_result_backend(
25+
MongoResultBackend(MONGO_URI, "my_app"),
26+
)
27+
28+
29+
@broker.task
30+
async def add_one(value: int) -> int:
31+
return value + 1
32+
33+
34+
async def main() -> None:
35+
await broker.startup()
36+
task = await add_one.kiq(1)
37+
result = await task.wait_result(timeout=5)
38+
print(result.return_value) # 2
39+
await broker.shutdown()
40+
41+
42+
if __name__ == "__main__":
43+
asyncio.run(main())
44+
```
45+
46+
Run a worker in one process and the script above in another:
47+
48+
```bash
49+
taskiq worker broker:broker -w 1
50+
python broker.py
51+
```
52+
53+
More runnable examples live in [`examples/`](examples/).
54+
55+
## Result backend
56+
57+
`MongoResultBackend` stores task results (and progress, see below) as documents in a collection, one per task id.
58+
59+
```python
60+
from taskiq_mongodb import MongoResultBackend
61+
62+
result_backend = MongoResultBackend(
63+
"mongodb://root:password@localhost:27017", # pragma: allowlist secret
64+
"my_app",
65+
collection_name="task_results", # default
66+
keep_results=True, # keep the document after get_result reads it
67+
ttl_seconds=3600, # auto-expire results after an hour; 0 disables expiry
68+
)
69+
```
70+
71+
It works standalone too, without a broker — useful when you just need durable storage for results or progress computed
72+
elsewhere. See [`examples/progress.py`](examples/progress.py).
73+
74+
## Broker
75+
76+
`MongoBroker` distributes tasks between worker processes. A worker claims the oldest pending message for a queue with
77+
an atomic `find_one_and_update`, and polls again when the queue is empty. A claimed message that isn't acknowledged
78+
within its queue's visibility timeout (worker crashed or hung) is returned to the queue automatically, or moved to the
79+
"dead" status once the queue's retry limit is exceeded.
80+
81+
```python
82+
from taskiq_mongodb import MongoBroker
83+
84+
broker = MongoBroker(
85+
"mongodb://root:password@localhost:27017", # pragma: allowlist secret
86+
"my_app",
87+
queues="taskiq", # default; see Multiple queues below
88+
collection_name="taskiq_messages", # default
89+
)
90+
```
91+
92+
Acknowledging a message deletes its document, so a healthy queue collection stays small.
93+
94+
### Multiple queues
95+
96+
Pass a queue name, a single queue configuration, or a sequence of either to `queues=`. Each queue is polled by its own
97+
background task, so different queues can have different poll intervals, visibility timeouts and retry limits without
98+
affecting each other:
99+
100+
```python
101+
from taskiq_mongodb import MongoBroker, MongoQueue
102+
103+
broker = MongoBroker(
104+
"mongodb://root:password@localhost:27017", # pragma: allowlist secret
105+
"my_app",
106+
queues=[
107+
"default",
108+
MongoQueue(name="critical", poll_interval=0.1, visibility_timeout=30),
109+
MongoQueue(name="reports", visibility_timeout=600, max_retries=3),
110+
],
111+
)
112+
```
113+
114+
The first queue (`"default"` above) is used for tasks that don't say otherwise. Route a task to a different queue with
115+
the `queue_name` label:
116+
117+
```python
118+
await add_one.kicker().with_labels(queue_name="critical").kiq(1)
119+
```
120+
121+
Kicking a task to a queue that isn't configured on the broker raises `UnknownQueueError` (from `broker.kick()` directly;
122+
through `.kiq()`, as above, TaskIQ wraps it in `SendTaskError` with `UnknownQueueError` as the cause).
123+
124+
### Priority
125+
126+
Set the `priority` label (an integer, default 0) to have a task claimed before lower-priority ones waiting in the same
127+
queue:
128+
129+
```python
130+
await add_one.kicker().with_labels(priority=10).kiq(1)
131+
```
132+
133+
### Delayed messages
134+
135+
Set the `delay` label (seconds) to make a task claimable only after that delay has passed:
136+
137+
```python
138+
await add_one.kicker().with_labels(delay=30).kiq(1) # claimable in 30 seconds
139+
```
140+
141+
### Retries and dead-lettering
142+
143+
`MongoBroker` has two independent, complementary retry mechanisms:
144+
145+
- **Crash recovery** — every queue has a `visibility_timeout` (default 300s) and `max_retries` (default 0).
146+
If a worker claims a message and crashes before acknowledging it, the message becomes claimable again once the timeout
147+
passes. After `max_retries` such claims it's moved to the "dead" status instead of being requeued. `max_retries=0`
148+
(the default) dead-letters after the very first unacknowledged attempt; there's no dedicated "unlimited" value — pass
149+
a very large number instead.
150+
- **Application-level retries** — for a task that raises an exception (as opposed to a worker that crashes),
151+
use TaskIQ's own `SimpleRetryMiddleware` to re-kick it a bounded number of times:
152+
153+
```python
154+
from taskiq.middlewares import SimpleRetryMiddleware
155+
156+
broker = broker.with_middlewares(SimpleRetryMiddleware(default_retry_count=3))
157+
158+
159+
@broker.task(retry_on_error=True, max_retries=3)
160+
async def flaky_task() -> None: ...
161+
```
162+
163+
See [`examples/dead_letter.py`](examples/dead_letter.py) for a full runnable example of the second kind.
164+
165+
## Progress reporting
166+
167+
Tasks can report their own progress while running, using TaskIQ's `ProgressTracker`, and a caller can poll it
168+
through the result backend while the task is still executing:
169+
170+
```python
171+
from taskiq import TaskiqDepends
172+
from taskiq.depends.progress_tracker import ProgressTracker, TaskState
173+
174+
175+
@broker.task
176+
async def process_batch(total: int, tracker: ProgressTracker[int] = TaskiqDepends()) -> str:
177+
for done in range(1, total + 1):
178+
...
179+
await tracker.set_progress(TaskState.STARTED, meta=done)
180+
return f"processed {total} items"
181+
182+
183+
task = await process_batch.kiq(total=5)
184+
progress = await task.get_progress() # -> TaskProgress(state=..., meta=...) or None
185+
```
186+
187+
See [`examples/progress.py`](examples/progress.py) for the full picture, including the polling loop.
188+
189+
## License
190+
191+
MIT — see [LICENSE](LICENSE).

‎examples/dead_letter.py‎

Lines changed: 34 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -1,60 +1,58 @@
11
"""
2-
Example showing how MongoBroker requeues and eventually dead-letters
3-
messages that are claimed but never acknowledged (e.g. a worker crashed
4-
mid-task).
2+
Example of automatic retries with MongoBroker and TaskIQ's SimpleRetryMiddleware.
3+
4+
A task that raises an exception is re-kicked by the middleware, up to the configured retry limit. A task that
5+
keeps failing exhausts its retry budget and taskiq gives up on it (the application-level analogue of a message
6+
landing in a dead-letter queue).
57
68
Requires a MongoDB instance, e.g. `make run_infra` from the repo root.
7-
This script drives the broker directly, without a `taskiq worker` process,
8-
so it can control acknowledgement itself.
99
1010
How to run:
11-
uv run examples/dead_letter.py
11+
1. Run worker: taskiq worker examples.dead_letter:broker -w 1
12+
2. Run client: uv run examples/dead_letter.py
1213
"""
1314

1415
import asyncio
15-
import contextlib
1616
import os
1717

18-
from taskiq import BrokerMessage
18+
from taskiq import Context, TaskiqDepends
19+
from taskiq.middlewares import SimpleRetryMiddleware
1920

20-
from taskiq_mongodb import MongoBroker
21+
from taskiq_mongodb import MongoBroker, MongoResultBackend
2122

2223

2324
MONGO_URI = os.environ.get("TASKIQ_MONGODB_URI", "mongodb://root:password@localhost:27017")
2425

26+
broker = (
27+
MongoBroker(MONGO_URI, "taskiq_example")
28+
.with_result_backend(MongoResultBackend(MONGO_URI, "taskiq_example"))
29+
.with_middlewares(SimpleRetryMiddleware(default_retry_count=3))
30+
)
2531

26-
async def main() -> None:
27-
# Small visibility_timeout/poll_interval and max_retries=2 so the example
28-
# finishes quickly; production values would be minutes, not seconds.
29-
broker = MongoBroker(
30-
MONGO_URI,
31-
"taskiq_example",
32-
poll_interval=0.2,
33-
visibility_timeout=1,
34-
max_retries=2,
35-
)
36-
await broker.startup()
3732

38-
await broker.kick(
39-
BrokerMessage(task_id="demo", task_name="demo:task", message=b"payload", labels={}),
40-
)
33+
@broker.task(retry_on_error=True, max_retries=3)
34+
async def flaky_task(fail_times: int, context: Context = TaskiqDepends()) -> str:
35+
attempt = int(context.message.labels.get("_retries", 0)) + 1
36+
if attempt <= fail_times:
37+
message = f"attempt {attempt} failed on purpose ({fail_times} failures configured)"
38+
raise RuntimeError(message)
39+
return f"succeeded on attempt {attempt}"
4140

42-
listener = broker.listen()
43-
for attempt in range(1, broker.max_retries + 1):
44-
await anext(listener)
45-
print(f"attempt {attempt}: claimed message, deliberately not acknowledging it")
46-
# Simulate a crashed worker: never call `await message.ack()`.
47-
await asyncio.sleep(broker.visibility_timeout + 0.5) # let the claim go stale
4841

49-
# One more poll cycle is what actually notices attempts >= max_retries
50-
# and flips the message to "dead"; nothing is left to claim afterwards.
51-
with contextlib.suppress(TimeoutError):
52-
await asyncio.wait_for(anext(listener), timeout=1)
42+
async def main() -> None:
43+
await broker.startup()
44+
45+
# Fails twice, then succeeds within the max_retries=3 budget.
46+
recovers = await flaky_task.kiq(fail_times=2)
47+
result = await recovers.wait_result(timeout=10)
48+
print(f"recovers after retries: is_err={result.is_err} value={result.return_value!r}")
5349

54-
doc = await broker.col.find_one({})
55-
print(f"final status after {broker.max_retries} unacknowledged attempts: {doc['status']}")
50+
# Fails more times than max_retries allows: retries are exhausted and taskiq gives up on the task, same idea
51+
# as dead-lettering a message.
52+
gives_up = await flaky_task.kiq(fail_times=10)
53+
result = await gives_up.wait_result(timeout=10)
54+
print(f"exhausts retries: is_err={result.is_err}")
5655

57-
await listener.aclose()
5856
await broker.shutdown()
5957

6058

‎examples/progress.py‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
"""
2+
Example of task progress reporting: a task updates its own progress via ProgressTracker while it runs, and the client
3+
polls it through MongoResultBackend while the worker is still executing the task.
4+
5+
Requires a MongoDB instance, e.g. `make run_infra` from the repo root.
6+
7+
How to run:
8+
1. Run worker: taskiq worker examples.progress:broker -w 1
9+
2. Run client: uv run examples/progress.py
10+
"""
11+
12+
import asyncio
13+
import os
14+
15+
from taskiq import TaskiqDepends
16+
from taskiq.depends.progress_tracker import ProgressTracker, TaskState
17+
18+
from taskiq_mongodb import MongoBroker, MongoResultBackend
19+
20+
21+
MONGO_URI = os.environ.get("TASKIQ_MONGODB_URI", "mongodb://root:password@localhost:27017")
22+
23+
broker = MongoBroker(MONGO_URI, "taskiq_example").with_result_backend(
24+
MongoResultBackend(MONGO_URI, "taskiq_example"),
25+
)
26+
27+
28+
@broker.task
29+
async def process_batch(total: int, tracker: ProgressTracker[int] = TaskiqDepends()) -> str:
30+
for done in range(1, total + 1):
31+
await asyncio.sleep(0.5)
32+
await tracker.set_progress(TaskState.STARTED, meta=done)
33+
return f"processed {total} items"
34+
35+
36+
async def main() -> None:
37+
await broker.startup()
38+
39+
task = await process_batch.kiq(total=5)
40+
41+
last_meta = None
42+
while not await task.is_ready():
43+
progress = await task.get_progress()
44+
if progress is not None and progress.meta != last_meta:
45+
last_meta = progress.meta
46+
print(f"progress: state={progress.state} meta={progress.meta}")
47+
await asyncio.sleep(0.2)
48+
49+
result = await task.get_result()
50+
print(f"result: {result.return_value}")
51+
52+
await broker.shutdown()
53+
54+
55+
if __name__ == "__main__":
56+
asyncio.run(main())

0 commit comments

Comments
 (0)