def test_currrent_message_middleware_exposes_the_current_message( stub_broker, stub_worker): # Given that I have a CurrentMessage middleware stub_broker.add_middleware(CurrentMessage()) # And an actor that accesses the current message sent_messages = [] received_messages = [] @dramatiq.actor def accessor(x): message_proxy = CurrentMessage.get_current_message() received_messages.append(message_proxy._message) # When I send it a couple messages sent_messages.append(accessor.send(1)) sent_messages.append(accessor.send(2)) # And wait for it to finish its work stub_broker.join(accessor.queue_name) # Then the sent messages and the received messages should be the same assert sorted(sent_messages) == sorted(received_messages) # When I try to access the current message from a non-worker thread # Then I should get back None assert CurrentMessage.get_current_message() is None
def test_async_actor(stub_broker): stub_broker.add_middleware(CurrentMessage()) sent_messages = [] received_messages = [] @dramatiq.actor async def async_accessor(x): message_proxy = CurrentMessage.get_current_message() received_messages.append(message_proxy._message) sent_messages.append(accessor.send(1)) sent_messages.append(accessor.send(2)) stub_broker.join(accessor.queue_name) assert sorted(sent_messages) == sorted(received_messages) assert CurrentMessage.get_current_message() is None
def __init__(self, middleware=None, settings=None): super().__init__(middleware=middleware) self.__settings = settings or Settings() self.__db_engine = db_engine(self.__settings) self.__shared_session = ContextVar("shared_session") self.__broker_id = uuid.uuid4() self.__queue_events = defaultdict(Event) # We have some actors using this, so it's always enabled. self.add_middleware(CurrentMessage()) # Enable special handling of actors with 'scheduled=True' in options. self.add_middleware( SchedulerMiddleware(self.__settings, self.__db_engine)) self.add_middleware(LocalNotifyMiddleware()) # This is postgres-specific, so... if "postgresql" in str(self.__db_engine.url): self.add_middleware(PostgresNotifyMiddleware(self.__db_engine))
def commit(publish_id: str, env: str, from_date: str) -> None: actor_msg_id = CurrentMessage.get_current_message().message_id commit_obj = Commit(publish_id, env, from_date, actor_msg_id) if not commit_obj.should_write(): return commit_obj.task.state = TaskStates.in_progress commit_obj.db.commit() try: commit_obj.write_publish_items() commit_obj.task.state = TaskStates.complete commit_obj.publish.state = PublishStates.committed commit_obj.db.commit() except Exception as exc_info: LOG.exception("Task %s encountered an error", commit_obj.task.id) commit_obj.rollback_publish_items(exc_info) commit_obj.task.state = TaskStates.failed commit_obj.publish.state = PublishStates.failed commit_obj.db.commit() return
from typing import Optional, List from .config import BaseConfig config = BaseConfig() import dramatiq from dramatiq.brokers.rabbitmq import RabbitmqBroker from dramatiq.middleware import CurrentMessage broker = RabbitmqBroker(host=config.RABBITMQ_HOST, middleware=[CurrentMessage()]) dramatiq.set_broker(broker) from fastapi import FastAPI from fastapi import Depends, FastAPI, HTTPException from fastapi.responses import HTMLResponse from fastapi.middleware.cors import CORSMiddleware from . import models, schemas from .database import session, engine, worker_session from .websockets import router models.BaseModel.metadata.create_all(bind=engine) app = FastAPI(docs_url='/') app.add_middleware( CORSMiddleware, allow_origins=["http://localhost:3000"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], )
def accessor(x): message_proxy = CurrentMessage.get_current_message() received_messages.append(message_proxy._message)
def process_task(task_request: dict): """ Main `drama` actor. Executes an arbitrary function defined by a task and updates its state. """ message = CurrentMessage.get_current_message() task_id = message.message_id # required attributes task_name = task_request["name"] task_module = task_request["module"] task_parent = task_request["parent"] # workflow id # optional attributes task_params = task_request["params"] task_inputs = task_request["inputs"] # task options task_opts = task_request["options"] force_interruption = task_opts["on_fail_force_interruption"] remove_local_dir = task_opts["on_fail_remove_local_dir"] # configure data file storage storage = get_available_storage() dfs = storage(bucket_name=task_parent, folder_name=task_name ) # bucket folder is shared across tasks in a workflow # create process task_process = Process( name=task_name, module=task_module, parent=task_parent, params=task_params, inputs=task_inputs, storage=dfs, ) task_process.debug(f"Running task {task_id} with name {task_name}") try: # import `execute` function from module task_process.debug(f"Importing function from {task_module}") func = get_process_func(task_module) # set process status to `running` process_running(message) # execute imported function data = func(**task_params, pcs=task_process) if data: if not isinstance(data, TaskResult): data = TaskResult(message=str(data)) if not data: data = TaskResult() except ImportError: task_process.error(traceback.format_exc()) task_process.close(force_interruption=force_interruption) raise ImportError( f"Module {task_module} from task {task_id} is not available") except StopIteration: task_process.error(traceback.format_exc()) task_process.close(force_interruption=force_interruption) raise StopIteration("Could not get data from upstream") except Exception: task_process.error("Unexpected unknown exception was raised by actor:") task_process.error(traceback.format_exc()) task_process.close(force_interruption=force_interruption, remove_local_dir=remove_local_dir) raise remote_logging_file = task_process.close() task_process.info(f"Task {task_id} successfully executed") data.log = remote_logging_file # result of this function *must be* JSON-encodable data_as_json = data.json() return data_as_json
from drama.config import settings from drama.database import get_db_connection from drama.logger import get_logger from drama.manager import TaskManager from drama.models.task import TaskResult, TaskStatus from drama.process import Process from drama.storage.helpers import get_available_storage from drama.worker.helpers import get_process_func logger = get_logger(__name__) # setup broker logger.debug("Setting up RabbitMQ broker") broker = RabbitmqBroker(url=settings.RABBIT_DNS) broker.add_middleware(CurrentMessage()) # set broker logger.debug("Attaching broker") dramatiq.set_broker(broker) @dramatiq.actor(**settings.DEFAULT_ACTOR_OPTS.dict()) def process_task(task_request: dict): """ Main `drama` actor. Executes an arbitrary function defined by a task and updates its state. """ message = CurrentMessage.get_current_message() task_id = message.message_id