Repository: scoutapp/scout_apm_python

Issue:
Add support for Dramatiq
Could we investigate the feasibility of adding support for the Dramatiq background job library?

https://dramatiq.io/

We have only had the one request for it as far as I am aware. But if is not a huge job then it would be beneficial to have going forward.


Discussion:
Will look into it!

Resolving patch:
diff --git a/src/scout_apm/dramatiq.py b/src/scout_apm/dramatiq.py
new file mode 100644
--- /dev/null
+++ b/src/scout_apm/dramatiq.py
@@ -0,0 +1,41 @@
+# coding=utf-8
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+import dramatiq
+
+import scout_apm.core
+from scout_apm.core.tracked_request import TrackedRequest
+
+
+class ScoutMiddleware(dramatiq.Middleware):
+    def __init__(self):
+        installed = scout_apm.core.install()
+        self._do_nothing = not installed
+
+    def before_process_message(self, broker, message):
+        if self._do_nothing:
+            return
+        tracked_request = TrackedRequest.instance()
+        tracked_request.tag("queue", message.queue_name)
+        tracked_request.tag("message_id", message.message_id)
+        tracked_request.start_span(operation="Job/" + message.actor_name)
+
+    def after_process_message(self, broker, message, result=None, exception=None):
+        if self._do_nothing:
+            return
+        tracked_request = TrackedRequest.instance()
+        tracked_request.mark_real_request()
+        if exception:
+            tracked_request.tag("error", "true")
+        tracked_request.stop_span()
+
+    def after_skip_message(self, broker, message):
+        """
+        The message was skipped by another middleware raising SkipMessage.
+        Stop the span and thus the request, it won't have been marked as real
+        so that's alright.
+        """
+        if self._do_nothing:
+            return
+        tracked_request = TrackedRequest.instance()
+        tracked_request.stop_span()


Test patch:
diff --git a/tests/integration/test_celery.py b/tests/integration/test_celery.py
--- a/tests/integration/test_celery.py
+++ b/tests/integration/test_celery.py
@@ -12,8 +12,7 @@
 @contextmanager
 def app_with_scout(config=None):
     """
-    Context manager that configures and installs the Scout plugin for Bottle.
-
+    Context manager that configures a Celery app with Scout installed.
     """
     # Enable Scout by default in tests.
     if config is None:
diff --git a/tests/integration/test_dramatiq.py b/tests/integration/test_dramatiq.py
new file mode 100644
--- /dev/null
+++ b/tests/integration/test_dramatiq.py
@@ -0,0 +1,115 @@
+# coding=utf-8
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+import sys
+from collections import namedtuple
+from contextlib import contextmanager
+
+import pytest
+
+if sys.version_info < (3, 5):  # Minimum version dramatiq should be installable on
+    pytest.skip(
+        "Dramatiq not installable on this Python version", allow_module_level=True
+    )
+else:
+    import dramatiq
+    from dramatiq.brokers.stub import StubBroker
+
+    from scout_apm.api import Config
+    from scout_apm.dramatiq import ScoutMiddleware
+
+
+@contextmanager
+def app_with_scout(config=None):
+    """
+    Context manager that configures a Dramatiq app with Scout middleware
+    installed.
+    """
+    # Enable Scout by default in tests.
+    if config is None:
+        config = {"monitor": True}
+
+    # Disable running the agent.
+    config["core_agent_launch"] = False
+
+    broker = StubBroker()
+    broker.emit_after("process_boot")
+    dramatiq.set_broker(broker)
+
+    @dramatiq.actor(max_retries=0)
+    def hello():
+        return "Hello World!"
+
+    @dramatiq.actor(max_retries=0)
+    def fail():
+        raise ValueError("BØØM!")  # non-ASCII
+
+    worker = dramatiq.Worker(broker, worker_timeout=0)
+
+    # Setup according to https://docs.scoutapm.com/#dramatiq
+    Config.set(**config)
+    broker.add_middleware(ScoutMiddleware(), before=broker.middleware[0].__class__)
+    worker.start()
+
+    App = namedtuple("App", ["broker", "worker", "hello", "fail"])
+    try:
+        yield App(broker=broker, worker=worker, hello=hello, fail=fail)
+    finally:
+        worker.stop()
+        # Reset Scout configuration.
+        Config.reset_all()
+
+
+def test_hello(tracked_requests):
+    with app_with_scout() as app:
+        app.hello.send()
+        app.broker.join(app.hello.queue_name)
+        app.worker.join()
+
+    assert len(tracked_requests) == 1
+    tracked_request = tracked_requests[0]
+    assert tracked_request.tags["queue"] == "default"
+    assert "message_id" in tracked_request.tags
+    assert len(tracked_request.complete_spans) == 1
+    span = tracked_request.complete_spans[0]
+    assert span.operation == "Job/hello"
+
+
+def test_fail(tracked_requests):
+    with app_with_scout() as app:
+        app.fail.send()
+        app.broker.join(app.fail.queue_name)
+        app.worker.join()
+
+    assert len(tracked_requests) == 1
+    tracked_request = tracked_requests[0]
+    assert tracked_request.tags["queue"] == "default"
+    assert "message_id" in tracked_request.tags
+    assert tracked_request.tags["error"] == "true"
+    assert len(tracked_request.complete_spans) == 1
+    span = tracked_request.complete_spans[0]
+    assert span.operation == "Job/fail"
+
+
+def test_not_installed(tracked_requests):
+    with app_with_scout() as app:
+        app.broker.middleware[0]._do_nothing = True
+        app.hello.send()
+        app.broker.join(app.fail.queue_name)
+        app.worker.join()
+
+    assert tracked_requests == []
+
+
+def test_skipped(tracked_requests):
+    class SkipMiddleware(dramatiq.Middleware):
+        def before_process_message(self, broker, message):
+            raise dramatiq.middleware.SkipMessage()
+
+    with app_with_scout() as app:
+        app.broker.add_middleware(SkipMiddleware(), after=ScoutMiddleware)
+        app.hello.send()
+        app.broker.join(app.hello.queue_name)
+        app.worker.join()
+
+    assert tracked_requests == []