sentry-sdk 0.18.0__py2.py3-none-any.whl → 2.46.0__py2.py3-none-any.whl

sentry_sdk/integrations/spark/spark_driver.py

@@ -1,15 +1,15 @@
-from sentry_sdk import configure_scope
-from sentry_sdk.hub import Hub
+import sentry_sdk
 from sentry_sdk.integrations import Integration
-from sentry_sdk.utils import capture_internal_exceptions
+from sentry_sdk.utils import capture_internal_exceptions, ensure_integration_enabled
 
-from sentry_sdk._types import MYPY
+from typing import TYPE_CHECKING
 
-if MYPY:
+if TYPE_CHECKING:
     from typing import Any
     from typing import Optional
 
     from sentry_sdk._types import Event, Hint
+    from pyspark import SparkContext
 
 
 class SparkIntegration(Integration):
@@ -18,7 +18,7 @@ class SparkIntegration(Integration):
     @staticmethod
     def setup_once():
         # type: () -> None
-        patch_spark_context_init()
+        _setup_sentry_tracing()
 
 
 def _set_app_properties():
@@ -31,14 +31,18 @@ def _set_app_properties():
 
     spark_context = SparkContext._active_spark_context
     if spark_context:
-        spark_context.setLocalProperty("sentry_app_name", spark_context.appName)
         spark_context.setLocalProperty(
-            "sentry_application_id", spark_context.applicationId
+            "sentry_app_name",
+            spark_context.appName,
+        )
+        spark_context.setLocalProperty(
+            "sentry_application_id",
+            spark_context.applicationId,
         )
 
 
 def _start_sentry_listener(sc):
-    # type: (Any) -> None
+    # type: (SparkContext) -> None
     """
     Start java gateway server to add custom `SparkListener`
     """
@@ -50,62 +54,77 @@ def _start_sentry_listener(sc):
     sc._jsc.sc().addSparkListener(listener)
 
 
-def patch_spark_context_init():
+def _add_event_processor(sc):
+    # type: (SparkContext) -> None
+    scope = sentry_sdk.get_isolation_scope()
+
+    @scope.add_event_processor
+    def process_event(event, hint):
+        # type: (Event, Hint) -> Optional[Event]
+        with capture_internal_exceptions():
+            if sentry_sdk.get_client().get_integration(SparkIntegration) is None:
+                return event
+
+            if sc._active_spark_context is None:
+                return event
+
+            event.setdefault("user", {}).setdefault("id", sc.sparkUser())
+
+            event.setdefault("tags", {}).setdefault(
+                "executor.id", sc._conf.get("spark.executor.id")
+            )
+            event["tags"].setdefault(
+                "spark-submit.deployMode",
+                sc._conf.get("spark.submit.deployMode"),
+            )
+            event["tags"].setdefault("driver.host", sc._conf.get("spark.driver.host"))
+            event["tags"].setdefault("driver.port", sc._conf.get("spark.driver.port"))
+            event["tags"].setdefault("spark_version", sc.version)
+            event["tags"].setdefault("app_name", sc.appName)
+            event["tags"].setdefault("application_id", sc.applicationId)
+            event["tags"].setdefault("master", sc.master)
+            event["tags"].setdefault("spark_home", sc.sparkHome)
+
+            event.setdefault("extra", {}).setdefault("web_url", sc.uiWebUrl)
+
+        return event
+
+
+def _activate_integration(sc):
+    # type: (SparkContext) -> None
+
+    _start_sentry_listener(sc)
+    _set_app_properties()
+    _add_event_processor(sc)
+
+
+def _patch_spark_context_init():
     # type: () -> None
     from pyspark import SparkContext
 
     spark_context_init = SparkContext._do_init
 
+    @ensure_integration_enabled(SparkIntegration, spark_context_init)
     def _sentry_patched_spark_context_init(self, *args, **kwargs):
         # type: (SparkContext, *Any, **Any) -> Optional[Any]
-        init = spark_context_init(self, *args, **kwargs)
-
-        if Hub.current.get_integration(SparkIntegration) is None:
-            return init
-
-        _start_sentry_listener(self)
-        _set_app_properties()
+        rv = spark_context_init(self, *args, **kwargs)
+        _activate_integration(self)
+        return rv
 
