add nerelic background rasks contextmanager

Author: dpastranak

agave/core/__init__.py

@@ -3,6 +3,7 @@
     accept_trace_from_queue,
     accept_trace_headers,
     add_custom_attribute,
+    background_task,
     get_trace_headers,
     inject_trace_headers,
     trace_attributes,
@@ -13,6 +14,7 @@
     'accept_trace_from_queue',
     'accept_trace_headers',
     'add_custom_attribute',
+    'background_task',
     'get_trace_headers',
     'inject_trace_headers',
     'trace_attributes',

agave/core/tracing.py

@@ -1,15 +1,32 @@
 from __future__ import annotations
 
 import inspect
+from contextlib import contextmanager
 from functools import wraps
-from typing import Any, Callable
+from typing import Any, Callable, Optional
 
 import newrelic.agent
 
 # trace headers key
 TRACE_HEADERS_KEY = "_nr_trace_headers"
 
 
+@contextmanager
+def background_task(
+    name: str,
+    group: str = "Task",
+    trace_headers: Optional[dict] = None,
+):
+    with newrelic.agent.BackgroundTask(
+        application=newrelic.agent.application(),
+        name=name,
+        group=group,
+    ):
+        if trace_headers:
+            accept_trace_headers(trace_headers, transport_type="Queue")
+        yield
+
+
 def get_trace_headers() -> dict:
     headers_list: list = []
     newrelic.agent.insert_distributed_trace_headers(headers_list)

agave/version.py

@@ -1 +1 @@
-__version__ = "1.5.3.dev00"
+__version__ = "1.5.3.dev01"

tests/core/test_tracing.py

@@ -9,6 +9,7 @@
     accept_trace_from_queue,
     accept_trace_headers,
     add_custom_attribute,
+    background_task,
     get_trace_headers,
     inject_trace_headers,
     trace_attributes,
@@ -416,3 +417,51 @@ def test_add_attributes_with_non_callable_non_string_extractor():
     with patch("agave.core.tracing.add_custom_attribute") as mock_add:
         _add_attributes(kwargs, extractors)
         mock_add.assert_called_once_with("result", None)
+
+
+def test_background_task_creates_transaction():
+    """Test that background_task creates a New Relic BackgroundTask."""
+    with patch("agave.core.tracing.newrelic.agent") as mock_agent:
+        mock_agent.application.return_value = "test_app"
+        mock_agent.BackgroundTask.return_value.__enter__ = lambda s: None
+        mock_agent.BackgroundTask.return_value.__exit__ = lambda s, *a: None
+
+        with background_task("my_task", "Redis/Queue"):
+            pass
+
+        mock_agent.BackgroundTask.assert_called_once_with(
+            application="test_app",
+            name="my_task",
+            group="Redis/Queue",
+        )
+
+
+def test_background_task_accepts_trace_headers():
+    """Test that background_task accepts trace headers inside transaction."""
+    with patch("agave.core.tracing.newrelic.agent") as mock_agent:
+        mock_agent.application.return_value = "test_app"
+        mock_agent.BackgroundTask.return_value.__enter__ = lambda s: None
+        mock_agent.BackgroundTask.return_value.__exit__ = lambda s, *a: None
+
+        with patch("agave.core.tracing.accept_trace_headers") as mock_accept:
+            headers = {"traceparent": "abc"}
+            with background_task("my_task", "SQS/Queue", headers):
+                pass
+
+            mock_accept.assert_called_once_with(
+                headers, transport_type="Queue"
+            )
+
+
+def test_background_task_does_not_accept_when_no_headers():
+    """Test that background_task skips accept when no trace headers."""
+    with patch("agave.core.tracing.newrelic.agent") as mock_agent:
+        mock_agent.application.return_value = "test_app"
+        mock_agent.BackgroundTask.return_value.__enter__ = lambda s: None
+        mock_agent.BackgroundTask.return_value.__exit__ = lambda s, *a: None
+
+        with patch("agave.core.tracing.accept_trace_headers") as mock_accept:
+            with background_task("my_task", "Task"):
+                pass
+
+            mock_accept.assert_not_called()