Merge pull request #266 from appsignal/fix-collector-mode-helpers

Fix collector mode helpers

Author: unflxw

.changesets/emit-action-name-in-collector-mode.md

@@ -0,0 +1,9 @@
+---
+bump: patch
+type: fix
+---
+
+Emit the `appsignal.action_name` attribute instead of `appsignal.root_name`
+when `set_root_name` is used in collector mode. The collector only reads the
+action name from `appsignal.action_name`, so root names set this way were
+previously ignored, falling back to the OpenTelemetry span name.

.changesets/sanitize-sql-body-in-collector-mode.md

@@ -0,0 +1,9 @@
+---
+bump: patch
+type: fix
+---
+
+Emit the `db.system.name` and `db.query.text` semantic-convention attributes
+instead of `appsignal.sql_body` when `set_sql_body` is used in collector
+mode. This fixes SQL query sanitization when using the experimental collector
+mode. 

.changesets/update-span-name-in-collector-mode.md

@@ -0,0 +1,9 @@
+---
+bump: patch
+type: fix
+---
+
+Update the OpenTelemetry span name directly when `set_name` is used in
+collector mode, instead of setting the `appsignal.name` attribute. The
+collector uses the span name as-is, so names set this way were previously
+ignored.

src/appsignal/tracing.py

@@ -44,6 +44,24 @@ def _set_prefixed_attribute(
         _set_attribute(f"{prefix}.{suffix}", value, span)
 
 
+def _update_span_name(name: str, span: Span | None = None) -> None:
+    span = span or trace.get_current_span()
+
+    if span is trace.INVALID_SPAN:
+        logger.debug("There is no active span, cannot set `name`")
+        return
+
+    span.update_name(name)
+
+
+def _use_collector() -> bool:
+    # Imported here to avoid an import cycle with the client module.
+    from .client import Client
+
+    config = Client.config()
+    return config is not None and config.should_use_collector()
+
+
 def set_params(params: Any, span: Span | None = None) -> None:
     _set_serialised_attribute("appsignal.request.parameters", params, span)
 
@@ -65,7 +83,13 @@ def set_header(header: str, value: Any, span: Span | None = None) -> None:
 
 
 def set_name(name: str, span: Span | None = None) -> None:
-    _set_attribute("appsignal.name", name, span)
+    # The collector uses the span name directly, and its enrichers cannot
+    # override it, so update the span name. The agent derives the name from
+    # semantic conventions, so set an attribute it honors as a hard override.
+    if _use_collector():
+        _update_span_name(name, span)
+    else:
+        _set_attribute("appsignal.name", name, span)
 
 
 def set_category(category: str, span: Span | None = None) -> None:
@@ -77,15 +101,27 @@ def set_body(body: str, span: Span | None = None) -> None:
 
 
 def set_sql_body(body: str, span: Span | None = None) -> None:
-    _set_attribute("appsignal.sql_body", body, span)
+    # The collector sanitizes SQL via the `db.query.text` attribute, gated on
+    # `db.system.name` being set to a recognized SQL system. The agent
+    # sanitizes via the `appsignal.sql_body` magic attribute.
+    if _use_collector():
+        _set_attribute("db.system.name", "other_sql", span)
+        _set_attribute("db.query.text", body, span)
+    else:
+        _set_attribute("appsignal.sql_body", body, span)
 
 
 def set_namespace(namespace: str, span: Span | None = None) -> None:
     _set_attribute("appsignal.namespace", namespace, span)
 
 
 def set_root_name(root_name: str, span: Span | None = None) -> None:
-    _set_attribute("appsignal.root_name", root_name, span)
+    # The collector reads the action name from `appsignal.action_name`, while
+    # the agent reads it from `appsignal.root_name`.
+    if _use_collector():
+        _set_attribute("appsignal.action_name", root_name, span)
+    else:
+        _set_attribute("appsignal.root_name", root_name, span)
 
 
 def set_error(error: Exception, span: Span | None = None) -> None:

tests/test_tracing.py

@@ -19,6 +19,7 @@
     set_sql_body,
     set_tag,
 )
+from appsignal.client import Client
 
 
 tracer = trace.get_tracer("appsignal/tests")
@@ -64,6 +65,55 @@ def test_set_attributes(spans):
     }
 
 
+def test_set_root_name_collector_mode(spans):
+    Client(
+        active=True,
+        name="MyApp",
+        push_api_key="0000-0000-0000-0000",
+        collector_endpoint="https://custom-endpoint.appsignal.com",
+    )
+
+    with tracer.start_as_current_span("span"):
+        set_root_name("Root name")
+
+    attributes = dict(spans()[0].attributes)
+    assert attributes["appsignal.action_name"] == "Root name"
+    assert "appsignal.root_name" not in attributes
+
+
+def test_set_name_collector_mode(spans):
+    Client(
+        active=True,
+        name="MyApp",
+        push_api_key="0000-0000-0000-0000",
+        collector_endpoint="https://custom-endpoint.appsignal.com",
+    )
+
+    with tracer.start_as_current_span("span"):
+        set_name("New name")
+
+    span = spans()[0]
+    assert span.name == "New name"
+    assert "appsignal.name" not in dict(span.attributes)
+
+
+def test_set_sql_body_collector_mode(spans):
+    Client(
+        active=True,
+        name="MyApp",
+        push_api_key="0000-0000-0000-0000",
+        collector_endpoint="https://custom-endpoint.appsignal.com",
+    )
+
+    with tracer.start_as_current_span("span"):
+        set_sql_body("SELECT * FROM users")
+
+    attributes = dict(spans()[0].attributes)
+    assert attributes["db.system.name"] == "other_sql"
+    assert attributes["db.query.text"] == "SELECT * FROM users"
+    assert "appsignal.sql_body" not in attributes
+
+
 def test_set_attributes_on_span(spans):
     with tracer.start_as_current_span("parent span") as span:
         with tracer.start_as_current_span("child span"):