sentry-sdk 3.0.0a2__py2.py3-none-any.whl → 3.0.0a4__py2.py3-none-any.whl

sentry_sdk/integrations/spark/spark_driver.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 import sentry_sdk
 from sentry_sdk.integrations import Integration
 from sentry_sdk.utils import capture_internal_exceptions, ensure_integration_enabled
@@ -16,13 +17,11 @@ class SparkIntegration(Integration):
     identifier = "spark"
 
     @staticmethod
-    def setup_once():
-        # type: () -> None
+    def setup_once() -> None:
         _setup_sentry_tracing()
 
 
-def _set_app_properties():
-    # type: () -> None
+def _set_app_properties() -> 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.
@@ -41,8 +40,7 @@ def _set_app_properties():
         )
 
 
-def _start_sentry_listener(sc):
-    # type: (SparkContext) -> None
+def _start_sentry_listener(sc: SparkContext) -> None:
     """
     Start java gateway server to add custom `SparkListener`
     """
@@ -54,13 +52,11 @@ def _start_sentry_listener(sc):
     sc._jsc.sc().addSparkListener(listener)
 
 
-def _add_event_processor(sc):
-    # type: (SparkContext) -> None
+def _add_event_processor(sc: SparkContext) -> None:
     scope = sentry_sdk.get_isolation_scope()
 
     @scope.add_event_processor
-    def process_event(event, hint):
-        # type: (Event, Hint) -> Optional[Event]
+    def process_event(event: Event, hint: Hint) -> Optional[Event]:
         with capture_internal_exceptions():
             if sentry_sdk.get_client().get_integration(SparkIntegration) is None:
                 return event
@@ -90,23 +86,22 @@ def _add_event_processor(sc):
         return event
 
 
-def _activate_integration(sc):
-    # type: (SparkContext) -> None
+def _activate_integration(sc: SparkContext) -> None:
 
     _start_sentry_listener(sc)
     _set_app_properties()
     _add_event_processor(sc)
 
 
-def _patch_spark_context_init():
-    # type: () -> None
+def _patch_spark_context_init() -> 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]
+    def _sentry_patched_spark_context_init(
+        self: SparkContext, *args: Any, **kwargs: Any
+    ) -> Optional[Any]:
         rv = spark_context_init(self, *args, **kwargs)
         _activate_integration(self)
         return rv
@@ -114,8 +109,7 @@ def _patch_spark_context_init():
     SparkContext._do_init = _sentry_patched_spark_context_init
 
 
-def _setup_sentry_tracing():
-    # type: () -> None
+def _setup_sentry_tracing() -> None:
     from pyspark import SparkContext
 
     if SparkContext._active_spark_context is not None:
@@ -125,102 +119,76 @@ def _setup_sentry_tracing():
 
 
 class SparkListener:
-    def onApplicationEnd(self, applicationEnd):  # noqa: N802,N803
-        # type: (Any) -> None
+    def onApplicationEnd(self, applicationEnd: Any) -> None:
         pass
 
-    def onApplicationStart(self, applicationStart):  # noqa: N802,N803
-        # type: (Any) -> None
+    def onApplicationStart(self, applicationStart: Any) -> None:
         pass
 
-    def onBlockManagerAdded(self, blockManagerAdded):  # noqa: N802,N803
-        # type: (Any) -> None
+    def onBlockManagerAdded(self, blockManagerAdded: Any) -> None:
         pass
 
-    def onBlockManagerRemoved(self, blockManagerRemoved):  # noqa: N802,N803
-        # type: (Any) -> None
+    def onBlockManagerRemoved(self, blockManagerRemoved: Any) -> None:
         pass
 
-    def onBlockUpdated(self, blockUpdated):  # noqa: N802,N803
-        # type: (Any) -> None
+    def onBlockUpdated(self, blockUpdated: Any) -> None:
         pass
 
-    def onEnvironmentUpdate(self, environmentUpdate):  # noqa: N802,N803
-        # type: (Any) -> None
+    def onEnvironmentUpdate(self, environmentUpdate: Any) -> None:
         pass
 
-    def onExecutorAdded(self, executorAdded):  # noqa: N802,N803
-        # type: (Any) -> None
+    def onExecutorAdded(self, executorAdded: Any) -> None:
         pass
 
-    def onExecutorBlacklisted(self, executorBlacklisted):  # noqa: N802,N803
-        # type: (Any) -> None
+    def onExecutorBlacklisted(self, executorBlacklisted: Any) -> None:
         pass
 
-    def onExecutorBlacklistedForStage(  # noqa: N802
-        self, executorBlacklistedForStage  # noqa: N803
-    ):
-        # type: (Any) -> None
+    def onExecutorBlacklistedForStage(self, executorBlacklistedForStage: Any) -> None:
         pass
 
-    def onExecutorMetricsUpdate(self, executorMetricsUpdate):  # noqa: N802,N803
-        # type: (Any) -> None
+    def onExecutorMetricsUpdate(self, executorMetricsUpdate: Any) -> None:
         pass
 
