Repository: scoutapp/scout_apm_python

Issue:
Instrument rq
RQ is a popular task queue: https://python-rq.org/

We should have an integration just like we do other task queues.


Discussion:

Resolving patch:
diff --git a/src/scout_apm/rq.py b/src/scout_apm/rq.py
new file mode 100644
--- /dev/null
+++ b/src/scout_apm/rq.py
@@ -0,0 +1,83 @@
+# coding=utf-8
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+import datetime as dt
+
+import wrapt
+from rq import SimpleWorker as RqSimpleWorker
+from rq import Worker as RqWorker
+from rq.job import Job
+
+import scout_apm.core
+from scout_apm.core.tracked_request import TrackedRequest
+
+# Custom worker class
+# rq worker --worker-class scout_apm.rq.Job
+# Or instantiate queues with `job_class='scout_apm.rq.Job'` or job_class=Job
+# or custom Job subclass
+# or add JobMixin to your own Job class
+# Config via environment variables or custom Job class with Config.set() call
+# in the same file
+# django-rq should allow using Django config since django.setup() should run as
+# part of the management command - need to test that
+
+install_attempted = False
+installed = None
+
+
+def ensure_scout_installed():
+    global install_attempted, installed
+
+    if not install_attempted:
+        install_attempted = True
+        installed = scout_apm.core.install()
+
+
+class WorkerMixin(object):
+    def __init__(self, *args, **kwargs):
+        global installed
+        ensure_scout_installed()
+        if installed:
+            ensure_job_instrumented()
+        super(WorkerMixin, self).__init__(*args, **kwargs)
+
+
+class Worker(WorkerMixin, RqWorker):
+    pass
+
+
+class SimpleWorker(WorkerMixin, RqSimpleWorker):
+    pass
+
+
+job_instrumented = False
+
+
+def ensure_job_instrumented():
+    global job_instrumented
+    if job_instrumented:
+        return
+    job_instrumented = True
+    Job.perform = wrap_perform(Job.perform)
+
+
+@wrapt.decorator
+def wrap_perform(wrapped, instance, args, kwargs):
+    global installed
+    if not installed:
+        return wrapped(*args, **kwargs)
+
+    tracked_request = TrackedRequest.instance()
+    tracked_request.is_real_request = True
+    tracked_request.tag("task_id", instance.get_id())
+    tracked_request.tag("queue", instance.origin)
+    queue_time = (dt.datetime.utcnow() - instance.enqueued_at).total_seconds()
+    tracked_request.tag("queue_time", queue_time)
+    tracked_request.start_span(operation="Job/{}".format(instance.func_name))
+    try:
+        return wrapped(*args, **kwargs)
+    except Exception:
+        tracked_request.tag("error", "true")
+        raise
+    finally:
+        tracked_request.stop_span()


Test patch:
diff --git a/tests/integration/test_rq.py b/tests/integration/test_rq.py
new file mode 100644
--- /dev/null
+++ b/tests/integration/test_rq.py
@@ -0,0 +1,98 @@
+# coding=utf-8
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+from collections import namedtuple
+from contextlib import contextmanager
+
+from fakeredis import FakeStrictRedis
+from rq import Queue
+
+import scout_apm.rq
+from scout_apm import compat
+from scout_apm.api import Config
+
+
+def hello():
+    return "Hello World!"
+
+
+def fail():
+    raise ValueError("BØØM!")  # non-ASCII
+
+
+@contextmanager
+def app_with_scout(scout_config=None):
+    """
+    Context manager that configures a Huey app with Scout installed.
+    """
+    # Enable Scout by default in tests.
+    if scout_config is None:
+        scout_config = {}
+    scout_config.setdefault("monitor", True)
+    scout_config["core_agent_launch"] = False
+
+    # Reset global state
+    scout_apm.rq.install_attempted = False
+    scout_apm.rq.installed = None
+
+    # Setup according to https://docs.scoutapm.com/#rq
+    # Using job_class argument to Queue
+    Config.set(**scout_config)
+    queue = Queue(name="myqueue", connection=FakeStrictRedis())
+    worker = scout_apm.rq.SimpleWorker([queue], connection=queue.connection)
+
+    App = namedtuple("App", ["queue", "worker"])
+    try:
+        yield App(queue=queue, worker=worker)
+    finally:
+        Config.reset_all()
+
+
+def test_hello(tracked_requests):
+    with app_with_scout() as app:
+        job = app.queue.enqueue(hello)
+        app.worker.work(burst=True)
+
+    assert job.is_finished
+    assert job.result == "Hello World!"
+    assert len(tracked_requests) == 1
+    tracked_request = tracked_requests[0]
+    task_id = tracked_request.tags["task_id"]
+    assert isinstance(task_id, compat.string_type) and len(task_id) == 36
+    assert tracked_request.tags["queue"] == "myqueue"
+    assert 0.0 < tracked_request.tags["queue_time"] < 60.0
+    assert len(tracked_request.complete_spans) == 2
+    assert tracked_request.complete_spans[0].operation == "Redis/PERSIST"
+    assert (
+        tracked_request.complete_spans[1].operation
+        == "Job/tests.integration.test_rq.hello"
+    )
+
+
+def test_fail(tracked_requests):
+    with app_with_scout() as app:
+        app.queue.enqueue(fail)
+        app.worker.work(burst=True)
+
+    tracked_request = tracked_requests[0]
+    task_id = tracked_request.tags["task_id"]
+    assert isinstance(task_id, compat.string_type) and len(task_id) == 36
+    assert tracked_request.tags["queue"] == "myqueue"
+    assert 0.0 < tracked_request.tags["queue_time"] < 60.0
+    assert tracked_request.tags["error"] == "true"
+    assert len(tracked_request.complete_spans) == 2
+    assert tracked_request.complete_spans[0].operation == "Redis/PERSIST"
+    assert (
+        tracked_request.complete_spans[1].operation
+        == "Job/tests.integration.test_rq.fail"
+    )
+
+
+def test_no_monitor(tracked_requests):
+    with app_with_scout(scout_config={"monitor": False}) as app:
+        job = app.queue.enqueue(hello)
+        app.worker.work(burst=True)
+
+    assert job.is_finished
+    assert job.result == "Hello World!"
+    assert tracked_requests == []