Merge branch 'potel-base' into potel-base-run-all-tests

Author: antonpirker

sentry_sdk/integrations/opentelemetry/integration.py

@@ -5,6 +5,7 @@
 """
 
 from sentry_sdk.integrations import DidNotEnable, Integration
+from sentry_sdk.integrations.opentelemetry.scope import setup_initial_scopes
 from sentry_sdk.integrations.opentelemetry.propagator import SentryPropagator
 from sentry_sdk.integrations.opentelemetry.span_processor import (
     SentrySpanProcessor,
@@ -74,6 +75,7 @@ def _setup_scope_context_management():
     import opentelemetry.context
 
     opentelemetry.context._RUNTIME_CONTEXT = SentryContextVarsRuntimeContext()
+    setup_initial_scopes()
 
 
 def _setup_sentry_tracing():

sentry_sdk/integrations/opentelemetry/propagator.py

@@ -99,7 +99,7 @@ def inject(self, carrier, context=None, setter=default_setter):
 
             # TODO-neel-potel check trace_propagation_targets
             # TODO-neel-potel test propagator works with twp
-            for (key, value) in current_scope.iter_trace_propagation_headers():
+            for key, value in current_scope.iter_trace_propagation_headers():
                 setter.set(carrier, key, value)
 
     @property

sentry_sdk/integrations/opentelemetry/sampler.py

@@ -17,7 +17,7 @@
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from typing import Optional, Sequence, Union
+    from typing import Any, Optional, Sequence, Union
     from opentelemetry.context import Context
     from opentelemetry.trace import Link, SpanKind
     from opentelemetry.trace.span import SpanContext
@@ -152,15 +152,9 @@ def should_sample(
         has_traces_sampler = callable(client.options.get("traces_sampler"))
 
         if is_root_span and has_traces_sampler:
-            sampling_context = {
-                "transaction_context": {
-                    "name": name,
-                    "op": attributes.get(SentrySpanAttribute.OP),
-                    "source": attributes.get(SentrySpanAttribute.SOURCE),
-                },
-                "parent_sampled": get_parent_sampled(parent_span_context, trace_id),
-            }
-            sampling_context.update(attributes)
+            sampling_context = create_sampling_context(
+                name, attributes, parent_span_context, trace_id
+            )
             sample_rate = client.options["traces_sampler"](sampling_context)
         else:
             # Check if there is a parent with a sampling decision
@@ -193,3 +187,19 @@ def should_sample(
 
     def get_description(self) -> str:
         return self.__class__.__name__
+
+
+def create_sampling_context(name, attributes, parent_span_context, trace_id):
+    # type: (str, Attributes, SpanContext, str) -> dict[str, Any]
+    sampling_context = {
+        "transaction_context": {
+            "name": name,
+            "op": attributes.get(SentrySpanAttribute.OP),
+            "source": attributes.get(SentrySpanAttribute.SOURCE),
+        },
+        "parent_sampled": get_parent_sampled(parent_span_context, trace_id),
+    }
+
+    sampling_context.update(attributes)
+
+    return sampling_context

sentry_sdk/integrations/opentelemetry/scope.py

@@ -1,7 +1,14 @@
 from typing import cast
 from contextlib import contextmanager
 
-from opentelemetry.context import get_value, set_value, attach, detach, get_current
+from opentelemetry.context import (
+    Context,
+    get_value,
+    set_value,
+    attach,
+    detach,
+    get_current,
+)
 from opentelemetry.trace import (
     SpanContext,
     NonRecordingSpan,
@@ -136,13 +143,13 @@ def start_span(self, **kwargs):
 _INITIAL_ISOLATION_SCOPE = None
 
 
-def _setup_initial_scopes():
+def setup_initial_scopes():
     global _INITIAL_CURRENT_SCOPE, _INITIAL_ISOLATION_SCOPE
     _INITIAL_CURRENT_SCOPE = PotelScope(ty=ScopeType.CURRENT)
     _INITIAL_ISOLATION_SCOPE = PotelScope(ty=ScopeType.ISOLATION)
 
-
-_setup_initial_scopes()
+    scopes = (_INITIAL_CURRENT_SCOPE, _INITIAL_ISOLATION_SCOPE)
+    attach(set_value(SENTRY_SCOPES_KEY, scopes))
 
 
 @contextmanager

sentry_sdk/integrations/opentelemetry/span_processor.py

@@ -20,6 +20,7 @@
     get_profiler_id,
 )
 from sentry_sdk.profiler.transaction_profiler import Profile
+from sentry_sdk.integrations.opentelemetry.sampler import create_sampling_context
 from sentry_sdk.integrations.opentelemetry.utils import (
     is_sentry_span,
     convert_from_otel_timestamp,
@@ -126,8 +127,10 @@ def _start_profile(self, span):
             # unix timestamp that is on span.start_time
             # setting it to 0 means the profiler will internally measure time on start
             profile = Profile(sampled, 0)
-            # TODO-neel-potel sampling context??
-            profile._set_initial_sampling_decision(sampling_context={})
+            sampling_context = create_sampling_context(
+                span.name, span.attributes, span.parent, span.context.trace_id
+            )
+            profile._set_initial_sampling_decision(sampling_context)
             profile.__enter__()
             set_sentry_meta(span, "profile", profile)
 

sentry_sdk/integrations/sqlalchemy.py

@@ -76,15 +76,15 @@ def _after_cursor_execute(conn, cursor, statement, parameters, context, *args):
         context, "_sentry_sql_span_manager", None
     )  # type: Optional[ContextManager[Any]]
 
-    if ctx_mgr is not None:
-        context._sentry_sql_span_manager = None
-        ctx_mgr.__exit__(None, None, None)
-
     span = getattr(context, "_sentry_sql_span", None)  # type: Optional[Span]
     if span is not None:
         with capture_internal_exceptions():
             add_query_source(span)
 
+    if ctx_mgr is not None:
+        context._sentry_sql_span_manager = None
+        ctx_mgr.__exit__(None, None, None)
+
 
 def _handle_error(context, *args):
     # type: (Any, *Any) -> None

sentry_sdk/tracing_utils.py

@@ -33,7 +33,6 @@
     from typing import Generator
     from typing import Optional
     from typing import Union
-
     from types import FrameType
 
 
@@ -148,6 +147,7 @@ def record_sql_queries(
         op=OP.DB,
         name=query,
         origin=span_origin,
+        only_if_parent=True,
     ) as span:
         for k, v in data.items():
             span.set_data(k, v)
@@ -731,6 +731,3 @@ def get_current_span(scope=None):
     LOW_QUALITY_TRANSACTION_SOURCES,
     SENTRY_TRACE_HEADER_NAME,
 )
-
-if TYPE_CHECKING:
-    from sentry_sdk.tracing import Span

tests/conftest.py

@@ -74,7 +74,7 @@ def clean_scopes():
     scope._isolation_scope.set(None)
     scope._current_scope.set(None)
 
-    potel_scope._setup_initial_scopes()
+    potel_scope.setup_initial_scopes()
 
 
 @pytest.fixture(autouse=True)

tests/integrations/clickhouse_driver/test_clickhouse_driver.py

@@ -1,7 +1,7 @@
 """
 Tests need a local clickhouse instance running, this can best be done using
 ```sh
-docker run -d -p 18123:8123 -p9000:9000 --name clickhouse-test --ulimit nofile=262144:262144 --rm clickhouse/clickhouse-server
+docker run -d -p 18123:8123 -p9000:9000 --name clickhouse-test --ulimit nofile=262144:262144 --rm clickhouse
 ```
 """
 

tests/integrations/pymongo/test_pymongo.py

@@ -1,6 +1,6 @@
 import re
 
-from sentry_sdk import capture_message, start_transaction
+import sentry_sdk
 from sentry_sdk.consts import SPANDATA
 from sentry_sdk.integrations.pymongo import PyMongoIntegration, _strip_pii
 
@@ -37,7 +37,7 @@ def test_transactions(sentry_init, capture_events, mongo_server, with_pii):
 
     connection = MongoClient(mongo_server.uri)
 
-    with start_transaction():
+    with sentry_sdk.start_span():
         list(
             connection["test_db"]["test_collection"].find({"foobar": 1})
         )  # force query execution
@@ -119,13 +119,9 @@ def test_breadcrumbs(
     list(
         connection["test_db"]["test_collection"].find({"foobar": 1})
     )  # force query execution
-    capture_message("hi")
-
-    if traces_sample_rate:
-        event = events[1]
-    else:
-        event = events[0]
+    sentry_sdk.capture_message("hi")
 
+    (event,) = events
     (crumb,) = event["breadcrumbs"]["values"]
 
     assert crumb["category"] == "query"
@@ -450,7 +446,7 @@ def test_span_origin(sentry_init, capture_events, mongo_server):
 
     connection = MongoClient(mongo_server.uri)
 
-    with start_transaction():
+    with sentry_sdk.start_span():
         list(
             connection["test_db"]["test_collection"].find({"foobar": 1})
         )  # force query execution