-    def onExecutorRemoved(self, executorRemoved):  # noqa: N802,N803
-        # type: (Any) -> None
+    def onExecutorRemoved(self, executorRemoved: Any) -> None:
         pass
 
-    def onJobEnd(self, jobEnd):  # noqa: N802,N803
-        # type: (Any) -> None
+    def onJobEnd(self, jobEnd: Any) -> None:
         pass
 
-    def onJobStart(self, jobStart):  # noqa: N802,N803
-        # type: (Any) -> None
+    def onJobStart(self, jobStart: Any) -> None:
         pass
 
-    def onNodeBlacklisted(self, nodeBlacklisted):  # noqa: N802,N803
-        # type: (Any) -> None
+    def onNodeBlacklisted(self, nodeBlacklisted: Any) -> None:
         pass
 
-    def onNodeBlacklistedForStage(self, nodeBlacklistedForStage):  # noqa: N802,N803
-        # type: (Any) -> None
+    def onNodeBlacklistedForStage(self, nodeBlacklistedForStage: Any) -> None:
         pass
 
-    def onNodeUnblacklisted(self, nodeUnblacklisted):  # noqa: N802,N803
-        # type: (Any) -> None
+    def onNodeUnblacklisted(self, nodeUnblacklisted: Any) -> None:
         pass
 
-    def onOtherEvent(self, event):  # noqa: N802,N803
-        # type: (Any) -> None
+    def onOtherEvent(self, event: Any) -> None:
         pass
 
-    def onSpeculativeTaskSubmitted(self, speculativeTask):  # noqa: N802,N803
-        # type: (Any) -> None
+    def onSpeculativeTaskSubmitted(self, speculativeTask: Any) -> None:
         pass
 
-    def onStageCompleted(self, stageCompleted):  # noqa: N802,N803
-        # type: (Any) -> None
+    def onStageCompleted(self, stageCompleted: Any) -> None:
         pass
 
-    def onStageSubmitted(self, stageSubmitted):  # noqa: N802,N803
-        # type: (Any) -> None
+    def onStageSubmitted(self, stageSubmitted: Any) -> None:
         pass
 
-    def onTaskEnd(self, taskEnd):  # noqa: N802,N803
-        # type: (Any) -> None
+    def onTaskEnd(self, taskEnd: Any) -> None:
         pass
 
-    def onTaskGettingResult(self, taskGettingResult):  # noqa: N802,N803
-        # type: (Any) -> None
+    def onTaskGettingResult(self, taskGettingResult: Any) -> None:
         pass
 
-    def onTaskStart(self, taskStart):  # noqa: N802,N803
-        # type: (Any) -> None
+    def onTaskStart(self, taskStart: Any) -> None:
         pass
 
-    def onUnpersistRDD(self, unpersistRDD):  # noqa: N802,N803
-        # type: (Any) -> None
+    def onUnpersistRDD(self, unpersistRDD: Any) -> None:
         pass
 
     class Java:
@@ -230,25 +198,22 @@ class SparkListener:
 class SentryListener(SparkListener):
     def _add_breadcrumb(
         self,
-        level,  # type: str
-        message,  # type: str
-        data=None,  # type: Optional[dict[str, Any]]
-    ):
-        # type: (...) -> None
+        level: str,
+        message: str,
+        data: Optional[dict[str, Any]] = None,
+    ) -> None:
         sentry_sdk.get_isolation_scope().add_breadcrumb(
             level=level, message=message, data=data
         )
 
