feat(tracing): Add sampling context from AWS and GCP (getsentry#916)

lobsterkatie · web-flow · commit a7ef7c05df66 · 2020-11-10T15:36:04.000-08:00
diff --git a/sentry_sdk/_compat.py b/sentry_sdk/_compat.py
@@ -7,7 +7,6 @@
     from typing import Tuple
     from typing import Any
     from typing import Type
-
     from typing import TypeVar
 
     T = TypeVar("T")
diff --git a/sentry_sdk/integrations/aws_lambda.py b/sentry_sdk/integrations/aws_lambda.py
@@ -65,7 +65,7 @@ def sentry_init_error(*args, **kwargs):
 
 def _wrap_handler(handler):
     # type: (F) -> F
-    def sentry_handler(aws_event, context, *args, **kwargs):
+    def sentry_handler(aws_event, aws_context, *args, **kwargs):
         # type: (Any, Any, *Any, **Any) -> Any
 
         # Per https://docs.aws.amazon.com/lambda/latest/dg/python-handler.html,
@@ -94,21 +94,23 @@ def sentry_handler(aws_event, context, *args, **kwargs):
         hub = Hub.current
         integration = hub.get_integration(AwsLambdaIntegration)
         if integration is None:
-            return handler(aws_event, context, *args, **kwargs)
+            return handler(aws_event, aws_context, *args, **kwargs)
 
         # If an integration is there, a client has to be there.
         client = hub.client  # type: Any
-        configured_time = context.get_remaining_time_in_millis()
+        configured_time = aws_context.get_remaining_time_in_millis()
 
         with hub.push_scope() as scope:
             with capture_internal_exceptions():
                 scope.clear_breadcrumbs()
                 scope.add_event_processor(
                     _make_request_event_processor(
-                        request_data, context, configured_time
+                        request_data, aws_context, configured_time
                     )
                 )
-                scope.set_tag("aws_region", context.invoked_function_arn.split(":")[3])
+                scope.set_tag(
+                    "aws_region", aws_context.invoked_function_arn.split(":")[3]
+                )
                 if batch_size > 1:
                     scope.set_tag("batch_request", True)
                     scope.set_tag("batch_size", batch_size)
@@ -134,11 +136,17 @@ def sentry_handler(aws_event, context, *args, **kwargs):
 
             headers = request_data.get("headers", {})
             transaction = Transaction.continue_from_headers(
-                headers, op="serverless.function", name=context.function_name
+                headers, op="serverless.function", name=aws_context.function_name
             )
-            with hub.start_transaction(transaction):
+            with hub.start_transaction(
+                transaction,
+                custom_sampling_context={
+                    "aws_event": aws_event,
+                    "aws_context": aws_context,
+                },
+            ):
                 try:
-                    return handler(aws_event, context, *args, **kwargs)
+                    return handler(aws_event, aws_context, *args, **kwargs)
                 except Exception:
                     exc_info = sys.exc_info()
                     sentry_event, hint = event_from_exception(
@@ -177,23 +185,8 @@ def __init__(self, timeout_warning=False):
     def setup_once():
         # type: () -> None
 
-        # Python 2.7: Everything is in `__main__`.
-        #
-        # Python 3.7: If the bootstrap module is *already imported*, it is the
-        # one we actually want to use (no idea what's in __main__)
-        #
-        # On Python 3.8 bootstrap is also importable, but will be the same file
-        # as __main__ imported under a different name:
-        #
-        #     sys.modules['__main__'].__file__ == sys.modules['bootstrap'].__file__
-        #     sys.modules['__main__'] is not sys.modules['bootstrap']
-        #
-        # Such a setup would then make all monkeypatches useless.
-        if "bootstrap" in sys.modules:
-            lambda_bootstrap = sys.modules["bootstrap"]  # type: Any
-        elif "__main__" in sys.modules:
-            lambda_bootstrap = sys.modules["__main__"]
-        else:
+        lambda_bootstrap = get_lambda_bootstrap()
+        if not lambda_bootstrap:
             logger.warning(
                 "Not running in AWS Lambda environment, "
                 "AwsLambdaIntegration disabled (could not find bootstrap module)"
@@ -280,6 +273,29 @@ def inner(*args, **kwargs):
             )
 
 
+def get_lambda_bootstrap():
+    # type: () -> Optional[Any]
+
+    # Python 2.7: Everything is in `__main__`.
+    #
+    # Python 3.7: If the bootstrap module is *already imported*, it is the
+    # one we actually want to use (no idea what's in __main__)
+    #
+    # On Python 3.8 bootstrap is also importable, but will be the same file
+    # as __main__ imported under a different name:
+    #
+    #     sys.modules['__main__'].__file__ == sys.modules['bootstrap'].__file__
+    #     sys.modules['__main__'] is not sys.modules['bootstrap']
+    #
+    # Such a setup would then make all monkeypatches useless.
+    if "bootstrap" in sys.modules:
+        return sys.modules["bootstrap"]
+    elif "__main__" in sys.modules:
+        return sys.modules["__main__"]
+    else:
+        return None
+
+
 def _make_request_event_processor(aws_event, aws_context, configured_timeout):
     # type: (Any, Any, Any) -> EventProcessor
     start_time = datetime.utcnow()
diff --git a/sentry_sdk/integrations/gcp.py b/sentry_sdk/integrations/gcp.py
@@ -34,13 +34,13 @@
 
 def _wrap_func(func):
     # type: (F) -> F
-    def sentry_func(functionhandler, event, *args, **kwargs):
+    def sentry_func(functionhandler, gcp_event, *args, **kwargs):
         # type: (Any, Any, *Any, **Any) -> Any
 
         hub = Hub.current
         integration = hub.get_integration(GcpIntegration)
         if integration is None:
-            return func(functionhandler, event, *args, **kwargs)
+            return func(functionhandler, gcp_event, *args, **kwargs)
 
         # If an integration is there, a client has to be there.
         client = hub.client  # type: Any
@@ -50,7 +50,7 @@ def sentry_func(functionhandler, event, *args, **kwargs):
             logger.debug(
                 "The configured timeout could not be fetched from Cloud Functions configuration."
             )
-            return func(functionhandler, event, *args, **kwargs)
+            return func(functionhandler, gcp_event, *args, **kwargs)
 
         configured_time = int(configured_time)
 
@@ -60,7 +60,9 @@ def sentry_func(functionhandler, event, *args, **kwargs):
             with capture_internal_exceptions():
                 scope.clear_breadcrumbs()
                 scope.add_event_processor(
-                    _make_request_event_processor(event, configured_time, initial_time)
+                    _make_request_event_processor(
+                        gcp_event, configured_time, initial_time
+                    )
                 )
                 scope.set_tag("gcp_region", environ.get("FUNCTION_REGION"))
                 timeout_thread = None
@@ -76,22 +78,34 @@ def sentry_func(functionhandler, event, *args, **kwargs):
                     timeout_thread.start()
 
             headers = {}
-            if hasattr(event, "headers"):
-                headers = event.headers
+            if hasattr(gcp_event, "headers"):
+                headers = gcp_event.headers
             transaction = Transaction.continue_from_headers(
                 headers, op="serverless.function", name=environ.get("FUNCTION_NAME", "")
             )
-            with hub.start_transaction(transaction):
+            sampling_context = {
+                "gcp_env": {
+                    "function_name": environ.get("FUNCTION_NAME"),
+                    "function_entry_point": environ.get("ENTRY_POINT"),
+                    "function_identity": environ.get("FUNCTION_IDENTITY"),
+                    "function_region": environ.get("FUNCTION_REGION"),
+                    "function_project": environ.get("GCP_PROJECT"),
+                },
+                "gcp_event": gcp_event,
+            }
+            with hub.start_transaction(
+                transaction, custom_sampling_context=sampling_context
+            ):
                 try:
-                    return func(functionhandler, event, *args, **kwargs)
+                    return func(functionhandler, gcp_event, *args, **kwargs)
                 except Exception:
                     exc_info = sys.exc_info()
-                    event, hint = event_from_exception(
+                    sentry_event, hint = event_from_exception(
                         exc_info,
                         client_options=client.options,
                         mechanism={"type": "gcp", "handled": False},
                     )
-                    hub.capture_event(event, hint=hint)
+                    hub.capture_event(sentry_event, hint=hint)
                     reraise(*exc_info)
                 finally:
                     if timeout_thread:
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -355,18 +355,60 @@ class StringContaining(object):
         def __init__(self, substring):
             self.substring = substring
 
+            try:
+                # unicode only exists in python 2
+                self.valid_types = (str, unicode)  # noqa
+            except NameError:
+                self.valid_types = (str,)
+
         def __eq__(self, test_string):
-            if not isinstance(test_string, str):
+            if not isinstance(test_string, self.valid_types):
                 return False
 
             if len(self.substring) > len(test_string):
                 return False
 
             return self.substring in test_string
 
+        def __ne__(self, test_string):
+            return not self.__eq__(test_string)
+
     return StringContaining
 
 
+def _safe_is_equal(x, y):
+    """
+    Compares two values, preferring to use the first's __eq__ method if it
+    exists and is implemented.
+
+    Accounts for py2/py3 differences (like ints in py2 not having a __eq__
+    method), as well as the incomparability of certain types exposed by using
+    raw __eq__ () rather than ==.
+    """
+
+    # Prefer using __eq__ directly to ensure that examples like
+    #
+    #   maisey = Dog()
+    #   maisey.name = "Maisey the Dog"
+    #   maisey == ObjectDescribedBy(attrs={"name": StringContaining("Maisey")})
+    #
+    # evaluate to True (in other words, examples where the values in self.attrs
+    # might also have custom __eq__ methods; this makes sure those methods get
+    # used if possible)
+    try:
+        is_equal = x.__eq__(y)
+    except AttributeError:
+        is_equal = NotImplemented
+
+    # this can happen on its own, too (i.e. without an AttributeError being
+    # thrown), which is why this is separate from the except block above
+    if is_equal == NotImplemented:
+        # using == smoothes out weird variations exposed by raw __eq__
+        return x == y
+
+    return is_equal
+
+
 @pytest.fixture(name="DictionaryContaining")
 def dictionary_containing_matcher():
     """
@@ -397,13 +439,19 @@ def __eq__(self, test_dict):
             if len(self.subdict) > len(test_dict):
                 return False
 
-            # Have to test self == other (rather than vice-versa) in case
-            # any of the values in self.subdict is another matcher with a custom
-            # __eq__ method (in LHS == RHS, LHS's __eq__ is tried before RHS's).
-            # In other words, this order is important so that examples like
-            # {"dogs": "are great"} == DictionaryContaining({"dogs": StringContaining("great")})
-            # evaluate to True
-            return all(self.subdict[key] == test_dict.get(key) for key in self.subdict)
+            for key, value in self.subdict.items():
+                try:
+                    test_value = test_dict[key]
+                except KeyError:  # missing key
+                    return False
+
+                if not _safe_is_equal(value, test_value):
+                    return False
+
+            return True
+
+        def __ne__(self, test_dict):
+            return not self.__eq__(test_dict)
 
     return DictionaryContaining
 
@@ -442,19 +490,19 @@ def __eq__(self, test_obj):
                 if not isinstance(test_obj, self.type):
                     return False
 
-            # all checks here done with getattr rather than comparing to
-            # __dict__ because __dict__ isn't guaranteed to exist
             if self.attrs:
-                # attributes must exist AND values must match
-                try:
-                    if any(
-                        getattr(test_obj, attr_name) != attr_value
-                        for attr_name, attr_value in self.attrs.items()
-                    ):
-                        return False  # wrong attribute value
-                except AttributeError:  # missing attribute
-                    return False
+                for attr_name, attr_value in self.attrs.items():
+                    try:
+                        test_value = getattr(test_obj, attr_name)
+                    except AttributeError:  # missing attribute
+                        return False
+
+                    if not _safe_is_equal(attr_value, test_value):
+                        return False
 
             return True
 
+        def __ne__(self, test_obj):
+            return not self.__eq__(test_obj)
+
     return ObjectDescribedBy
diff --git a/tests/integrations/aws_lambda/client.py b/tests/integrations/aws_lambda/client.py
@@ -49,6 +49,13 @@ def run_lambda_function(
             **subprocess_kwargs
         )
 
+        subprocess.check_call(
+            "pip install mock==3.0.0 funcsigs -t .",
+            cwd=tmpdir,
+            shell=True,
+            **subprocess_kwargs
+        )
+
         # https://docs.aws.amazon.com/lambda/latest/dg/lambda-python-how-to-create-deployment-package.html
         subprocess.check_call(
             "pip install ../*.tar.gz -t .", cwd=tmpdir, shell=True, **subprocess_kwargs
@@ -69,9 +76,19 @@ def run_lambda_function(
             )
 
         @add_finalizer
-        def delete_function():
+        def clean_up():
             client.delete_function(FunctionName=fn_name)
 
+            # this closes the web socket so we don't get a
+            #   ResourceWarning: unclosed <ssl.SSLSocket ... >
+            # warning on every test
+            # based on https://github.com/boto/botocore/pull/1810
+            # (if that's ever merged, this can just become client.close())
+            session = client._endpoint.http_session
+            managers = [session._manager] + list(session._proxy_managers.values())
+            for manager in managers:
+                manager.clear()
+
         response = client.invoke(
             FunctionName=fn_name,
             InvocationType="RequestResponse",
diff --git a/tests/integrations/aws_lambda/test_aws.py b/tests/integrations/aws_lambda/test_aws.py
diff --git a/tests/integrations/gcp/test_gcp.py b/tests/integrations/gcp/test_gcp.py
