Documentation Example: Task Polling

[skwzrd]

Since a lot of users might be coming here from a Celery background where it's common to submit a task and poll it's results, it'd be good to demonstrate this with the following, somewhere in the docs.

import asyncio
import time

import dramatiq
from dramatiq.brokers.redis import RedisBroker
from dramatiq.middleware import AsyncIO
from dramatiq.results import Results
from dramatiq.results.backend import DEFAULT_TIMEOUT
from dramatiq.results.backends.redis import RedisBackend
from dramatiq.results.errors import ResultMissing, ResultTimeout
from quart import Quart, jsonify


class RedisBackendExt(RedisBackend):
    def get_result_from_message_key(self, message_key, *, block=False, timeout=None):
        if timeout is None:
            timeout = DEFAULT_TIMEOUT

        if block:
            timeout = int(timeout / 1000)
            if timeout == 0:
                data = self.client.rpoplpush(message_key, message_key)
            else:
                data = self.client.brpoplpush(message_key, message_key, timeout)

            if data is None:
                raise ResultTimeout(message_key)

        else:
            data = self.client.lindex(message_key, 0)
            if data is None:
                raise ResultMissing(message_key)

        return self.unwrap_result(self.encoder.decode(data))


backend = RedisBackendExt()
broker = RedisBroker(host='localhost', port=6379)
broker.add_middleware(Results(backend=backend))
broker.add_middleware(AsyncIO())
dramatiq.set_broker(broker)

app = Quart(__name__) # totally decoupled from dramatiq


@dramatiq.actor(store_results=True)
async def cut_granite_with_toothpick(enqueue_time: float):
    print(f'{enqueue_time=}')
    await asyncio.sleep(4)
    return 'finally, the granite is cut...'


@app.get('/message/<string:message_key>')
async def get_result(message_key):
    result = backend.get_result_from_message_key(message_key)
    return jsonify({'result': result}), 200


@app.get('/')
async def index():
    enqueue_time = time.time()
    message = cut_granite_with_toothpick.send(enqueue_time)
    message_key = backend.build_message_key(message)
    print(f'http://127.0.0.1:8000/message/{message_key}')
    return jsonify({'message_key': message_key}), 202


if __name__ == '__main__':
    app.run(host='127.0.0.1', port=8000, debug=True)

    # run this with python3.13 test.py
    # also run dramatiq --processes 2 --threads 2 test