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

sentry_sdk/integrations/spark/spark_driver.py

@@ -0,0 +1,316 @@
+import sentry_sdk
+from sentry_sdk.integrations import Integration
+from sentry_sdk.utils import capture_internal_exceptions, ensure_integration_enabled
+
+from typing import TYPE_CHECKING
+
+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):
+    identifier = "spark"
+
+    @staticmethod
+    def setup_once():
+        # type: () -> None
+        _setup_sentry_tracing()
+
+
+def _set_app_properties():
+    # type: () -> None
+    """
+    Set properties in driver that propagate to worker processes, allowing for workers to have access to those properties.
+    This allows worker integration to have access to app_name and application_id.
+    """
+    from pyspark import SparkContext
+
+    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,
+        )
+
+
+def _start_sentry_listener(sc):
+    # type: (SparkContext) -> None
+    """
+    Start java gateway server to add custom `SparkListener`
+    """
+    from pyspark.java_gateway import ensure_callback_server_started
+
+    gw = sc._gateway
+    ensure_callback_server_started(gw)
+    listener = SentryListener()
+    sc._jsc.sc().addSparkListener(listener)
+
+
+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]
+        rv = spark_context_init(self, *args, **kwargs)
+        _activate_integration(self)
+        return rv
+
+    SparkContext._do_init = _sentry_patched_spark_context_init
+
+
+def _setup_sentry_tracing():
+    # type: () -> None
+    from pyspark import SparkContext
+
+    if SparkContext._active_spark_context is not None:
+        _activate_integration(SparkContext._active_spark_context)
+        return
+    _patch_spark_context_init()
+
+
+class SparkListener:
+    def onApplicationEnd(self, applicationEnd):  # noqa: N802,N803
+        # type: (Any) -> None
+        pass
+
+    def onApplicationStart(self, applicationStart):  # noqa: N802,N803
+        # type: (Any) -> None
+        pass
+
+    def onBlockManagerAdded(self, blockManagerAdded):  # noqa: N802,N803
+        # type: (Any) -> None
+        pass
+
+    def onBlockManagerRemoved(self, blockManagerRemoved):  # noqa: N802,N803
+        # type: (Any) -> None
+        pass
+
+    def onBlockUpdated(self, blockUpdated):  # noqa: N802,N803
+        # type: (Any) -> None
+        pass
+
+    def onEnvironmentUpdate(self, environmentUpdate):  # noqa: N802,N803
+        # type: (Any) -> None
+        pass
+
+    def onExecutorAdded(self, executorAdded):  # noqa: N802,N803
+        # type: (Any) -> None
+        pass
+
+    def onExecutorBlacklisted(self, executorBlacklisted):  # noqa: N802,N803
+        # type: (Any) -> None
+        pass
+
+    def onExecutorBlacklistedForStage(  # noqa: N802
+        self,
+        executorBlacklistedForStage,  # noqa: N803
+    ):
+        # type: (Any) -> None
+        pass
+
+    def onExecutorMetricsUpdate(self, executorMetricsUpdate):  # noqa: N802,N803
+        # type: (Any) -> None
+        pass
+
+    def onExecutorRemoved(self, executorRemoved):  # noqa: N802,N803
+        # type: (Any) -> None
+        pass
+
+    def onJobEnd(self, jobEnd):  # noqa: N802,N803
+        # type: (Any) -> None
+        pass
+
+    def onJobStart(self, jobStart):  # noqa: N802,N803
+        # type: (Any) -> None
+        pass
+
+    def onNodeBlacklisted(self, nodeBlacklisted):  # noqa: N802,N803
+        # type: (Any) -> None
+        pass
+
+    def onNodeBlacklistedForStage(self, nodeBlacklistedForStage):  # noqa: N802,N803
+        # type: (Any) -> None
+        pass
+
+    def onNodeUnblacklisted(self, nodeUnblacklisted):  # noqa: N802,N803
+        # type: (Any) -> None
+        pass
+
+    def onOtherEvent(self, event):  # noqa: N802,N803
+        # type: (Any) -> None
+        pass
+
+    def onSpeculativeTaskSubmitted(self, speculativeTask):  # noqa: N802,N803
+        # type: (Any) -> None
+        pass
+
+    def onStageCompleted(self, stageCompleted):  # noqa: N802,N803
+        # type: (Any) -> None
+        pass
+
+    def onStageSubmitted(self, stageSubmitted):  # noqa: N802,N803
+        # type: (Any) -> None
+        pass
+
+    def onTaskEnd(self, taskEnd):  # noqa: N802,N803
+        # type: (Any) -> None
+        pass
+
+    def onTaskGettingResult(self, taskGettingResult):  # noqa: N802,N803
+        # type: (Any) -> None
+        pass
+
+    def onTaskStart(self, taskStart):  # noqa: N802,N803
+        # type: (Any) -> None
+        pass
+
+    def onUnpersistRDD(self, unpersistRDD):  # noqa: N802,N803
+        # type: (Any) -> None
+        pass
+
+    class Java:
+        implements = ["org.apache.spark.scheduler.SparkListenerInterface"]
+
+
+class SentryListener(SparkListener):
+    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._add_breadcrumb(level="info", message=message)
+        _set_app_properties()
+
+    def onJobEnd(self, jobEnd):  # noqa: N802,N803
+        # type: (Any) -> None
+        level = ""
+        message = ""
+        data = {"result": jobEnd.jobResult().toString()}
+
+        if jobEnd.jobResult().toString() == "JobSucceeded":
+            level = "info"
+            message = "Job {} Ended".format(jobEnd.jobId())
+        else:
+            level = "warning"
+            message = "Job {} Failed".format(jobEnd.jobId())
+
+        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 = {"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
+        # type: (Any) -> None
+        from py4j.protocol import Py4JJavaError  # type: ignore
+
+        stage_info = stageCompleted.stageInfo()
+        message = ""
+        level = ""
+
+        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:
+            data["reason"] = stage_info.failureReason().get()
+            message = "Stage {} Failed".format(stage_info.stageId())
+            level = "warning"
+        except Py4JJavaError:
+            message = "Stage {} Completed".format(stage_info.stageId())
+            level = "info"
+
+        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

@@ -0,0 +1,116 @@
+import sys
+
+import sentry_sdk
+from sentry_sdk.integrations import Integration
+from sentry_sdk.utils import (
+    capture_internal_exceptions,
+    exc_info_from_error,
+    single_exception_from_error_tuple,
+    walk_exception_chain,
+    event_hint_with_exc_info,
+)
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any
+    from typing import Optional
+
+    from sentry_sdk._types import ExcInfo, Event, Hint
+
+
+class SparkWorkerIntegration(Integration):
+    identifier = "spark_worker"
+
+    @staticmethod
+    def setup_once():
+        # type: () -> None
+        import pyspark.daemon as original_daemon
+
+        original_daemon.worker_main = _sentry_worker_main
+
+
+def _capture_exception(exc_info):
+    # type: (ExcInfo) -> None
+    client = sentry_sdk.get_client()
+
+    mechanism = {"type": "spark", "handled": False}
+
+    exc_info = exc_info_from_error(exc_info)
+
+    exc_type, exc_value, tb = exc_info
+    rv = []
+
+    # On Exception worker will call sys.exit(-1), so we can ignore SystemExit and similar errors
+    for exc_type, exc_value, tb in walk_exception_chain(exc_info):
+        if exc_type not in (SystemExit, EOFError, ConnectionResetError):
+            rv.append(
+                single_exception_from_error_tuple(
+                    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}}  # type: Event
+
+        _tag_task_context()
+
+        sentry_sdk.capture_event(event, hint=hint)
+
+
+def _tag_task_context():
+    # type: () -> None
+    from pyspark.taskcontext import TaskContext
+
+    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 = sentry_sdk.get_client().get_integration(
+                SparkWorkerIntegration
+            )
+            task_context = TaskContext.get()
+
+            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()))
+
+            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"]
+                    )
+
+        return event
+
+
+def _sentry_worker_main(*args, **kwargs):
+    # type: (*Optional[Any], **Optional[Any]) -> None
+    import pyspark.worker as original_worker
+
+    try:
+        original_worker.main(*args, **kwargs)
+    except SystemExit:
+        if sentry_sdk.get_client().get_integration(SparkWorkerIntegration) is not None:
+            exc_info = sys.exc_info()
+            with capture_internal_exceptions():
+                _capture_exception(exc_info)

