feat(metrics): add global tags to sentry_sdk tags as well (#102795)

This PR adds the option to also configure tags set in the metrics
backend via `add_global_tags` to also be set in the sentry_sdk.

---------

Co-authored-by: Markus Unterwaditzer <[email protected]>

Author: bmckerry

src/sentry/consumers/__init__.py

@@ -667,7 +667,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)
 

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(
@@ -52,45 +53,57 @@ 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"):
             stack = _THREAD_LOCAL_TAGS.stack = []
         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:

src/sentry/runner/commands/run.py

@@ -580,7 +580,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(

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 ""})

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"]
 

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:

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}
 

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