-    def onJobStart(self, jobStart):  # noqa: N802,N803
-        # type: (Any) -> None
+    def onJobStart(self, jobStart: 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
+    def onJobEnd(self, jobEnd: Any) -> None:
         level = ""
         message = ""
         data = {"result": jobEnd.jobResult().toString()}
@@ -262,8 +227,7 @@ class SentryListener(SparkListener):
 
         self._add_breadcrumb(level=level, message=message, data=data)
 
-    def onStageSubmitted(self, stageSubmitted):  # noqa: N802,N803
-        # type: (Any) -> None
+    def onStageSubmitted(self, stageSubmitted: Any) -> None:
         stage_info = stageSubmitted.stageInfo()
         message = "Stage {} Submitted".format(stage_info.stageId())
 
@@ -275,8 +239,7 @@ class SentryListener(SparkListener):
         self._add_breadcrumb(level="info", message=message, data=data)
         _set_app_properties()
 
-    def onStageCompleted(self, stageCompleted):  # noqa: N802,N803
-        # type: (Any) -> None
+    def onStageCompleted(self, stageCompleted: Any) -> None:
         from py4j.protocol import Py4JJavaError  # type: ignore
 
         stage_info = stageCompleted.stageInfo()
@@ -300,8 +263,7 @@ class SentryListener(SparkListener):
         self._add_breadcrumb(level=level, message=message, data=data)
 
 
-def _get_attempt_id(stage_info):
-    # type: (Any) -> Optional[int]
+def _get_attempt_id(stage_info: Any) -> Optional[int]:
     try:
         return stage_info.attemptId()
     except Exception:

sentry_sdk/integrations/spark/spark_worker.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 import sys
 
 import sentry_sdk
@@ -23,15 +24,13 @@ class SparkWorkerIntegration(Integration):
     identifier = "spark_worker"
 
     @staticmethod
-    def setup_once():
-        # type: () -> None
+    def setup_once() -> None:
         import pyspark.daemon as original_daemon
 
         original_daemon.worker_main = _sentry_worker_main
 
 
-def _capture_exception(exc_info):
-    # type: (ExcInfo) -> None
+def _capture_exception(exc_info: ExcInfo) -> None:
     client = sentry_sdk.get_client()
 
     mechanism = {"type": "spark", "handled": False}
@@ -53,22 +52,20 @@ def _capture_exception(exc_info):
     if rv:
         rv.reverse()
         hint = event_hint_with_exc_info(exc_info)
-        event = {"level": "error", "exception": {"values": rv}}  # type: Event
+        event: Event = {"level": "error", "exception": {"values": rv}}
 
         _tag_task_context()
 
         sentry_sdk.capture_event(event, hint=hint)
 
 
-def _tag_task_context():
-    # type: () -> None
+def _tag_task_context() -> 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]
+    def process_event(event: Event, hint: Hint) -> Optional[Event]:
         with capture_internal_exceptions():
             integration = sentry_sdk.get_client().get_integration(
                 SparkWorkerIntegration
@@ -103,8 +100,7 @@ def _tag_task_context():
         return event
 
 
-def _sentry_worker_main(*args, **kwargs):
-    # type: (*Optional[Any], **Optional[Any]) -> None
+def _sentry_worker_main(*args: Optional[Any], **kwargs: Optional[Any]) -> None:
     import pyspark.worker as original_worker
 
     try:

sentry_sdk/integrations/sqlalchemy.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 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
@@ -29,8 +30,7 @@ class SqlalchemyIntegration(Integration):
     origin = f"auto.db.{identifier}"
 
     @staticmethod
-    def setup_once():
-        # type: () -> None
+    def setup_once() -> None:
         version = parse_version(SQLALCHEMY_VERSION)
         _check_minimum_version(SqlalchemyIntegration, version)
 
@@ -41,9 +41,14 @@ class SqlalchemyIntegration(Integration):
 
 @ensure_integration_enabled(SqlalchemyIntegration)
 def _before_cursor_execute(
-    conn, cursor, statement, parameters, context, executemany, *args
-):
-    # type: (Any, Any, Any, Any, Any, bool, *Any) -> None
+    conn: Any,
+    cursor: Any,
+    statement: Any,
+    parameters: Any,
+    context: Any,
+    executemany: bool,
+    *args: Any,
+) -> None:
     ctx_mgr = record_sql_queries(
         cursor,
         statement,
@@ -62,13 +67,14 @@ def _before_cursor_execute(
 
 
 @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(
+def _after_cursor_execute(
+    conn: Any, cursor: Any, statement: Any, parameters: Any, context: Any, *args: Any
+) -> None:
+    ctx_mgr: Optional[ContextManager[Any]] = getattr(
         context, "_sentry_sql_span_manager", None
-    )  # type: Optional[ContextManager[Any]]
+    )
 
-    span = getattr(context, "_sentry_sql_span", None)  # type: Optional[Span]
+    span: Optional[Span] = getattr(context, "_sentry_sql_span", None)
     if span is not None:
         with capture_internal_exceptions():
             add_query_source(span)
@@ -78,13 +84,12 @@ def _after_cursor_execute(conn, cursor, statement, parameters, context, *args):
         ctx_mgr.__exit__(None, None, None)
 
 
-def _handle_error(context, *args):
-    # type: (Any, *Any) -> None
+def _handle_error(context: Any, *args: Any) -> None:
     execution_context = context.execution_context
     if execution_context is None:
         return
 
-    span = getattr(execution_context, "_sentry_sql_span", None)  # type: Optional[Span]
+    span: Optional[Span] = getattr(execution_context, "_sentry_sql_span", None)
 
     if span is not None:
         span.set_status(SPANSTATUS.INTERNAL_ERROR)
@@ -92,9 +97,9 @@ def _handle_error(context, *args):
     # _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(
+    ctx_mgr: Optional[ContextManager[Any]] = 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
@@ -102,8 +107,7 @@ def _handle_error(context, *args):
 
 
 # See: https://docs.sqlalchemy.org/en/20/dialects/index.html
-def _get_db_system(name):
-    # type: (str) -> Optional[str]
+def _get_db_system(name: str) -> Optional[str]:
     name = str(name)
 
     if "sqlite" in name:
@@ -124,8 +128,7 @@ def _get_db_system(name):
     return None
 
 
-def _set_db_data(span, conn):
-    # type: (Span, Any) -> None
+def _set_db_data(span: Span, conn: Any) -> None:
     db_system = _get_db_system(conn.engine.name)
     if db_system is not None:
         span.set_attribute(SPANDATA.DB_SYSTEM, db_system)