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