Skip to content
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 20 additions & 2 deletions sentry_sdk/integrations/spark/spark_driver.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,8 @@
from sentry_sdk._types import Event, Hint
from pyspark import SparkContext

_spark_context_class = None


class SparkIntegration(Integration):
identifier = "spark"
Expand Down Expand Up @@ -100,23 +102,39 @@ def _activate_integration(sc):

def _patch_spark_context_init():
# type: () -> None
from pyspark import SparkContext
global _spark_context_class
if _spark_context_class is None:
from pyspark import SparkContext

_spark_context_class = SparkContext
else:
SparkContext = _spark_context_class

spark_context_init = SparkContext._do_init

if getattr(spark_context_init, "_sentry_patched", False):
return

@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

_sentry_patched_spark_context_init._sentry_patched = True
SparkContext._do_init = _sentry_patched_spark_context_init


def _setup_sentry_tracing():
# type: () -> None
from pyspark import SparkContext
global _spark_context_class
if _spark_context_class is None:
from pyspark import SparkContext

_spark_context_class = SparkContext
else:
SparkContext = _spark_context_class

if SparkContext._active_spark_context is not None:
_activate_integration(SparkContext._active_spark_context)
Expand Down