Add distributed tracing module for New Relic integration#180
Add distributed tracing module for New Relic integration#180dpastranak wants to merge 6 commits intomainfrom
Conversation
WalkthroughAdds a new agave/core/tracing.py module implementing New Relic distributed tracing integration (public API, decorators, helpers, and a background_task context manager), test coverage in tests/core/test_tracing.py, an examples/tasks/tracing_example.py example, a tracing extras requirement in setup.py, newrelic added to requirements-test.txt, and a version bump in agave/version.py to "1.5.3.dev01". No other modules were modified. Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
Comment |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #180 +/- ##
==========================================
Coverage 100.00% 100.00%
==========================================
Files 19 20 +1
Lines 711 814 +103
Branches 74 88 +14
==========================================
+ Hits 711 814 +103
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Fix all issues with AI agents
In `@agave/core/tracing.py`:
- Line 106: The function signature for trace_attributes uses the Python 3.10
union syntax (Callable | str) which breaks 3.9; update the signature in
trace_attributes to use typing.Union (e.g., Union[Callable, str]) and add the
appropriate import (from typing import Callable, Union) at the module top, or
alternatively enable postponed evaluation by adding from __future__ import
annotations at the top of the file and keep the current signature—pick one
approach and apply it consistently for trace_attributes.
- Around line 17-29: The type annotation in accept_trace_headers uses Python
3.10+ PEP 604 syntax (dict | None) which breaks on 3.9; update the signature to
use typing.Optional (import Optional from typing) or enable postponed evaluation
(from __future__ import annotations) and replace "headers: dict | None" with
"headers: Optional[dict]" so accept_trace_headers retains the same behavior
while being compatible with Python 3.9.
In `@tests/core/test_tracing.py`:
- Around line 103-114: Update test function signatures that use the Python 3.10
union type syntax (e.g., in test_inject_trace_headers_async_keyword_arg and
other tests using inject_trace_headers/my_func/get_trace_headers) so they are
compatible with Python 3.9: replace occurrences of "dict | None" with
"Optional[dict]" and add "from typing import Optional" to the test module,
applying the same change to all affected test functions (lines referenced in the
review: the functions around the test_inject_trace_headers_async_keyword_arg and
the other listed tests).
🧹 Nitpick comments (2)
agave/core/tracing.py (2)
37-60: Consider adding async support for queue consumers.Unlike
inject_trace_headersandtrace_attributes, this decorator only handles sync functions. If an async function is decorated, it won't be properly awaited. Consider adding async support for consistency:♻️ Proposed enhancement
def accept_trace_from_queue(func: Callable) -> Callable: ... + if inspect.iscoroutinefunction(func): + `@wraps`(func) + async def async_wrapper(*args, **kwargs): + trace_headers = kwargs.pop(TRACE_HEADERS_KEY, None) + if trace_headers: + accept_trace_headers(trace_headers, transport_type="Queue") + return await func(*args, **kwargs) + return async_wrapper + `@wraps`(func) def wrapper(*args, **kwargs): trace_headers = kwargs.pop(TRACE_HEADERS_KEY, None) if trace_headers: accept_trace_headers(trace_headers, transport_type="Queue") return func(*args, **kwargs) return wrapper
174-186: Consider logging extraction failures at debug level for observability.The silent
except Exception: passmakes debugging difficult when attribute extraction fails unexpectedly. While the intent to not disrupt the primary logic is valid, adding debug-level logging would aid troubleshooting without impacting performance.♻️ Optional improvement
+import logging + +logger = logging.getLogger(__name__) + def _add_attributes(kwargs: dict, extractors: dict) -> None: ... except Exception: - pass # Silent exception + logger.debug( + "Failed to extract trace attribute %s", + attr_name, + exc_info=True, + ) # we don't want to fail if unable to extract an attribute
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@agave/core/tracing.py`:
- Around line 1-6: The file imports newrelic.agent at module import time;
instead, add a lazy loader function (e.g., _get_newrelic()) that tries to import
newrelic.agent and on ImportError raises a clear RuntimeError explaining tracing
requires the "tracing" extra, then replace all direct references to
newrelic.agent in this module (e.g., inside any decorator or helper functions in
agave/core/tracing.py) to call _get_newrelic() and use the returned module so
users who don't have New Relic installed can import agave.core without failing
until a tracing function is actually invoked.
🧹 Nitpick comments (2)
agave/core/tracing.py (2)
79-87: Validateparam_nameexists or**kwargsis accepted.If
param_nameisn’t in the target signature and the function doesn’t accept**kwargs, the wrapper will fail later with an unexpectedTypeError. Consider failing fast with a clear error at decoration time.♻️ Proposed fix
def decorator(func: Callable) -> Callable: sig = inspect.signature(func) + if param_name not in sig.parameters and not any( + p.kind == inspect.Parameter.VAR_KEYWORD + for p in sig.parameters.values() + ): + raise ValueError( + f"inject_trace_headers param_name '{param_name}' is not in " + f"{func.__name__} signature and the function does not accept **kwargs." + )
180-191: Avoid fully silent exception swallowing.Swallowing all exceptions makes debugging attribute extraction failures hard. Consider at least logging at debug level (or narrowing the exception set) to preserve observability.
♻️ Proposed fix (debug logging)
+import logging @@ +logger = logging.getLogger(__name__) @@ - except Exception: - pass # Silent exception + except Exception as exc: + logger.debug( + "Trace attribute extraction failed for %s", + attr_name, + exc_info=exc, + )
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@agave/core/tracing.py`:
- Around line 39-62: The decorator accept_trace_from_queue only supports
synchronous functions; detect async functions (use inspect.iscoroutinefunction
or asyncio.iscoroutinefunction on func) and return an async wrapper that pops
TRACE_HEADERS_KEY from kwargs, calls accept_trace_headers(...,
transport_type="Queue") if present, then awaits func(*args, **kwargs); for sync
funcs keep the existing wrapper behavior. Ensure both wrappers are decorated
with functools.wraps(func) and reference TRACE_HEADERS_KEY and
accept_trace_headers exactly as in the diff so the fix preserves behavior and
signatures.
🧹 Nitpick comments (1)
agave/core/tracing.py (1)
190-191: Consider logging exceptions instead of silently swallowing them.While the intent to not fail on extraction errors is correct, completely silent failures make debugging difficult. Consider logging at DEBUG level to aid troubleshooting without impacting production behavior.
♻️ Proposed refactor
+import logging + +logger = logging.getLogger(__name__) + # In _add_attributes: except Exception: - pass # Silent exception - # we don't want to fail if unable to extract an attribute + logger.debug( + "Failed to extract trace attribute '%s'", + attr_name, + exc_info=True, + )
Introduces agave.core.tracing with utilities to propagate and accept distributed trace headers across HTTP calls and queue messages: - get_trace_headers(): extract current trace headers - accept_trace_headers(): continue trace from incoming headers - add_custom_attribute(): add attributes to current trace - @accept_trace_from_queue: decorator for queue consumers - @inject_trace_headers: decorator for HTTP client calls - @trace_attributes: decorator to add custom trace attributes Also adds tracing optional dependency (newrelic>=7.0.0,<12.0.0) and bumps version to 1.5.3.
dpastranak
force-pushed
the
distributed-tracing
branch
from
f381e8a to
d8a5cd3
Compare
- Update accept_trace_from_queue decorator to handle both sync and async functions - Add tests for async queue handler support
requirements-test.txt
Outdated
| pytest-vcr==1.0.2 | ||
| pytest-asyncio==0.18.* | ||
| typing_extensions==4.12.2 | ||
| newrelic==11.2.0 No newline at end of file |
There was a problem hiding this comment.
Agregar una linea al final del archivo
agave/core/tracing.py
Outdated
| from functools import wraps | ||
| from typing import Any, Callable, Optional | ||
|
|
||
| import newrelic.agent |
There was a problem hiding this comment.
Creo que esto puede fallar, dado que newrelic es opcional, puedes hacer unas pruebas? Quizá podemos hacer algo como aquí:
agave/agave/tools/asyncio/sqs_client.py
Lines 7 to 14 in e7aa7f5
agave/core/tracing.py
Outdated
| yield | ||
|
|
||
|
|
||
| def get_trace_headers() -> dict: |
There was a problem hiding this comment.
Aquí podemos ser más específicos dict[str, str]
tests/core/test_tracing.py
Outdated
| @@ -0,0 +1,467 @@ | |||
| from __future__ import annotations | |||
There was a problem hiding this comment.
Que opinas si para mantener el estilo actual del proyecto ya que tiene soporte para 3.9 usemos Optional?
agave/core/__init__.py
Outdated
| from .tracing import ( | ||
| TRACE_HEADERS_KEY, | ||
| accept_trace_from_queue, | ||
| accept_trace_headers, | ||
| add_custom_attribute, | ||
| background_task, | ||
| get_trace_headers, | ||
| inject_trace_headers, | ||
| trace_attributes, | ||
| ) |
There was a problem hiding this comment.
Sugiero quitarlos de aqui.
Para el modulo agave.core no dependa de newrelic
- Make newrelic import explicit with clear error message for missing dependency - Remove tracing re-exports from agave.core to avoid forcing newrelic installation - Use Optional and Union for Python 3.9 compatibility instead of PEP 604 syntax and future module - Add specific type hints dict[str, str] for trace headers - Add trailing newline to requirements-test.txt
dpastranak
force-pushed
the
distributed-tracing
branch
from
da7b2d7 to
2002fbd
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@agave/core/tracing.py`:
- Around line 6-12: The ImportError caught when trying to import newrelic.agent
in agave/core/tracing.py must be re-raised while preserving the original
exception context; update the except ImportError handler that currently raises a
new ImportError to chain the original exception using "raise ... from e" (i.e.,
capture the caught exception variable and re-raise the new ImportError with
"from" the original) so the ImportError cause is preserved for debugging.
| return dict(headers_list) | ||
|
|
||
|
|
||
| def accept_trace_headers( |
There was a problem hiding this comment.
Esta función no es necesaria.
Realmente solo hace una linea y el transport_type lo usas como Queue
There was a problem hiding this comment.
En este caso la función también incluye la validación if not headers: return que evita errores cuando los headers son None. Además, el transport_type por default es "HTTP" y se puede usar en contexto HHTP no solo como Queue.
Sin esta función cada que se use tendría que repetirse la validación también. Creo que la función aporta valor.
| def get_trace_headers() -> dict[str, str]: | ||
| headers_list: list = [] | ||
| newrelic.agent.insert_distributed_trace_headers(headers_list) | ||
| return dict(headers_list) |
There was a problem hiding this comment.
Esta funcion es un poco extraña, inserta una lista vacía y el resultado siempre es el mismo.
Creo que es mejor quitar la funcion y hacer directo el insert_distributed_trace_headers
There was a problem hiding this comment.
es que más bien el API de newrelic está extraña. Pide pasar una lista vacía que muta internamente donde regresa los headers y esta función encapsula esta complejidad.
sin esta función se tendría que hacer
headers_list = []
newrelic.agent.insert_distributed_trace_headers(headers_list)
# Después de llamar, headers_list ya tiene datos:
# [('newrelic', 'eyJ2IjpbMCwxXS...'), ('traceparent', '00-abc123...')]
headers = dict(headers_list) # lo paso a un diccionario para pasarlo entre servicio.
con la función solo se hace
headers = get_trace_headers()
o por qué dices que siempre regresa lo mismo?
Creo que justo esta función hace la interfaz más clara y simple.
| def test_get_trace_headers_returns_dict_with_headers(): | ||
| with patch("newrelic.agent") as mock_agent: | ||
| mock_agent.insert_distributed_trace_headers.side_effect = ( | ||
| lambda h: h.extend([("traceparent", "abc"), ("tracestate", "xyz")]) | ||
| ) | ||
| result = get_trace_headers() | ||
| assert result == {"traceparent": "abc", "tracestate": "xyz"} |
There was a problem hiding this comment.
En general los test aunque cumplen la covertura, sería mejor tener test con los examples.
@celery_sqs.task
@accept_trace_from_queue
def example_task(...):
# extract trace headers
return {...headers}Y el test debe ejecutar la task y evaluar el resultado
- Remove docstrings from tracing module now examples serves as docs - Add examples/tasks/tracing_example.py with usage examples - Update accept_trace_from_queue to extract headers from body dict - Add integration tests for tracing with SQS tasks
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@examples/tasks/tracing_example.py`:
- Around line 1-4: Add "from __future__ import annotations" at the top of the
module to enable postponed evaluation of annotations so PEP 604 union types like
"dict | None" used later (e.g., in the tracing_example.py annotations and any
function signatures such as those involving accept_trace_from_queue) work on
Python 3.9; insert this single import as the very first line before other
imports (above BaseModel and agave.core.tracing imports) so runtime parsing of
"dict | None" is deferred.
🧹 Nitpick comments (1)
agave/core/tracing.py (1)
156-169: Consider logging extraction failures for debugging.The silent
except Exception: passis intentional to avoid failing traces due to attribute extraction errors, which is a reasonable design choice. However, adding a debug-level log would help diagnose issues without impacting production.♻️ Optional: Add debug logging
+import logging + +logger = logging.getLogger(__name__) + def _add_attributes(kwargs: dict, extractors: dict) -> None: for attr_name, extractor in extractors.items(): try: if callable(extractor): value = extractor(kwargs) elif isinstance(extractor, str): value = _get_nested_value(kwargs, extractor) else: value = None add_custom_attribute(attr_name, value) - except Exception: - pass # Silent exception - # we don't want to fail if unable to extract an attribute + except Exception: + logger.debug( + "Failed to extract trace attribute %s", attr_name, exc_info=True + )
| from pydantic import BaseModel | ||
|
|
||
| from agave.core.tracing import ( | ||
| accept_trace_from_queue, |
There was a problem hiding this comment.
Add from __future__ import annotations for Python 3.9 compatibility.
Line 37 uses dict | None (PEP 604 union syntax) which requires Python 3.10+ at runtime. Adding the future import enables postponed annotation evaluation, making this compatible with Python 3.9.
🐛 Proposed fix
+from __future__ import annotations
+
from pydantic import BaseModel
from agave.core.tracing import (🤖 Prompt for AI Agents
In `@examples/tasks/tracing_example.py` around lines 1 - 4, Add "from __future__
import annotations" at the top of the module to enable postponed evaluation of
annotations so PEP 604 union types like "dict | None" used later (e.g., in the
tracing_example.py annotations and any function signatures such as those
involving accept_trace_from_queue) work on Python 3.9; insert this single import
as the very first line before other imports (above BaseModel and
agave.core.tracing imports) so runtime parsing of "dict | None" is deferred.
There was a problem hiding this comment.
Pull request overview
This PR adds a New Relic–based distributed tracing module and integrates it with Agave’s task queues, plus an example and tests. It introduces utilities and decorators for propagating and accepting distributed trace headers across HTTP calls and SQS-based tasks, along with an optional tracing extra and version bump.
Changes:
- Added
agave.core.tracingwith helpers (get_trace_headers,accept_trace_headers,add_custom_attribute) and decorators (accept_trace_from_queue,inject_trace_headers,trace_attributes) plus abackground_taskcontext manager. - Introduced a new
tracingextras group fornewrelicand adjusted test requirements to include the New Relic agent. - Added comprehensive tests for tracing behavior and an example task module demonstrating tracing and queue integration, and bumped the library version to
1.5.3.dev01.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
agave/core/tracing.py |
Implements New Relic tracing primitives (header extraction/acceptance, custom attributes, queue integration, and background task context) used by both HTTP and SQS flows. |
tests/core/test_tracing.py |
Adds unit and integration-style tests covering sync/async decorators, attribute extraction, SQS task integration, and background task behavior; a few tests currently use an invalid patch target and will fail. |
setup.py |
Declares a new tracing optional dependency (newrelic>=7.0.0,<12.0.0) and already targets Python >=3.9. |
requirements-test.txt |
Pins newrelic==11.2.0 so the tracing tests can run in CI. |
examples/tasks/tracing_example.py |
Provides an example of using the new tracing utilities with SQS tasks and background tasks; one type annotation currently uses syntax that is invalid on Python 3.9. |
agave/version.py |
Bumps the library version to 1.5.3.dev01, which slightly differs from the PR description text. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| @@ -1 +1 @@ | |||
| __version__ = '1.5.2' | |||
| __version__ = "1.5.3.dev01" | |||
There was a problem hiding this comment.
The PR description mentions bumping the version to 1.5.3, but the code sets __version__ to "1.5.3.dev01". Please either update the version string here or adjust the PR description so they are consistent.
| 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: |
There was a problem hiding this comment.
The patch target "agave.core.tracing.newrelic.agent" is not a valid module path for unittest.mock.patch (it tries to import a module named agave.core.tracing.newrelic, which does not exist), so this test will raise ModuleNotFoundError instead of patching New Relic. To patch the newrelic.agent object actually used by agave.core.tracing, you should target the real module path (for example, "newrelic.agent" as in the other tests in this file).
| 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: |
There was a problem hiding this comment.
As in the previous background_task test, the patch target "agave.core.tracing.newrelic.agent" is not a real module and will cause unittest.mock.patch to fail with ModuleNotFoundError, so this test will never exercise the background_task logic. Please switch the target to the actual newrelic.agent module used by the tracing code (e.g. "newrelic.agent") so the New Relic agent calls are correctly mocked.
| 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: |
There was a problem hiding this comment.
This test also patches "agave.core.tracing.newrelic.agent", which is not a valid module path, so patch will fail before the test runs and background_task will not be exercised. Align this with the other tests by patching the real newrelic.agent module instead (for example, using "newrelic.agent" as the patch target).
| method: str, endpoint: str, trace_headers: dict | None = None | ||
| ) -> None: |
There was a problem hiding this comment.
The type annotation trace_headers: dict | None = None uses the PEP 604 union operator, which is only valid syntax in Python 3.10+, but this project declares python_requires='>=3.9' in setup.py. To keep the examples compatible with Python 3.9, please use Optional[dict] (with the appropriate import) or dict[str, Any] | None only after dropping 3.9 support.
|
se aprueba sin ver! ❤️ |
4 participants
Introduces agave.core.tracing with utilities to propagate and accept distributed trace headers across HTTP calls and queue messages:
Also adds 'tracing' optional dependency (newrelic>=7.0.0,<12.0.0) and bumps version to 1.5.3.
Summary by CodeRabbit
New Features
Tests
Chores
Examples
✏️ Tip: You can customize this high-level summary in your review settings.