-        with configure_scope() as scope:
-
-            @scope.add_event_processor
-            def process_event(event, hint):
-                # type: (Event, Hint) -> Optional[Event]
-                with capture_internal_exceptions():
-                    if Hub.current.get_integration(SparkIntegration) is None:
-                        return event
-
-                    event.setdefault("user", {}).setdefault("id", self.sparkUser())
-
-                    event.setdefault("tags", {}).setdefault(
-                        "executor.id", self._conf.get("spark.executor.id")
-                    )
-                    event["tags"].setdefault(
-                        "spark-submit.deployMode",
-                        self._conf.get("spark.submit.deployMode"),
-                    )
-                    event["tags"].setdefault(
-                        "driver.host", self._conf.get("spark.driver.host")
-                    )
-                    event["tags"].setdefault(
-                        "driver.port", self._conf.get("spark.driver.port")
-                    )
-                    event["tags"].setdefault("spark_version", self.version)
-                    event["tags"].setdefault("app_name", self.appName)
-                    event["tags"].setdefault("application_id", self.applicationId)
-                    event["tags"].setdefault("master", self.master)
-                    event["tags"].setdefault("spark_home", self.sparkHome)
-
-                    event.setdefault("extra", {}).setdefault("web_url", self.uiWebUrl)
+    SparkContext._do_init = _sentry_patched_spark_context_init
 
-                return event
 
-        return init
+def _setup_sentry_tracing():
+    # type: () -> None
+    from pyspark import SparkContext
 
-    SparkContext._do_init = _sentry_patched_spark_context_init
+    if SparkContext._active_spark_context is not None:
+        _activate_integration(SparkContext._active_spark_context)
+        return
+    _patch_spark_context_init()
 
 
-class SparkListener(object):
+class SparkListener:
     def onApplicationEnd(self, applicationEnd):  # noqa: N802,N803
         # type: (Any) -> None
         pass
@@ -139,7 +158,8 @@ class SparkListener(object):
         pass
 
     def onExecutorBlacklistedForStage(  # noqa: N802
-        self, executorBlacklistedForStage  # noqa: N803
+        self,
+        executorBlacklistedForStage,  # noqa: N803
     ):
         # type: (Any) -> None
         pass
@@ -209,14 +229,23 @@ class SparkListener(object):
 
 
 class SentryListener(SparkListener):
-    def __init__(self):
-        # type: () -> None
-        self.hub = Hub.current
+    def _add_breadcrumb(
+        self,
+        level,  # type: str
+        message,  # type: str
+        data=None,  # type: Optional[dict[str, Any]]
+    ):
+        # type: (...) -> None
+        sentry_sdk.get_isolation_scope().add_breadcrumb(
+            level=level, message=message, data=data
+        )
 
     def onJobStart(self, jobStart):  # noqa: N802,N803
         # type: (Any) -> None
+        sentry_sdk.get_isolation_scope().clear_breadcrumbs()
+
         message = "Job {} Started".format(jobStart.jobId())
-        self.hub.add_breadcrumb(level="info", message=message)
+        self._add_breadcrumb(level="info", message=message)
         _set_app_properties()
 
     def onJobEnd(self, jobEnd):  # noqa: N802,N803
@@ -232,14 +261,19 @@ class SentryListener(SparkListener):
             level = "warning"
             message = "Job {} Failed".format(jobEnd.jobId())
 
-        self.hub.add_breadcrumb(level=level, message=message, data=data)
+        self._add_breadcrumb(level=level, message=message, data=data)
 
     def onStageSubmitted(self, stageSubmitted):  # noqa: N802,N803
         # type: (Any) -> None
         stage_info = stageSubmitted.stageInfo()
         message = "Stage {} Submitted".format(stage_info.stageId())
-        data = {"attemptId": stage_info.attemptId(), "name": stage_info.name()}
-        self.hub.add_breadcrumb(level="info", message=message, data=data)
+
+        data = {"name": stage_info.name()}
+        attempt_id = _get_attempt_id(stage_info)
+        if attempt_id is not None:
+            data["attemptId"] = attempt_id
+
+        self._add_breadcrumb(level="info", message=message, data=data)
         _set_app_properties()
 
     def onStageCompleted(self, stageCompleted):  # noqa: N802,N803
