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 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
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