Repository: scoutapp/scout_apm_python

Issue:
Huey Django integration
Thank so much for adding Huey integration (#362). When using Django the Huey task worker is started with a management command which is part of the Huey package. We are unable to modify the managements command. I suggest that a Django app is added to Scout which will connect to the Huey signals during the Django startup process.


Discussion:
Hi @cburza ,

It looks like you can set [`HUEY` in your Django settings](https://huey.readthedocs.io/en/latest/django.html) to a `Huey()` instance. If you do this, you should be able to call `attach_scout` on that instance, and the integration should work. You can configure Scout through your Django settings.

I'll look into testing this further and adding a documentation change if nothing else.

Thanks,

Adam
@cburza I've just given this a test and it looks harder than I thought. Our `attach_scout` function will only work in a Django-less Huey setup. I will look into adding a Django-startup-patch-huey function as you suggested.

Resolving patch:
diff --git a/src/scout_apm/django/apps.py b/src/scout_apm/django/apps.py
--- a/src/scout_apm/django/apps.py
+++ b/src/scout_apm/django/apps.py
@@ -7,6 +7,7 @@
 
 import scout_apm.core
 from scout_apm.core.config import scout_config
+from scout_apm.django.instruments.huey import ensure_huey_instrumented
 from scout_apm.django.instruments.sql import ensure_sql_instrumented
 from scout_apm.django.instruments.template import ensure_templates_instrumented
 
@@ -28,6 +29,7 @@ def ready(self):
         self.install_middleware()
 
         # Setup Instruments
+        ensure_huey_instrumented()
         ensure_sql_instrumented()
         ensure_templates_instrumented()
 
diff --git a/src/scout_apm/django/instruments/huey.py b/src/scout_apm/django/instruments/huey.py
new file mode 100644
--- /dev/null
+++ b/src/scout_apm/django/instruments/huey.py
@@ -0,0 +1,31 @@
+# coding=utf-8
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+import logging
+
+from django.conf import settings
+
+logger = logging.getLogger(__name__)
+
+huey_instrumented = False
+
+
+def ensure_huey_instrumented():
+    global huey_instrumented
+    if huey_instrumented:
+        return
+    huey_instrumented = True
+
+    # Avoid importing if not installed
+    if "huey.contrib.djhuey" not in settings.INSTALLED_APPS:  # pragma: no cover
+        return
+
+    try:
+        from huey.contrib.djhuey import HUEY
+    except ImportError:  # pragma: no cover
+        return
+
+    from scout_apm.huey import attach_scout_handlers
+
+    attach_scout_handlers(HUEY)
+    logger.debug("Instrumented huey.contrib.djhuey")
diff --git a/src/scout_apm/huey.py b/src/scout_apm/huey.py
--- a/src/scout_apm/huey.py
+++ b/src/scout_apm/huey.py
@@ -15,9 +15,13 @@
 def attach_scout(huey):
     installed = scout_apm.core.install()
     if installed:
-        huey.pre_execute()(scout_on_pre_execute)
-        huey.post_execute()(scout_on_post_execute)
-        huey.signal(SIGNAL_CANCELED)(scout_on_cancelled)
+        attach_scout_handlers(huey)
+
+
+def attach_scout_handlers(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):


Test patch:
diff --git a/tests/integration/django_app.py b/tests/integration/django_app.py
--- a/tests/integration/django_app.py
+++ b/tests/integration/django_app.py
@@ -35,8 +35,11 @@
         "django.contrib.contenttypes",
         "django.contrib.messages",
         "django.contrib.sessions",
+        "huey.contrib.djhuey",
         "scout_apm.django",
     ],
+    # https://huey.readthedocs.io/en/latest/django.html
+    "HUEY": {"backend_class": "huey.MemoryHuey", "immediate": True},
 }
 
 if django.VERSION > (1, 10):
diff --git a/tests/integration/test_django.py b/tests/integration/test_django.py
--- a/tests/integration/test_django.py
+++ b/tests/integration/test_django.py
@@ -17,6 +17,7 @@
 
 from scout_apm.compat import datetime_to_timestamp
 from scout_apm.core.config import scout_config
+from scout_apm.django.instruments.huey import ensure_huey_instrumented
 from scout_apm.django.instruments.sql import ensure_sql_instrumented
 from scout_apm.django.instruments.template import ensure_templates_instrumented
 from tests.compat import mock
@@ -88,18 +89,15 @@ def make_admin_user():
     return user
 
 
-def test_ensure_sql_instrumented_again():
-    """
-    Check second call doesn't crash (should be a no-op)
-    """
-    ensure_sql_instrumented()
-
-
-def test_ensure_templates_instrumented_again():
+@pytest.mark.parametrize(
+    "func",
+    [ensure_huey_instrumented, ensure_sql_instrumented, ensure_templates_instrumented],
+)
+def test_instruments_idempotent(func):
     """
     Check second call doesn't crash (should be a no-op)
     """
-    ensure_templates_instrumented()
+    func()
 
 
 def test_on_setting_changed_application_root():
@@ -920,3 +918,23 @@ def test_old_style_view_middleware_deleted(tracked_requests):
 
     assert response.status_int == 200
     assert len(tracked_requests) == 0
+
+
+def test_huey_basic_task(tracked_requests):
+    with app_with_scout():
+        from huey.contrib.djhuey import task
+
+        @task()
+        def hello():
+            return "Hello World!"
+
+        result = 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_django.hello"