|
12 | 12 | # See the License for the specific language governing permissions and |
13 | 13 | # limitations under the License. |
14 | 14 |
|
15 | | -from typing import Callable, Optional |
| 15 | +from typing import Callable, Optional, Sequence, Union |
16 | 16 |
|
17 | 17 | from opentelemetry.baggage import get_all as get_all_baggage |
18 | 18 | from opentelemetry.context import Context |
|
21 | 21 |
|
22 | 22 | # A BaggageKeyPredicate is a function that takes a baggage key and returns a boolean |
23 | 23 | BaggageKeyPredicateT = Callable[[str], bool] |
| 24 | +BaggageKeyPredicates = Union[ |
| 25 | + BaggageKeyPredicateT, Sequence[BaggageKeyPredicateT] |
| 26 | +] |
24 | 27 |
|
25 | 28 | # A BaggageKeyPredicate that always returns True, allowing all baggage keys to be added to spans |
26 | 29 | ALLOW_ALL_BAGGAGE_KEYS: BaggageKeyPredicateT = lambda _: True # noqa: E731 |
@@ -50,13 +53,16 @@ class BaggageSpanProcessor(SpanProcessor): |
50 | 53 |
|
51 | 54 | """ |
52 | 55 |
|
53 | | - def __init__(self, baggage_key_predicate: BaggageKeyPredicateT) -> None: |
54 | | - self._baggage_key_predicate = baggage_key_predicate |
| 56 | + def __init__(self, baggage_key_predicate: BaggageKeyPredicates) -> None: |
| 57 | + if callable(baggage_key_predicate): |
| 58 | + self._predicates = [baggage_key_predicate] |
| 59 | + else: |
| 60 | + self._predicates = list(baggage_key_predicate) |
55 | 61 |
|
56 | 62 | def on_start( |
57 | 63 | self, span: "Span", parent_context: Optional[Context] = None |
58 | 64 | ) -> None: |
59 | 65 | baggage = get_all_baggage(parent_context) |
60 | 66 | for key, value in baggage.items(): |
61 | | - if self._baggage_key_predicate(key): |
| 67 | + if any(predicate(key) for predicate in self._predicates): |
62 | 68 | span.set_attribute(key, value) |
0 commit comments