def enqueue_using_session(self, db, message, delay=None):
# Given a dramatiq message, saves it to the queue in the DB.
queue_name = message.queue_name
if delay is not None:
queue_name = dq_name(queue_name)
message.options["eta"] = current_millis() + delay
db_message = DramatiqMessage(id=message.message_id,
actor=message.actor_name,
queue=queue_name)
message_dict = message.asdict()
# Drop these so we're not storing them in two places.
del message_dict["message_id"]
del message_dict["queue_name"]
del message_dict["actor_name"]
db_message.body = message_dict
# Use merge rather than add since the message may already exist;
# for instance this happens in case of retry.
db_message = db.merge(db_message)
# Explicitly wipe out the consumer, since if we've updated an existing
# message it'll have to be consumed again.
db_message.consumer_id = None
def __ensure_enqueued(self, broker, actor, *args, **kwargs):
msg = actor.message(*args, **kwargs)
# Use a fixed message ID for this actor; this ensures there's only one
# scheduler message in the system for this actor.
msg = msg.copy(message_id=actor.options["scheduled_message_id"])
session = Session(bind=self.__db_engine)
try:
broker.set_session(session)
# Enqueue ourselves
broker.enqueue(msg, delay=Settings().scheduler_delay * 60 * 1000)
# Clean any other messages to same actor & queue which are NOT this one
queues = [actor.queue_name, dq_name(actor.queue_name)]
session.query(DramatiqMessage).filter(
DramatiqMessage.actor == actor.actor_name,
DramatiqMessage.queue.in_(queues),
DramatiqMessage.id != msg.message_id,
).delete(synchronize_session=False)
# And commit
session.commit()
finally:
session.close()
broker.set_session(None)
def tab_from_q_name(name):
if name == dq_name(name):
return "delayed"
elif name == xq_name(name):
return "failed"
else:
return "standard"
def declare_queue(self, queue_name):
if queue_name not in self.queues:
self.emit_before("declare_queue", queue_name)
self.queues[queue_name] = None
self.emit_after("declare_queue", queue_name)
delayed_name = dq_name(queue_name)
self.queues[delayed_name] = None
self.delay_queues.add(delayed_name)
self.emit_after("declare_delay_queue", delayed_name)
def test_redis_requeues_unhandled_delay_messages_on_shutdown(redis_broker):
# Given that I have an actor that takes its time
@dramatiq.actor
def do_work():
pass
# If I send it a delayed message
message = do_work.send_with_options(delay=10000)
# Then start a worker and subsequently shut it down
with worker(redis_broker, worker_threads=1):
pass
# I expect it to have re-enqueued the message
messages = redis_broker.client.lrange("dramatiq:%s" % dq_name(do_work.queue_name), 0, 10)
assert message.options["redis_message_id"].encode("utf-8") in messages
def test_dq_name_returns_canonical_delay_names(given, expected):
assert dq_name(given) == expected
def rabbitmq_random_queue(rabbitmq_broker):
queue_name = "rabbit-queue-%s" % uuid.uuid4()
yield queue_name
rabbitmq_broker.channel.queue_delete(queue_name)
rabbitmq_broker.channel.queue_delete(dq_name(queue_name))
rabbitmq_broker.channel.queue_delete(xq_name(queue_name))
def queue_for_tab(name, tab):
return {
"standard": name,
"delayed": dq_name(name),
"failed": xq_name(name),
}[tab]
