bababuah

Osmar Coronel · Osmar Coronel · commit 7ed9f7dc864d · 2019-08-02T11:08:08.000-07:00
diff --git a/sentry_sdk/integrations/beam.py b/sentry_sdk/integrations/beam.py
@@ -2,6 +2,7 @@
 
 import sys
 import types
+import linecache
 from functools import wraps
 
 from sentry_sdk.hub import Hub
@@ -17,6 +18,18 @@
 class BeamIntegration(Integration):
     identifier = "beam"
 
+    def __init__(self, send_source=False):
+        self.cached_source = None
+        self.cached_file = None
+        if send_source:
+            import __main__ as main
+
+            filename = main.__file__
+            if filename not in linecache.cache:
+                linecache.getlines(filename)
+            self.cached_source = linecache.cache[filename]
+            self.cached_file = filename
+
     @staticmethod
     def setup_once():
         # type: () -> None
@@ -130,6 +143,7 @@ def raiseException(client):
         hub.bind_client(client)
     exc_info = sys.exc_info()
     with capture_internal_exceptions():
+        setup_file(hub)
         _capture_exception(exc_info, hub)
     reraise(*exc_info)
 
@@ -145,3 +159,12 @@ def _wrap_generator_call(gen, client):
             break
         except Exception:
             raiseException(client)
+
+
+def setup_file(hub):
+    integration = hub.get_integration(BeamIntegration)
+    if integration:
+        cached_source = integration.cached_source
+        cached_file = integration.cached_file
+        if cached_source is not None and cached_file not in linecache.cache:
+            linecache.cache[cached_file] = cached_source
diff --git a/tests/integrations/beam/test_beam.py b/tests/integrations/beam/test_beam.py
@@ -163,8 +163,11 @@ def process_outputs(self, windowed_input_element, results):
 
 @pytest.fixture
 def init_beam(sentry_init):
-    def inner(fn):
-        sentry_init(default_integrations=False, integrations=[BeamIntegration()])
+    def inner(fn, send_source=False):
+        sentry_init(
+            default_integrations=False,
+            integrations=[BeamIntegration(send_source=send_source)],
+        )
         # Little hack to avoid having to run the whole pipeline.
         pardo = ParDo(fn)
         signature = pardo._signature
@@ -176,9 +179,19 @@ def inner(fn):
     return inner
 
 
-@pytest.mark.parametrize("fn", [test_simple, test_callable, test_place_holder])
-def test_invoker_normal(init_beam, fn):
-    invoker = init_beam(fn)
+@pytest.mark.parametrize(
+    "fn,send_source",
+    [
+        [test_simple, False],
+        [test_simple, True],
+        [test_callable, False],
+        [test_callable, True],
+        [test_place_holder, False],
+        [test_place_holder, True]
+    ],
+)
+def test_invoker_normal(init_beam, fn, send_source):
+    invoker = init_beam(fn, send_source)
     print("Normal testing {} with {} invoker.".format(fn, invoker))
     windowed_value = WindowedValue(False, 0, [None])
     invoker.invoke_process(windowed_value)
