Add key predicate to baggage span processor (#2535)

* Add key predicate to baggage span processor

* add changelog entry

* fix linting

* more linter fixes

---------

Co-authored-by: Diego Hurtado <ocelotl@users.noreply.github.com>
This commit is contained in:
Mike Goldsmith
2024-05-28 19:49:00 +01:00
committed by GitHub
parent 88111d0a83
commit 59a737c285
5 changed files with 126 additions and 10 deletions

View File

@ -12,12 +12,16 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import re
import unittest
from opentelemetry.baggage import get_all as get_all_baggage
from opentelemetry.baggage import set_baggage
from opentelemetry.context import attach, detach
from opentelemetry.processor.baggage import BaggageSpanProcessor
from opentelemetry.processor.baggage import (
ALLOW_ALL_BAGGAGE_KEYS,
BaggageSpanProcessor,
)
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SpanProcessor
from opentelemetry.trace import Span, Tracer
@ -25,13 +29,77 @@ from opentelemetry.trace import Span, Tracer
class BaggageSpanProcessorTest(unittest.TestCase):
def test_check_the_baggage(self):
self.assertIsInstance(BaggageSpanProcessor(), SpanProcessor)
self.assertIsInstance(
BaggageSpanProcessor(ALLOW_ALL_BAGGAGE_KEYS), SpanProcessor
)
def test_set_baggage_attaches_to_child_spans_and_detaches_properly_with_context(
self,
):
tracer_provider = TracerProvider()
tracer_provider.add_span_processor(BaggageSpanProcessor())
tracer_provider.add_span_processor(
BaggageSpanProcessor(ALLOW_ALL_BAGGAGE_KEYS)
)
# tracer has no baggage to start
tracer = tracer_provider.get_tracer("my-tracer")
self.assertIsInstance(tracer, Tracer)
self.assertEqual(get_all_baggage(), {})
# set baggage in context
ctx = set_baggage("queen", "bee")
with tracer.start_as_current_span(
name="bumble", context=ctx
) as bumble_span:
# span should have baggage key-value pair in context
self.assertEqual(get_all_baggage(ctx), {"queen": "bee"})
# span should have baggage key-value pair in attribute
self.assertEqual(bumble_span._attributes["queen"], "bee")
with tracer.start_as_current_span(
name="child_span", context=ctx
) as child_span:
self.assertIsInstance(child_span, Span)
# child span should have baggage key-value pair in context
self.assertEqual(get_all_baggage(ctx), {"queen": "bee"})
# child span should have baggage key-value pair in attribute
self.assertEqual(child_span._attributes["queen"], "bee")
def test_baggage_span_processor_with_string_prefix(
self,
):
tracer_provider = TracerProvider()
tracer_provider.add_span_processor(
BaggageSpanProcessor(self.has_prefix)
)
# tracer has no baggage to start
tracer = tracer_provider.get_tracer("my-tracer")
self.assertIsInstance(tracer, Tracer)
self.assertEqual(get_all_baggage(), {})
# set baggage in context
ctx = set_baggage("queen", "bee")
with tracer.start_as_current_span(
name="bumble", context=ctx
) as bumble_span:
# span should have baggage key-value pair in context
self.assertEqual(get_all_baggage(ctx), {"queen": "bee"})
# span should have baggage key-value pair in attribute
self.assertEqual(bumble_span._attributes["queen"], "bee")
with tracer.start_as_current_span(
name="child_span", context=ctx
) as child_span:
self.assertIsInstance(child_span, Span)
# child span should have baggage key-value pair in context
self.assertEqual(get_all_baggage(ctx), {"queen": "bee"})
# child span should have baggage key-value pair in attribute
self.assertEqual(child_span._attributes["queen"], "bee")
def test_baggage_span_processor_with_regex(
self,
):
tracer_provider = TracerProvider()
tracer_provider.add_span_processor(
BaggageSpanProcessor(self.matches_regex)
)
# tracer has no baggage to start
tracer = tracer_provider.get_tracer("my-tracer")
@ -59,7 +127,9 @@ class BaggageSpanProcessorTest(unittest.TestCase):
self,
):
tracer_provider = TracerProvider()
tracer_provider.add_span_processor(BaggageSpanProcessor())
tracer_provider.add_span_processor(
BaggageSpanProcessor(ALLOW_ALL_BAGGAGE_KEYS)
)
# tracer has no baggage to start
tracer = tracer_provider.get_tracer("my-tracer")
@ -87,3 +157,11 @@ class BaggageSpanProcessorTest(unittest.TestCase):
detach(moar_token)
detach(honey_token)
self.assertEqual(get_all_baggage(), {})
@staticmethod
def has_prefix(baggage_key: str) -> bool:
return baggage_key.startswith("que")
@staticmethod
def matches_regex(baggage_key: str) -> bool:
return re.match(r"que.*", baggage_key) is not None