@@ -249,7 +283,11 @@ class SentryListener(SparkListener):
         stage_info = stageCompleted.stageInfo()
         message = ""
         level = ""
-        data = {"attemptId": stage_info.attemptId(), "name": stage_info.name()}
+
+        data = {"name": stage_info.name()}
+        attempt_id = _get_attempt_id(stage_info)
+        if attempt_id is not None:
+            data["attemptId"] = attempt_id
 
         # Have to Try Except because stageInfo.failureReason() is typed with Scala Option
         try:
@@ -260,4 +298,19 @@ class SentryListener(SparkListener):
             message = "Stage {} Completed".format(stage_info.stageId())
             level = "info"
 
-        self.hub.add_breadcrumb(level=level, message=message, data=data)
+        self._add_breadcrumb(level=level, message=message, data=data)
+
+
+def _get_attempt_id(stage_info):
+    # type: (Any) -> Optional[int]
+    try:
+        return stage_info.attemptId()
+    except Exception:
+        pass
+
+    try:
+        return stage_info.attemptNumber()
+    except Exception:
+        pass
+
+    return None

sentry_sdk/integrations/spark/spark_worker.py

@@ -1,9 +1,6 @@
-from __future__ import absolute_import
-
 import sys
 
-from sentry_sdk import configure_scope
-from sentry_sdk.hub import Hub
+import sentry_sdk
 from sentry_sdk.integrations import Integration
 from sentry_sdk.utils import (
     capture_internal_exceptions,
@@ -13,9 +10,9 @@ from sentry_sdk.utils import (
     event_hint_with_exc_info,
 )
 
-from sentry_sdk._types import MYPY
+from typing import TYPE_CHECKING
 
-if MYPY:
+if TYPE_CHECKING:
     from typing import Any
     from typing import Optional
 
@@ -33,11 +30,9 @@ class SparkWorkerIntegration(Integration):
         original_daemon.worker_main = _sentry_worker_main
 
 
-def _capture_exception(exc_info, hub):
-    # type: (ExcInfo, Hub) -> None
-    client = hub.client
-
-    client_options = client.options  # type: ignore
+def _capture_exception(exc_info):
+    # type: (ExcInfo) -> None
+    client = sentry_sdk.get_client()
 
     mechanism = {"type": "spark", "handled": False}
 
@@ -51,63 +46,61 @@ def _capture_exception(exc_info, hub):
         if exc_type not in (SystemExit, EOFError, ConnectionResetError):
             rv.append(
                 single_exception_from_error_tuple(
-                    exc_type, exc_value, tb, client_options, mechanism
+                    exc_type, exc_value, tb, client.options, mechanism
                 )
             )
 
     if rv:
         rv.reverse()
         hint = event_hint_with_exc_info(exc_info)
-        event = {"level": "error", "exception": {"values": rv}}
+        event = {"level": "error", "exception": {"values": rv}}  # type: Event
 
         _tag_task_context()
 
-        hub.capture_event(event, hint=hint)
+        sentry_sdk.capture_event(event, hint=hint)
 
 
 def _tag_task_context():
     # type: () -> None
     from pyspark.taskcontext import TaskContext
 
-    with configure_scope() as scope:
+    scope = sentry_sdk.get_isolation_scope()
 
-        @scope.add_event_processor
-        def process_event(event, hint):
-            # type: (Event, Hint) -> Optional[Event]
-            with capture_internal_exceptions():
-                integration = Hub.current.get_integration(SparkWorkerIntegration)
-                task_context = TaskContext.get()
+    @scope.add_event_processor
+    def process_event(event, hint):
+        # type: (Event, Hint) -> Optional[Event]
+        with capture_internal_exceptions():
+            integration = sentry_sdk.get_client().get_integration(
+                SparkWorkerIntegration
+            )
+            task_context = TaskContext.get()
 
-                if integration is None or task_context is None:
-                    return event
+            if integration is None or task_context is None:
+                return event
 
-                event.setdefault("tags", {}).setdefault(
-                    "stageId", str(task_context.stageId())
-                )
-                event["tags"].setdefault("partitionId", str(task_context.partitionId()))
-                event["tags"].setdefault(
-                    "attemptNumber", str(task_context.attemptNumber())
-                )
-                event["tags"].setdefault(
-                    "taskAttemptId", str(task_context.taskAttemptId())
-                )
+            event.setdefault("tags", {}).setdefault(
+                "stageId", str(task_context.stageId())
+            )
+            event["tags"].setdefault("partitionId", str(task_context.partitionId()))
+            event["tags"].setdefault("attemptNumber", str(task_context.attemptNumber()))
+            event["tags"].setdefault("taskAttemptId", str(task_context.taskAttemptId()))
 
-                if task_context._localProperties:
-                    if "sentry_app_name" in task_context._localProperties:
-                        event["tags"].setdefault(
-                            "app_name", task_context._localProperties["sentry_app_name"]
-                        )
-                        event["tags"].setdefault(
-                            "application_id",
-                            task_context._localProperties["sentry_application_id"],
-                        )
+            if task_context._localProperties:
+                if "sentry_app_name" in task_context._localProperties:
+                    event["tags"].setdefault(
+                        "app_name", task_context._localProperties["sentry_app_name"]
+                    )
+                    event["tags"].setdefault(
+                        "application_id",
+                        task_context._localProperties["sentry_application_id"],
+                    )
 
-                    if "callSite.short" in task_context._localProperties:
-                        event.setdefault("extra", {}).setdefault(
-                            "callSite", task_context._localProperties["callSite.short"]
-                        )
+                if "callSite.short" in task_context._localProperties:
+                    event.setdefault("extra", {}).setdefault(
+                        "callSite", task_context._localProperties["callSite.short"]
+                    )
 
-            return event
+        return event
 
 
 def _sentry_worker_main(*args, **kwargs):
@@ -117,8 +110,7 @@ def _sentry_worker_main(*args, **kwargs):
     try:
         original_worker.main(*args, **kwargs)
     except SystemExit:
-        if Hub.current.get_integration(SparkWorkerIntegration) is not None:
-            hub = Hub.current
+        if sentry_sdk.get_client().get_integration(SparkWorkerIntegration) is not None:
             exc_info = sys.exc_info()
             with capture_internal_exceptions():
-                _capture_exception(exc_info, hub)
+                _capture_exception(exc_info)

sentry_sdk/integrations/sqlalchemy.py

@@ -1,9 +1,11 @@
-from __future__ import absolute_import
-
-from sentry_sdk._types import MYPY
-from sentry_sdk.hub import Hub
-from sentry_sdk.integrations import Integration, DidNotEnable
-from sentry_sdk.tracing import record_sql_queries
+from sentry_sdk.consts import SPANSTATUS, SPANDATA
+from sentry_sdk.integrations import _check_minimum_version, Integration, DidNotEnable
+from sentry_sdk.tracing_utils import add_query_source, record_sql_queries
+from sentry_sdk.utils import (
+    capture_internal_exceptions,
+    ensure_integration_enabled,
+    parse_version,
+)
 
 try:
     from sqlalchemy.engine import Engine  # type: ignore
@@ -12,7 +14,9 @@ try:
 except ImportError:
     raise DidNotEnable("SQLAlchemy not installed.")
 
-if MYPY:
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
     from typing import Any
     from typing import ContextManager
     from typing import Optional
@@ -22,76 +26,117 @@ if MYPY:
 
 class SqlalchemyIntegration(Integration):
     identifier = "sqlalchemy"
+    origin = f"auto.db.{identifier}"
 
     @staticmethod
     def setup_once():
         # type: () -> None
-
-        try:
-            version = tuple(map(int, SQLALCHEMY_VERSION.split("b")[0].split(".")))
-        except (TypeError, ValueError):
-            raise DidNotEnable(
-                "Unparsable SQLAlchemy version: {}".format(SQLALCHEMY_VERSION)
-            )
-
-        if version < (1, 2):
-            raise DidNotEnable("SQLAlchemy 1.2 or newer required.")
+        version = parse_version(SQLALCHEMY_VERSION)
+        _check_minimum_version(SqlalchemyIntegration, version)
 
         listen(Engine, "before_cursor_execute", _before_cursor_execute)
         listen(Engine, "after_cursor_execute", _after_cursor_execute)
         listen(Engine, "handle_error", _handle_error)
 
 
+@ensure_integration_enabled(SqlalchemyIntegration)
 def _before_cursor_execute(
     conn, cursor, statement, parameters, context, executemany, *args
 ):
     # type: (Any, Any, Any, Any, Any, bool, *Any) -> None
-    hub = Hub.current
-    if hub.get_integration(SqlalchemyIntegration) is None:
-        return
-
     ctx_mgr = record_sql_queries(
-        hub,
         cursor,
         statement,
         parameters,
         paramstyle=context and context.dialect and context.dialect.paramstyle or None,
         executemany=executemany,
+        span_origin=SqlalchemyIntegration.origin,
     )
-    conn._sentry_sql_span_manager = ctx_mgr
+    context._sentry_sql_span_manager = ctx_mgr
 
     span = ctx_mgr.__enter__()
 
     if span is not None:
-        conn._sentry_sql_span = span
+        _set_db_data(span, conn)
+        context._sentry_sql_span = span
 
 
-def _after_cursor_execute(conn, cursor, statement, *args):
-    # type: (Any, Any, Any, *Any) -> None
-    ctx_mgr = getattr(
-        conn, "_sentry_sql_span_manager", None
-    )  # type: ContextManager[Any]
+@ensure_integration_enabled(SqlalchemyIntegration)
+def _after_cursor_execute(conn, cursor, statement, parameters, context, *args):
+    # type: (Any, Any, Any, Any, Any, *Any) -> None
+    ctx_mgr = getattr(context, "_sentry_sql_span_manager", None)  # type: Optional[ContextManager[Any]]
 
     if ctx_mgr is not None:
-        conn._sentry_sql_span_manager = 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)
+
 
 def _handle_error(context, *args):
     # type: (Any, *Any) -> None
