diff --git a/posthog/test/tracing/test_config.py b/posthog/test/tracing/test_config.py index 29fee059..2bc1af5e 100644 --- a/posthog/test/tracing/test_config.py +++ b/posthog/test/tracing/test_config.py @@ -229,3 +229,51 @@ def test_an_unusable_value_falls_back_rather_than_dropping_every_span(self, valu assert resolved.max_attributes_per_span == DEFAULT_MAX_ATTRIBUTES_PER_SPAN assert resolved.max_events_per_span == DEFAULT_MAX_EVENTS_PER_SPAN assert resolved.max_attribute_value_length == DEFAULT_MAX_ATTRIBUTE_VALUE_LENGTH + + +class TestBeforeSpanSendConfig: + def test_accepts_one_hook_or_a_list(self): + def hook(span): + return span + + assert resolve_traces_config({"before_span_send": hook}).before_span_send == ( + hook, + ) + assert resolve_traces_config( + {"before_span_send": [hook, hook]} + ).before_span_send == (hook, hook) + + def test_defaults_to_no_hooks(self): + assert resolve_traces_config({}).before_span_send == () + + def test_skips_falsy_entries_silently(self, caplog): + caplog.set_level("WARNING", logger="posthog") + + def hook(span): + return span + + resolved = resolve_traces_config({"before_span_send": [None, False, hook]}) + assert resolved.before_span_send == (hook,) + assert not caplog.records + + def test_rejects_an_entry_that_is_not_callable_rather_than_exporting_unhooked( + self, + ): + def hook(span): + return span + + with pytest.raises(ValueError, match="not callable"): + resolve_traces_config({"before_span_send": ["scrub", hook]}) + + def test_a_hook_whose_truthiness_raises_is_still_resolved(self): + class Hook: + def __bool__(self): + raise RuntimeError("no") + + def __call__(self, span): + return span + + hook = Hook() + assert resolve_traces_config({"before_span_send": hook}).before_span_send == ( + hook, + ) diff --git a/posthog/test/tracing/test_limits.py b/posthog/test/tracing/test_limits.py index 123c57fc..7799c5a7 100644 --- a/posthog/test/tracing/test_limits.py +++ b/posthog/test/tracing/test_limits.py @@ -4,12 +4,14 @@ from posthog.tracing import _limits as limits_module from posthog.tracing._limits import ( + apply_span_limits, bound_attributes, truncate_attribute_value, truncate_attributes, ) from posthog.tracing._otlp import ( CIRCULAR_VALUE, + SpanRecord, MAX_VALUE_ITEMS, MAX_VALUE_NODES, TRUNCATED_VALUE, @@ -181,3 +183,18 @@ def test_stops_walking_a_mapping_at_the_encoders_item_cap(self): value = {f"k{i}": "v" * 50 for i in range(MAX_VALUE_ITEMS + 50)} bounded = truncate_attribute_value(value, 4) assert len(bounded) == MAX_VALUE_ITEMS + + +class TestApplySpanLimits: + def test_rebounds_a_container_the_hook_grew_in_place(self): + grown = ["x" * 3] + record = SpanRecord("t", "s", "n", 1, 2, attributes={"k": grown}) + grown.append("y" * 20) + apply_span_limits(record, frozenset(), 128, 128, 128, 8) + assert record.attributes["k"] == ["xxx", "y" * 8] + + def test_an_empty_key_spends_no_slot(self): + record = SpanRecord("t", "s", "n", 1, 2, attributes={"": 1, "a": 2, "b": 3}) + apply_span_limits(record, frozenset(), 2, 128, 128, 8) + assert record.attributes == {"a": 2, "b": 3} + assert record.dropped_attributes_count == 0 diff --git a/posthog/test/tracing/test_pipeline.py b/posthog/test/tracing/test_pipeline.py index 7920cb6b..da6dbab4 100644 --- a/posthog/test/tracing/test_pipeline.py +++ b/posthog/test/tracing/test_pipeline.py @@ -2,6 +2,7 @@ import logging import threading import time +import warnings import weakref from types import SimpleNamespace from unittest import mock @@ -12,6 +13,7 @@ SPAN_ID, TRACE_ID, FakeSender, + FakeTimer, clock, fake_timers, make, @@ -20,9 +22,10 @@ ) from posthog.tracing import _pipeline as pipeline_module from posthog.tracing import _span as span_module +from posthog.tracing._config import resolve_traces_config from posthog.tracing._drops import DropLog -from posthog.tracing._transport import SendOutcome from posthog.tracing._span import NOOP_SPAN, PassThroughSpan, RecordingSpan +from posthog.tracing._transport import SendOutcome __all__ = ["clock", "fake_timers"] @@ -626,3 +629,395 @@ def test_reports_a_spans_limit_drops_once_at_debug(self, caplog): 'Span limits discarded data from "capped": 1 attributes, 1 events, ' "0 event attributes" ) + + +class TestBeforeSpanSend: + def test_a_hook_returning_none_drops_the_span_quietly(self, caplog): + caplog.set_level("DEBUG", logger="posthog") + pipeline, _, _ = make(before_span_send=lambda span: None) + pipeline.start_span("a").end() + pipeline.close() + assert queued(pipeline) == [] + assert not [r for r in caplog.records if r.levelname == "WARNING"] + assert "before_span_send dropped the span" in caplog.text + + def test_a_raising_hook_drops_the_span_rather_than_exporting_it(self, caplog): + caplog.set_level("WARNING", logger="posthog") + + def broken(span): + raise RuntimeError("scrubber bug") + + pipeline, _, _ = make(before_span_send=broken) + pipeline.start_span("a", attributes={"password": "hunter2"}).end() + assert queued(pipeline) == [] + assert any("before_span_send failed" in r.getMessage() for r in caplog.records) + + def test_a_raising_hook_warns_with_its_traceback_once_per_interval( + self, caplog, clock + ): + caplog.set_level("WARNING", logger="posthog") + + def broken(span): + raise RuntimeError("scrubber bug") + + pipeline, _, _ = make(before_span_send=broken, flush_interval=5) + for _ in range(3): + pipeline.start_span("a").end() + clock["now"] += 6 + pipeline.start_span("a").end() + warned = [r for r in caplog.records if "before_span_send raised" in r.message] + assert len(warned) == 2 + assert "RuntimeError: scrubber bug" in caplog.text + + def test_an_async_hook_is_named_as_unsupported(self, caplog): + caplog.set_level("WARNING", logger="posthog") + + async def hook(span): + return span + + pipeline, _, _ = make(before_span_send=hook) + pipeline.start_span("a").end() + pipeline.close() + assert "before_span_send is async, which is not supported" in caplog.text + + def test_the_hook_sees_plain_values_not_the_wire_encoding(self): + seen = {} + + def hook(span): + seen.update(span) + return span + + pipeline, _, _ = make(before_span_send=hook, context={"distinct_id": "u1"}) + pipeline.start_span("checkout", kind="server", attributes={"userId": 42}).end() + assert seen["attributes"] == {"posthogDistinctId": "u1", "userId": 42} + assert seen["name"] == "checkout" + assert seen["kind"] == "server" + assert seen["status"] is None + assert len(seen["trace_id"]) == 32 and len(seen["span_id"]) == 16 + assert seen["parent_span_id"] is None + assert isinstance(seen["start_time_ns"], int) + assert seen["end_time_ns"] >= seen["start_time_ns"] + + def test_redacts_in_place(self): + def scrub(span): + span["attributes"].pop("http.request.header.authorization", None) + return span + + pipeline, _, _ = make(before_span_send=scrub) + pipeline.start_span( + "a", attributes={"http.request.header.authorization": "Bearer x", "ok": 1} + ).end() + assert queued(pipeline)[0].attributes == {"ok": 1} + + def test_identity_fields_are_readable_but_not_writable(self, caplog): + caplog.set_level("DEBUG", logger="posthog") + + def forge(span): + span["trace_id"] = "f" * 32 + span["span_id"] = "e" * 16 + return span + + pipeline, _, _ = make(before_span_send=forge) + span = pipeline.start_span("a", parent=f"00-{TRACE_ID}-{SPAN_ID}-00") + span_id = span._span_id + span.end() + record = queued(pipeline)[0] + assert (record.trace_id, record.span_id) == (TRACE_ID, span_id) + assert record.parent_span_id == SPAN_ID + assert any("identity field" in r.getMessage() for r in caplog.records) + + def test_propagation_state_survives_a_hook_that_rebuilds_the_dict(self): + def rebuild(span): + return {k: v for k, v in span.items() if k != "trace_id"} + + pipeline, _, _ = make(before_span_send=rebuild) + pipeline.start_span( + "a", parent=f"00-{TRACE_ID}-{SPAN_ID}-00", tracestate="vendor=abc" + ).end() + record = queued(pipeline)[0] + assert record.trace_id == TRACE_ID + assert record.trace_flags == "00" + assert record.parent_is_remote is True + assert record.trace_state == "vendor=abc" + + def test_runs_a_list_left_to_right_and_the_first_none_stops_it(self): + calls = [] + + def first(span): + calls.append("first") + span["attributes"]["step"] = 1 + return span + + def second(span): + calls.append("second") + return None + + def third(span): + calls.append("third") + return span + + pipeline, _, _ = make(before_span_send=[first, second, third]) + pipeline.start_span("a").end() + assert calls == ["first", "second"] + assert queued(pipeline) == [] + + def test_a_hook_can_remove_the_stacktrace(self): + def strip_stacks(span): + for event in span["events"]: + event["attributes"].pop("exception.stacktrace", None) + return span + + pipeline, _, _ = make(before_span_send=strip_stacks) + with pytest.raises(ValueError): + with pipeline.start_span("job"): + raise ValueError("boom") + (event,) = queued(pipeline)[0].events + assert event.attributes == { + "exception.type": "ValueError", + "exception.message": "boom", + } + + def test_reapplies_the_caps_to_what_the_hook_added(self): + def enrich(span): + for i in range(10): + span["attributes"][f"extra{i}"] = "x" * 50 + span["events"].extend({"name": f"e{i}"} for i in range(5)) + return span + + pipeline, _, _ = make( + before_span_send=enrich, + context={"distinct_id": "u1"}, + max_attributes_per_span=3, + max_events_per_span=2, + max_attribute_value_length=10, + ) + pipeline.start_span("a", attributes={"first": 1}).end() + record = queued(pipeline)[0] + assert list(record.attributes) == [ + "posthogDistinctId", + "first", + "extra0", + "extra1", + ] + assert record.attributes["extra0"] == "x" * 10 + assert record.dropped_attributes_count == 8 + assert [e.name for e in record.events] == ["e0", "e1"] + assert record.dropped_events_count == 3 + + def test_keeps_the_counts_the_span_already_dropped(self): + pipeline, _, _ = make( + before_span_send=lambda span: span, max_attributes_per_span=1 + ) + span = pipeline.start_span("a", attributes={"a": 1, "b": 2}) + span.add_event("e", {"k1": 1}) + span.end() + assert queued(pipeline)[0].dropped_attributes_count == 1 + + @pytest.mark.parametrize( + "result", + [ + "not a dict", + 42, + {"name": "a"}, + {"attributes": {}, "events": "no"}, + ], + ) + def test_drops_a_value_that_is_not_a_span_dict(self, result): + pipeline, _, _ = make(before_span_send=lambda span: result) + pipeline.start_span("a").end() + assert queued(pipeline) == [] + + def test_drops_the_span_from_an_async_hook_without_a_never_awaited_warning(self): + async def hook(span): + return span + + pipeline, _, _ = make(before_span_send=hook) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + pipeline.start_span("a").end() + gc.collect() + assert queued(pipeline) == [] + assert not [w for w in caught if "never awaited" in str(w.message)] + + def test_resanitizes_names_times_status_and_events(self, caplog): + caplog.set_level("DEBUG", logger="posthog") + + def mangle(span): + span["name"] = "" + span["start_time_ns"] = -5 + span["end_time_ns"] = "later" + span["status"] = {"code": "maybe"} + span["events"].append({"name": None, "timestamp_ns": 2**70}) + span["events"].append(None) + return span + + pipeline, _, _ = make(before_span_send=mangle) + span = pipeline.start_span("a", start_time=1_700_000_000) + span.set_status("error", "real failure") + span.end(end_time=1_700_000_005) + record = queued(pipeline)[0] + assert record.name == "unknown" + assert record.start_ns == 1_700_000_000 * 10**9 + assert record.end_ns == 1_700_000_005 * 10**9 + assert record.status.code == "error" + assert record.status.message == "real failure" + (event,) = record.events + assert event.name == "unknown" + assert event.timestamp_ns == record.start_ns + assert any( + "not an epoch-nanosecond int" in r.getMessage() for r in caplog.records + ) + + def test_clears_the_status_when_the_hook_sets_none(self): + def clear(span): + span["status"] = None + return span + + pipeline, _, _ = make(before_span_send=clear) + span = pipeline.start_span("a") + span.set_status("error", "x") + span.end() + assert queued(pipeline)[0].status is None + + def test_a_hook_that_ends_a_span_of_its_own_does_not_deadlock(self): + pipeline, _, _ = make() + seen = [] + + def nested(span): + if span["name"] == "outer": + pipeline.start_span("from-hook").end() + seen.append(span["name"]) + return span + + pipeline._config = resolve_traces_config({"before_span_send": nested}) + pipeline.start_span("outer").end() + assert sorted(seen) == ["from-hook", "outer"] + assert sorted(r.name for r in queued(pipeline)) == ["from-hook", "outer"] + + def test_is_not_called_for_a_span_whose_client_was_disabled(self): + client = SimpleNamespace(disabled=False, send=True) + hook = mock.Mock(side_effect=lambda span: span) + pipeline, _, _ = make(client=client, before_span_send=hook) + span = pipeline.start_span("a") + client.disabled = True + span.end() + assert not hook.called + + def test_a_span_whose_hook_outlives_close_is_not_queued(self): + # shutdown() closes tracing while another thread's hook is running: + # that span must not reach a queue nothing will flush. + pipeline, _, _ = make_traces() + + def close_mid_hook(span): + pipeline.close() + return span + + pipeline._config = resolve_traces_config({"before_span_send": close_mid_hook}) + timers_before = len(FakeTimer.instances) + pipeline.start_span("a").end() + assert queued(pipeline) == [] + assert len(FakeTimer.instances) == timers_before + + def test_the_limit_report_is_logged_with_no_tracing_lock_held(self): + pipeline, _, _ = make_traces(max_attributes_per_span=1) + held = [] + original = pipeline_module._report_limit_drops + + def spy(record): + held.append((pipeline._lock.locked(), pipeline._exporter._lock.locked())) + original(record) + + with mock.patch.object(pipeline_module, "_report_limit_drops", spy): + pipeline.start_span("a", attributes={"x": 1, "y": 2}).end() + assert held == [(False, False)] + + +class TestBeforeSpanSendBounds: + def test_bounds_a_status_message_the_hook_sets(self): + def long_message(span): + span["status"] = {"code": "error", "message": "m" * 100} + return span + + pipeline, _, _ = make( + before_span_send=long_message, max_attribute_value_length=5 + ) + pipeline.start_span("a").end() + assert queued(pipeline)[0].status.message == "mmmmm" + + def test_keeps_a_falsy_status_message_the_hook_sets(self): + def zero_message(span): + span["status"] = {"code": "error", "message": 0} + return span + + pipeline, _, _ = make(before_span_send=zero_message) + pipeline.start_span("a").end() + assert queued(pipeline)[0].status.message == "0" + + def test_a_status_message_whose_str_raises_keeps_the_span(self): + class Hostile: + def __str__(self): + raise RuntimeError("no") + + def hostile_message(span): + span["status"] = {"code": "error", "message": Hostile()} + return span + + pipeline, _, _ = make(before_span_send=hostile_message) + pipeline.start_span("a").end() + assert queued(pipeline)[0].status.message == "[Unserializable]" + + def test_a_hook_can_scrub_the_auto_context_keys(self): + def scrub(span): + span["attributes"].pop("posthogDistinctId") + span["attributes"].pop("sessionId") + return span + + pipeline, _, _ = make( + before_span_send=scrub, context={"distinct_id": "u", "session_id": "s"} + ) + pipeline.start_span("a").end() + assert queued(pipeline)[0].attributes == {} + + def test_recaps_event_attributes_the_hook_widens(self): + def widen(span): + span["events"][0]["attributes"].update({f"k{i}": i for i in range(200)}) + return span + + pipeline, _, _ = make(before_span_send=widen) + span = pipeline.start_span("a") + span.add_event("e", {"first": 1}) + span.end() + (event,) = queued(pipeline)[0].events + assert len(event.attributes) == 128 + assert event.dropped_attributes_count == 73 + + def test_keeps_an_events_dropped_count_when_the_hook_rebuilds_the_events(self): + def rebuild(span): + span["events"] = [dict(event) for event in span["events"]] + return span + + pipeline, _, _ = make(before_span_send=rebuild) + span = pipeline.start_span("a") + span.add_event("wide", {f"k{i}": i for i in range(130)}) + span.end() + assert queued(pipeline)[0].events[0].dropped_attributes_count == 2 + + def test_a_field_the_hook_leaves_out_keeps_the_originals_value(self): + def allowlist(span): + return {"attributes": {"kept": 1}, "events": []} + + pipeline, _, _ = make(before_span_send=allowlist) + pipeline.start_span("a", kind="server").end() + record = queued(pipeline)[0] + assert record.name == "a" + assert record.kind == "server" + assert record.start_ns <= record.end_ns + assert record.attributes == {"kept": 1} + + def test_drops_a_dict_without_attributes_and_events(self, caplog): + caplog.set_level("WARNING", logger="posthog") + pipeline, _, _ = make(before_span_send=lambda span: {"name": "x"}) + pipeline.start_span("a").end() + pipeline.close() + assert queued(pipeline) == [] + assert "unusable record" in caplog.text diff --git a/posthog/tracing/_before_span_send.py b/posthog/tracing/_before_span_send.py new file mode 100644 index 00000000..577fe454 --- /dev/null +++ b/posthog/tracing/_before_span_send.py @@ -0,0 +1,208 @@ +"""The ``before_span_send`` hook chain, where a finished span is scrubbed or dropped. + +The hook receives a plain dict, like the events ``before_send`` hook. It is the +scrubbing point, so a hook that raises drops the span rather than letting the +unscrubbed record through. +""" + +import inspect +import logging +from typing import Any, Dict, Mapping, Optional + +from ._config import MAX_ATTRIBUTES_PER_EVENT, ResolvedTracesConfig +from ._drops import DropLog +from ._limits import apply_span_limits +from ._otlp import SpanEventRecord, SpanRecord, SpanStatus, non_negative_count +from ._sanitize import ( + MAX_TIMESTAMP_NS, + MIN_TIMESTAMP_NS, + clamp_end_ns, + copy_user_attributes, + safe_str, + sanitize_name, +) + +log = logging.getLogger("posthog") + +_READ_ONLY_ID_KEYS = ("trace_id", "span_id", "parent_span_id") + + +def run_before_span_send( + record: SpanRecord, config: ResolvedTracesConfig, drops: DropLog +) -> Optional[SpanRecord]: + """Run the hook chain, returning the span to queue or ``None`` to drop it. + + The ids are read-only: they are restored after every hook, since a + rewritten id would orphan children already sent. + """ + if not config.before_span_send: + return record + + identity = (record.trace_id, record.span_id, record.parent_span_id) + # The span's own write order, so the caps keep the earliest-set entries. + keys_before_hook = list(record.attributes) + data = _hook_view(record) + try: + for hook in config.before_span_send: + result: Any = hook(data) + if result is None: + # The documented way to filter, so not a drop worth warning about. + log.debug("before_span_send dropped the span") + return None + if inspect.iscoroutine(result): + # Not awaited; closed so it does not warn. + result.close() + drops.record(1, "before_span_send is async, which is not supported") + return None + if not isinstance(result, Mapping): + log.debug( + "before_span_send did not return a span dict; dropping the span" + ) + drops.record(1, "before_span_send returned an unusable record") + return None + data = _keep_identity(result, identity) + rebuilt = _rebuild(record, data, config) + if rebuilt is None: + log.debug("before_span_send did not return a span dict; dropping the span") + drops.record(1, "before_span_send returned an unusable record") + return None + apply_span_limits( + rebuilt, + record.auto_attribute_keys, + config.max_attributes_per_span, + config.max_events_per_span, + MAX_ATTRIBUTES_PER_EVENT, + config.max_attribute_value_length, + keys_before_hook, + ) + return rebuilt + except Exception: + drops.warn_failure( + "before_span_send raised; dropping the span rather than exporting it " + "unscrubbed" + ) + drops.record(1, "before_span_send failed") + return None + + +def _hook_view(record: SpanRecord) -> Dict[str, Any]: + return { + "trace_id": record.trace_id, + "span_id": record.span_id, + "parent_span_id": record.parent_span_id, + "name": record.name, + "kind": record.kind, + "status": ( + {"code": record.status.code, "message": record.status.message} + if record.status is not None + else None + ), + "attributes": dict(record.attributes), + "events": [ + { + "name": event.name, + "timestamp_ns": event.timestamp_ns, + "attributes": dict(event.attributes or {}), + "dropped_attributes_count": event.dropped_attributes_count, + } + for event in record.events + ], + "start_time_ns": record.start_ns, + "end_time_ns": record.end_ns, + } + + +def _keep_identity(result: Mapping, identity: tuple) -> Dict[str, Any]: + data = result if isinstance(result, dict) else dict(result) + if any(data.get(key) != value for key, value in zip(_READ_ONLY_ID_KEYS, identity)): + log.debug( + "before_span_send changed a span identity field; keeping the original ids" + ) + data.update(zip(_READ_ONLY_ID_KEYS, identity)) + return data + + +def _rebuild( + record: SpanRecord, data: Mapping, config: ResolvedTracesConfig +) -> Optional[SpanRecord]: + """A span record from what the chain returned, sanitized as ``end()`` would. + + ``None`` when it is not a span dict, rather than exporting a span of + fallbacks joinable to nothing. A field the hook left out or made unusable + keeps the original's value. + """ + attributes = data.get("attributes") + events = data.get("events") + if not isinstance(attributes, Mapping) or not isinstance(events, (list, tuple)): + return None + + max_length = config.max_attribute_value_length + start_ns = _valid_ns(data.get("start_time_ns"), record.start_ns) + end_ns = clamp_end_ns(_valid_ns(data.get("end_time_ns"), record.end_ns), start_ns) + kind = data.get("kind") + + rebuilt_events = [] + for event in events: + try: + event_attributes = event.get("attributes") + rebuilt_events.append( + SpanEventRecord( + name=sanitize_name( + event.get("name"), "Span event name", max_length + ), + timestamp_ns=_valid_ns(event.get("timestamp_ns"), start_ns), + attributes=copy_user_attributes({}, event_attributes) + if event_attributes is not None + else None, + dropped_attributes_count=non_negative_count( + event.get("dropped_attributes_count") + ), + ) + ) + except Exception: + log.debug("before_span_send left an unreadable span event; dropping it") + + return SpanRecord( + trace_id=record.trace_id, + span_id=record.span_id, + parent_span_id=record.parent_span_id, + # SDK bookkeeping the hook is not shown, so it cannot erase it. + trace_state=record.trace_state, + trace_flags=record.trace_flags, + parent_is_remote=record.parent_is_remote, + dropped_attributes_count=record.dropped_attributes_count, + dropped_events_count=record.dropped_events_count, + auto_attribute_keys=record.auto_attribute_keys, + name=sanitize_name(data.get("name", record.name), "Span name", max_length), + kind=kind if isinstance(kind, str) else record.kind, + status=_hook_status(data.get("status"), record.status), + attributes=copy_user_attributes({}, attributes), + events=rebuilt_events, + start_ns=start_ns, + end_ns=end_ns, + ) + + +def _hook_status(value: Any, original: Optional[SpanStatus]) -> Optional[SpanStatus]: + if value is None: + return None + if isinstance(value, Mapping) and value.get("code") in ("ok", "error"): + message = value.get("message") + text = None if message is None else safe_str(message) + return SpanStatus(value["code"], text or None) + # An unknown code would lose an error the span really had. + log.debug("before_span_send set an unknown span status; keeping the original") + return original + + +def _valid_ns(value: Any, fallback: int) -> int: + if ( + isinstance(value, int) + and not isinstance(value, bool) + and MIN_TIMESTAMP_NS <= value <= MAX_TIMESTAMP_NS + ): + return value + log.debug( + "before_span_send set a time that is not an epoch-nanosecond int; keeping the original" + ) + return fallback diff --git a/posthog/tracing/_config.py b/posthog/tracing/_config.py index 71574d63..579d9ff4 100644 --- a/posthog/tracing/_config.py +++ b/posthog/tracing/_config.py @@ -6,8 +6,9 @@ import logging import math from dataclasses import dataclass, field, fields -from typing import Any, Dict, Mapping, Optional +from typing import Any, Callable, Dict, Mapping, Optional, Tuple +from ..types import BeforeSpanSendCallback from ._sanitize import attribute_key log = logging.getLogger("posthog") @@ -47,6 +48,7 @@ class ResolvedTracesConfig: max_attributes_per_span: int = DEFAULT_MAX_ATTRIBUTES_PER_SPAN max_events_per_span: int = DEFAULT_MAX_EVENTS_PER_SPAN max_attribute_value_length: int = DEFAULT_MAX_ATTRIBUTE_VALUE_LENGTH + before_span_send: Tuple[BeforeSpanSendCallback, ...] = () _KNOWN_KEYS = frozenset(field.name for field in fields(ResolvedTracesConfig)) @@ -120,6 +122,27 @@ def _usable_resource_attributes(value: Any) -> Dict[str, Any]: return attributes +def _resolve_before_span_send(value: Any) -> Tuple[Callable, ...]: + """The hooks, in order. Raises for an entry that is not callable. + + The hook is the scrubbing point, so a broken entry turns tracing off + rather than exporting spans the entry was meant to redact. The client + reports the error and leaves tracing off for the life of the client. + """ + if value is None: + return () + supplied = list(value) if isinstance(value, (list, tuple)) else [value] + # `[enabled and scrub]` yields None or False: no hook, rather than a broken one. + supplied = [hook for hook in supplied if hook is not None and hook is not False] + for hook in supplied: + if not callable(hook): + raise ValueError( + "traces before_span_send entry {!r} is not callable; tracing is off " + "rather than exporting spans it was meant to redact".format(hook) + ) + return tuple(supplied) + + def resolve_traces_config( config: Any, host_resource_attributes: Optional[Mapping[str, str]] = None ) -> ResolvedTracesConfig: @@ -182,4 +205,5 @@ def resolve_traces_config( max_attribute_value_length=_positive_int( config, "max_attribute_value_length", DEFAULT_MAX_ATTRIBUTE_VALUE_LENGTH ), + before_span_send=_resolve_before_span_send(config.get("before_span_send")), ) diff --git a/posthog/tracing/_drops.py b/posthog/tracing/_drops.py index 947055e1..33cbcd3c 100644 --- a/posthog/tracing/_drops.py +++ b/posthog/tracing/_drops.py @@ -22,6 +22,7 @@ def __init__(self, interval: float) -> None: # A dict for its order: reasons are named in the order they happened. self._reasons: Dict[str, None] = {} self._last_warning_at = 0.0 + self._last_failure_at = 0.0 def record(self, count: int, reason: str) -> None: with self._lock: @@ -47,8 +48,25 @@ def warn_if_due(self, force: bool = False) -> None: # A raising logging handler must not surface through span creation. pass + def warn_failure(self, message: str) -> None: + """Warn with the current traceback, at most once per interval; debug otherwise. + + For a failure in application code, where the aggregate count alone + gives nothing to act on. + """ + with self._lock: + now = time.monotonic() + due = now - self._last_failure_at >= self._interval + if due: + self._last_failure_at = now + try: + log.log(logging.WARNING if due else logging.DEBUG, message, exc_info=True) + except Exception: + pass + def reinit_after_fork(self) -> None: self._lock = threading.Lock() self._count = 0 self._reasons.clear() self._last_warning_at = 0.0 + self._last_failure_at = 0.0 diff --git a/posthog/tracing/_limits.py b/posthog/tracing/_limits.py index cdec65ad..faeb5542 100644 --- a/posthog/tracing/_limits.py +++ b/posthog/tracing/_limits.py @@ -7,7 +7,7 @@ """ from datetime import date -from typing import Any, Dict, Mapping, Tuple +from typing import AbstractSet, Any, Dict, List, Mapping, Sequence, Tuple from ._otlp import ( CIRCULAR_VALUE, @@ -15,6 +15,9 @@ MAX_VALUE_ITEMS, MAX_VALUE_NODES, TRUNCATED_VALUE, + SpanRecord, + SpanStatus, + non_negative_count, ) from ._sanitize import UNSERIALIZABLE_VALUE, attribute_key @@ -141,3 +144,72 @@ def truncate_attributes(attributes: Mapping, max_length: int) -> Dict[str, Any]: key: truncate_attribute_value(value, max_length) for key, value in attributes.items() } + + +def _ordered_keys(attributes: Mapping, keys_before_hook: Sequence[str]) -> List[Any]: + """The keys the span set first, in its order, so the earliest-set entries win.""" + keys = list(attributes.keys()) + present = set(keys) + before = [key for key in keys_before_hook if key in present] + seen = set(before) + return before + [key for key in keys if key not in seen] + + +def apply_span_limits( + record: SpanRecord, + auto_keys: AbstractSet[str], + max_attributes: int, + max_events: int, + max_attributes_per_event: int, + max_length: int, + keys_before_hook: Sequence[str] = (), +) -> None: + """Re-apply the per-span caps after a ``before_span_send`` hook, which + bypasses the span's own writer. Counts add to what the span already dropped. + + Every container is walked again: the hook holds the same nested objects + the span stored, so one it grew in place has the same identity.""" + attributes: Dict[str, Any] = {} + kept = 0 + dropped_attributes = 0 + for key in _ordered_keys(record.attributes, keys_before_hook): + value = record.attributes[key] + if value is None or not key: + continue + if key not in auto_keys: + if kept >= max_attributes: + dropped_attributes += 1 + continue + kept += 1 + attributes[key] = truncate_attribute_value(value, max_length) + record.attributes = attributes + if dropped_attributes: + record.dropped_attributes_count = ( + non_negative_count(record.dropped_attributes_count) + dropped_attributes + ) + + kept_events: list = [] + dropped_events = 0 + for event in record.events: + if len(kept_events) >= max_events: + dropped_events += 1 + continue + if event.attributes: + event.attributes, dropped = bound_attributes( + event.attributes, max_attributes_per_event, max_length + ) + if dropped: + event.dropped_attributes_count = ( + non_negative_count(event.dropped_attributes_count) + dropped + ) + kept_events.append(event) + record.events = kept_events + if dropped_events: + record.dropped_events_count = ( + non_negative_count(record.dropped_events_count) + dropped_events + ) + + if record.status is not None and record.status.message: + record.status = SpanStatus( + record.status.code, truncate_string(record.status.message, max_length) + ) diff --git a/posthog/tracing/_otlp.py b/posthog/tracing/_otlp.py index 735ced0b..d5e0745b 100644 --- a/posthog/tracing/_otlp.py +++ b/posthog/tracing/_otlp.py @@ -11,7 +11,7 @@ import platform from dataclasses import dataclass, field from datetime import date, datetime -from typing import Any, Dict, List, Mapping, Optional, Tuple +from typing import Any, Dict, FrozenSet, List, Mapping, Optional, Tuple from ..version import VERSION from ._sanitize import FUNCTION_VALUE, UNSERIALIZABLE_VALUE, attribute_key, safe_str @@ -80,6 +80,9 @@ class SpanRecord: # User attributes and events the per-span caps refused. dropped_attributes_count: int = 0 dropped_events_count: int = 0 + # Keys the SDK attached itself, exempt from the attribute cap when it is + # re-applied after before_span_send. + auto_attribute_keys: FrozenSet[str] = frozenset() MAX_UINT32 = 0xFFFFFFFF diff --git a/posthog/tracing/_pipeline.py b/posthog/tracing/_pipeline.py index 65bebd83..3e8cd735 100644 --- a/posthog/tracing/_pipeline.py +++ b/posthog/tracing/_pipeline.py @@ -10,6 +10,7 @@ from contextvars import ContextVar from typing import Any, Callable, Dict, List, Mapping, Optional, Protocol +from ._before_span_send import run_before_span_send from ._config import ResolvedTracesConfig from ._drops import DropLog from ._ids import new_span_id, new_trace_id @@ -301,8 +302,11 @@ def _on_span_end(self, record: SpanRecord) -> None: if disabled: self._drops.record(1, "the client is disabled") else: - _report_limit_drops(record) - self._exporter.enqueue(record) + # The hook is application code, so it runs with no lock held. + hooked = run_before_span_send(record, self._config, self._drops) + if hooked is not None: + _report_limit_drops(hooked) + self._exporter.enqueue(hooked) self._drops.warn_if_due() diff --git a/posthog/tracing/_span.py b/posthog/tracing/_span.py index ed22c4db..f78ae464 100644 --- a/posthog/tracing/_span.py +++ b/posthog/tracing/_span.py @@ -452,6 +452,7 @@ def end(self, end_time: Optional[SpanTimeInput] = None) -> None: end_ns=clamp_end_ns(resolved, self._start_ns), dropped_attributes_count=dropped_attributes, dropped_events_count=dropped_events, + auto_attribute_keys=self._auto_keys, ) try: self._on_end(record) diff --git a/posthog/types.py b/posthog/types.py index 8e8d7a63..b0a4bd1e 100644 --- a/posthog/types.py +++ b/posthog/types.py @@ -8,6 +8,10 @@ # Takes an event dictionary and returns the modified event or None to drop it BeforeSendCallback = Callable[[dict[str, Any]], Optional[dict[str, Any]]] +# Type alias for the traces before_span_send callback function +# Takes a span dictionary and returns the modified span or None to drop it +BeforeSpanSendCallback = Callable[[dict[str, Any]], Optional[dict[str, Any]]] + # Type alias for the send_feature_flags parameter class SendFeatureFlagsOptions(TypedDict, total=False): diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index bc64b740..1b6089d3 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -888,6 +888,7 @@ attribute posthog.send = True attribute posthog.super_properties = None attribute posthog.sync_mode = False attribute posthog.types.BeforeSendCallback = Callable[[dict[str, Any]], Optional[dict[str, Any]]] +attribute posthog.types.BeforeSpanSendCallback = Callable[[dict[str, Any]], Optional[dict[str, Any]]] attribute posthog.types.FeatureFlag.enabled: bool attribute posthog.types.FeatureFlag.key: str attribute posthog.types.FeatureFlag.metadata: Union[FlagMetadata, LegacyFlagMetadata]