requests api should be traced as a whole (#203)

* requests api should be traced as a whole

Author: saartochner-lumigo

src/lumigo_tracer/spans_container.py

@@ -143,12 +143,14 @@ def start_timeout_timer(self, context=None) -> None:
                 return
             TimeoutMechanism.start(remaining_time - buffer, self.handle_timeout)
 
-    def add_span(self, span: dict):
+    def add_span(self, span: dict) -> dict:
         """
         This function parses an request event and add it to the span.
         """
-        self.spans.append(recursive_json_join(span, self.base_msg))
+        new_span = recursive_json_join(span, self.base_msg)
+        self.spans.append(new_span)
         self.span_ids_to_send.add(span["id"])
+        return new_span
 
     def get_last_span(self) -> Optional[dict]:
         if not self.spans:
@@ -185,8 +187,9 @@ def update_event_times(
             self.spans[-1]["started"] = int(start_timestamp * 1000)
             self.spans[-1]["ended"] = int(end_timestamp * 1000)
 
+    @staticmethod
     def _create_exception_event(
-        self, exc_type: str, message: str, stacktrace: str = "", frames: Optional[List[dict]] = None
+        exc_type: str, message: str, stacktrace: str = "", frames: Optional[List[dict]] = None
     ):
         return {
             "type": exc_type,
@@ -195,19 +198,25 @@ def _create_exception_event(
             "frames": frames or [],
         }
 
+    @staticmethod
+    def add_exception_to_span(
+        span: dict, exception: Exception, frames_infos: List[inspect.FrameInfo]
+    ):
+        message = exception.args[0] if exception.args else None
+        if not isinstance(message, str):
+            message = str(message)
+        span["error"] = SpansContainer._create_exception_event(
+            exc_type=exception.__class__.__name__,
+            message=message,
+            stacktrace=get_stacktrace(exception),
+            frames=format_frames(frames_infos) if Configuration.verbose else [],
+        )
+
     def add_exception_event(
         self, exception: Exception, frames_infos: List[inspect.FrameInfo]
     ) -> None:
         if self.function_span:
-            message = exception.args[0] if exception.args else None
-            if not isinstance(message, str):
-                message = str(message)
-            self.function_span["error"] = self._create_exception_event(
-                exc_type=exception.__class__.__name__,
-                message=message,
-                stacktrace=get_stacktrace(exception),
-                frames=format_frames(frames_infos) if Configuration.verbose else [],
-            )
+            self.add_exception_to_span(self.function_span, exception, frames_infos)
 
     def add_step_end_event(self, ret_val):
         message_id = str(uuid.uuid4())

src/lumigo_tracer/wrappers/http/sync_http_wrappers.py

@@ -2,7 +2,7 @@
 import http.client
 from io import BytesIO
 import importlib.util
-from typing import Optional
+from typing import Optional, Dict
 
 from lumigo_tracer.libs.wrapt import wrap_function_wrapper
 from lumigo_tracer.parsing_utils import safe_get_list, recursive_json_join
@@ -36,16 +36,17 @@ def is_lumigo_edge(host: Optional[str]) -> bool:
     return False
 
 
-def add_request_event(parse_params: HttpRequest):
+def add_request_event(parse_params: HttpRequest) -> Dict:
     """
     This function parses an request event and add it to the span.
     """
     if is_lumigo_edge(parse_params.host):
-        return
+        return {}
     parser = get_parser(parse_params.host)()
     msg = parser.parse_request(parse_params)
     HttpState.previous_request = parse_params
-    SpansContainer.get_span().add_span(msg)
+    new_span = SpansContainer.get_span().add_span(msg)
+    return new_span
 
 
 def add_unparsed_request(parse_params: HttpRequest):
@@ -204,9 +205,29 @@ def _requests_wrapper(func, instance, args, kwargs):
     This is the wrapper of the function `requests.request`.
     This function is being wrapped specifically because it initializes the connection by itself and parses the response,
         which creates a gap from the traditional http.client wrapping.
+    Moreover, these "extra" steps may raise exceptions. We should attach the error to the http span.
     """
     start_time = datetime.now()
-    ret_val = func(*args, **kwargs)
+    try:
+        ret_val = func(*args, **kwargs)
+    except Exception as exception:
+        with lumigo_safe_execute("requests wrapper exception occurred"):
+            method = safe_get_list(args, 0, kwargs.get("method", "")).upper()
+            url = safe_get_list(args, 1, kwargs.get("url"))
+            if HttpState.previous_request and HttpState.previous_request.host in url:
+                span = SpansContainer.get_span().get_last_span()
+            else:
+                span = add_request_event(
+                    HttpRequest(
+                        host=url,
+                        method=method,
+                        uri=url,
+                        body=kwargs.get("data"),
+                        headers=kwargs.get("headers"),
+                    )
+                )
+            SpansContainer.add_exception_to_span(span, exception, [])
+        raise
     with lumigo_safe_execute("requests wrapper time updates"):
         SpansContainer.get_span().update_event_times(start_time=start_time)
     return ret_val
@@ -276,3 +297,4 @@ def wrap_http_calls():
             )
         if importlib.util.find_spec("requests"):
             wrap_function_wrapper("requests.api", "request", _requests_wrapper)
+            wrap_function_wrapper("requests", "request", _requests_wrapper)

src/test/unit/wrappers/http/test_sync_http_wrappers.py

@@ -265,6 +265,55 @@ def delayed_getaddrinfo(*args, **kwargs):
     assert span["started"] - start_time < 100
 
 
+@pytest.mark.parametrize(
+    "func_to_patch",
+    [
+        (socket, "getaddrinfo"),  # this function is being called before the request
+        (http.client.HTTPConnection, "getresponse"),  # after the request
+    ],
+)
+def test_requests_failure_before_http_call(monkeypatch, context, func_to_patch):
+    @lumigo_tracer.lumigo_tracer()
+    def lambda_test_function(event, context):
+        try:
+            requests.post("https://www.google.com", data=b"123", headers={"a": "b"})
+        except ZeroDivisionError:
+            return True
+        return False
+
+    # requests executes this function before/after the http call
+    monkeypatch.setattr(*func_to_patch, lambda *args, **kwargs: 1 / 0)
+
+    assert lambda_test_function({}, context) is True
+
+    assert len(SpansContainer.get_span().spans) == 1
+    span = SpansContainer.get_span().spans[0]
+    assert span["error"]["message"] == "division by zero"
+    assert span["info"]["httpInfo"]["request"]["method"] == "POST"
+    assert span["info"]["httpInfo"]["request"]["body"] == '"123"'
+    assert span["info"]["httpInfo"]["request"]["headers"]
+
+
+def test_requests_failure_with_kwargs(monkeypatch, context):
+    @lumigo_tracer.lumigo_tracer()
+    def lambda_test_function(event, context):
+        try:
+            requests.request(
+                method="POST", url="https://www.google.com", data=b"123", headers={"a": "b"}
+            )
+        except ZeroDivisionError:
+            return True
+        return False
+
+    monkeypatch.setattr(socket, "getaddrinfo", lambda *args, **kwargs: 1 / 0)
+
+    assert lambda_test_function({}, context) is True
+
+    assert len(SpansContainer.get_span().spans) == 1
+    span = SpansContainer.get_span().spans[0]
+    assert span["info"]["httpInfo"]["request"]["method"] == "POST"
+
+
 def test_wrapping_with_tags_for_api_gw_headers(monkeypatch, context):
     monkeypatch.setattr(auto_tag_event, "AUTO_TAG_API_GW_HEADERS", ["Accept"])