Repository: scoutapp/scout_apm_python

Issue:
Add Huey integration
[Huey](https://huey.readthedocs.io/en/latest/) is a relatively popular task queue, and we've had at least one enquiry for it. We should support it.


Discussion:
Since it's gevent driven, we'd also have to  add official support with #28.
@adamchainz You can run Huey with processes, threads or greenlets. We haven't run it with greenlet, so would be great to have a version that works with processes or threads.
Ah right. Will look into it then :)

Resolving patch:
diff --git a/src/scout_apm/huey.py b/src/scout_apm/huey.py
new file mode 100644
--- /dev/null
+++ b/src/scout_apm/huey.py
@@ -0,0 +1,47 @@
+# coding=utf-8
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+from huey.exceptions import RetryTask, TaskLockedException
+from huey.signals import SIGNAL_CANCELED
+
+from scout_apm.core.tracked_request import TrackedRequest
+
+# Because neither hooks nor signals are called in *all* cases, we need to use
+# both in order to capture every case. See source:
+# https://github.com/coleifer/huey/blob/e6710bd6a9f581ebc728e24f5923d26eb0047750/huey/api.py#L331  # noqa
+
+
+def attach_scout(huey):
+    huey.pre_execute()(scout_on_pre_execute)
+    huey.post_execute()(scout_on_post_execute)
+    huey.signal(SIGNAL_CANCELED)(scout_on_cancelled)
+
+
+def scout_on_pre_execute(task):
+    tracked_request = TrackedRequest.instance()
+
+    tracked_request.tag("task_id", task.id)
+
+    operation = "Job/{}.{}".format(task.__module__, task.__class__.__name__)
+    tracked_request.start_span(operation=operation)
+
+
+def scout_on_post_execute(task, task_value, exception):
+    tracked_request = TrackedRequest.instance()
+    if exception is None:
+        tracked_request.is_real_request = True
+    elif isinstance(exception, TaskLockedException):
+        pass
+    elif isinstance(exception, RetryTask):
+        tracked_request.is_real_request = True
+        tracked_request.tag("retrying", True)
+    else:
+        tracked_request.is_real_request = True
+        tracked_request.tag("error", "true")
+    tracked_request.stop_span()
+
+
+def scout_on_cancelled(signal, task, exc=None):
+    # In the case of a cancelled signal, Huey doesn't run the post_execute
+    # handler, so we need to tidy up
+    TrackedRequest.instance().stop_span()


Test patch:
diff --git a/tests/integration/test_huey.py b/tests/integration/test_huey.py
new file mode 100644
--- /dev/null
+++ b/tests/integration/test_huey.py
@@ -0,0 +1,126 @@
+# coding=utf-8
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+from collections import namedtuple
+from contextlib import contextmanager
+
+import pytest
+from huey import MemoryHuey
+from huey.exceptions import CancelExecution, HueyException, RetryTask, TaskException
+
+from scout_apm.api import Config
+from scout_apm.huey import attach_scout
+
+
+@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
+
+    huey = MemoryHuey(immediate=True)
+
+    @huey.task()
+    @huey.lock_task("hello")
+    def hello():
+        return "Hello World!"
+
+    @huey.task()
+    def retry_once():
+        if not retry_once._did_retry:
+            retry_once._did_retry = True
+            raise RetryTask()
+        return "Done."
+
+    retry_once._did_retry = False
+
+    @huey.task()
+    def fail():
+        raise ValueError("BØØM!")  # non-ASCII
+
+    # Setup according to https://docs.scoutapm.com/#huey
+    Config.set(**scout_config)
+    attach_scout(huey)
+
+    App = namedtuple("App", ["huey", "hello", "retry_once", "fail"])
+    try:
+        yield App(huey=huey, hello=hello, retry_once=retry_once, fail=fail)
+    finally:
+        Config.reset_all()
+
+
+def test_hello(tracked_requests):
+    with app_with_scout() as app:
+        result = app.hello()
+        value = result(blocking=True, timeout=1)
+
+    assert value == "Hello World!"
+    assert len(tracked_requests) == 1
+    tracked_request = tracked_requests[0]
+    assert "task_id" in tracked_request.tags
+    assert len(tracked_request.complete_spans) == 1
+    span = tracked_request.complete_spans[0]
+    assert span.operation == "Job/tests.integration.test_huey.hello"
+
+
+def test_retry_once(tracked_requests):
+    with app_with_scout() as app:
+        result = app.retry_once()
+        value = result(blocking=True, timeout=1)
+
+    assert value == "Done."
+    assert len(tracked_requests) == 2
+    tracked_request = tracked_requests[0]
+    assert "task_id" in tracked_request.tags
+    assert tracked_request.tags["retrying"] is True
+    assert len(tracked_request.complete_spans) == 1
+    span = tracked_request.complete_spans[0]
+    assert span.operation == "Job/tests.integration.test_huey.retry_once"
+
+
+def test_fail(tracked_requests):
+    with app_with_scout() as app:
+        result = app.fail()
+        with pytest.raises(TaskException):
+            result(blocking=True, timeout=1)
+
+    assert len(tracked_requests) == 1
+    tracked_request = tracked_requests[0]
+    assert tracked_request.tags["error"] == "true"
+    assert "task_id" in tracked_request.tags
+    assert len(tracked_request.complete_spans) == 1
+    span = tracked_request.complete_spans[0]
+    assert span.operation == "Job/tests.integration.test_huey.fail"
+
+
+def test_cancelled(tracked_requests):
+    with app_with_scout() as app:
+
+        @app.huey.pre_execute()
+        def cancel_all_tasks(task):
+            raise CancelExecution()
+
+        result = app.hello()
+
+        with pytest.raises(HueyException) as excinfo:
+            result(blocking=True, timeout=0.1)
+
+    assert excinfo.value.args == ("timed out waiting for result",)
+    assert tracked_requests == []
+
+
+def test_locked(tracked_requests):
+    with app_with_scout() as app:
+        with app.huey.lock_task("hello"):
+            result = app.hello()
+
+            with pytest.raises(TaskException) as excinfo:
+                result(blocking=True, timeout=0.1)
+
+    assert excinfo.value.metadata["error"].startswith("TaskLockedException")
+    assert tracked_requests == []