From 47a0f5d03bbd66d4ff945dc7deebbb20f53b5721 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Alex=20Szab=C3=B3?= Date: Mon, 16 Dec 2024 12:05:15 +0100 Subject: [PATCH] Implement defensive programming techniques for setting spans in ContextVars to avoid accidentally trying to merge a None with a tuple of Spans Related to #2056 --- elasticapm/context/contextvars.py | 2 +- tests/context/test_context.py | 12 ++++++++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/elasticapm/context/contextvars.py b/elasticapm/context/contextvars.py index e9ac2ea82..14b7f871e 100644 --- a/elasticapm/context/contextvars.py +++ b/elasticapm/context/contextvars.py @@ -78,7 +78,7 @@ def set_span(self, span: "elasticapm.traces.Span") -> None: The previously-activated span will be saved to be re-activated later. """ - spans = self.elasticapm_spans_var.get() + spans: tuple = self.elasticapm_spans_var.get() or () self.elasticapm_spans_var.set(spans + (span,)) def unset_span(self, clear_all: bool = False) -> "elasticapm.traces.Span": diff --git a/tests/context/test_context.py b/tests/context/test_context.py index 058d60bab..a74fed362 100644 --- a/tests/context/test_context.py +++ b/tests/context/test_context.py @@ -30,10 +30,13 @@ import sys +import mock import pytest import elasticapm.context +from elasticapm.context.contextvars import ContextVarsContext from elasticapm.context.threadlocal import ThreadLocalContext +from elasticapm.traces import Span def test_execution_context_backing(): @@ -63,3 +66,12 @@ def test_execution_context_monkeypatched(monkeypatch): # Should always use ThreadLocalContext when thread local is monkey patched assert isinstance(execution_context, ThreadLocalContext) + + +def test_none_spans_should_not_raise_a_type_error_on_set_span(): + context = ContextVarsContext() + context.elasticapm_spans_var.set(None) + + context.set_span(mock.MagicMock(spec=Span)) + + assert context.get_span() is not None