def import_broker(value):
module, broker_or_callable = import_object(value)
if broker_or_callable is None:
return module, get_broker()
if callable(broker_or_callable):
broker_or_callable()
return module, get_broker()
if not isinstance(broker_or_callable, Broker):
raise ImportError("%r is not a Broker." % value)
return module, broker_or_callable
def _enqueue_correlated(ctx: _Context):
"""Enqueues a correlated deploy for further processing by orchestrator.
"""
dramatiq.get_broker().enqueue(
dramatiq.Message(
queue_name="orchestration.engine.step",
actor_name="on_step_deploy_finalized",
args=([
encoder.encode(ctx.deploy_execution_ctx),
encoder.encode(ctx.node_id), ctx.block_hash, ctx.deploy_hash
]),
kwargs=dict(),
options=dict(),
))
def __init__(self, children, *, broker=None):
self.broker = broker or dramatiq.get_broker()
self.messages = []
messages = []
for child in children:
if isinstance(child, pipeline):
messages.extend(message for message in child.messages)
else:
messages.append(child)
# The msg.copy() is important here. It serves several functions.
# 1. It prevents addions to message from altering the source messages.
# 2. Those unaltered messages are the ones that need to be used for
# pipe_target and pipe_source in order to prevent circular references
# when executing the pipeline.
prev_msg = None
for n, message in enumerate(msg.copy() for msg in messages):
with suppress(IndexError):
next_msg = messages[n + 1]
message.options['pipe_target'] = next_msg.asdict()
if prev_msg:
message.options['pipe_source'] = prev_msg.asdict()
self.messages.append(message)
prev_msg = messages[n]
def sqlite_broker_in_tests(sqlite_in_tests):
"""Reset dramatiq broker to a new instance pointing at sqlite DB for duration of
tests.
This is required because the dramatiq design is such that a broker needs to be
installed at import-time, and actors declared using the decorator will be pointing
at that broker. Since that happens too early for us to set up our test fixtures,
we need to reinstall another broker pointing at our sqlite DB after we've set
that up.
"""
old_broker = dramatiq.get_broker()
new_broker = Broker()
dramatiq.set_broker(new_broker)
# All actors are currently pointing at old_broker which is not using sqlite.
# This will break a call to .send() on those actors.
# Point them towards new_broker instead which will allow them to work.
actors = []
for actor in old_broker.actors.values():
actors.append(actor)
actor.broker = new_broker
new_broker.declare_actor(actor)
# Everything now points at the sqlite-enabled broker, so proceed with test
yield
# Now roll back our changes
for actor in actors:
actor.broker = old_broker
dramatiq.set_broker(old_broker)
def distribute_city_council_objects_to_sync(payload):
"""Recebe o payload e dispara uma task para cada registro.
O webservice da Câmara retorna uma lista de ações (inserção,
atualização e deleção) e os registros que sofreram cada uma
delas. Essa task executa cada uma de maneira separada para que,
caso tenham algum erro, possam ser tratados de maneira separada.
"""
action_methods = {
"inclusoesContrato": add_citycouncil_contract,
"alteracoesContrato": update_citycouncil_contract,
"exclusoesContrato": remove_citycouncil_contract,
"inclusoesLicitacao": add_citycouncil_bid,
"alteracoesLicitacao": update_citycouncil_bid,
"exclusoesLicitacao": remove_citycouncil_bid,
"inclusoesReceita": add_citycouncil_revenue,
"alteracoesReceita": update_citycouncil_revenue,
"exclusoesReceita": remove_citycouncil_revenue,
"inclusoesDespesa": add_citycouncil_expense,
"alteracoesDespesa": update_citycouncil_expense,
"exclusoesDespesa": remove_citycouncil_expense,
}
broker = get_broker()
for action_name, records in payload.items():
task = action_methods.get(action_name)
for record in records:
broker.enqueue(task.message(record))
def send(self, *messages):
"""Send the first unsent message bound to this context. If :messages: are provided, they are bound to the
context before a message is sent. Returns True if a message is sent, otherwise False."""
# Bind the messages
self.bind(*messages)
if self.status != self.Status.running:
return
# Get the current state of messages and sent
sent_key = self._key_for('sent')
messages, sent, errors = self.messages, self.sent, self.errors
for msg in messages:
if msg.message_id not in sent\
and msg.message_id not in errors\
and self.redis.execute_command('ZADD', sent_key, 'NX', time.time(), msg.message_id):
broker = dramatiq.get_broker()
broker.enqueue(msg)
self.count_message(msg)
return
completed, total = self.progress()
if completed == total:
self.status = self.Status.complete
def build_message_key(self, message):
broker = dramatiq.get_broker()
actor = broker.get_actor(message.actor_name)
signature = inspect.signature(actor)
bound_args = signature.bind(*message.args, **message.kwargs)
message_key = cache.signature_key(bound_args, message.actor_name)
return f'{self.key_prefix}:{message_key}'
def abort(
message_id: str,
middleware: Optional[Abortable] = None,
abort_ttl: Optional[int] = None,
mode: AbortMode = AbortMode.ABORT,
) -> None:
"""Abort a pending or running message given its ``message_id``.
:param message_id: Message to abort. Use the return value of ``actor.send``
or ``actor.send_with_options`` to then use its ``.message_id`` attribute.
:param middleware: :class:`Abortable` middleware used by the workers and
broker used to signal termination. If set to ``None``, use the default broker
from ``dramatiq.get_broker()`` and retrieve the configured :class:`Abortable`
middleware. If no :class:`Abortable` middleware is set on the broker and
``middleware`` is ``None``, raises a :class:`RuntimeError`.
:type middleware: :class:`Abortable`
:param abort_ttl: Change default abort TTL value, optional argument. If set to
``None`` default value from :class:`Abortable` is used.
:param mode: "AbortMode.ABORT" or "AbortMode.CANCEL".In "cancel" mode,
only pending message will be aborted,
running message will also be aborted additionally in "abort" mode.
"""
if not middleware:
broker = get_broker()
for middleware in broker.middleware:
if isinstance(middleware, Abortable):
break
else:
raise RuntimeError("The default broker doesn't have an abortable backend.")
middleware.abort(message_id, abort_ttl, mode)
def setup_broker(settings):
"""Setup all brokers."""
from dramatiq.brokers.redis import RedisBroker
dramatiq.set_broker(RedisBroker(url=settings['app.broker.url']))
dramatiq.get_broker().add_middleware(ConfigMiddleware())
dramatiq.get_broker().add_middleware(DBSessionMiddleware())
dramatiq.get_broker().add_middleware(ElasticsearchMiddleware())
dramatiq.get_broker().add_middleware(RateLimiterMiddleware())
def stub_worker():
broker = dramatiq.get_broker()
worker = dramatiq.Worker(broker, worker_timeout=1000)
worker.start()
try:
yield worker
finally:
worker.stop()
def _pre_setup(self):
super()._pre_setup()
self.broker = get_broker()
self.broker.flush_all()
self.worker = Worker(self.broker, worker_timeout=100)
self.worker.start()
def import_broker(value):
module, broker = import_object(value)
if broker is None:
return module, get_broker()
if not isinstance(broker, Broker):
raise ImportError("%r is not a Broker." % value)
return module, broker
def test_broker_uses_rabbitmq_if_not_set():
# Given that no global broker is set
dramatiq.broker.global_broker = None
# If I try to get the global broker
broker = dramatiq.get_broker()
# I expect it to be a RabbitmqBroker instance
assert isinstance(broker, RabbitmqBroker)
def make_wsgi_middleware(prefix, broker=None):
broker = broker or dramatiq.get_broker()
if not isinstance(broker, RedisBroker):
raise RuntimeError("broker must be a RedisBroker")
def middleware(app):
return DashboardMiddleware(app, broker, prefix)
return middleware
def send2broker(cls, *messages, delay=None, run_at=None):
"""Send all the messages with the delay
:param _*messages: message instance list
:param delay: delay before send
:param run_at: datetime when the process must be executed
"""
if delay is None and run_at:
delay = (run_at - datetime.now()).seconds * 1000
for message in messages:
broker = message.broker or get_broker()
broker.enqueue(message, delay=delay)
def build_message_key(self, message) -> str:
broker = dramatiq.get_broker()
actor = broker.get_actor(message.actor_name)
signature = inspect.signature(actor.fn)
# even though .bind() will sort any declared keyword arguments, if the
# actor.fn accepts any **kwargs then this ensures that those are sorted
# in the final bound args as well
kwargs = {k: v for k, v in sorted(message.kwargs.items())}
try:
bound_args = signature.bind(*message.args, **kwargs)
except TypeError:
raise TypeError('Cannot cache partial messages')
message_key = signature_key(bound_args, message.actor_name)
return message_key
def main(args):
broker = dramatiq.get_broker()
broker.emit_after("process_boot")
worker = dramatiq.Worker(broker, worker_threads=1)
worker.start()
while True:
try:
time.sleep(1)
except KeyboardInterrupt:
break
worker.stop()
broker.close()
def __call__(self, *args, **kwargs):
message = self.message(*args, **kwargs)
with suppress(dramatiq.results.errors.ResultMissing):
return message.get_result()
result = super().__call__(*args, **kwargs)
broker = dramatiq.get_broker()
for middleware in broker.middleware:
if not isinstance(middleware, dramatiq.results.Results):
continue
middleware.backend
result_ttl = self.options.get('result_ttl', middleware.result_ttl)
middleware.backend.store_result(message, result, result_ttl)
return result
def _declare_actor_for(ActorCls,
method,
queue_name="default",
priority=0,
**options):
if not hasattr(method, '__self__'):
raise AnyBlokActorException(
"The method %r must be declared as a classmethod" % method)
Model = method.__self__
db_name = Model.registry.db_name
registry_name = Model.__registry_name__
actor_name = db_name + ':' + registry_name + '=>' + method.__name__
logger.info("Declare the actor : %r", actor_name)
if not _queue_name_re.fullmatch(queue_name):
raise AnyBlokActorException(
"Queue names must start with a letter or an underscore followed "
"by any number of letters, digits, dashes or underscores.")
broker = dramatiq.get_broker()
invalid_options = set(options) - broker.actor_options
if invalid_options:
raise AnyBlokActorException(
("The following actor options are undefined: "
"{%s}. Did you forget to add a middleware "
"to your Broker?") % ', '.join(invalid_options))
if isinstance(method, ActorCls):
raise AnyBlokActorException(
"The actor %r is declared two time as an actor" % actor_name)
def fn(*a, **kw):
return method(*a, **kw)
setattr(fn, 'registry', Model.registry)
setattr(fn, 'method', method)
actor = ActorCls(
fn,
actor_name=actor_name,
queue_name=queue_name,
priority=priority,
broker=broker,
options=options,
)
setattr(fn, 'actor', actor)
setattr(Model, method.__name__, actor)
logger.debug('declare actor on "%s:%s"', Model.__registry_name__,
method.__name__)
def import_broker(value):
modname, varname = value, None
if ":" in value:
modname, varname = value.split(":", 1)
module = importlib.import_module(modname)
if varname is not None:
if not hasattr(module, varname):
raise ImportError("Module %r does not define a %r variable." % (modname, varname))
broker = getattr(module, varname)
if not isinstance(broker, Broker):
raise ImportError("Variable %r from module %r is not a Broker." % (varname, modname))
return module, broker
return module, get_broker()
def setUp(self):
super().setUp()
self.TEST_USER_NAME = 'testuser'
self.TEST_USER_PASSWORD = '******'
# Create a user
test_user1 = User.objects.create_user(
username=self.TEST_USER_NAME,
password=self.TEST_USER_PASSWORD,
)
test_user1.save()
print(f'TEST USER {test_user1.username}')
print(f'TEST USER {test_user1.id}')
self.test_user = test_user1
self.broker = dramatiq.get_broker()
self.worker = dramatiq.Worker(self.broker, worker_timeout=100)
self.worker.start()
def import_broker(value):
modname, varname = value, None
if ":" in value:
modname, varname = value.split(":", 1)
module = importlib.import_module(modname)
if varname is not None:
if not hasattr(module, varname):
raise ImportError(
f"Module {modname!r} does not define a {varname!r} variable.")
broker = getattr(module, varname)
if not isinstance(broker, Broker):
raise ImportError(
f"Variable {varname!r} from module {modname!r} is not a Broker."
)
return module, broker
return module, get_broker()
def import_broker(value):
modname, varname = value, None
if ":" in value:
modname, varname = value.split(":", 1)
module = importlib.import_module(modname)
if varname is not None:
varnames = varname.split('.')
try:
broker = functools.reduce(getattr, varnames, module)
except AttributeError:
raise ImportError("Module %r does not define a %r variable." %
(modname, varname))
if not isinstance(broker, Broker):
raise ImportError("Variable %r from module %r is not a Broker." %
(varname, modname))
return module, broker
return module, get_broker()
def stub_worker(monkeypatch, settings, _runner):
if settings.JOEFLOW_TASK_RUNNER == "joeflow.runner.celery.task_runner":
yield mock.Mock()
else:
import dramatiq
broker = dramatiq.get_broker()
broker.emit_after("process_boot")
broker.flush_all()
worker = dramatiq.Worker(broker, worker_timeout=100)
worker.start()
class Meta:
@staticmethod
def wait():
broker.join(settings.JOEFLOW_CELERY_QUEUE_NAME, timeout=60000)
worker.join()
yield Meta
worker.stop()
def abort(message_id: str, middleware: Optional[Abortable] = None) -> None:
"""Abort a pending or running message given its ``message_id``.
:param message_id: Message to abort. Use the return value of ``actor.send``
or ``actor.send_with_options`` to then use its ``.message_id`` attribute.
:param middleware: :class:`Abortable` middleware used by the workers and
broker used to signal termination. If set to ``None``, use the default broker
from ``dramatiq.get_broker()`` and retrieve the configured :class:`Abortable`
middleware. If no :class:`Abortable` middleware is set on the broker and
``middleware`` is ``None``, raises a :class:`RuntimeError`.
:type middleware: :class:`Abortable`
"""
if not middleware:
broker = get_broker()
for middleware in broker.middleware:
if isinstance(middleware, Abortable):
break
else:
raise RuntimeError(
"The default broker doesn't have an abortable backend.")
middleware.abort(message_id)
async def db_session(request: Request, call_next):
"""Maintain a DB session around each request, which is also shared
with the dramatiq broker.
An implicit commit occurs if and only if the request succeeds.
"""
request.state.db = new_db_session(app.state.db_engine)
# Any dramatiq operations should also make use of this session.
broker = dramatiq.get_broker()
broker.set_session(request.state.db)
try:
response = await call_next(request)
if response.status_code >= 200 and response.status_code < 300:
await run_in_threadpool(request.state.db.commit)
finally:
broker.set_session(None)
await run_in_threadpool(request.state.db.close)
request.state.db = None
return response
def execute_task(task_request: TaskRequest, workflow_id: str) -> Task:
"""
Send task request to main `drama` actor.
"""
db = get_db_connection()
# appends workflow id to task
task_dict = task_request.dict()
task_dict["parent"] = workflow_id
# triggers actor execution
_message = process_task.message_with_options(
args=(task_dict, ),
on_failure=process_failure,
on_success=process_succeeded,
)
_message = _message.copy(
queue_name=settings.DEFAULT_ACTOR_OPTS.queue_name, )
broker = dramatiq.get_broker()
message = broker.enqueue(_message)
# creates task on database
task = TaskManager(db).create_or_update_from_id(
message.message_id,
name=task_request.name,
module=task_request.module,
parent=workflow_id,
params=task_request.params,
inputs=task_request.inputs,
labels=task_request.labels,
metadata=task_request.metadata,
status=TaskStatus.STATUS_PENDING,
created_at=datetime.now(),
)
return task
def broker():
broker = dramatiq.get_broker()
broker.flush_all()
return broker
def _get_actor_schema(actor_name: str) -> Optional[mapper.BoundSchema]:
with suppress(AttributeError):
broker = dramatiq.get_broker()
actor = broker.get_actor(actor_name)
return actor.fn.schema
def broker():
return dramatiq.get_broker()
