feat: add flag to avoid collecting the HTTP's request and response body (#247)

* feat: add flag to avoid collecting the HTTP's request and response body

Co-authored-by: CircleCI <[email protected]>

Author: saartochner-lumigo

src/lumigo_tracer/lumigo_utils.py

@@ -83,6 +83,7 @@
 STACKTRACE_LINE_TO_DROP = "lumigo_tracer/tracer.py"
 Container = TypeVar("Container", dict, list)
 DEFAULT_AUTO_TAG_KEY = "LUMIGO_AUTO_TAG"
+SKIP_COLLECTING_HTTP_BODY_KEY = "LUMIGO_SKIP_COLLECTING_HTTP_BODY"
 
 _logger: Dict[str, logging.Logger] = {}
 
@@ -141,6 +142,7 @@ class Configuration:
     should_scrub_known_services: bool = False
     is_sync_tracer: bool = False
     auto_tag: List[str] = []
+    skip_collecting_http_body: bool = False
 
     @staticmethod
     def get_max_entry_size(has_error: bool = False) -> int:
@@ -165,6 +167,7 @@ def config(
     edge_kinesis_aws_access_key_id: Optional[str] = None,
     edge_kinesis_aws_secret_access_key: Optional[str] = None,
     auto_tag: Optional[List[str]] = None,
+    skip_collecting_http_body: bool = False,
 ) -> None:
     """
     This function configure the lumigo wrapper.
@@ -185,6 +188,7 @@ def config(
     :param edge_kinesis_aws_access_key_id: The credentials to push to the Kinesis in China region
     :param edge_kinesis_aws_secret_access_key: The credentials to push to the Kinesis in China region
     :param auto_tag: The keys from the event that should be used as execution tags.
+    :param skip_collecting_http_body: Should we not collect the HTTP request and response bodies.
     """
 
     Configuration.token = token or os.environ.get(LUMIGO_TOKEN_KEY, "")
@@ -232,7 +236,7 @@ def config(
     Configuration.max_entry_size = int(os.environ.get("LUMIGO_MAX_ENTRY_SIZE", max_entry_size))
     Configuration.edge_kinesis_stream_name = (
         edge_kinesis_stream_name
-        or os.environ.get("LUMIGO_EDGE_KINESIS_STREAM_NAME")  # noqa`
+        or os.environ.get("LUMIGO_EDGE_KINESIS_STREAM_NAME")  # noqa
         or EDGE_KINESIS_STREAM_NAME  # noqa
     )
     Configuration.edge_kinesis_aws_access_key_id = edge_kinesis_aws_access_key_id or os.environ.get(
@@ -249,6 +253,11 @@ def config(
     Configuration.auto_tag = auto_tag or os.environ.get(
         "LUMIGO_AUTO_TAG", DEFAULT_AUTO_TAG_KEY
     ).split(",")
+    Configuration.skip_collecting_http_body = (
+        not Configuration.verbose
+        or skip_collecting_http_body  # noqa: W503
+        or os.environ.get(SKIP_COLLECTING_HTTP_BODY_KEY, "false").lower() == "true"  # noqa: W503
+    )
 
 
 def _is_span_has_error(span: dict) -> bool:

src/lumigo_tracer/wrappers/http/http_parser.py

@@ -48,7 +48,7 @@ def parse_request(self, parse_params: HttpRequest) -> dict:
             additional_info = {
                 "headers": lumigo_dumps(parse_params.headers),
                 "body": lumigo_dumps(parse_params.body, omit_skip_path=HttpState.omit_skip_path)
-                if parse_params.body
+                if parse_params.body and not Configuration.skip_collecting_http_body
                 else "",
                 "method": parse_params.method,
                 "uri": parse_params.uri,
@@ -78,7 +78,9 @@ def parse_response(self, url: str, status_code: int, headers: dict, body: bytes)
         if Configuration.verbose and not should_scrub_domain(url):
             additional_info = {
                 "headers": lumigo_dumps(headers, max_size),
-                "body": lumigo_dumps(body, max_size) if body else "",
+                "body": lumigo_dumps(body, max_size)
+                if body and not Configuration.skip_collecting_http_body
+                else "",
                 "statusCode": status_code,
             }
         else:

src/lumigo_tracer/wrappers/http/sync_http_wrappers.py

@@ -74,8 +74,12 @@ def add_unparsed_request(span_id: Optional[str], parse_params: HttpRequest) -> O
             if http_info.get("host") == parse_params.host:
                 if "response" not in http_info:
                     SpansContainer.get_span().get_span_by_id(span_id)
-                    http_info["request"]["body"] = concat_old_body_to_new(
-                        http_info.get("request", {}).get("body"), parse_params.body
+                    http_info["request"]["body"] = (
+                        concat_old_body_to_new(
+                            http_info.get("request", {}).get("body"), parse_params.body
+                        )
+                        if not Configuration.skip_collecting_http_body
+                        else ""
                     )
                     if HttpState.previous_span_id == span_id and HttpState.previous_request:
                         HttpState.previous_request.body += parse_params.body
@@ -130,7 +134,7 @@ def _update_request_data_increased_size_limit(http_info: dict, max_size: int) ->
                 max_size,
                 omit_skip_path=HttpState.omit_skip_path,
             )
-            if HttpState.previous_request.body
+            if HttpState.previous_request.body and not Configuration.skip_collecting_http_body
             else "",
             "headers": lumigo_dumps(HttpState.previous_request.headers, max_size),
         }

src/test/unit/test_tracer.py

@@ -22,6 +22,7 @@
     EXECUTION_TAGS_KEY,
     report_json,
     EDGE_KINESIS_STREAM_NAME,
+    SKIP_COLLECTING_HTTP_BODY_KEY,
 )
 
 from lumigo_tracer.spans_container import SpansContainer
@@ -203,6 +204,30 @@ def lambda_test_function(event, context):
         assert any(line == "hello" for line in capturer.get_lines())
 
 
+@pytest.mark.parametrize("is_verbose", [True, False])
+def test_skip_collecting_http_parts(monkeypatch, context, is_verbose):
+    if is_verbose:
+        monkeypatch.setenv("LUMIGO_VERBOSE", "false")
+    else:
+        monkeypatch.setenv(SKIP_COLLECTING_HTTP_BODY_KEY, "true")
+
+    @lumigo_tracer()
+    def lambda_test_function(event, context):
+        conn = http.client.HTTPConnection("www.google.com")
+        conn.request("POST", "/", json.dumps({"a": "b"}))
+        return {"hello": "world"}
+
+    lambda_test_function({}, context)
+    http_spans = list(SpansContainer.get_span().spans.values())
+    assert http_spans[0]["info"]["httpInfo"]["request"]["body"] == ""
+    if is_verbose:
+        assert "uri" not in http_spans[0]["info"]["httpInfo"]["request"]
+        assert "headers" not in http_spans[0]["info"]["httpInfo"]["request"]
+    else:
+        assert http_spans[0]["info"]["httpInfo"]["request"]["uri"] == "www.google.com/"
+        assert http_spans[0]["info"]["httpInfo"]["request"]["headers"]
+
+
 def test_lumigo_chalice(context):
     class App:
         @property