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 <[email protected]>

Author: MikeGoldsmith

CHANGELOG.md

@@ -36,6 +36,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   ([#2397](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2397)))
 - `opentelemetry-processor-baggage` Initial release
   ([#2436](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2436))
+- `opentelemetry-processor-baggage` Add baggage key predicate
+  ([#2535](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2535))
 
 ### Fixed
 

processor/opentelemetry-processor-baggage/README.rst

@@ -20,3 +20,32 @@ Do not put sensitive information in Baggage.
 
 To repeat: a consequence of adding data to Baggage is that the keys and
 values will appear in all outgoing HTTP headers from the application.
+
+## Usage
+
+Add the span processor when configuring the tracer provider.
+
+To configure the span processor to copy all baggage entries during configuration:
+
+```python
+from opentelemetry.processor.baggage import BaggageSpanProcessor, ALLOW_ALL_BAGGAGE_KEYS
+
+tracer_provider = TracerProvider()
+tracer_provider.add_span_processor(BaggageSpanProcessor(ALLOW_ALL_BAGGAGE_KEYS))
+```
+
+Alternatively, you can provide a custom baggage key predicate to select which baggage keys you want to copy.
+
+For example, to only copy baggage entries that start with `my-key`:
+
+```python
+starts_with_predicate = lambda baggage_key: baggage_key.startswith("my-key")
+tracer_provider.add_span_processor(BaggageSpanProcessor(starts_with_predicate))
+```
+
+For example, to only copy baggage entries that match the regex `^key.+`:
+
+```python
+regex_predicate = lambda baggage_key: baggage_key.startswith("^key.+")
+tracer_provider.add_span_processor(BaggageSpanProcessor(regex_predicate))
+```

processor/opentelemetry-processor-baggage/src/opentelemetry/processor/baggage/__init__.py

@@ -14,7 +14,7 @@
 
 # pylint: disable=import-error
 
-from .processor import BaggageSpanProcessor
+from .processor import ALLOW_ALL_BAGGAGE_KEYS, BaggageSpanProcessor
 from .version import __version__
 
-__all__ = ["BaggageSpanProcessor", "__version__"]
+__all__ = ["ALLOW_ALL_BAGGAGE_KEYS", "BaggageSpanProcessor", "__version__"]

processor/opentelemetry-processor-baggage/src/opentelemetry/processor/baggage/processor.py

@@ -12,13 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
+from typing import Callable, Optional
 
 from opentelemetry.baggage import get_all as get_all_baggage
 from opentelemetry.context import Context
 from opentelemetry.sdk.trace.export import SpanProcessor
 from opentelemetry.trace import Span
 
+# A BaggageKeyPredicate is a function that takes a baggage key and returns a boolean
+BaggageKeyPredicateT = Callable[[str], bool]
+
+# A BaggageKeyPredicate that always returns True, allowing all baggage keys to be added to spans
+ALLOW_ALL_BAGGAGE_KEYS: BaggageKeyPredicateT = lambda _: True
+
 
 class BaggageSpanProcessor(SpanProcessor):
     """
@@ -44,12 +50,13 @@ class BaggageSpanProcessor(SpanProcessor):
 
     """
 
-    def __init__(self) -> None:
-        pass
+    def __init__(self, baggage_key_predicate: BaggageKeyPredicateT) -> None:
+        self._baggage_key_predicate = baggage_key_predicate
 
     def on_start(
         self, span: "Span", parent_context: Optional[Context] = None
     ) -> None:
         baggage = get_all_baggage(parent_context)
         for key, value in baggage.items():
-            span.set_attribute(key, value)
+            if self._baggage_key_predicate(key):
+                span.set_attribute(key, value)

processor/opentelemetry-processor-baggage/tests/test_baggage_processor.py

@@ -12,26 +12,94 @@
 # 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
 
 
 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 @@ def test_set_baggage_attaches_to_child_spans_and_detaches_properly_with_token(
         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 @@ def test_set_baggage_attaches_to_child_spans_and_detaches_properly_with_token(
             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