sentry_sdk/integrations/sqlalchemy.py

@@ -0,0 +1,142 @@
+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
+    from sqlalchemy.event import listen  # type: ignore
+    from sqlalchemy import __version__ as SQLALCHEMY_VERSION  # type: ignore
+except ImportError:
+    raise DidNotEnable("SQLAlchemy not installed.")
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any
+    from typing import ContextManager
+    from typing import Optional
+
+    from sentry_sdk.tracing import Span
+
+
+class SqlalchemyIntegration(Integration):
+    identifier = "sqlalchemy"
+    origin = f"auto.db.{identifier}"
+
+    @staticmethod
+    def setup_once():
+        # type: () -> None
+        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
+    ctx_mgr = record_sql_queries(
+        cursor,
+        statement,
+        parameters,
+        paramstyle=context and context.dialect and context.dialect.paramstyle or None,
+        executemany=executemany,
+        span_origin=SqlalchemyIntegration.origin,
+    )
+    context._sentry_sql_span_manager = ctx_mgr
+
+    span = ctx_mgr.__enter__()
+
+    if span is not None:
+        _set_db_data(span, conn)
+        context._sentry_sql_span = span
+
+
+@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:
+        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
+    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(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(execution_context, "_sentry_sql_span_manager", None)  # type: Optional[ContextManager[Any]]
+
+    if ctx_mgr is not 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)