diff --git a/src/sentry/consumers/__init__.py b/src/sentry/consumers/__init__.py index 0246a48374ec49..65aeb9a9ee3420 100644 --- a/src/sentry/consumers/__init__.py +++ b/src/sentry/consumers/__init__.py @@ -672,7 +672,7 @@ def create_with_partitions(self, commit, partitions): # Update the min_partition global tag based on current partition assignment if partitions: min_partition = min(p.index for p in partitions) - add_global_tags(min_partition=str(min_partition)) + add_global_tags(tags={"min_partition": str(min_partition)}) return self.inner.create_with_partitions(commit, partitions) diff --git a/src/sentry/metrics/middleware.py b/src/sentry/metrics/middleware.py index c3116a6f74c4b8..4a4b115b69d34e 100644 --- a/src/sentry/metrics/middleware.py +++ b/src/sentry/metrics/middleware.py @@ -2,9 +2,10 @@ from contextlib import contextmanager from threading import local +import sentry_sdk from django.conf import settings -from sentry.metrics.base import MetricsBackend, MutableTags, Tags, TagValue +from sentry.metrics.base import MetricsBackend, MutableTags, Tags _BAD_TAGS = frozenset(["event", "project", "group"]) _NOT_BAD_TAGS = frozenset(["use_case_id", "consumer_member_id", "broker_id", "kafka_slice_id"]) @@ -50,8 +51,10 @@ def _filter_tags(key: str, tags: MutableTags) -> MutableTags: _GLOBAL_TAGS: list[Tags] = [] -def _add_global_tags(_all_threads: bool = False, **tags: TagValue) -> list[Tags]: - if _all_threads: +def _add_global_tags( + all_threads: bool = False, set_sentry_tags: bool = False, tags: Tags | None = None +) -> list[Tags]: + if all_threads: stack = _GLOBAL_TAGS else: if not hasattr(_THREAD_LOCAL_TAGS, "stack"): @@ -59,36 +62,46 @@ def _add_global_tags(_all_threads: bool = False, **tags: TagValue) -> list[Tags] else: stack = _THREAD_LOCAL_TAGS.stack + if tags is None: + tags = {} stack.append(tags) + if set_sentry_tags: + sentry_sdk.set_tags(tags) return stack -def add_global_tags(_all_threads: bool = False, **tags: TagValue) -> None: +def add_global_tags( + all_threads: bool = False, set_sentry_tags: bool = False, tags: Tags | None = None +) -> None: """ Set multiple metric tags onto the global or thread-local stack which then apply to all metrics. + If `set_sentry_tags` is True, also sets the given tags in sentry_sdk. When used in combination with the `global_tags` context manager, - `add_global_tags` is reverted in any wrapping invocaation of `global_tags`. + `add_global_tags` is reverted in any wrapping invocation of `global_tags`. + However, tags set in the current sentry_sdk instance will remain set there. For example:: - with global_tags(tag_a=123): - add_global_tags(tag_b=123) + with global_tags(tags={"tag_a": 123}): + add_global_tags(tags={"tag_b": 123}) # tag_b is no longer visible """ - _add_global_tags(_all_threads=_all_threads, **tags) + _add_global_tags(all_threads=all_threads, set_sentry_tags=set_sentry_tags, tags=tags) @contextmanager -def global_tags(_all_threads: bool = False, **tags: TagValue) -> Generator[None]: +def global_tags( + all_threads: bool = False, set_sentry_tags: bool = False, tags: Tags | None = None +) -> Generator[None]: """ The context manager version of `add_global_tags` that reverts all tag changes upon exit. See docstring of `add_global_tags` for how those two methods interact. """ - stack = _add_global_tags(_all_threads=_all_threads, **tags) + stack = _add_global_tags(all_threads=all_threads, set_sentry_tags=set_sentry_tags, tags=tags) old_len = len(stack) - 1 try: diff --git a/src/sentry/runner/commands/run.py b/src/sentry/runner/commands/run.py index ff24c4c9f696d1..f156de071f3321 100644 --- a/src/sentry/runner/commands/run.py +++ b/src/sentry/runner/commands/run.py @@ -593,7 +593,12 @@ def basic_consumer( logging.getLogger("arroyo").setLevel(log_level.upper()) add_global_tags( - kafka_topic=topic, consumer_group=options["group_id"], kafka_slice_id=kafka_slice_id + set_sentry_tags=True, + tags={ + "kafka_topic": topic, + "consumer_group": options["group_id"], + "kafka_slice_id": kafka_slice_id, + }, ) processor = get_stream_processor( consumer_name, diff --git a/src/sentry/sentry_metrics/configuration.py b/src/sentry/sentry_metrics/configuration.py index 8c8b63752f26f5..9e9aec0514e981 100644 --- a/src/sentry/sentry_metrics/configuration.py +++ b/src/sentry/sentry_metrics/configuration.py @@ -167,6 +167,4 @@ def initialize_main_process_state(config: MetricsIngestConfiguration) -> None: from sentry.utils.metrics import add_global_tags - global_tag_map = {"pipeline": config.internal_metrics_tag or ""} - - add_global_tags(_all_threads=True, **global_tag_map) + add_global_tags(all_threads=True, tags={"pipeline": config.internal_metrics_tag or ""}) diff --git a/src/sentry/tasks/store.py b/src/sentry/tasks/store.py index cf840981f5738a..c0fb9294e91418 100644 --- a/src/sentry/tasks/store.py +++ b/src/sentry/tasks/store.py @@ -533,7 +533,7 @@ def _do_save_event( event_data=data, ) - with metrics.global_tags(event_type=event_type): + with metrics.global_tags(tags={"event_type": event_type}): if event_id is None and data is not None: event_id = data["event_id"] diff --git a/src/sentry/utils/arroyo.py b/src/sentry/utils/arroyo.py index c27c4b0cbe6970..5391eccc8ffb58 100644 --- a/src/sentry/utils/arroyo.py +++ b/src/sentry/utils/arroyo.py @@ -126,7 +126,7 @@ def _initialize_arroyo_subprocess(initializer: Callable[[], None] | None, tags: from sentry.metrics.middleware import add_global_tags # Inherit global tags from the parent process - add_global_tags(_all_threads=True, **tags) + add_global_tags(all_threads=True, tags=tags) def initialize_arroyo_main() -> None: diff --git a/tests/sentry/metrics/test_middleware.py b/tests/sentry/metrics/test_middleware.py index e860169b2ed2a7..dede56da490a00 100644 --- a/tests/sentry/metrics/test_middleware.py +++ b/tests/sentry/metrics/test_middleware.py @@ -31,9 +31,9 @@ def test_filter_tags_prod() -> None: def test_global() -> None: assert get_current_global_tags() == {} - with global_tags(tag_a=123): + with global_tags(tags={"tag_a": 123}): assert get_current_global_tags() == {"tag_a": 123} - add_global_tags(tag_b=123) + add_global_tags(tags={"tag_b": 123}) assert get_current_global_tags() == {"tag_a": 123, "tag_b": 123} diff --git a/tests/sentry/sentry_metrics/test_parallel_indexer.py b/tests/sentry/sentry_metrics/test_parallel_indexer.py index 516484c2341791..6d91f565968af0 100644 --- a/tests/sentry/sentry_metrics/test_parallel_indexer.py +++ b/tests/sentry/sentry_metrics/test_parallel_indexer.py @@ -28,7 +28,7 @@ def reset_global_metrics_state(): # running a MetricsConsumerStrategyFactory has a side-effect of mutating # global metrics tags - with global_tags(_all_threads=True): + with global_tags(all_threads=True): yield