-    conn = context.connection
-    span = getattr(conn, "_sentry_sql_span", None)  # type: Optional[Span]
+    execution_context = context.execution_context
+    if execution_context is None:
+        return
+
+    span = getattr(execution_context, "_sentry_sql_span", None)  # type: Optional[Span]
 
     if span is not None:
-        span.set_status("internal_error")
+        span.set_status(SPANSTATUS.INTERNAL_ERROR)
 
     # _after_cursor_execute does not get called for crashing SQL stmts. Judging
     # from SQLAlchemy codebase it does seem like any error coming into this
     # handler is going to be fatal.
-    ctx_mgr = getattr(
-        conn, "_sentry_sql_span_manager", None
-    )  # type: ContextManager[Any]
+    ctx_mgr = getattr(execution_context, "_sentry_sql_span_manager", None)  # type: Optional[ContextManager[Any]]
 
     if ctx_mgr is not None:
-        conn._sentry_sql_span_manager = None
+        execution_context._sentry_sql_span_manager = None
         ctx_mgr.__exit__(None, None, None)
+
+
+# See: https://docs.sqlalchemy.org/en/20/dialects/index.html
+def _get_db_system(name):
+    # type: (str) -> Optional[str]
+    name = str(name)
+
+    if "sqlite" in name:
+        return "sqlite"
+
+    if "postgres" in name:
+        return "postgresql"
+
+    if "mariadb" in name:
+        return "mariadb"
+
+    if "mysql" in name:
+        return "mysql"
+
+    if "oracle" in name:
+        return "oracle"
+
+    return None
+
+
+def _set_db_data(span, conn):
+    # type: (Span, Any) -> None
+    db_system = _get_db_system(conn.engine.name)
+    if db_system is not None:
+        span.set_data(SPANDATA.DB_SYSTEM, db_system)
+
+    if conn.engine.url is None:
+        return
+
+    db_name = conn.engine.url.database
+    if db_name is not None:
+        span.set_data(SPANDATA.DB_NAME, db_name)
+
+    server_address = conn.engine.url.host
+    if server_address is not None:
+        span.set_data(SPANDATA.SERVER_ADDRESS, server_address)
+
+    server_port = conn.engine.url.port
+    if server_port is not None:
+        span.set_data(SPANDATA.SERVER_PORT, server_port)