From 60bc2ec337fb8a809f05591e227b8fd6690d5faa Mon Sep 17 00:00:00 2001 From: saartochner Date: Sun, 4 Oct 2020 17:13:36 +0300 Subject: [PATCH 01/24] retry in case of error with edge --- src/lumigo_tracer/utils.py | 27 ++++++++++++++++----------- src/test/unit/test_main_utils.py | 18 ++++++++++++++++++ 2 files changed, 34 insertions(+), 11 deletions(-) diff --git a/src/lumigo_tracer/utils.py b/src/lumigo_tracer/utils.py index 2c98e056..eb10daeb 100644 --- a/src/lumigo_tracer/utils.py +++ b/src/lumigo_tracer/utils.py @@ -215,23 +215,26 @@ def prepare_host(host): return host -def report_json(region: Union[None, str], msgs: List[dict]) -> int: +def report_json(region: Union[None, str], msgs: List[dict], should_retry: bool = True) -> int: """ This function sends the information back to the edge. :param region: The region to use as default if not configured otherwise. :param msgs: the message to send. + :param should_retry: False to disable the default retry on unsuccessful sending :return: The duration of reporting (in milliseconds), or 0 if we didn't send (due to configuration or fail). """ global edge_connection get_logger().info(f"reporting the messages: {msgs[:10]}") - host = prepare_host(Configuration.host or EDGE_HOST.format(region=region)) - duration = 0 - if not edge_connection or edge_connection.host != host: - edge_connection = establish_connection(host) - if not edge_connection: - return duration + host = None + with lumigo_safe_execute("report json: establish connection"): + host = prepare_host(Configuration.host or EDGE_HOST.format(region=region)) + duration = 0 + if not edge_connection or edge_connection.host != host: + edge_connection = establish_connection(host) + if not edge_connection: + return duration if Configuration.should_report: try: prune_trace: bool = not os.environ.get("LUMIGO_PRUNE_TRACE_OFF", "").lower() == "true" @@ -245,10 +248,12 @@ def report_json(region: Union[None, str], msgs: List[dict]) -> int: duration = int((time.time() - start_time) * 1000) get_logger().info(f"successful reporting, code: {getattr(response, 'code', 'unknown')}") except Exception as e: - get_logger().exception( - f"Could not report json to {host}. Retrying to establish connection.", exc_info=e - ) - edge_connection = establish_connection(host) + if should_retry: + get_logger().exception(f"Could not report to {host}. Retrying.", exc_info=e) + edge_connection = establish_connection(host) + report_json(region, msgs, should_retry=False) + else: + get_logger().exception("Could not report: A span was lost.", exc_info=e) return duration diff --git a/src/test/unit/test_main_utils.py b/src/test/unit/test_main_utils.py index 1bbaa13c..69ba9919 100644 --- a/src/test/unit/test_main_utils.py +++ b/src/test/unit/test_main_utils.py @@ -2,6 +2,8 @@ from collections import OrderedDict from decimal import Decimal import datetime +import http.client +from mock import Mock import pytest from lumigo_tracer.utils import ( @@ -26,6 +28,7 @@ lumigo_dumps, prepare_host, EDGE_PATH, + report_json, ) import json @@ -362,3 +365,18 @@ def test_get_timeout_buffer(remaining_time, conf, expected): ) def test_prepare_host(arg, host): assert prepare_host(arg) == host + + +@pytest.mark.parametrize( + "errors, final_log", [(ValueError, "ERROR"), ([ValueError, Mock()], "INFO")] +) +def test_report_json_retry(monkeypatch, reporter_mock, caplog, errors, final_log): + reporter_mock.side_effect = report_json + monkeypatch.setattr(Configuration, "host", "force_reconnect") + monkeypatch.setattr(Configuration, "should_report", True) + monkeypatch.setattr(http.client, "HTTPSConnection", Mock()) + http.client.HTTPSConnection("force_reconnect").getresponse.side_effect = errors + + report_json(None, [{"a": "b"}]) + + assert caplog.records[-1].levelname == final_log From e471345669d11194986c4a3b5bbb552b8f6d06aa Mon Sep 17 00:00:00 2001 From: saartochner Date: Sun, 4 Oct 2020 19:46:46 +0300 Subject: [PATCH 02/24] split event and parsers --- src/lumigo_tracer/auto_tag/auto_tag_event.py | 4 +- src/lumigo_tracer/event/__init__.py | 0 .../event_parser.py => event/event_dumper.py} | 20 +- src/lumigo_tracer/event/event_trigger.py | 185 ++++++++++++++++++ .../{utils.py => lumigo_utils.py} | 2 +- .../parsers/{parser.py => http_parser.py} | 4 +- .../{parsers/utils.py => parsing_utils.py} | 184 +---------------- src/lumigo_tracer/spans_container.py | 32 ++- src/lumigo_tracer/sync_http/sync_hook.py | 4 +- src/lumigo_tracer/user_utils.py | 2 +- src/test/component/test_component.py | 2 +- src/test/conftest.py | 14 +- src/test/unit/auto_tag/test_auto_tag_event.py | 2 +- .../test_event_dumper.py} | 31 ++- .../test_event_trigger.py} | 146 +------------- .../{test_parser.py => test_http_parser.py} | 2 +- src/test/unit/sync_http/test_sync_hook.py | 8 +- ...est_main_utils.py => test_lumigo_utils.py} | 2 +- src/test/unit/test_parsing_utils.py | 154 +++++++++++++++ src/test/unit/test_spans_container.py | 4 +- src/test/unit/test_user_utils.py | 2 +- 21 files changed, 407 insertions(+), 397 deletions(-) create mode 100644 src/lumigo_tracer/event/__init__.py rename src/lumigo_tracer/{parsers/event_parser.py => event/event_dumper.py} (89%) create mode 100644 src/lumigo_tracer/event/event_trigger.py rename src/lumigo_tracer/{utils.py => lumigo_utils.py} (99%) rename src/lumigo_tracer/parsers/{parser.py => http_parser.py} (98%) rename src/lumigo_tracer/{parsers/utils.py => parsing_utils.py} (51%) rename src/test/unit/{parsers/test_event_parser.py => event/test_event_dumper.py} (95%) rename src/test/unit/{parsers/test_utils.py => event/test_event_trigger.py} (72%) rename src/test/unit/parsers/{test_parser.py => test_http_parser.py} (98%) rename src/test/unit/{test_main_utils.py => test_lumigo_utils.py} (99%) create mode 100644 src/test/unit/test_parsing_utils.py diff --git a/src/lumigo_tracer/auto_tag/auto_tag_event.py b/src/lumigo_tracer/auto_tag/auto_tag_event.py index f1a15ce2..df9cb9fd 100644 --- a/src/lumigo_tracer/auto_tag/auto_tag_event.py +++ b/src/lumigo_tracer/auto_tag/auto_tag_event.py @@ -3,8 +3,8 @@ from typing import Dict, List, Optional from lumigo_tracer.user_utils import add_execution_tag -from lumigo_tracer.parsers.utils import str_to_list -from lumigo_tracer.utils import get_logger, is_api_gw_event +from lumigo_tracer.parsing_utils import str_to_list +from lumigo_tracer.lumigo_utils import get_logger, is_api_gw_event AUTO_TAG_API_GW_HEADERS: Optional[List[str]] = ( str_to_list(os.environ.get("LUMIGO_AUTO_TAG_API_GW_HEADERS", "")) or [] diff --git a/src/lumigo_tracer/event/__init__.py b/src/lumigo_tracer/event/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/lumigo_tracer/parsers/event_parser.py b/src/lumigo_tracer/event/event_dumper.py similarity index 89% rename from src/lumigo_tracer/parsers/event_parser.py rename to src/lumigo_tracer/event/event_dumper.py index 57c3868b..3b6ff6e3 100644 --- a/src/lumigo_tracer/parsers/event_parser.py +++ b/src/lumigo_tracer/event/event_dumper.py @@ -4,8 +4,8 @@ from collections import OrderedDict from typing import Dict, List -from lumigo_tracer.parsers.utils import str_to_list, safe_get -from lumigo_tracer.utils import get_logger, is_api_gw_event +from lumigo_tracer.parsing_utils import str_to_list, safe_get +from lumigo_tracer.lumigo_utils import get_logger, is_api_gw_event, lumigo_dumps API_GW_KEYS_ORDER = str_to_list(os.environ.get("LUMIGO_API_GW_KEYS_ORDER", "")) or [ "version", @@ -53,7 +53,7 @@ def is_supported(event) -> bool: @staticmethod @abstractmethod - def parse(event) -> Dict: + def parse(event) -> OrderedDict: raise NotImplementedError() @@ -65,7 +65,7 @@ def is_supported(event) -> bool: return False @staticmethod - def parse(event) -> Dict: + def parse(event) -> OrderedDict: new_event: OrderedDict = OrderedDict() # Add order keys for order_key in API_GW_KEYS_ORDER: @@ -94,7 +94,7 @@ def is_supported(event) -> bool: return safe_get(event, ["Records", 0, "EventSource"]) == "aws:sns" @staticmethod - def parse(event) -> Dict: + def parse(event) -> OrderedDict: new_sns_event: OrderedDict = OrderedDict() new_sns_event["Records"] = [] # Add order keys @@ -113,7 +113,7 @@ def is_supported(event) -> bool: return safe_get(event, ["Records", 0, "eventSource"]) == "aws:sqs" @staticmethod - def parse(event) -> Dict: + def parse(event) -> OrderedDict: new_sqs_event: OrderedDict = OrderedDict() new_sqs_event["Records"] = [] # Add order keys @@ -126,17 +126,17 @@ def parse(event) -> Dict: return new_sqs_event -class EventParser: +class EventDumper: @staticmethod - def parse_event(event: Dict, handlers: List[EventParseHandler] = None): + def dump_event(event: Dict, handlers: List[EventParseHandler] = None) -> str: handlers = handlers or [ApiGWHandler(), SNSHandler(), SQSHandler()] for handler in handlers: try: if handler.is_supported(event): - return handler.parse(event) + return lumigo_dumps(handler.parse(event)) except Exception as e: get_logger().debug( f"Error while trying to parse with handler {handler.__class__.__name__} event {event}", exc_info=e, ) - return event + return lumigo_dumps(event) diff --git a/src/lumigo_tracer/event/event_trigger.py b/src/lumigo_tracer/event/event_trigger.py new file mode 100644 index 00000000..ceb52bdf --- /dev/null +++ b/src/lumigo_tracer/event/event_trigger.py @@ -0,0 +1,185 @@ +from typing import Union, List, Dict + +from lumigo_tracer.parsing_utils import recursive_get_key, safe_get +from lumigo_tracer.lumigo_utils import ( + lumigo_safe_execute, + Configuration, + STEP_FUNCTION_UID_KEY, + LUMIGO_EVENT_KEY, + md5hash, +) + +TRIGGER_CREATION_TIME_KEY = "approxEventCreationTime" +MESSAGE_ID_KEY = "messageId" +MESSAGE_IDS_KEY = "messageIds" + + +def parse_triggered_by(event: dict): + """ + This function parses the event and build the dictionary that describes the given event. + + The current possible values are: + * {triggeredBy: unknown} + * {triggeredBy: apigw, api: , resource: <>, httpMethod: <>, stage: <>, identity: <>, referer: <>} + """ + with lumigo_safe_execute("triggered by"): + if not isinstance(event, dict): + if _is_step_function(event): + return _parse_step_function(event) + return None + if _is_supported_http_method(event): + return _parse_http_method(event) + elif _is_supported_sns(event): + return _parse_sns(event) + elif _is_supported_streams(event): + return _parse_streams(event) + elif _is_supported_cw(event): + return _parse_cw(event) + elif _is_step_function(event): + return _parse_step_function(event) + + return _parse_unknown(event) + + +def _parse_unknown(event: dict): + result = {"triggeredBy": "unknown"} + return result + + +def _is_step_function(event: Union[List, Dict]): + return ( + Configuration.is_step_function + and isinstance(event, (list, dict)) # noqa + and STEP_FUNCTION_UID_KEY in recursive_get_key(event, LUMIGO_EVENT_KEY, default={}) # noqa + ) + + +def _parse_step_function(event: dict): + result = { + "triggeredBy": "stepFunction", + "messageId": recursive_get_key(event, LUMIGO_EVENT_KEY)[STEP_FUNCTION_UID_KEY], + } + return result + + +def _is_supported_http_method(event: dict): + return ( + "httpMethod" in event # noqa + and "headers" in event # noqa + and "requestContext" in event # noqa + and event.get("requestContext", {}).get("elb") is None # noqa + ) or ( # noqa + event.get("version", "") == "2.0" and "headers" in event # noqa + ) # noqa # noqa + + +def _parse_http_method(event: dict): + version = event.get("version") + if version and version.startswith("2.0"): + return _parse_http_method_v2(event) + return _parse_http_method_v1(event) + + +def _parse_http_method_v1(event: dict): + result = { + "triggeredBy": "apigw", + "httpMethod": event.get("httpMethod", ""), + "resource": event.get("resource", ""), + "messageId": event.get("requestContext", {}).get("requestId", ""), + } + if isinstance(event.get("headers"), dict): + result["api"] = event["headers"].get("Host", "unknown.unknown.unknown") + if isinstance(event.get("requestContext"), dict): + result["stage"] = event["requestContext"].get("stage", "unknown") + return result + + +def _parse_http_method_v2(event: dict): + result = { + "triggeredBy": "apigw", + "httpMethod": event.get("requestContext", {}).get("http", {}).get("method"), + "resource": event.get("requestContext", {}).get("http", {}).get("path"), + "messageId": event.get("requestContext", {}).get("requestId", ""), + "api": event.get("requestContext", {}).get("domainName", ""), + "stage": event.get("requestContext", {}).get("stage", "unknown"), + } + return result + + +def _is_supported_sns(event: dict): + return event.get("Records", [{}])[0].get("EventSource") == "aws:sns" + + +def _parse_sns(event: dict): + return { + "triggeredBy": "sns", + "arn": event["Records"][0]["Sns"]["TopicArn"], + "messageId": event["Records"][0]["Sns"].get("MessageId"), + } + + +def _is_supported_cw(event: dict): + return event.get("detail-type") == "Scheduled Event" and "source" in event and "time" in event + + +def _parse_cw(event: dict): + resource = event.get("resources", ["/unknown"])[0].split("/")[1] + return { + "triggeredBy": "cloudwatch", + "resource": resource, + "region": event.get("region"), + "detailType": event.get("detail-type"), + } + + +def _is_supported_streams(event: dict): + return event.get("Records", [{}])[0].get("eventSource") in [ + "aws:kinesis", + "aws:dynamodb", + "aws:sqs", + "aws:s3", + ] + + +def _parse_streams(event: dict) -> Dict[str, str]: + """ + :return: {"triggeredBy": str, "arn": str} + If has messageId, return also: {"messageId": str} + """ + triggered_by = event["Records"][0]["eventSource"].split(":")[1] + result = {"triggeredBy": triggered_by} + if triggered_by == "s3": + result["arn"] = event["Records"][0]["s3"]["bucket"]["arn"] + result["messageId"] = ( + event["Records"][0].get("responseElements", {}).get("x-amz-request-id") + ) + else: + result["arn"] = event["Records"][0]["eventSourceARN"] + if triggered_by == "sqs": + result.update(_parse_sqs_event(event)) + elif triggered_by == "kinesis": + result["messageId"] = safe_get(event, ["Records", 0, "kinesis", "sequenceNumber"]) + elif triggered_by == "dynamodb": + result.update(_parse_dynamomdb_event(event)) + return result + + +def _get_ddb_approx_creation_time_ms(event) -> int: + return event["Records"][0].get("dynamodb", {}).get("ApproximateCreationDateTime", 0) * 1000 + + +def _parse_dynamomdb_event(event) -> Dict[str, Union[int, List[str]]]: + creation_time = _get_ddb_approx_creation_time_ms(event) + mids = [] + for record in event["Records"]: + event_name = record.get("eventName") + if event_name in ("MODIFY", "REMOVE") and record.get("dynamodb", {}).get("Keys"): + mids.append(md5hash(record["dynamodb"]["Keys"])) + elif event_name == "INSERT" and record.get("dynamodb", {}).get("NewImage"): + mids.append(md5hash(record["dynamodb"]["NewImage"])) + return {MESSAGE_IDS_KEY: mids, TRIGGER_CREATION_TIME_KEY: creation_time} + + +def _parse_sqs_event(event) -> Dict[str, Union[int, List[str]]]: + mids = [record["messageId"] for record in event["Records"] if record.get("messageId")] + return {MESSAGE_IDS_KEY: mids} if len(mids) > 1 else {MESSAGE_ID_KEY: mids[0]} diff --git a/src/lumigo_tracer/utils.py b/src/lumigo_tracer/lumigo_utils.py similarity index 99% rename from src/lumigo_tracer/utils.py rename to src/lumigo_tracer/lumigo_utils.py index eb10daeb..04af01fc 100644 --- a/src/lumigo_tracer/utils.py +++ b/src/lumigo_tracer/lumigo_utils.py @@ -467,7 +467,7 @@ def lumigo_dumps( regexes: Optional[Pattern[str]] = None, enforce_jsonify: bool = False, decimal_safe=False, -): +) -> str: regexes = regexes or get_omitting_regex() max_size = max_size if max_size is not None else Configuration.max_entry_size is_truncated = False diff --git a/src/lumigo_tracer/parsers/parser.py b/src/lumigo_tracer/parsers/http_parser.py similarity index 98% rename from src/lumigo_tracer/parsers/parser.py rename to src/lumigo_tracer/parsers/http_parser.py index 4ed565fd..85f13bd8 100644 --- a/src/lumigo_tracer/parsers/parser.py +++ b/src/lumigo_tracer/parsers/http_parser.py @@ -3,7 +3,7 @@ from typing import Type, Optional import time -from lumigo_tracer.parsers.utils import ( +from lumigo_tracer.parsing_utils import ( safe_split_get, safe_key_from_json, safe_key_from_xml, @@ -12,7 +12,7 @@ safe_get, should_scrub_domain, ) -from lumigo_tracer.utils import Configuration, lumigo_dumps, md5hash, get_logger +from lumigo_tracer.lumigo_utils import Configuration, lumigo_dumps, md5hash, get_logger from lumigo_tracer.parsers.http_data_classes import HttpRequest HTTP_TYPE = "http" diff --git a/src/lumigo_tracer/parsers/utils.py b/src/lumigo_tracer/parsing_utils.py similarity index 51% rename from src/lumigo_tracer/parsers/utils.py rename to src/lumigo_tracer/parsing_utils.py index 74ef1994..c41368ec 100644 --- a/src/lumigo_tracer/parsers/utils.py +++ b/src/lumigo_tracer/parsing_utils.py @@ -8,18 +8,7 @@ import itertools from collections.abc import Iterable -from lumigo_tracer.utils import ( - Configuration, - LUMIGO_EVENT_KEY, - STEP_FUNCTION_UID_KEY, - lumigo_safe_execute, - get_logger, - md5hash, -) - -MESSAGE_ID_KEY = "messageId" -MESSAGE_IDS_KEY = "messageIds" -TRIGGER_CREATION_TIME_KEY = "approxEventCreationTime" +from lumigo_tracer.lumigo_utils import Configuration, get_logger def safe_get(d: Union[dict, list], keys: List[Union[str, int]], default: Any = None) -> Any: @@ -153,177 +142,6 @@ def recursive_json_join(d1: Optional[dict], d2: Optional[dict]): return d -def parse_triggered_by(event: dict): - """ - This function parses the event and build the dictionary that describes the given event. - - The current possible values are: - * {triggeredBy: unknown} - * {triggeredBy: apigw, api: , resource: <>, httpMethod: <>, stage: <>, identity: <>, referer: <>} - """ - with lumigo_safe_execute("triggered by"): - if not isinstance(event, dict): - if _is_step_function(event): - return _parse_step_function(event) - return None - if _is_supported_http_method(event): - return parse_http_method(event) - elif _is_supported_sns(event): - return _parse_sns(event) - elif _is_supported_streams(event): - return _parse_streams(event) - elif _is_supported_cw(event): - return _parse_cw(event) - elif _is_step_function(event): - return _parse_step_function(event) - - return _parse_unknown(event) - - -def _parse_unknown(event: dict): - result = {"triggeredBy": "unknown"} - return result - - -def _is_step_function(event: Union[List, Dict]): - return ( - Configuration.is_step_function - and isinstance(event, (list, dict)) # noqa - and STEP_FUNCTION_UID_KEY in recursive_get_key(event, LUMIGO_EVENT_KEY, default={}) # noqa - ) - - -def _parse_step_function(event: dict): - result = { - "triggeredBy": "stepFunction", - "messageId": recursive_get_key(event, LUMIGO_EVENT_KEY)[STEP_FUNCTION_UID_KEY], - } - return result - - -def _is_supported_http_method(event: dict): - return ( - "httpMethod" in event # noqa - and "headers" in event # noqa - and "requestContext" in event # noqa - and event.get("requestContext", {}).get("elb") is None # noqa - ) or ( # noqa - event.get("version", "") == "2.0" and "headers" in event # noqa - ) # noqa # noqa - - -def parse_http_method(event: dict): - version = event.get("version") - if version and version.startswith("2.0"): - return _parse_http_method_v2(event) - return _parse_http_method_v1(event) - - -def _parse_http_method_v1(event: dict): - result = { - "triggeredBy": "apigw", - "httpMethod": event.get("httpMethod", ""), - "resource": event.get("resource", ""), - "messageId": event.get("requestContext", {}).get("requestId", ""), - } - if isinstance(event.get("headers"), dict): - result["api"] = event["headers"].get("Host", "unknown.unknown.unknown") - if isinstance(event.get("requestContext"), dict): - result["stage"] = event["requestContext"].get("stage", "unknown") - return result - - -def _parse_http_method_v2(event: dict): - result = { - "triggeredBy": "apigw", - "httpMethod": event.get("requestContext", {}).get("http", {}).get("method"), - "resource": event.get("requestContext", {}).get("http", {}).get("path"), - "messageId": event.get("requestContext", {}).get("requestId", ""), - "api": event.get("requestContext", {}).get("domainName", ""), - "stage": event.get("requestContext", {}).get("stage", "unknown"), - } - return result - - -def _is_supported_sns(event: dict): - return event.get("Records", [{}])[0].get("EventSource") == "aws:sns" - - -def _parse_sns(event: dict): - return { - "triggeredBy": "sns", - "arn": event["Records"][0]["Sns"]["TopicArn"], - "messageId": event["Records"][0]["Sns"].get("MessageId"), - } - - -def _is_supported_cw(event: dict): - return event.get("detail-type") == "Scheduled Event" and "source" in event and "time" in event - - -def _parse_cw(event: dict): - resource = event.get("resources", ["/unknown"])[0].split("/")[1] - return { - "triggeredBy": "cloudwatch", - "resource": resource, - "region": event.get("region"), - "detailType": event.get("detail-type"), - } - - -def _is_supported_streams(event: dict): - return event.get("Records", [{}])[0].get("eventSource") in [ - "aws:kinesis", - "aws:dynamodb", - "aws:sqs", - "aws:s3", - ] - - -def _parse_streams(event: dict) -> Dict[str, str]: - """ - :return: {"triggeredBy": str, "arn": str} - If has messageId, return also: {"messageId": str} - """ - triggered_by = event["Records"][0]["eventSource"].split(":")[1] - result = {"triggeredBy": triggered_by} - if triggered_by == "s3": - result["arn"] = event["Records"][0]["s3"]["bucket"]["arn"] - result["messageId"] = ( - event["Records"][0].get("responseElements", {}).get("x-amz-request-id") - ) - else: - result["arn"] = event["Records"][0]["eventSourceARN"] - if triggered_by == "sqs": - result.update(_parse_sqs_event(event)) - elif triggered_by == "kinesis": - result["messageId"] = safe_get(event, ["Records", 0, "kinesis", "sequenceNumber"]) - elif triggered_by == "dynamodb": - result.update(_parse_dynamomdb_event(event)) - return result - - -def _get_ddb_approx_creation_time_ms(event) -> int: - return event["Records"][0].get("dynamodb", {}).get("ApproximateCreationDateTime", 0) * 1000 - - -def _parse_dynamomdb_event(event) -> Dict[str, Union[int, List[str]]]: - creation_time = _get_ddb_approx_creation_time_ms(event) - mids = [] - for record in event["Records"]: - event_name = record.get("eventName") - if event_name in ("MODIFY", "REMOVE") and record.get("dynamodb", {}).get("Keys"): - mids.append(md5hash(record["dynamodb"]["Keys"])) - elif event_name == "INSERT" and record.get("dynamodb", {}).get("NewImage"): - mids.append(md5hash(record["dynamodb"]["NewImage"])) - return {MESSAGE_IDS_KEY: mids, TRIGGER_CREATION_TIME_KEY: creation_time} - - -def _parse_sqs_event(event) -> Dict[str, Union[int, List[str]]]: - mids = [record["messageId"] for record in event["Records"] if record.get("messageId")] - return {MESSAGE_IDS_KEY: mids} if len(mids) > 1 else {MESSAGE_ID_KEY: mids[0]} - - def should_scrub_domain(url: str) -> bool: if url and Configuration.domains_scrubber: for regex in Configuration.domains_scrubber: diff --git a/src/lumigo_tracer/spans_container.py b/src/lumigo_tracer/spans_container.py index e3e824a7..96762b76 100644 --- a/src/lumigo_tracer/spans_container.py +++ b/src/lumigo_tracer/spans_container.py @@ -9,8 +9,8 @@ import traceback from typing import List, Dict, Optional, Callable, Set -from lumigo_tracer.parsers.event_parser import EventParser -from lumigo_tracer.utils import ( +from lumigo_tracer.event.event_dumper import EventDumper +from lumigo_tracer.lumigo_utils import ( Configuration, LUMIGO_EVENT_KEY, STEP_FUNCTION_UID_KEY, @@ -18,17 +18,14 @@ lumigo_dumps, EXECUTION_TAGS_KEY, get_timeout_buffer, + get_logger, + _is_span_has_error, ) -from lumigo_tracer import utils -from lumigo_tracer.parsers.parser import get_parser, HTTP_TYPE, StepFunctionParser -from lumigo_tracer.parsers.utils import ( - parse_trace_id, - safe_split_get, - recursive_json_join, - parse_triggered_by, -) -from lumigo_tracer.utils import get_logger, _is_span_has_error -from .parsers.http_data_classes import HttpRequest +from lumigo_tracer import lumigo_utils +from lumigo_tracer.parsers.http_parser import get_parser, HTTP_TYPE, StepFunctionParser +from lumigo_tracer.parsing_utils import parse_trace_id, safe_split_get, recursive_json_join +from lumigo_tracer.event.event_trigger import parse_triggered_by +from lumigo_tracer.parsers.http_data_classes import HttpRequest _VERSION_PATH = os.path.join(os.path.dirname(__file__), "VERSION") MAX_LAMBDA_TIME = 15 * 60 * 1000 @@ -111,7 +108,7 @@ def _generate_start_span(self) -> dict: def start(self, event=None, context=None): to_send = self._generate_start_span() if not Configuration.send_only_if_error: - report_duration = utils.report_json(region=self.region, msgs=[to_send]) + report_duration = lumigo_utils.report_json(region=self.region, msgs=[to_send]) self.function_span["reporter_rtt"] = report_duration self.http_spans = [] else: @@ -123,7 +120,7 @@ def handle_timeout(self, *args): to_send = [s for s in self.http_spans if s["id"] in self.http_span_ids_to_send] if Configuration.send_only_if_error: to_send.append(self._generate_start_span()) - utils.report_json(region=self.region, msgs=to_send) + lumigo_utils.report_json(region=self.region, msgs=to_send) self.http_span_ids_to_send.clear() def start_timeout_timer(self, context=None) -> None: @@ -279,7 +276,7 @@ def end(self, ret_val=None) -> Optional[int]: to_send = [self.function_span] + [ s for s in self.http_spans if s["id"] in self.http_span_ids_to_send ] - reported_rtt = utils.report_json(region=self.region, msgs=to_send) + reported_rtt = lumigo_utils.report_json(region=self.region, msgs=to_send) else: get_logger().debug( "No Spans were sent, `Configuration.send_only_if_error` is on and no span has error" @@ -311,10 +308,7 @@ def create_span(cls, event=None, context=None, force=False) -> "SpansContainer": additional_info = {} if Configuration.verbose: additional_info.update( - { - "event": lumigo_dumps(EventParser.parse_event(event)), - "envs": lumigo_dumps(dict(os.environ)), - } + {"event": EventDumper.dump_event(event), "envs": lumigo_dumps(dict(os.environ))} ) trace_root, transaction_id, suffix = parse_trace_id(os.environ.get("_X_AMZN_TRACE_ID", "")) diff --git a/src/lumigo_tracer/sync_http/sync_hook.py b/src/lumigo_tracer/sync_http/sync_hook.py index 6939af89..a08003f2 100644 --- a/src/lumigo_tracer/sync_http/sync_hook.py +++ b/src/lumigo_tracer/sync_http/sync_hook.py @@ -11,8 +11,8 @@ from lumigo_tracer.auto_tag.auto_tag_event import AutoTagEvent from lumigo_tracer.libs.wrapt import wrap_function_wrapper -from lumigo_tracer.parsers.utils import safe_get_list -from lumigo_tracer.utils import ( +from lumigo_tracer.parsing_utils import safe_get_list +from lumigo_tracer.lumigo_utils import ( config, Configuration, get_logger, diff --git a/src/lumigo_tracer/user_utils.py b/src/lumigo_tracer/user_utils.py index c4d2550a..93e8f090 100644 --- a/src/lumigo_tracer/user_utils.py +++ b/src/lumigo_tracer/user_utils.py @@ -1,5 +1,5 @@ from lumigo_tracer.spans_container import SpansContainer -from lumigo_tracer.utils import Configuration, warn_client +from lumigo_tracer.lumigo_utils import Configuration, warn_client LUMIGO_REPORT_ERROR_STRING = "[LUMIGO_LOG]" MAX_TAGS = 50 diff --git a/src/test/component/test_component.py b/src/test/component/test_component.py index 12ab1d11..b447fe2d 100644 --- a/src/test/component/test_component.py +++ b/src/test/component/test_component.py @@ -7,7 +7,7 @@ from lumigo_tracer.sync_http.sync_hook import lumigo_tracer from lumigo_tracer.spans_container import SpansContainer -from lumigo_tracer.utils import md5hash +from lumigo_tracer.lumigo_utils import md5hash DEFAULT_USER = "cicd" diff --git a/src/test/conftest.py b/src/test/conftest.py index 383c63d1..804fb502 100644 --- a/src/test/conftest.py +++ b/src/test/conftest.py @@ -1,20 +1,20 @@ import builtins import logging -from lumigo_tracer import utils +from lumigo_tracer import lumigo_utils from lumigo_tracer.spans_container import SpansContainer import mock import pytest -from lumigo_tracer.utils import Configuration, get_omitting_regex +from lumigo_tracer.lumigo_utils import Configuration, get_omitting_regex @pytest.fixture(autouse=True) def reporter_mock(monkeypatch): - utils.Configuration.should_report = False - reporter_mock = mock.Mock(utils.report_json) + lumigo_utils.Configuration.should_report = False + reporter_mock = mock.Mock(lumigo_utils.report_json) reporter_mock.return_value = 123 - monkeypatch.setattr(utils, "report_json", reporter_mock) + monkeypatch.setattr(lumigo_utils, "report_json", reporter_mock) return reporter_mock @@ -52,8 +52,8 @@ def verbose_logger(): """ This fixture make sure that we will see all the log in the tests. """ - utils.get_logger().setLevel(logging.DEBUG) - utils.config(should_report=False, verbose=True) + lumigo_utils.get_logger().setLevel(logging.DEBUG) + lumigo_utils.config(should_report=False, verbose=True) def pytest_addoption(parser): diff --git a/src/test/unit/auto_tag/test_auto_tag_event.py b/src/test/unit/auto_tag/test_auto_tag_event.py index e9583666..9ad440d0 100644 --- a/src/test/unit/auto_tag/test_auto_tag_event.py +++ b/src/test/unit/auto_tag/test_auto_tag_event.py @@ -1,7 +1,7 @@ from lumigo_tracer.auto_tag import auto_tag_event from lumigo_tracer.auto_tag.auto_tag_event import EventAutoTagHandler, AutoTagEvent from lumigo_tracer.spans_container import SpansContainer -from lumigo_tracer.utils import EXECUTION_TAGS_KEY +from lumigo_tracer.lumigo_utils import EXECUTION_TAGS_KEY class ExceptionHandler(EventAutoTagHandler): diff --git a/src/test/unit/parsers/test_event_parser.py b/src/test/unit/event/test_event_dumper.py similarity index 95% rename from src/test/unit/parsers/test_event_parser.py rename to src/test/unit/event/test_event_dumper.py index a59e86b1..601cb1c6 100644 --- a/src/test/unit/parsers/test_event_parser.py +++ b/src/test/unit/event/test_event_dumper.py @@ -1,8 +1,7 @@ -import json from collections import OrderedDict -from typing import Dict -from lumigo_tracer.parsers.event_parser import EventParser, EventParseHandler +from lumigo_tracer.event.event_dumper import EventDumper, EventParseHandler +from lumigo_tracer.lumigo_utils import lumigo_dumps class ExceptionHandler(EventParseHandler): @@ -11,22 +10,22 @@ def is_supported(event) -> bool: raise Exception() @staticmethod - def parse(event) -> Dict: + def parse(event) -> OrderedDict: raise Exception() def test_parse_event_not_api_gw_none_check(): - new_event = EventParser.parse_event(event=None) + new_event = EventDumper.dump_event(event=None) - assert new_event is None + assert new_event == lumigo_dumps(None) def test_parse_event_not_api_gw(): event = {"a": 1} - new_event = EventParser.parse_event(event=event, handlers=[ExceptionHandler()]) + new_event = EventDumper.dump_event(event=event, handlers=[ExceptionHandler()]) - assert new_event == event + assert new_event == lumigo_dumps(event) def test_parse_event_api_gw_v1(): @@ -145,9 +144,9 @@ def test_parse_event_api_gw_v1(): "isBase64Encoded": False, } - order_api_gw_event = EventParser.parse_event(event=not_order_api_gw_event) + order_api_gw_event = EventDumper.dump_event(event=not_order_api_gw_event) - assert json.dumps(order_api_gw_event) == json.dumps( + assert order_api_gw_event == lumigo_dumps( OrderedDict( { "resource": "/add-user", @@ -236,9 +235,9 @@ def test_parse_event_api_gw_v2(): "isBase64Encoded": True, } - order_api_gw_event = EventParser.parse_event(event=not_order_api_gw_event) + order_api_gw_event = EventDumper.dump_event(event=not_order_api_gw_event) - assert json.dumps(order_api_gw_event) == json.dumps( + assert order_api_gw_event == lumigo_dumps( OrderedDict( { "version": "2.0", @@ -314,9 +313,9 @@ def test_parse_event_sns(): ] } - order_sns_event = EventParser.parse_event(event=not_order_sns_event) + order_sns_event = EventDumper.dump_event(event=not_order_sns_event) - assert json.dumps(order_sns_event) == json.dumps( + assert order_sns_event == lumigo_dumps( OrderedDict( { "Records": [ @@ -384,9 +383,9 @@ def test_parse_event_sqs(): ] } - order_sns_event = EventParser.parse_event(event=not_order_sns_event) + order_sns_event = EventDumper.dump_event(event=not_order_sns_event) - assert json.dumps(order_sns_event) == json.dumps( + assert order_sns_event == lumigo_dumps( OrderedDict( { "Records": [ diff --git a/src/test/unit/parsers/test_utils.py b/src/test/unit/event/test_event_trigger.py similarity index 72% rename from src/test/unit/parsers/test_utils.py rename to src/test/unit/event/test_event_trigger.py index cbbcfd11..a2d019e0 100644 --- a/src/test/unit/parsers/test_utils.py +++ b/src/test/unit/event/test_event_trigger.py @@ -1,89 +1,7 @@ -import os - import pytest -from lumigo_tracer.parsers import utils -from lumigo_tracer.parsers.utils import should_scrub_domain, str_to_list, str_to_tuple -from lumigo_tracer.utils import config, Configuration - - -@pytest.mark.parametrize( - ("input_params", "expected_output"), - [ - (("a.b.c", ".", 0), "a"), # happy flow - (("a.b.c", ".", 1), "b"), - (("a.b.c", ".", 5, "d"), "d"), # return the default - ], -) -def test_safe_split_get(input_params, expected_output): - assert utils.safe_split_get(*input_params) == expected_output - - -@pytest.mark.parametrize( - ("input_params", "expected_output"), - [ - ((b'{"a": "b"}', "a"), "b"), # happy flow - ((b'{"a": "b"}', "c"), None), # return the default - ((b"b", "c"), None), # not a json - ], -) -def test_key_from_json(input_params, expected_output): - assert utils.safe_key_from_json(*input_params) == expected_output - - -@pytest.mark.parametrize( - ("input_params", "expected_output"), - [ - ((b"b", "a"), "b"), # happy flow - one parameter - ((b"c", "a/b"), "c"), # happy flow - longer path - ((b"b", "c"), None), # not existing key - ((b"c", "a/e"), None), # not existing sub-key - ((b'{"a": "b"}', "c"), None), # not an xml - ], -) -def test_key_from_xml(input_params, expected_output): - assert utils.safe_key_from_xml(*input_params) == expected_output - - -@pytest.mark.parametrize( - ("input_params", "expected_output"), - [ - ((b"a=b", "a"), "b"), # happy flow - one parameter - ((b"a=b&c=d", "c"), "d"), # happy flow - multiple parameters - ((b"a=b&c=d", "e"), None), # not existing key - ((b'{"a": "b"}', "c"), None), # not an query, no '&' - ((b"a&b", "a"), None), # not an query, with '&' - ], -) -def test_key_from_query(input_params, expected_output): - assert utils.safe_key_from_query(*input_params) == expected_output - - -@pytest.mark.parametrize( - ("trace_id", "result"), - [ - ("Root=1-2-3;Parent=34;Sampled=0", ("1-2-3", "3", ";Parent=34;Sampled=0")), # happy flow - ("Root=1-2-3;", ("1-2-3", "3", ";")), - ("Root=1-2;", ("1-2", "", ";")), - ("a;1", ("", "", ";1")), - ("123", ("", "", "123")), - ], -) -def test_parse_trace_id(trace_id, result): - assert utils.parse_trace_id(trace_id) == result - - -@pytest.mark.parametrize( - ("d1", "d2", "result"), - [ - ({1: 2}, {3: 4}, {1: 2, 3: 4}), # happy flow - ({1: 2}, {1: 3}, {1: 2}), # same key twice - ({1: {2: 3}}, {4: 5}, {1: {2: 3}, 4: 5}), # dictionary in d1 and nothing in d2 - ({1: {2: 3}}, {1: {4: 5}}, {1: {2: 3, 4: 5}}), # merge two inner dictionaries - ], -) -def test_recursive_json_join(d1, d2, result): - assert utils.recursive_json_join(d1, d2) == result +from lumigo_tracer.event.event_trigger import parse_triggered_by +from lumigo_tracer.lumigo_utils import Configuration @pytest.mark.parametrize( @@ -398,62 +316,4 @@ def test_recursive_json_join(d1, d2, result): ) def test_parse_triggered_by(event, output): Configuration.is_step_function = True - assert utils.parse_triggered_by(event) == output - - -def test_config_with_verbose_param_with_no_env_verbose_verbose_is_false(): - config(verbose=False) - - assert Configuration.verbose is False - - -def test_config_no_verbose_param_and_no_env_verbose_is_true(): - config() - - assert Configuration.verbose - - -def test_config_no_verbose_param_and_with_env_verbose_equals_to_false_verbose_is_false(monkeypatch): - monkeypatch.setattr(os, "environ", {"LUMIGO_VERBOSE": "FALSE"}) - config() - - assert Configuration.verbose is False - - -@pytest.mark.parametrize( - ("d", "keys", "result_value", "default"), - [ - ({"k": ["a", "b"]}, ["k", 1], "b", None), # Happy flow. - ({"k": ["a"]}, ["k", 1], "default", "default"), # List index out of range. - ({"k": "a"}, ["b"], "default", "default"), # Key doesn't exist. - ({"k": "a"}, [1], "default", "default"), # Wrong key type. - ({"k": "a"}, ["k", 0, 1], "default", "default"), # Wrong keys length. - ], -) -def test_safe_get(d, keys, result_value, default): - assert utils.safe_get(d, keys, default) == result_value - - -@pytest.mark.parametrize( - ("regexes", "url", "expected"), - [(["secret.*"], "lumigo.io", False), (["not-relevant", "secret.*"], "secret.aws.com", True)], -) -def test_should_scrub_domain(regexes, url, expected): - config(domains_scrubber=regexes) - assert should_scrub_domain(url) == expected - - -def test_str_to_list(): - assert str_to_list("a,b,c,d") == ["a", "b", "c", "d"] - - -def test_str_to_list_exception(): - assert str_to_list("") is None - - -def test_str_to_tuple(): - assert str_to_tuple("a,b,c,d") == ("a", "b", "c", "d") - - -def test_str_to_tuple_exception(): - assert str_to_tuple([]) is None + assert parse_triggered_by(event) == output diff --git a/src/test/unit/parsers/test_parser.py b/src/test/unit/parsers/test_http_parser.py similarity index 98% rename from src/test/unit/parsers/test_parser.py rename to src/test/unit/parsers/test_http_parser.py index a43ef7f6..cf6c4d70 100644 --- a/src/test/unit/parsers/test_parser.py +++ b/src/test/unit/parsers/test_http_parser.py @@ -3,7 +3,7 @@ import pytest from lumigo_tracer.parsers.http_data_classes import HttpRequest -from lumigo_tracer.parsers.parser import ( +from lumigo_tracer.parsers.http_parser import ( ServerlessAWSParser, Parser, get_parser, diff --git a/src/test/unit/sync_http/test_sync_hook.py b/src/test/unit/sync_http/test_sync_hook.py index e6c42fbb..12b99dec 100644 --- a/src/test/unit/sync_http/test_sync_hook.py +++ b/src/test/unit/sync_http/test_sync_hook.py @@ -19,10 +19,10 @@ import requests from capturer import CaptureOutput -from lumigo_tracer import lumigo_tracer, LumigoChalice, utils, add_execution_tag +from lumigo_tracer import lumigo_tracer, LumigoChalice, lumigo_utils, add_execution_tag from lumigo_tracer.auto_tag import auto_tag_event -from lumigo_tracer.parsers.parser import Parser -from lumigo_tracer.utils import ( +from lumigo_tracer.parsers.http_parser import Parser +from lumigo_tracer.lumigo_utils import ( Configuration, STEP_FUNCTION_UID_KEY, LUMIGO_EVENT_KEY, @@ -260,7 +260,7 @@ def lambda_test_function(): pass lambda_test_function() - assert utils.Configuration.should_report is False + assert lumigo_utils.Configuration.should_report is False def test_domains_scrubber_happy_flow(monkeypatch): diff --git a/src/test/unit/test_main_utils.py b/src/test/unit/test_lumigo_utils.py similarity index 99% rename from src/test/unit/test_main_utils.py rename to src/test/unit/test_lumigo_utils.py index 69ba9919..5f7233f4 100644 --- a/src/test/unit/test_main_utils.py +++ b/src/test/unit/test_lumigo_utils.py @@ -6,7 +6,7 @@ from mock import Mock import pytest -from lumigo_tracer.utils import ( +from lumigo_tracer.lumigo_utils import ( _create_request_body, _is_span_has_error, _get_event_base64_size, diff --git a/src/test/unit/test_parsing_utils.py b/src/test/unit/test_parsing_utils.py new file mode 100644 index 00000000..00e8b894 --- /dev/null +++ b/src/test/unit/test_parsing_utils.py @@ -0,0 +1,154 @@ +import os + +import pytest + +from lumigo_tracer.parsing_utils import ( + should_scrub_domain, + str_to_list, + str_to_tuple, + safe_split_get, + safe_key_from_json, + safe_get, + recursive_json_join, + parse_trace_id, + safe_key_from_query, + safe_key_from_xml, +) +from lumigo_tracer.lumigo_utils import config, Configuration + + +@pytest.mark.parametrize( + ("input_params", "expected_output"), + [ + (("a.b.c", ".", 0), "a"), # happy flow + (("a.b.c", ".", 1), "b"), + (("a.b.c", ".", 5, "d"), "d"), # return the default + ], +) +def test_safe_split_get(input_params, expected_output): + assert safe_split_get(*input_params) == expected_output + + +@pytest.mark.parametrize( + ("input_params", "expected_output"), + [ + ((b'{"a": "b"}', "a"), "b"), # happy flow + ((b'{"a": "b"}', "c"), None), # return the default + ((b"b", "c"), None), # not a json + ], +) +def test_key_from_json(input_params, expected_output): + assert safe_key_from_json(*input_params) == expected_output + + +@pytest.mark.parametrize( + ("input_params", "expected_output"), + [ + ((b"b", "a"), "b"), # happy flow - one parameter + ((b"c", "a/b"), "c"), # happy flow - longer path + ((b"b", "c"), None), # not existing key + ((b"c", "a/e"), None), # not existing sub-key + ((b'{"a": "b"}', "c"), None), # not an xml + ], +) +def test_key_from_xml(input_params, expected_output): + assert safe_key_from_xml(*input_params) == expected_output + + +@pytest.mark.parametrize( + ("input_params", "expected_output"), + [ + ((b"a=b", "a"), "b"), # happy flow - one parameter + ((b"a=b&c=d", "c"), "d"), # happy flow - multiple parameters + ((b"a=b&c=d", "e"), None), # not existing key + ((b'{"a": "b"}', "c"), None), # not an query, no '&' + ((b"a&b", "a"), None), # not an query, with '&' + ], +) +def test_key_from_query(input_params, expected_output): + assert safe_key_from_query(*input_params) == expected_output + + +@pytest.mark.parametrize( + ("trace_id", "result"), + [ + ("Root=1-2-3;Parent=34;Sampled=0", ("1-2-3", "3", ";Parent=34;Sampled=0")), # happy flow + ("Root=1-2-3;", ("1-2-3", "3", ";")), + ("Root=1-2;", ("1-2", "", ";")), + ("a;1", ("", "", ";1")), + ("123", ("", "", "123")), + ], +) +def test_parse_trace_id(trace_id, result): + assert parse_trace_id(trace_id) == result + + +@pytest.mark.parametrize( + ("d1", "d2", "result"), + [ + ({1: 2}, {3: 4}, {1: 2, 3: 4}), # happy flow + ({1: 2}, {1: 3}, {1: 2}), # same key twice + ({1: {2: 3}}, {4: 5}, {1: {2: 3}, 4: 5}), # dictionary in d1 and nothing in d2 + ({1: {2: 3}}, {1: {4: 5}}, {1: {2: 3, 4: 5}}), # merge two inner dictionaries + ], +) +def test_recursive_json_join(d1, d2, result): + assert recursive_json_join(d1, d2) == result + + +def test_config_with_verbose_param_with_no_env_verbose_verbose_is_false(): + config(verbose=False) + + assert Configuration.verbose is False + + +def test_config_no_verbose_param_and_no_env_verbose_is_true(): + config() + + assert Configuration.verbose + + +def test_config_no_verbose_param_and_with_env_verbose_equals_to_false_verbose_is_false(monkeypatch): + monkeypatch.setattr(os, "environ", {"LUMIGO_VERBOSE": "FALSE"}) + config() + + assert Configuration.verbose is False + + +@pytest.mark.parametrize( + ("d", "keys", "result_value", "default"), + [ + ({"k": ["a", "b"]}, ["k", 1], "b", None), # Happy flow. + ({"k": ["a"]}, ["k", 1], "default", "default"), # List index out of range. + ({"k": "a"}, ["b"], "default", "default"), # Key doesn't exist. + ({"k": "a"}, [1], "default", "default"), # Wrong key type. + ({"k": "a"}, ["k", 0, 1], "default", "default"), # Wrong keys length. + ], +) +def test_safe_get(d, keys, result_value, default): + assert safe_get(d, keys, default) == result_value + + +@pytest.mark.parametrize( + ("regexes", "url", "expected"), + [(["secret.*"], "lumigo.io", False), (["not-relevant", "secret.*"], "secret.aws.com", True)], +) +def test_should_scrub_domain(regexes, url, expected): + config(domains_scrubber=regexes) + assert should_scrub_domain(url) == expected + + +def test_str_to_list(): + assert str_to_list("a,b,c,d") == ["a", "b", "c", "d"] + + +def test_str_to_list_exception(): + assert str_to_list("") is None + + +def test_str_to_tuple(): + assert str_to_tuple("a,b,c,d") == ("a", "b", "c", "d") + + +def test_str_to_tuple_exception(): + assert str_to_tuple([]) is None diff --git a/src/test/unit/test_spans_container.py b/src/test/unit/test_spans_container.py index 78eee4c3..fb87c3b7 100644 --- a/src/test/unit/test_spans_container.py +++ b/src/test/unit/test_spans_container.py @@ -3,9 +3,9 @@ import pytest from lumigo_tracer.parsers.http_data_classes import HttpRequest -from lumigo_tracer.parsers.parser import HTTP_TYPE +from lumigo_tracer.parsers.http_parser import HTTP_TYPE from lumigo_tracer.spans_container import SpansContainer, TimeoutMechanism, FUNCTION_TYPE -from lumigo_tracer.utils import Configuration, EXECUTION_TAGS_KEY, DEFAULT_MAX_ENTRY_SIZE +from lumigo_tracer.lumigo_utils import Configuration, EXECUTION_TAGS_KEY, DEFAULT_MAX_ENTRY_SIZE @pytest.fixture() diff --git a/src/test/unit/test_user_utils.py b/src/test/unit/test_user_utils.py index fe410790..03d86534 100644 --- a/src/test/unit/test_user_utils.py +++ b/src/test/unit/test_user_utils.py @@ -7,7 +7,7 @@ MAX_TAG_VALUE_LEN, MAX_TAGS, ) -from lumigo_tracer.utils import Configuration, EXECUTION_TAGS_KEY +from lumigo_tracer.lumigo_utils import Configuration, EXECUTION_TAGS_KEY def test_report_error_with_enhance_print(capsys): From 8383ea1400ab0d0e566e2e6159055ec999656f70 Mon Sep 17 00:00:00 2001 From: saartochner Date: Mon, 5 Oct 2020 11:35:54 +0300 Subject: [PATCH 03/24] split sync_hook to tracer and http hooks --- src/lumigo_tracer/__init__.py | 4 +- .../handler.py => auto_instrument_handler.py} | 0 src/lumigo_tracer/lumigo_utils.py | 28 +- src/lumigo_tracer/sync_http/sync_hook.py | 170 +------ src/lumigo_tracer/tracer.py | 163 +++++++ src/test/component/test_component.py | 2 +- src/test/conftest.py | 6 +- src/test/unit/sync_http/test_sync_hook.py | 455 ++---------------- ...ler.py => test_auto_instrument_handler.py} | 2 +- src/test/unit/test_lumigo_utils.py | 8 + src/test/unit/test_tracer.py | 366 ++++++++++++++ 11 files changed, 618 insertions(+), 586 deletions(-) rename src/lumigo_tracer/{sync_http/handler.py => auto_instrument_handler.py} (100%) create mode 100644 src/lumigo_tracer/tracer.py rename src/test/unit/{test_handler.py => test_auto_instrument_handler.py} (95%) create mode 100644 src/test/unit/test_tracer.py diff --git a/src/lumigo_tracer/__init__.py b/src/lumigo_tracer/__init__.py index d1cdd53c..3bf463ef 100644 --- a/src/lumigo_tracer/__init__.py +++ b/src/lumigo_tracer/__init__.py @@ -1,3 +1,3 @@ -from .sync_http.sync_hook import lumigo_tracer, LumigoChalice # noqa +from .tracer import lumigo_tracer, LumigoChalice # noqa from .user_utils import report_error, add_execution_tag # noqa -from .sync_http.handler import _handler # noqa +from .auto_instrument_handler import _handler # noqa diff --git a/src/lumigo_tracer/sync_http/handler.py b/src/lumigo_tracer/auto_instrument_handler.py similarity index 100% rename from src/lumigo_tracer/sync_http/handler.py rename to src/lumigo_tracer/auto_instrument_handler.py diff --git a/src/lumigo_tracer/lumigo_utils.py b/src/lumigo_tracer/lumigo_utils.py index 04af01fc..50daeb50 100644 --- a/src/lumigo_tracer/lumigo_utils.py +++ b/src/lumigo_tracer/lumigo_utils.py @@ -52,8 +52,10 @@ TRUNCATE_SUFFIX = "...[too long]" NUMBER_OF_SPANS_IN_REPORT_OPTIMIZATION = 200 DEFAULT_KEY_DEPTH = 4 +LUMIGO_TOKEN_KEY = "LUMIGO_TRACER_TOKEN" +KILL_SWITCH = "LUMIGO_SWITCH_OFF" -_logger: Union[logging.Logger, None] = None +_logger: Dict[str, logging.Logger] = {} edge_connection = None try: @@ -116,7 +118,7 @@ def config( elif not is_aws_environment(): Configuration.should_report = False Configuration.host = edge_host or os.environ.get("LUMIGO_TRACER_HOST", "") - Configuration.token = token or os.environ.get("LUMIGO_TRACER_TOKEN", "") + Configuration.token = token or os.environ.get(LUMIGO_TOKEN_KEY, "") Configuration.enhanced_print = ( enhance_print or os.environ.get("LUMIGO_ENHANCED_PRINT", "").lower() == "true" ) @@ -234,6 +236,7 @@ def report_json(region: Union[None, str], msgs: List[dict], should_retry: bool = if not edge_connection or edge_connection.host != host: edge_connection = establish_connection(host) if not edge_connection: + get_logger().warning("Can not establish connection. Skip sending span.") return duration if Configuration.should_report: try: @@ -257,7 +260,7 @@ def report_json(region: Union[None, str], msgs: List[dict], should_retry: bool = return duration -def get_logger(): +def get_logger(logger_name="lumigo"): """ This function returns lumigo's logger. The logger streams the logs to the stderr in format the explicitly say that those are lumigo's logs. @@ -266,16 +269,17 @@ def get_logger(): Add the environment variable `LUMIGO_DEBUG=true` to activate it. """ global _logger - if not _logger: - _logger = logging.getLogger("lumigo") + if logger_name not in _logger: + _logger[logger_name] = logging.getLogger(logger_name) + _logger[logger_name].propagate = False handler = logging.StreamHandler() handler.setFormatter(logging.Formatter(LOG_FORMAT)) if os.environ.get("LUMIGO_DEBUG", "").lower() == "true": - _logger.setLevel(logging.DEBUG) + _logger[logger_name].setLevel(logging.DEBUG) else: - _logger.setLevel(logging.CRITICAL) - _logger.addHandler(handler) - return _logger + _logger[logger_name].setLevel(logging.CRITICAL) + _logger[logger_name].addHandler(handler) + return _logger[logger_name] @contextmanager @@ -290,7 +294,7 @@ def is_aws_environment(): """ :return: heuristically determine rather we're running on an aws environment. """ - return bool(os.environ.get("LAMBDA_RUNTIME_DIR")) + return bool(os.environ.get("AWS_LAMBDA_FUNCTION_VERSION")) def ensure_str(s: Union[str, bytes]) -> str: @@ -500,3 +504,7 @@ def lumigo_dumps( return ( (retval[:max_size] + TRUNCATE_SUFFIX) if len(retval) >= max_size or is_truncated else retval ) + + +def is_kill_switch_on(): + return str(os.environ.get(KILL_SWITCH, "")).lower() == "true" diff --git a/src/lumigo_tracer/sync_http/sync_hook.py b/src/lumigo_tracer/sync_http/sync_hook.py index a08003f2..b62a0597 100644 --- a/src/lumigo_tracer/sync_http/sync_hook.py +++ b/src/lumigo_tracer/sync_http/sync_hook.py @@ -1,33 +1,17 @@ from datetime import datetime -import inspect -import logging import http.client from io import BytesIO -import os -import builtins -from functools import wraps import importlib.util -import botocore.awsrequest # noqa: F401 -from lumigo_tracer.auto_tag.auto_tag_event import AutoTagEvent from lumigo_tracer.libs.wrapt import wrap_function_wrapper from lumigo_tracer.parsing_utils import safe_get_list -from lumigo_tracer.lumigo_utils import ( - config, - Configuration, - get_logger, - lumigo_safe_execute, - is_aws_environment, - ensure_str, -) -from lumigo_tracer.spans_container import SpansContainer, TimeoutMechanism +from lumigo_tracer.lumigo_utils import get_logger, lumigo_safe_execute, ensure_str +from lumigo_tracer.spans_container import SpansContainer from lumigo_tracer.parsers.http_data_classes import HttpRequest from collections import namedtuple _BODY_HEADER_SPLITTER = b"\r\n\r\n" _FLAGS_HEADER_SPLITTER = b"\r\n" -_KILL_SWITCH = "LUMIGO_SWITCH_OFF" -CONTEXT_WRAPPED_BY_LUMIGO_KEY = "_wrapped_by_lumigo" MAX_READ_SIZE = 1024 already_wrapped = False LUMIGO_HEADERS_HOOK_KEY = "_lumigo_headers_hook" @@ -176,151 +160,6 @@ def _putheader_wrapper(func, instance, args, kwargs): return ret_val -def _is_context_already_wrapped(*args) -> bool: - """ - This function is here in order to validate that we didn't already wrap this lambda - (using the sls plugin / auto instrumentation / etc.) - """ - return len(args) >= 2 and hasattr(args[1], CONTEXT_WRAPPED_BY_LUMIGO_KEY) - - -def _add_wrap_flag_to_context(*args): - """ - This function is here in order to validate that we didn't already wrap this invocation - (using the sls plugin / auto instrumentation / etc.). - We are adding lumigo's flag to the context, and check it's value in _is_context_already_wrapped. - """ - if len(args) >= 2: - with lumigo_safe_execute("wrap context"): - setattr(args[1], CONTEXT_WRAPPED_BY_LUMIGO_KEY, True) - - -def _lumigo_tracer(func): - @wraps(func) - def lambda_wrapper(*args, **kwargs): - if str(os.environ.get(_KILL_SWITCH, "")).lower() == "true": - return func(*args, **kwargs) - - if _is_context_already_wrapped(*args): - return func(*args, **kwargs) - _add_wrap_flag_to_context(*args) - executed = False - ret_val = None - local_print = print - local_logging_format = logging.Formatter.format - try: - if Configuration.enhanced_print: - _enhance_output(args, local_print, local_logging_format) - SpansContainer.create_span(*args, force=True) - with lumigo_safe_execute("auto tag"): - AutoTagEvent.auto_tag_event(args[0]) - SpansContainer.get_span().start(*args) - wrap_http_calls() - try: - executed = True - ret_val = func(*args, **kwargs) - except Exception as e: - with lumigo_safe_execute("Customer's exception"): - SpansContainer.get_span().add_exception_event(e, inspect.trace()) - raise - finally: - SpansContainer.get_span().end(ret_val) - if Configuration.enhanced_print: - builtins.print = local_print - logging.Formatter.format = local_logging_format - return ret_val - except Exception: - # The case where our wrapping raised an exception - if not executed: - TimeoutMechanism.stop() - get_logger().exception("exception in the wrapper", exc_info=True) - return func(*args, **kwargs) - else: - raise - - return lambda_wrapper - - -def _enhance_output(args, local_print, local_logging_format): - if len(args) < 2: - return - request_id = getattr(args[1], "aws_request_id", "") - prefix = f"RequestId: {request_id}" - builtins.print = lambda *args, **kwargs: local_print( - *[_add_prefix_for_each_line(prefix, str(arg)) for arg in args], **kwargs - ) - logging.Formatter.format = lambda self, record: _add_prefix_for_each_line( - prefix, local_logging_format(self, record) - ) - - -def _add_prefix_for_each_line(prefix: str, text: str): - enhanced_lines = [] - for line in text.split("\n"): - if line and not line.startswith(prefix): - line = prefix + " " + line - enhanced_lines.append(line) - return "\n".join(enhanced_lines) - - -def lumigo_tracer(*args, **kwargs): - """ - This function should be used as wrapper to your lambda function. - It will trace your HTTP calls and send it to our backend, which will help you understand it better. - - If the kill switch is activated (env variable `LUMIGO_SWITCH_OFF` set to 1), this function does nothing. - - You can pass to this decorator more configurations to configure the interface to lumigo, - See `lumigo_tracer.reporter.config` for more details on the available configuration. - """ - config(*args, **kwargs) - return _lumigo_tracer - - -class LumigoChalice: - DECORATORS_OF_NEW_HANDLERS = [ - "on_s3_event", - "on_sns_message", - "on_sqs_message", - "schedule", - "authorizer", - "lambda_function", - ] - - def __init__(self, app, *args, **kwargs): - self.lumigo_conf_args = args - self.lumigo_conf_kwargs = kwargs - self.app = app - self.original_app_attr_getter = app.__getattribute__ - self.lumigo_app = lumigo_tracer(*self.lumigo_conf_args, **self.lumigo_conf_kwargs)(app) - - def __getattr__(self, item): - original_attr = self.original_app_attr_getter(item) - if is_aws_environment() and item in self.DECORATORS_OF_NEW_HANDLERS: - - def get_decorator(*args, **kwargs): - # calling the annotation, example `app.authorizer(THIS)` - chalice_actual_decorator = original_attr(*args, **kwargs) - - def wrapper2(func): - user_func_wrapped_by_chalice = chalice_actual_decorator(func) - return LumigoChalice( - user_func_wrapped_by_chalice, - *self.lumigo_conf_args, - **self.lumigo_conf_kwargs, - ) - - return wrapper2 - - return get_decorator - return original_attr - - def __call__(self, *args, **kwargs): - if len(args) < 2 and "context" not in kwargs: - kwargs["context"] = getattr(getattr(self.app, "current_request", None), "context", None) - return self.lumigo_app(*args, **kwargs) - - def wrap_http_calls(): global already_wrapped if not already_wrapped: @@ -330,7 +169,10 @@ def wrap_http_calls(): wrap_function_wrapper( "http.client", "HTTPConnection.request", _headers_reminder_wrapper ) - wrap_function_wrapper("botocore.awsrequest", "AWSRequest.__init__", _putheader_wrapper) + if importlib.util.find_spec("botocore"): + wrap_function_wrapper( + "botocore.awsrequest", "AWSRequest.__init__", _putheader_wrapper + ) wrap_function_wrapper("http.client", "HTTPConnection.getresponse", _response_wrapper) wrap_function_wrapper("http.client", "HTTPResponse.read", _read_wrapper) if importlib.util.find_spec("urllib3"): diff --git a/src/lumigo_tracer/tracer.py b/src/lumigo_tracer/tracer.py new file mode 100644 index 00000000..b1f012bd --- /dev/null +++ b/src/lumigo_tracer/tracer.py @@ -0,0 +1,163 @@ +import inspect +import logging +import builtins +from functools import wraps + +from lumigo_tracer.auto_tag.auto_tag_event import AutoTagEvent +from lumigo_tracer.lumigo_utils import ( + config, + Configuration, + get_logger, + lumigo_safe_execute, + is_aws_environment, + is_kill_switch_on, +) +from lumigo_tracer.spans_container import SpansContainer, TimeoutMechanism +from lumigo_tracer.sync_http.sync_hook import wrap_http_calls + +CONTEXT_WRAPPED_BY_LUMIGO_KEY = "_wrapped_by_lumigo" + + +def _is_context_already_wrapped(*args) -> bool: + """ + This function is here in order to validate that we didn't already wrap this lambda + (using the sls plugin / auto instrumentation / etc.) + """ + return len(args) >= 2 and hasattr(args[1], CONTEXT_WRAPPED_BY_LUMIGO_KEY) + + +def _add_wrap_flag_to_context(*args): + """ + This function is here in order to validate that we didn't already wrap this invocation + (using the sls plugin / auto instrumentation / etc.). + We are adding lumigo's flag to the context, and check it's value in _is_context_already_wrapped. + """ + if len(args) >= 2: + with lumigo_safe_execute("wrap context"): + setattr(args[1], CONTEXT_WRAPPED_BY_LUMIGO_KEY, True) + + +def _lumigo_tracer(func): + @wraps(func) + def lambda_wrapper(*args, **kwargs): + if is_kill_switch_on(): + return func(*args, **kwargs) + + if _is_context_already_wrapped(*args): + return func(*args, **kwargs) + _add_wrap_flag_to_context(*args) + executed = False + ret_val = None + local_print = print + local_logging_format = logging.Formatter.format + try: + if Configuration.enhanced_print: + _enhance_output(args, local_print, local_logging_format) + SpansContainer.create_span(*args, force=True) + with lumigo_safe_execute("auto tag"): + AutoTagEvent.auto_tag_event(args[0]) + SpansContainer.get_span().start(*args) + wrap_http_calls() + try: + executed = True + ret_val = func(*args, **kwargs) + except Exception as e: + with lumigo_safe_execute("Customer's exception"): + SpansContainer.get_span().add_exception_event(e, inspect.trace()) + raise + finally: + SpansContainer.get_span().end(ret_val) + if Configuration.enhanced_print: + builtins.print = local_print + logging.Formatter.format = local_logging_format + return ret_val + except Exception: + # The case where our wrapping raised an exception + if not executed: + TimeoutMechanism.stop() + get_logger().exception("exception in the wrapper", exc_info=True) + return func(*args, **kwargs) + else: + raise + + return lambda_wrapper + + +def _enhance_output(args, local_print, local_logging_format): + if len(args) < 2: + return + request_id = getattr(args[1], "aws_request_id", "") + prefix = f"RequestId: {request_id}" + builtins.print = lambda *args, **kwargs: local_print( + *[_add_prefix_for_each_line(prefix, str(arg)) for arg in args], **kwargs + ) + logging.Formatter.format = lambda self, record: _add_prefix_for_each_line( + prefix, local_logging_format(self, record) + ) + + +def _add_prefix_for_each_line(prefix: str, text: str): + enhanced_lines = [] + for line in text.split("\n"): + if line and not line.startswith(prefix): + line = prefix + " " + line + enhanced_lines.append(line) + return "\n".join(enhanced_lines) + + +def lumigo_tracer(*args, **kwargs): + """ + This function should be used as wrapper to your lambda function. + It will trace your HTTP calls and send it to our backend, which will help you understand it better. + + If the kill switch is activated (env variable `LUMIGO_SWITCH_OFF` set to 1), this function does nothing. + + You can pass to this decorator more configurations to configure the interface to lumigo, + See `lumigo_tracer.reporter.config` for more details on the available configuration. + """ + config(*args, **kwargs) + return _lumigo_tracer + + +class LumigoChalice: + DECORATORS_OF_NEW_HANDLERS = [ + "on_s3_event", + "on_sns_message", + "on_sqs_message", + "schedule", + "authorizer", + "lambda_function", + ] + + def __init__(self, app, *args, **kwargs): + self.lumigo_conf_args = args + self.lumigo_conf_kwargs = kwargs + self.app = app + self.original_app_attr_getter = app.__getattribute__ + self.lumigo_app = lumigo_tracer(*self.lumigo_conf_args, **self.lumigo_conf_kwargs)(app) + + def __getattr__(self, item): + original_attr = self.original_app_attr_getter(item) + if is_aws_environment() and item in self.DECORATORS_OF_NEW_HANDLERS: + + def get_decorator(*args, **kwargs): + # calling the annotation, example `app.authorizer(THIS)` + chalice_actual_decorator = original_attr(*args, **kwargs) + + def wrapper2(func): + user_func_wrapped_by_chalice = chalice_actual_decorator(func) + return LumigoChalice( + user_func_wrapped_by_chalice, + *self.lumigo_conf_args, + **self.lumigo_conf_kwargs, + ) + + return wrapper2 + + return get_decorator + return original_attr + + def __call__(self, *args, **kwargs): + if len(args) < 2 and "context" not in kwargs: + kwargs["context"] = getattr(getattr(self.app, "current_request", None), "context", None) + return self.lumigo_app(*args, **kwargs) diff --git a/src/test/component/test_component.py b/src/test/component/test_component.py index b447fe2d..063258be 100644 --- a/src/test/component/test_component.py +++ b/src/test/component/test_component.py @@ -5,7 +5,7 @@ import os -from lumigo_tracer.sync_http.sync_hook import lumigo_tracer +from lumigo_tracer.tracer import lumigo_tracer from lumigo_tracer.spans_container import SpansContainer from lumigo_tracer.lumigo_utils import md5hash diff --git a/src/test/conftest.py b/src/test/conftest.py index 804fb502..05ece2c2 100644 --- a/src/test/conftest.py +++ b/src/test/conftest.py @@ -1,12 +1,13 @@ import builtins import logging +from types import SimpleNamespace from lumigo_tracer import lumigo_utils from lumigo_tracer.spans_container import SpansContainer import mock import pytest -from lumigo_tracer.lumigo_utils import Configuration, get_omitting_regex +from lumigo_tracer.lumigo_utils import Configuration, get_omitting_regex, get_logger @pytest.fixture(autouse=True) @@ -72,8 +73,9 @@ def pytest_collection_modifyitems(config, items): @pytest.fixture(autouse=True) def capture_all_logs(caplog): caplog.set_level(logging.DEBUG, logger="lumigo") + get_logger().propagate = True @pytest.fixture def context(): - return mock.Mock(get_remaining_time_in_millis=lambda: 1000 * 2) + return SimpleNamespace(aws_request_id="1234", get_remaining_time_in_millis=lambda: 1000 * 2) diff --git a/src/test/unit/sync_http/test_sync_hook.py b/src/test/unit/sync_http/test_sync_hook.py index 12b99dec..5e36b637 100644 --- a/src/test/unit/sync_http/test_sync_hook.py +++ b/src/test/unit/sync_http/test_sync_hook.py @@ -1,98 +1,31 @@ import copy -import datetime import json -from decimal import Decimal import time import http.client -import os -import sys import urllib -from functools import wraps from io import BytesIO from types import SimpleNamespace -import logging from typing import Dict import socket import pytest import urllib3 import requests -from capturer import CaptureOutput -from lumigo_tracer import lumigo_tracer, LumigoChalice, lumigo_utils, add_execution_tag +from lumigo_tracer import lumigo_tracer, lumigo_utils from lumigo_tracer.auto_tag import auto_tag_event +from lumigo_tracer.lumigo_utils import EXECUTION_TAGS_KEY from lumigo_tracer.parsers.http_parser import Parser -from lumigo_tracer.lumigo_utils import ( - Configuration, - STEP_FUNCTION_UID_KEY, - LUMIGO_EVENT_KEY, - _create_request_body, - EXECUTION_TAGS_KEY, -) - from lumigo_tracer.spans_container import SpansContainer -def test_lambda_wrapper_basic_events(reporter_mock): - """ - This test checks that the basic events (start and end messages) has been sent. - """ - - @lumigo_tracer(token="123") - def lambda_test_function(): - pass - - lambda_test_function() - function_span = SpansContainer.get_span().function_span - assert not SpansContainer.get_span().http_spans - assert "started" in function_span - assert "ended" in function_span - assert reporter_mock.call_count == 2 - first_send = reporter_mock.call_args_list[0][1]["msgs"] - assert len(first_send) == 1 - assert first_send[0]["id"].endswith("_started") - assert first_send[0]["maxFinishTime"] - - -@pytest.mark.parametrize("exc", [ValueError("Oh no"), ValueError(), ValueError(Exception())]) -def test_lambda_wrapper_exception(exc): +def test_lambda_wrapper_http(context): @lumigo_tracer(token="123") - def lambda_test_function(): - a = "A" # noqa - raise exc - - try: - lambda_test_function() - except ValueError: - pass - else: - assert False - - function_span = SpansContainer.get_span().function_span - assert not SpansContainer.get_span().http_spans - assert function_span.get("error", {}).get("type") == "ValueError" - # Make sure no lumigo_tracer - assert len(function_span["error"]["frames"]) == 1 - assert function_span["error"]["frames"][0].pop("lineno") > 0 - assert function_span["error"]["frames"][0] == { - "function": "lambda_test_function", - "fileName": __file__, - "variables": {"a": '"A"', "exc": f'"{str(exc)}"'}, - } - assert not function_span["id"].endswith("_started") - assert "reporter_rtt" in function_span - assert "maxFinishTime" not in function_span - # Test that we can create an output message out of this span - assert _create_request_body([function_span], prune_size_flag=False) - - -def test_lambda_wrapper_http(): - @lumigo_tracer(token="123") - def lambda_test_function(): + def lambda_test_function(event, context): time.sleep(0.01) http.client.HTTPConnection("www.google.com").request("POST", "/") - lambda_test_function() + lambda_test_function({}, context) http_spans = SpansContainer.get_span().http_spans assert http_spans assert http_spans[0].get("info", {}).get("httpInfo", {}).get("host") == "www.google.com" @@ -102,71 +35,71 @@ def lambda_test_function(): assert "content-length" in http_spans[0]["info"]["httpInfo"]["request"]["headers"] -def test_lambda_wrapper_query_with_http_params(): +def test_lambda_wrapper_query_with_http_params(context): @lumigo_tracer(token="123") - def lambda_test_function(): + def lambda_test_function(event, context): http.client.HTTPConnection("www.google.com").request("GET", "/?q=123") - lambda_test_function() + lambda_test_function({}, context) http_spans = SpansContainer.get_span().http_spans assert http_spans assert http_spans[0]["info"]["httpInfo"]["request"]["uri"] == "www.google.com/?q=123" -def test_uri_requests(): +def test_uri_requests(context): @lumigo_tracer(token="123") - def lambda_test_function(): + def lambda_test_function(event, context): conn = http.client.HTTPConnection("www.google.com") conn.request("POST", "/?q=123", b"123") conn.send(BytesIO(b"456")) - lambda_test_function() + lambda_test_function({}, context) http_spans = SpansContainer.get_span().http_spans assert http_spans assert http_spans[0]["info"]["httpInfo"]["request"]["uri"] == "www.google.com/?q=123" -def test_lambda_wrapper_get_response(): +def test_lambda_wrapper_get_response(context): @lumigo_tracer(token="123") - def lambda_test_function(): + def lambda_test_function(event, context): conn = http.client.HTTPConnection("www.google.com") conn.request("GET", "") conn.getresponse() - lambda_test_function() + lambda_test_function({}, context) http_spans = SpansContainer.get_span().http_spans assert http_spans assert http_spans[0]["info"]["httpInfo"]["response"]["statusCode"] == 200 -def test_lambda_wrapper_http_splitted_send(): +def test_lambda_wrapper_http_splitted_send(context): """ This is a test for the specific case of requests, where they split the http requests into headers and body. We didn't use directly the package requests in order to keep the dependencies small. """ @lumigo_tracer(token="123") - def lambda_test_function(): + def lambda_test_function(event, context): conn = http.client.HTTPConnection("www.google.com") conn.request("POST", "/", b"123") conn.send(BytesIO(b"456")) - lambda_test_function() + lambda_test_function({}, context) http_spans = SpansContainer.get_span().http_spans assert http_spans assert http_spans[0]["info"]["httpInfo"]["request"]["body"] == '"123456"' assert "content-length" in http_spans[0]["info"]["httpInfo"]["request"]["headers"] -def test_lambda_wrapper_no_headers(): +def test_lambda_wrapper_no_headers(context): @lumigo_tracer(token="123") - def lambda_test_function(): + def lambda_test_function(event, context): http.client.HTTPConnection("www.google.com").send(BytesIO(b"123")) - lambda_test_function() + lambda_test_function({}, context) http_events = SpansContainer.get_span().http_spans assert len(http_events) == 1 assert http_events[0].get("info", {}).get("httpInfo", {}).get("host") == "www.google.com" @@ -174,28 +107,18 @@ def lambda_test_function(): assert "ended" in http_events[0] -def test_lambda_wrapper_http_non_splitted_send(): +def test_lambda_wrapper_http_non_splitted_send(context): @lumigo_tracer(token="123") - def lambda_test_function(): + def lambda_test_function(event, context): http.client.HTTPConnection("www.google.com").request("POST", "/") http.client.HTTPConnection("www.github.com").send(BytesIO(b"123")) - lambda_test_function() + lambda_test_function({}, context) http_events = SpansContainer.get_span().http_spans assert len(http_events) == 2 -def test_lambda_wrapper_return_decimal(): - @lumigo_tracer(token="123") - def lambda_test_function(): - return {"a": [Decimal(1)]} - - lambda_test_function() - span = SpansContainer.get_span().function_span - assert span["return_value"] == '{"a": [1.0]}' - - -def test_catch_file_like_object_sent_on_http(): +def test_catch_file_like_object_sent_on_http(context): class A: def seek(self, where): pass @@ -207,68 +130,37 @@ def read(self, amount=None): return b"body" @lumigo_tracer(token="123") - def lambda_test_function(): + def lambda_test_function(event, context): try: http.client.HTTPConnection("www.github.com").send(A()) except Exception: # We don't care about errors pass - lambda_test_function() + lambda_test_function({}, context) http_events = SpansContainer.get_span().http_spans assert len(http_events) == 1 span = SpansContainer.get_span().http_spans[0] assert span["info"]["httpInfo"]["request"]["body"] == '"body"' -def test_kill_switch(monkeypatch): - monkeypatch.setattr(os, "environ", {"LUMIGO_SWITCH_OFF": "true"}) - - @lumigo_tracer(token="123") - def lambda_test_function(): - return 1 - - assert lambda_test_function() == 1 - assert not SpansContainer._span - - -def test_wrapping_exception(monkeypatch): - monkeypatch.setattr(SpansContainer, "create_span", lambda x: 1 / 0) - - @lumigo_tracer(token="123") - def lambda_test_function(): - return 1 - - assert lambda_test_function() == 1 - assert not SpansContainer._span - - -def test_wrapping_with_parameters(): - @lumigo_tracer(should_report="123") - def lambda_test_function(): - return 1 - - assert lambda_test_function() == 1 - assert Configuration.should_report == "123" - - -def test_bad_domains_scrubber(monkeypatch): +def test_bad_domains_scrubber(monkeypatch, context): monkeypatch.setenv("LUMIGO_DOMAINS_SCRUBBER", '["bad json') @lumigo_tracer(token="123", should_report=True) - def lambda_test_function(): + def lambda_test_function(event, context): pass - lambda_test_function() + lambda_test_function({}, context) assert lumigo_utils.Configuration.should_report is False -def test_domains_scrubber_happy_flow(monkeypatch): +def test_domains_scrubber_happy_flow(monkeypatch, context): @lumigo_tracer(token="123", domains_scrubber=[".*google.*"]) - def lambda_test_function(): + def lambda_test_function(event, context): return http.client.HTTPConnection(host="www.google.com").send(b"\r\n") - lambda_test_function() + lambda_test_function({}, context) http_events = SpansContainer.get_span().http_spans assert len(http_events) == 1 assert http_events[0].get("info", {}).get("httpInfo", {}).get("host") == "www.google.com" @@ -276,51 +168,26 @@ def lambda_test_function(): assert http_events[0]["info"]["httpInfo"]["request"]["body"] == "The data is not available" -def test_domains_scrubber_override_allows_default_domains(monkeypatch): +def test_domains_scrubber_override_allows_default_domains(monkeypatch, context): ssm_url = "www.ssm.123.amazonaws.com" @lumigo_tracer(token="123", domains_scrubber=[".*google.*"]) - def lambda_test_function(): + def lambda_test_function(event, context): try: return http.client.HTTPConnection(host=ssm_url).send(b"\r\n") except Exception: return - lambda_test_function() + lambda_test_function({}, context) http_events = SpansContainer.get_span().http_spans assert len(http_events) == 1 assert http_events[0].get("info", {}).get("httpInfo", {}).get("host") == ssm_url assert http_events[0]["info"]["httpInfo"]["request"]["headers"] -def test_wrapping_with_print_override(): - @lumigo_tracer(enhance_print=True) - def lambda_test_function(event, context): - print("hello\nworld") - return 1 - - with CaptureOutput() as capturer: - assert lambda_test_function({}, SimpleNamespace(aws_request_id="1234")) == 1 - assert Configuration.enhanced_print is True - assert "RequestId: 1234 hello" in capturer.get_lines() - assert "RequestId: 1234 world" in capturer.get_lines() - - -def test_wrapping_without_print_override(): +def test_wrapping_json_request(context): @lumigo_tracer() def lambda_test_function(event, context): - print("hello") - return 1 - - with CaptureOutput() as capturer: - assert lambda_test_function({}, SimpleNamespace(aws_request_id="1234")) == 1 - assert Configuration.enhanced_print is False - assert any(line == "hello" for line in capturer.get_lines()) - - -def test_wrapping_json_request(): - @lumigo_tracer() - def lambda_test_function(): urllib.request.urlopen( urllib.request.Request( "http://api.github.com", b"{}", headers={"Content-Type": "application/json"} @@ -328,7 +195,7 @@ def lambda_test_function(): ) return 1 - assert lambda_test_function() == 1 + assert lambda_test_function({}, context) == 1 http_events = SpansContainer.get_span().http_spans assert any( '"content-type": "application/json"' @@ -337,155 +204,18 @@ def lambda_test_function(): ) -def test_exception_in_parsers(monkeypatch, caplog): +def test_exception_in_parsers(monkeypatch, caplog, context): monkeypatch.setattr(Parser, "parse_request", Exception) @lumigo_tracer(token="123") - def lambda_test_function(): + def lambda_test_function(event, context): return http.client.HTTPConnection(host="www.google.com").send(b"\r\n") - lambda_test_function() + lambda_test_function({}, context) assert caplog.records[-1].msg == "An exception occurred in lumigo's code add request event" -def test_lumigo_chalice(): - class App: - @property - def a(self): - return "a" - - def b(self): - return "b" - - def __call__(self, *args, **kwargs): - return "c" - - app = App() - app = LumigoChalice(app) - - # should not use lumigo's wrapper - assert app.a == "a" - assert app.b() == "b" - assert not SpansContainer._span - - # should create a new span (but return the original value) - assert app() == "c" - assert SpansContainer._span - - -def test_lumigo_chalice_create_extra_lambdas(monkeypatch): - # mimic aws env - monkeypatch.setitem(os.environ, "LAMBDA_RUNTIME_DIR", "true") - - class Chalice: - """ - This class in a mimic of chalice. - """ - - touched = False - - @staticmethod - def on_s3_event(**kwargs): - Chalice.touched = True # represents chalice's global analysis (in the deploy process) - - def _create_registration_function(func): - @wraps(func) - def user_lambda_handler(*args, **kwargs): - return func(*args, **kwargs) - - return user_lambda_handler - - return _create_registration_function - - app = Chalice() - app = LumigoChalice(app) - - @app.on_s3_event(name="test") - def handler(event, context): - return "hello world" - - # should run the outer code before lambda execution, but not create span (in global execution) - assert app.touched - assert not SpansContainer._span - - # should create a new span (but return the original value) - assert handler({}, {}) == "hello world" - assert SpansContainer._span - - -def test_wrapping_with_logging_override_default_usage(caplog): - @lumigo_tracer(enhance_print=True) - def lambda_test_function(event, context): - logging.warning("hello\nworld") - return 1 - - assert lambda_test_function({}, SimpleNamespace(aws_request_id="1234")) == 1 - assert Configuration.enhanced_print is True - assert any("RequestId: 1234" in line and "hello" in line for line in caplog.text.split("\n")) - assert any("RequestId: 1234" in line and "world" in line for line in caplog.text.split("\n")) - - -def test_wrapping_with_logging_exception(caplog): - @lumigo_tracer(enhance_print=True) - def lambda_test_function(event, context): - logger = logging.getLogger("logger_name") - handler = logging.StreamHandler() - logger.addHandler(handler) - - try: - 1 / 0 - except Exception: # You must call the logging.exception method just inside the except part. - logger.exception("hello") - return 1 - - assert lambda_test_function({}, SimpleNamespace(aws_request_id="1234")) == 1 - # Check all lines have exactly one RequestId. - for line in caplog.text.splitlines(): - assert line.startswith("RequestId: 1234") and line.count("RequestId: 1234") == 1 - # Check the message was logged. - test_message = [line for line in caplog.text.splitlines() if line.endswith("hello")][0].replace( - " ", "" - ) - assert "ERROR" in test_message and "hello" in test_message - - -def test_wrapping_with_logging_override_complex_usage(): - @lumigo_tracer(enhance_print=True) - def lambda_test_function(event, context): - handler = logging.StreamHandler(sys.stdout) - formatter = logging.Formatter("%(name)s [%(levelname)s] %(message)s") # Format of a client. - handler.setFormatter(formatter) - logger = logging.getLogger("my_test") - logger.handlers = [handler] - logger.setLevel("INFO") - - logger.info("hello\nworld") - return 1 - - with CaptureOutput() as capturer: - assert lambda_test_function({}, SimpleNamespace(aws_request_id="1234")) == 1 - assert Configuration.enhanced_print is True - assert "RequestId: 1234 my_test [INFO] hello" in capturer.get_lines() - assert "RequestId: 1234 world" in capturer.get_lines() - - -def test_wrapping_without_logging_override(caplog): - @lumigo_tracer() - def lambda_test_function(event, context): - logging.warning("hello\nworld") - return 1 - - assert lambda_test_function({}, SimpleNamespace(aws_request_id="1234")) == 1 - assert Configuration.enhanced_print is False - assert any( - "RequestId: 1234" not in line and "world" in line for line in caplog.text.split("\n") - ) - assert any( - "RequestId: 1234" not in line and "hello" in line for line in caplog.text.split("\n") - ) - - -def test_wrapping_urlib_stream_get(): +def test_wrapping_urlib_stream_get(context): """ This is the same case as the one of `requests.get`. """ @@ -495,7 +225,7 @@ def lambda_test_function(event, context): r = urllib3.PoolManager().urlopen("GET", "https://www.google.com", preload_content=False) return b"".join(r.stream(32)) - lambda_test_function({}, None) + lambda_test_function({}, context) assert len(SpansContainer.get_span().http_spans) == 1 event = SpansContainer.get_span().http_spans[0] assert event["info"]["httpInfo"]["response"]["body"] @@ -503,7 +233,7 @@ def lambda_test_function(event, context): assert event["info"]["httpInfo"]["host"] == "www.google.com" -def test_wrapping_requests_times(monkeypatch): +def test_wrapping_requests_times(monkeypatch, context): @lumigo_tracer() def lambda_test_function(event, context): start_time = time.time() * 1000 @@ -520,87 +250,19 @@ def delayed_getaddrinfo(*args, **kwargs): monkeypatch.setattr(socket, "getaddrinfo", delayed_getaddrinfo) # validate that the added delay didn't affect the start time - start_time = lambda_test_function({}, None) + start_time = lambda_test_function({}, context) span = SpansContainer.get_span().http_spans[0] assert span["started"] - start_time < 100 -@pytest.mark.parametrize( - "event, expected_triggered_by, expected_message_id", - [ - ({}, "unknown", None), - ({"result": 1, LUMIGO_EVENT_KEY: {STEP_FUNCTION_UID_KEY: "123"}}, "stepFunction", "123"), - ], -) -def test_wrapping_step_function(event, expected_triggered_by, expected_message_id): - @lumigo_tracer(step_function=True) - def lambda_test_function(event, context): - return {"result": 1} - - lambda_test_function(event, None) - span = SpansContainer.get_span() - assert len(span.http_spans) == 1 - assert span.function_span["info"]["triggeredBy"] == expected_triggered_by - assert span.function_span["info"].get("messageId") == expected_message_id - return_value = json.loads(span.function_span["return_value"]) - assert return_value["result"] == 1 - assert return_value[LUMIGO_EVENT_KEY][STEP_FUNCTION_UID_KEY] - assert span.http_spans[0]["info"]["httpInfo"]["host"] == "StepFunction" - - -def test_omitting_keys(): - @lumigo_tracer() - def lambda_test_function(event, context): - d = {"a": "b", "myPassword": "123"} - conn = http.client.HTTPConnection("www.google.com") - conn.request("POST", "/", json.dumps(d)) - return {"secret_password": "lumigo rulz"} - - lambda_test_function({"key": "24"}, None) - span = SpansContainer.get_span() - assert span.function_span["return_value"] == '{"secret_password": "****"}' - assert span.function_span["event"] == '{"key": "****"}' - spans = json.loads(_create_request_body(SpansContainer.get_span().http_spans, True)) - assert spans[0]["info"]["httpInfo"]["request"]["body"] == json.dumps( - {"a": "b", "myPassword": "****"} - ) - - -def test_can_not_wrap_twice(reporter_mock): - @lumigo_tracer() - @lumigo_tracer() - def lambda_test_function(event, context): - return "ret_value" - - result = lambda_test_function({}, SimpleNamespace(aws_request_id="1234")) - assert result == "ret_value" - assert reporter_mock.call_count == 2 - - -def test_wrapping_with_tags(): - key = "my_key" - value = "my_value" +def test_wrapping_with_tags_for_api_gw_headers(monkeypatch, context): + monkeypatch.setattr(auto_tag_event, "AUTO_TAG_API_GW_HEADERS", ["Accept"]) @lumigo_tracer() def lambda_test_function(event, context): - add_execution_tag(key, value) return "ret_value" - result = lambda_test_function({}, SimpleNamespace(aws_request_id="1234")) - assert result == "ret_value" - assert SpansContainer.get_span().function_span[EXECUTION_TAGS_KEY] == [ - {"key": key, "value": value} - ] - - -def test_wrapping_with_tags_for_api_gw_headers(monkeypatch): - set_header_key(monkeypatch, "Accept") - - @lumigo_tracer() - def lambda_test_function(event, context): - return "ret_value" - - result = lambda_test_function(api_gw_event(), SimpleNamespace(aws_request_id="1234")) + result = lambda_test_function(api_gw_event(), context) assert result == "ret_value" assert SpansContainer.get_span().function_span[EXECUTION_TAGS_KEY] == [ @@ -608,22 +270,7 @@ def lambda_test_function(event, context): ] -def test_not_jsonable_return(monkeypatch): - @lumigo_tracer() - def lambda_test_function(event, context): - return {"a": datetime.datetime.now()} - - lambda_test_function(api_gw_event(), SimpleNamespace(aws_request_id="1234")) - - function_span = SpansContainer.get_span().function_span - assert function_span["return_value"] is None - assert function_span["error"]["type"] == "ReturnValueError" - # following python's runtime: runtime/lambda_runtime_marshaller.py:27 - expected_message = 'The lambda will probably fail due to bad return value. Original message: "Object of type datetime is not JSON serializable"' - assert function_span["error"]["message"] == expected_message - - -def test_correct_headers_of_send_after_request(): +def test_correct_headers_of_send_after_request(context): @lumigo_tracer() def lambda_test_function(event, context): d = {"a": "b", "myPassword": "123"} @@ -632,16 +279,12 @@ def lambda_test_function(event, context): conn.send(b"GET\r\nc: d\r\n\r\nbody") return {"lumigo": "rulz"} - lambda_test_function({"key": "24"}, None) + lambda_test_function({"key": "24"}, context) spans = SpansContainer.get_span().http_spans assert spans[0]["info"]["httpInfo"]["request"]["headers"] == json.dumps({"a": "b"}) assert spans[1]["info"]["httpInfo"]["request"]["headers"] == json.dumps({"c": "d"}) -def set_header_key(monkeypatch, header: str): - monkeypatch.setattr(auto_tag_event, "AUTO_TAG_API_GW_HEADERS", [header]) - - def api_gw_event() -> Dict: return { "resource": "/add-user", diff --git a/src/test/unit/test_handler.py b/src/test/unit/test_auto_instrument_handler.py similarity index 95% rename from src/test/unit/test_handler.py rename to src/test/unit/test_auto_instrument_handler.py index 9abd0a05..c606f9c9 100644 --- a/src/test/unit/test_handler.py +++ b/src/test/unit/test_auto_instrument_handler.py @@ -6,7 +6,7 @@ import os import pytest -from lumigo_tracer.sync_http.handler import _handler, ORIGINAL_HANDLER_KEY +from lumigo_tracer.auto_instrument_handler import _handler, ORIGINAL_HANDLER_KEY def test_happy_flow(monkeypatch): diff --git a/src/test/unit/test_lumigo_utils.py b/src/test/unit/test_lumigo_utils.py index 5f7233f4..5f7371c9 100644 --- a/src/test/unit/test_lumigo_utils.py +++ b/src/test/unit/test_lumigo_utils.py @@ -29,6 +29,8 @@ prepare_host, EDGE_PATH, report_json, + is_kill_switch_on, + KILL_SWITCH, ) import json @@ -380,3 +382,9 @@ def test_report_json_retry(monkeypatch, reporter_mock, caplog, errors, final_log report_json(None, [{"a": "b"}]) assert caplog.records[-1].levelname == final_log + + +@pytest.mark.parametrize("env, expected", [("True", True), ("other", False), (123, False)]) +def test_is_kill_switch_on(monkeypatch, env, expected): + monkeypatch.setenv(KILL_SWITCH, env) + assert is_kill_switch_on() == expected diff --git a/src/test/unit/test_tracer.py b/src/test/unit/test_tracer.py new file mode 100644 index 00000000..876c4ae2 --- /dev/null +++ b/src/test/unit/test_tracer.py @@ -0,0 +1,366 @@ +import datetime +import json +from decimal import Decimal +import http.client +import os +import sys +from functools import wraps +import logging + +import pytest +from capturer import CaptureOutput + +from lumigo_tracer import lumigo_tracer, LumigoChalice, add_execution_tag +from lumigo_tracer.lumigo_utils import ( + Configuration, + STEP_FUNCTION_UID_KEY, + LUMIGO_EVENT_KEY, + _create_request_body, + EXECUTION_TAGS_KEY, +) + +from lumigo_tracer.spans_container import SpansContainer + + +def test_lambda_wrapper_basic_events(reporter_mock, context): + """ + This test checks that the basic events (start and end messages) has been sent. + """ + + @lumigo_tracer(token="123") + def lambda_test_function(event, context): + pass + + lambda_test_function({}, context) + function_span = SpansContainer.get_span().function_span + assert not SpansContainer.get_span().http_spans + assert "started" in function_span + assert "ended" in function_span + assert reporter_mock.call_count == 2 + first_send = reporter_mock.call_args_list[0][1]["msgs"] + assert len(first_send) == 1 + assert first_send[0]["id"].endswith("_started") + assert first_send[0]["maxFinishTime"] + + +@pytest.mark.parametrize("exc", [ValueError("Oh no"), ValueError(), ValueError(Exception())]) +def test_lambda_wrapper_exception(exc, context): + @lumigo_tracer(token="123") + def lambda_test_function(event, context): + a = "A" # noqa + raise exc + + try: + lambda_test_function({}, context) + except ValueError: + pass + else: + assert False + + function_span = SpansContainer.get_span().function_span + assert not SpansContainer.get_span().http_spans + assert function_span.get("error", {}).get("type") == "ValueError" + # Make sure no lumigo_tracer + assert len(function_span["error"]["frames"]) == 1 + assert function_span["error"]["frames"][0].pop("lineno") > 0 + assert function_span["error"]["frames"][0] == { + "function": "lambda_test_function", + "fileName": __file__, + "variables": { + "a": '"A"', + "context": f'"{str(context)}"', + "event": "{}", + "exc": f'"{str(exc)}"', + }, + } + assert not function_span["id"].endswith("_started") + assert "reporter_rtt" in function_span + assert "maxFinishTime" not in function_span + # Test that we can create an output message out of this span + assert _create_request_body([function_span], prune_size_flag=False) + + +def test_lambda_wrapper_return_decimal(context): + @lumigo_tracer(token="123") + def lambda_test_function(event, context): + return {"a": [Decimal(1)]} + + lambda_test_function({}, context) + span = SpansContainer.get_span().function_span + assert span["return_value"] == '{"a": [1.0]}' + + +def test_kill_switch(monkeypatch, context): + monkeypatch.setattr(os, "environ", {"LUMIGO_SWITCH_OFF": "true"}) + + @lumigo_tracer(token="123") + def lambda_test_function(event, context): + return 1 + + assert lambda_test_function({}, context) == 1 + assert not SpansContainer._span + + +def test_wrapping_exception(monkeypatch, context): + monkeypatch.setattr(SpansContainer, "create_span", lambda *args, **kwargs: 1 / 0) + + @lumigo_tracer(token="123") + def lambda_test_function(event, context): + return 1 + + assert lambda_test_function({}, context) == 1 + assert not SpansContainer._span + + +def test_wrapping_with_parameters(context): + @lumigo_tracer(should_report="123") + def lambda_test_function(event, context): + return 1 + + assert lambda_test_function({}, context) == 1 + assert Configuration.should_report == "123" + + +def test_wrapping_with_print_override(context): + @lumigo_tracer(enhance_print=True) + def lambda_test_function(event, context): + print("hello\nworld") + return 1 + + with CaptureOutput() as capturer: + assert lambda_test_function({}, context) == 1 + assert Configuration.enhanced_print is True + assert "RequestId: 1234 hello" in capturer.get_lines() + assert "RequestId: 1234 world" in capturer.get_lines() + + +def test_wrapping_without_print_override(context): + @lumigo_tracer() + def lambda_test_function(event, context): + print("hello") + return 1 + + with CaptureOutput() as capturer: + assert lambda_test_function({}, context) == 1 + assert Configuration.enhanced_print is False + assert any(line == "hello" for line in capturer.get_lines()) + + +def test_lumigo_chalice(context): + class App: + @property + def a(self): + return "a" + + def b(self): + return "b" + + def __call__(self, *args, **kwargs): + return "c" + + app = App() + app = LumigoChalice(app) + + # should not use lumigo's wrapper + assert app.a == "a" + assert app.b() == "b" + assert not SpansContainer._span + + # should create a new span (but return the original value) + assert app({}, context) == "c" + assert SpansContainer._span + + +def test_lumigo_chalice_create_extra_lambdas(monkeypatch, context): + # mimic aws env + monkeypatch.setitem(os.environ, "AWS_LAMBDA_FUNCTION_VERSION", "true") + + class Chalice: + """ + This class in a mimic of chalice. + """ + + touched = False + + @staticmethod + def on_s3_event(**kwargs): + Chalice.touched = True # represents chalice's global analysis (in the deploy process) + + def _create_registration_function(func): + @wraps(func) + def user_lambda_handler(*args, **kwargs): + return func(*args, **kwargs) + + return user_lambda_handler + + return _create_registration_function + + app = Chalice() + app = LumigoChalice(app) + + @app.on_s3_event(name="test") + def handler(event, context): + return "hello world" + + # should run the outer code before lambda execution, but not create span (in global execution) + assert app.touched + assert not SpansContainer._span + + # should create a new span (but return the original value) + assert handler({}, context) == "hello world" + assert SpansContainer._span + + +def test_wrapping_with_logging_override_default_usage(caplog, context): + @lumigo_tracer(enhance_print=True) + def lambda_test_function(event, context): + logging.warning("hello\nworld") + return 1 + + assert lambda_test_function({}, context) == 1 + assert Configuration.enhanced_print is True + assert any("RequestId: 1234" in line and "hello" in line for line in caplog.text.split("\n")) + assert any("RequestId: 1234" in line and "world" in line for line in caplog.text.split("\n")) + + +def test_wrapping_with_logging_exception(caplog, context): + @lumigo_tracer(enhance_print=True) + def lambda_test_function(event, context): + logger = logging.getLogger("logger_name") + handler = logging.StreamHandler() + logger.addHandler(handler) + + try: + 1 / 0 + except Exception: # You must call the logging.exception method just inside the except part. + logger.exception("hello") + return 1 + + assert lambda_test_function({}, context) == 1 + # Check all lines have exactly one RequestId. + for line in caplog.text.splitlines(): + assert line.startswith("RequestId: 1234") and line.count("RequestId: 1234") == 1 + # Check the message was logged. + test_message = [line for line in caplog.text.splitlines() if line.endswith("hello")][0].replace( + " ", "" + ) + assert "ERROR" in test_message and "hello" in test_message + + +def test_wrapping_with_logging_override_complex_usage(context): + @lumigo_tracer(enhance_print=True) + def lambda_test_function(event, context): + handler = logging.StreamHandler(sys.stdout) + formatter = logging.Formatter("%(name)s [%(levelname)s] %(message)s") # Format of a client. + handler.setFormatter(formatter) + logger = logging.getLogger("my_test") + logger.handlers = [handler] + logger.setLevel("INFO") + + logger.info("hello\nworld") + return 1 + + with CaptureOutput() as capturer: + assert lambda_test_function({}, context) == 1 + assert Configuration.enhanced_print is True + assert "RequestId: 1234 my_test [INFO] hello" in capturer.get_lines() + assert "RequestId: 1234 world" in capturer.get_lines() + + +def test_wrapping_without_logging_override(caplog, context): + @lumigo_tracer() + def lambda_test_function(event, context): + logging.warning("hello\nworld") + return 1 + + assert lambda_test_function({}, context) == 1 + assert Configuration.enhanced_print is False + assert any( + "RequestId: 1234" not in line and "world" in line for line in caplog.text.split("\n") + ) + assert any( + "RequestId: 1234" not in line and "hello" in line for line in caplog.text.split("\n") + ) + + +@pytest.mark.parametrize( + "event, expected_triggered_by, expected_message_id", + [ + ({}, "unknown", None), + ({"result": 1, LUMIGO_EVENT_KEY: {STEP_FUNCTION_UID_KEY: "123"}}, "stepFunction", "123"), + ], +) +def test_wrapping_step_function(event, expected_triggered_by, expected_message_id, context): + @lumigo_tracer(step_function=True) + def lambda_test_function(event, context): + return {"result": 1} + + lambda_test_function(event, context) + span = SpansContainer.get_span() + assert len(span.http_spans) == 1 + assert span.function_span["info"]["triggeredBy"] == expected_triggered_by + assert span.function_span["info"].get("messageId") == expected_message_id + return_value = json.loads(span.function_span["return_value"]) + assert return_value["result"] == 1 + assert return_value[LUMIGO_EVENT_KEY][STEP_FUNCTION_UID_KEY] + assert span.http_spans[0]["info"]["httpInfo"]["host"] == "StepFunction" + + +def test_omitting_keys(context): + @lumigo_tracer() + def lambda_test_function(event, context): + d = {"a": "b", "myPassword": "123"} + conn = http.client.HTTPConnection("www.google.com") + conn.request("POST", "/", json.dumps(d)) + return {"secret_password": "lumigo rulz"} + + lambda_test_function({"key": "24"}, context) + span = SpansContainer.get_span() + assert span.function_span["return_value"] == '{"secret_password": "****"}' + assert span.function_span["event"] == '{"key": "****"}' + spans = json.loads(_create_request_body(SpansContainer.get_span().http_spans, True)) + assert spans[0]["info"]["httpInfo"]["request"]["body"] == json.dumps( + {"a": "b", "myPassword": "****"} + ) + + +def test_can_not_wrap_twice(reporter_mock, context): + @lumigo_tracer() + @lumigo_tracer() + def lambda_test_function(event, context): + return "ret_value" + + result = lambda_test_function({}, context) + assert result == "ret_value" + assert reporter_mock.call_count == 2 + + +def test_wrapping_with_tags(context): + key = "my_key" + value = "my_value" + + @lumigo_tracer() + def lambda_test_function(event, context): + add_execution_tag(key, value) + return "ret_value" + + result = lambda_test_function({}, context) + assert result == "ret_value" + assert SpansContainer.get_span().function_span[EXECUTION_TAGS_KEY] == [ + {"key": key, "value": value} + ] + + +def test_not_jsonable_return(monkeypatch, context): + @lumigo_tracer() + def lambda_test_function(event, context): + return {"a": datetime.datetime.now()} + + lambda_test_function({}, context) + + function_span = SpansContainer.get_span().function_span + assert function_span["return_value"] is None + assert function_span["error"]["type"] == "ReturnValueError" + # following python's runtime: runtime/lambda_runtime_marshaller.py:27 + expected_message = 'The lambda will probably fail due to bad return value. Original message: "Object of type datetime is not JSON serializable"' + assert function_span["error"]["message"] == expected_message From f56b82611291f32a45caab7e27151dff7e0167b6 Mon Sep 17 00:00:00 2001 From: saartochner Date: Mon, 5 Oct 2020 16:20:00 +0300 Subject: [PATCH 04/24] refactor http hooks to be a wrapper --- src/lumigo_tracer/spans_container.py | 92 ++++++------------- src/lumigo_tracer/sync_http/__init__.py | 0 src/lumigo_tracer/tracer.py | 4 +- src/lumigo_tracer/wrappers/__init__.py | 5 + .../{parsers => wrappers/http}/__init__.py | 0 .../http}/http_data_classes.py | 11 +++ .../{parsers => wrappers/http}/http_parser.py | 2 +- .../http/sync_http_wrappers.py} | 83 ++++++++++++++--- src/test/component/test_component.py | 14 +-- src/test/conftest.py | 2 + src/test/unit/test_lumigo_utils.py | 2 +- src/test/unit/test_spans_container.py | 83 +++++------------ src/test/unit/test_tracer.py | 10 +- .../http}/test_http_parser.py | 4 +- .../http/test_sync_http_wrappers.py} | 56 +++++++---- 15 files changed, 193 insertions(+), 175 deletions(-) delete mode 100644 src/lumigo_tracer/sync_http/__init__.py create mode 100644 src/lumigo_tracer/wrappers/__init__.py rename src/lumigo_tracer/{parsers => wrappers/http}/__init__.py (100%) rename src/lumigo_tracer/{parsers => wrappers/http}/http_data_classes.py (67%) rename src/lumigo_tracer/{parsers => wrappers/http}/http_parser.py (99%) rename src/lumigo_tracer/{sync_http/sync_hook.py => wrappers/http/sync_http_wrappers.py} (68%) rename src/test/unit/{parsers => wrappers/http}/test_http_parser.py (97%) rename src/test/unit/{sync_http/test_sync_hook.py => wrappers/http/test_sync_http_wrappers.py} (89%) diff --git a/src/lumigo_tracer/spans_container.py b/src/lumigo_tracer/spans_container.py index 96762b76..6602f1e1 100644 --- a/src/lumigo_tracer/spans_container.py +++ b/src/lumigo_tracer/spans_container.py @@ -22,10 +22,8 @@ _is_span_has_error, ) from lumigo_tracer import lumigo_utils -from lumigo_tracer.parsers.http_parser import get_parser, HTTP_TYPE, StepFunctionParser from lumigo_tracer.parsing_utils import parse_trace_id, safe_split_get, recursive_json_join from lumigo_tracer.event.event_trigger import parse_triggered_by -from lumigo_tracer.parsers.http_data_classes import HttpRequest _VERSION_PATH = os.path.join(os.path.dirname(__file__), "VERSION") MAX_LAMBDA_TIME = 15 * 60 * 1000 @@ -92,10 +90,8 @@ def __init__( }, self.base_msg, ) - self.previous_request: Optional[HttpRequest] = None - self.previous_response_body: bytes = b"" - self.http_span_ids_to_send: Set[str] = set() - self.http_spans: List[Dict] = [] + self.span_ids_to_send: Set[str] = set() + self.spans: List[Dict] = [] SpansContainer.is_cold = False def _generate_start_span(self) -> dict: @@ -110,18 +106,18 @@ def start(self, event=None, context=None): if not Configuration.send_only_if_error: report_duration = lumigo_utils.report_json(region=self.region, msgs=[to_send]) self.function_span["reporter_rtt"] = report_duration - self.http_spans = [] + self.spans = [] else: get_logger().debug("Skip sending start because tracer in 'send only if error' mode .") self.start_timeout_timer(context) def handle_timeout(self, *args): get_logger().info("The tracer reached the end of the timeout timer") - to_send = [s for s in self.http_spans if s["id"] in self.http_span_ids_to_send] + to_send = [s for s in self.spans if s["id"] in self.span_ids_to_send] if Configuration.send_only_if_error: to_send.append(self._generate_start_span()) lumigo_utils.report_json(region=self.region, msgs=to_send) - self.http_span_ids_to_send.clear() + self.span_ids_to_send.clear() def start_timeout_timer(self, context=None) -> None: if Configuration.timeout_timer: @@ -135,41 +131,27 @@ def start_timeout_timer(self, context=None) -> None: return TimeoutMechanism.start(remaining_time - buffer, self.handle_timeout) - def add_request_event(self, parse_params: HttpRequest): + def add_span(self, span: dict): """ This function parses an request event and add it to the span. """ - parser = get_parser(parse_params.host)() - msg = parser.parse_request(parse_params) - self.previous_request = parse_params - self.http_spans.append(recursive_json_join(msg, self.base_msg)) - self.http_span_ids_to_send.add(msg["id"]) + self.spans.append(recursive_json_join(span, self.base_msg)) + self.span_ids_to_send.add(span["id"]) - def add_unparsed_request(self, parse_params: HttpRequest): - """ - This function handle the case where we got a request the is not fully formatted as we expected, - I.e. there isn't '\r\n' in the request data that logically splits the headers from the body. + def get_last_span(self) -> Optional[dict]: + if not self.spans: + return None + return self.spans[-1] - In that case, we will consider it as a continuance of the previous request if they got the same url, - and we didn't get any answer yet. - """ - if self.http_spans: - last_event = self.http_spans[-1] - if last_event and last_event.get("type") == HTTP_TYPE and self.previous_request: - if last_event.get("info", {}).get("httpInfo", {}).get("host") == parse_params.host: - if "response" not in last_event["info"]["httpInfo"]: - self.http_spans.pop() - body = (self.previous_request.body + parse_params.body)[:MAX_BODY_SIZE] - self.add_request_event(self.previous_request.clone(body=body)) - return - self.add_request_event(parse_params.clone(headers=None)) + def remove_last_span(self) -> Optional[dict]: + return self.spans.pop() if self.spans else None def update_event_end_time(self) -> None: """ This function assumes synchronous execution - we update the last http event. """ - if self.http_spans: - self.http_spans[-1]["ended"] = int(time.time() * 1000) + if self.spans: + self.spans[-1]["ended"] = int(time.time() * 1000) def update_event_times( self, start_time: Optional[datetime] = None, end_time: Optional[datetime] = None @@ -177,36 +159,11 @@ def update_event_times( """ This function assumes synchronous execution - we update the last http event. """ - if self.http_spans: + if self.spans: start_timestamp = start_time.timestamp() if start_time else time.time() end_timestamp = end_time.timestamp() if end_time else time.time() - self.http_spans[-1]["started"] = int(start_timestamp * 1000) - self.http_spans[-1]["ended"] = int(end_timestamp * 1000) - - def update_event_response( - self, host: Optional[str], status_code: int, headers: dict, body: bytes - ) -> None: - """ - :param host: If None, use the host from the last span, otherwise this is the first chuck and we can empty - the aggregated response body - This function assumes synchronous execution - we update the last http event. - """ - if self.http_spans: - last_event = self.http_spans.pop() - if not host: - host = last_event.get("info", {}).get("httpInfo", {}).get("host", "unknown") - else: - self.previous_response_body = b"" - - headers = {k.lower(): v for k, v in headers.items()} if headers else {} - parser = get_parser(host, headers)() # type: ignore - if len(self.previous_response_body) < Configuration.max_entry_size: - self.previous_response_body += body - update = parser.parse_response( # type: ignore - host, status_code, headers, self.previous_response_body # type: ignore - ) - self.http_spans.append(recursive_json_join(update, last_event)) - self.http_span_ids_to_send.add(update.get("id") or last_event["id"]) + self.spans[-1]["started"] = int(start_timestamp * 1000) + self.spans[-1]["ended"] = int(end_timestamp * 1000) def _create_exception_event( self, exc_type: str, message: str, stacktrace: str = "", frames: Optional[List[dict]] = None @@ -233,10 +190,13 @@ def add_exception_event( ) def add_step_end_event(self, ret_val): + # TODO + from lumigo_tracer.wrappers.http.http_parser import StepFunctionParser + message_id = str(uuid.uuid4()) step_function_span = StepFunctionParser().create_span(message_id) - self.http_spans.append(recursive_json_join(step_function_span, self.base_msg)) - self.http_span_ids_to_send.add(step_function_span["id"]) + self.spans.append(recursive_json_join(step_function_span, self.base_msg)) + self.span_ids_to_send.add(step_function_span["id"]) if isinstance(ret_val, dict): ret_val[LUMIGO_EVENT_KEY] = {STEP_FUNCTION_UID_KEY: message_id} get_logger().debug(f"Added key {LUMIGO_EVENT_KEY} to the user's return value") @@ -269,12 +229,12 @@ def end(self, ret_val=None) -> Optional[int]: ) self.function_span.update({"return_value": parsed_ret_val}) spans_contain_errors: bool = any( - _is_span_has_error(s) for s in self.http_spans + [self.function_span] + _is_span_has_error(s) for s in self.spans + [self.function_span] ) if (not Configuration.send_only_if_error) or spans_contain_errors: to_send = [self.function_span] + [ - s for s in self.http_spans if s["id"] in self.http_span_ids_to_send + s for s in self.spans if s["id"] in self.span_ids_to_send ] reported_rtt = lumigo_utils.report_json(region=self.region, msgs=to_send) else: diff --git a/src/lumigo_tracer/sync_http/__init__.py b/src/lumigo_tracer/sync_http/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/src/lumigo_tracer/tracer.py b/src/lumigo_tracer/tracer.py index b1f012bd..c2f98e4e 100644 --- a/src/lumigo_tracer/tracer.py +++ b/src/lumigo_tracer/tracer.py @@ -13,7 +13,7 @@ is_kill_switch_on, ) from lumigo_tracer.spans_container import SpansContainer, TimeoutMechanism -from lumigo_tracer.sync_http.sync_hook import wrap_http_calls +from lumigo_tracer.wrappers import wrap CONTEXT_WRAPPED_BY_LUMIGO_KEY = "_wrapped_by_lumigo" @@ -57,7 +57,7 @@ def lambda_wrapper(*args, **kwargs): with lumigo_safe_execute("auto tag"): AutoTagEvent.auto_tag_event(args[0]) SpansContainer.get_span().start(*args) - wrap_http_calls() + wrap() try: executed = True ret_val = func(*args, **kwargs) diff --git a/src/lumigo_tracer/wrappers/__init__.py b/src/lumigo_tracer/wrappers/__init__.py new file mode 100644 index 00000000..3617073d --- /dev/null +++ b/src/lumigo_tracer/wrappers/__init__.py @@ -0,0 +1,5 @@ +from .http.sync_http_wrappers import wrap_http_calls + + +def wrap(): + wrap_http_calls() diff --git a/src/lumigo_tracer/parsers/__init__.py b/src/lumigo_tracer/wrappers/http/__init__.py similarity index 100% rename from src/lumigo_tracer/parsers/__init__.py rename to src/lumigo_tracer/wrappers/http/__init__.py diff --git a/src/lumigo_tracer/parsers/http_data_classes.py b/src/lumigo_tracer/wrappers/http/http_data_classes.py similarity index 67% rename from src/lumigo_tracer/parsers/http_data_classes.py rename to src/lumigo_tracer/wrappers/http/http_data_classes.py index 49b92c0e..fec471aa 100644 --- a/src/lumigo_tracer/parsers/http_data_classes.py +++ b/src/lumigo_tracer/wrappers/http/http_data_classes.py @@ -1,4 +1,5 @@ from copy import deepcopy +from typing import Optional class HttpRequest: @@ -20,3 +21,13 @@ def clone(self, **kwargs): for k, v in kwargs.items(): setattr(clone_obj, k, v) return clone_obj + + +class HttpState: + previous_request: Optional[HttpRequest] = None + previous_response_body: bytes = b"" + + @staticmethod + def clear(): + HttpState.previous_request = None + HttpState.previous_response_body = b"" diff --git a/src/lumigo_tracer/parsers/http_parser.py b/src/lumigo_tracer/wrappers/http/http_parser.py similarity index 99% rename from src/lumigo_tracer/parsers/http_parser.py rename to src/lumigo_tracer/wrappers/http/http_parser.py index 85f13bd8..06e9534e 100644 --- a/src/lumigo_tracer/parsers/http_parser.py +++ b/src/lumigo_tracer/wrappers/http/http_parser.py @@ -13,7 +13,7 @@ should_scrub_domain, ) from lumigo_tracer.lumigo_utils import Configuration, lumigo_dumps, md5hash, get_logger -from lumigo_tracer.parsers.http_data_classes import HttpRequest +from lumigo_tracer.wrappers.http.http_data_classes import HttpRequest HTTP_TYPE = "http" diff --git a/src/lumigo_tracer/sync_http/sync_hook.py b/src/lumigo_tracer/wrappers/http/sync_http_wrappers.py similarity index 68% rename from src/lumigo_tracer/sync_http/sync_hook.py rename to src/lumigo_tracer/wrappers/http/sync_http_wrappers.py index b62a0597..8d1a3ac6 100644 --- a/src/lumigo_tracer/sync_http/sync_hook.py +++ b/src/lumigo_tracer/wrappers/http/sync_http_wrappers.py @@ -2,14 +2,17 @@ import http.client from io import BytesIO import importlib.util +from typing import Optional from lumigo_tracer.libs.wrapt import wrap_function_wrapper -from lumigo_tracer.parsing_utils import safe_get_list -from lumigo_tracer.lumigo_utils import get_logger, lumigo_safe_execute, ensure_str -from lumigo_tracer.spans_container import SpansContainer -from lumigo_tracer.parsers.http_data_classes import HttpRequest +from lumigo_tracer.parsing_utils import safe_get_list, recursive_json_join +from lumigo_tracer.lumigo_utils import get_logger, lumigo_safe_execute, ensure_str, Configuration +from lumigo_tracer.spans_container import SpansContainer, MAX_BODY_SIZE +from lumigo_tracer.wrappers.http.http_data_classes import HttpRequest, HttpState from collections import namedtuple +from lumigo_tracer.wrappers.http.http_parser import get_parser, HTTP_TYPE + _BODY_HEADER_SPLITTER = b"\r\n\r\n" _FLAGS_HEADER_SPLITTER = b"\r\n" MAX_READ_SIZE = 1024 @@ -20,6 +23,64 @@ HookedData = namedtuple("HookedData", ["headers", "path"]) +def add_request_event(parse_params: HttpRequest): + """ + This function parses an request event and add it to the span. + """ + parser = get_parser(parse_params.host)() + msg = parser.parse_request(parse_params) + HttpState.previous_request = parse_params + SpansContainer.get_span().add_span(msg) + + +def add_unparsed_request(parse_params: HttpRequest): + """ + This function handle the case where we got a request the is not fully formatted as we expected, + I.e. there isn't '\r\n' in the request data that logically splits the headers from the body. + + In that case, we will consider it as a continuance of the previous request if they got the same url, + and we didn't get any answer yet. + """ + last_event = SpansContainer.get_span().get_last_span() + if last_event: + if last_event and last_event.get("type") == HTTP_TYPE and HttpState.previous_request: + if last_event.get("info", {}).get("httpInfo", {}).get("host") == parse_params.host: + if "response" not in last_event["info"]["httpInfo"]: + SpansContainer.get_span().remove_last_span() + body = (HttpState.previous_request.body + parse_params.body)[:MAX_BODY_SIZE] + add_request_event(HttpState.previous_request.clone(body=body)) + return + add_request_event(parse_params.clone(headers=None)) + + +def update_event_response( + host: Optional[str], status_code: int, headers: dict, body: bytes +) -> None: + """ + :param host: If None, use the host from the last span, otherwise this is the first chuck and we can empty + the aggregated response body + This function assumes synchronous execution - we update the last http event. + """ + last_event = SpansContainer.get_span().remove_last_span() + if last_event: + if not host: + host = last_event.get("info", {}).get("httpInfo", {}).get("host", "unknown") + else: + HttpState.previous_response_body = b"" + + headers = {k.lower(): v for k, v in headers.items()} if headers else {} + parser = get_parser(host, headers)() # type: ignore + if len(HttpState.previous_response_body) < Configuration.max_entry_size: + HttpState.previous_response_body += body + update = parser.parse_response( # type: ignore + host, status_code, headers, HttpState.previous_response_body # type: ignore + ) + SpansContainer.get_span().add_span(recursive_json_join(update, last_event)) + + +# Wrappers # + + def _http_send_wrapper(func, instance, args, kwargs): """ This is the wrapper of the requests. it parses the http's message to conclude the url, headers, and body. @@ -69,13 +130,11 @@ def _http_send_wrapper(func, instance, args, kwargs): with lumigo_safe_execute("add request event"): if headers: - SpansContainer.get_span().add_request_event( + add_request_event( HttpRequest(host=host, method=method, uri=uri, headers=headers, body=body) ) else: - SpansContainer.get_span().add_unparsed_request( - HttpRequest(host=host, method=method, uri=uri, body=data) - ) + add_unparsed_request(HttpRequest(host=host, method=method, uri=uri, body=data)) ret_val = func(*args, **kwargs) with lumigo_safe_execute("add response event"): @@ -119,7 +178,7 @@ def _response_wrapper(func, instance, args, kwargs): with lumigo_safe_execute("parse response"): headers = dict(ret_val.headers.items()) status_code = ret_val.code - SpansContainer.get_span().update_event_response(instance.host, status_code, headers, b"") + update_event_response(instance.host, status_code, headers, b"") return ret_val @@ -130,9 +189,7 @@ def _read_wrapper(func, instance, args, kwargs): ret_val = func(*args, **kwargs) if ret_val: with lumigo_safe_execute("parse response.read"): - SpansContainer.get_span().update_event_response( - None, instance.code, dict(instance.headers.items()), ret_val - ) + update_event_response(None, instance.code, dict(instance.headers.items()), ret_val) return ret_val @@ -144,7 +201,7 @@ def _read_stream_wrapper(func, instance, args, kwargs): def _read_stream_wrapper_generator(stream_generator, instance): for partial_response in stream_generator: with lumigo_safe_execute("parse response.read_chunked"): - SpansContainer.get_span().update_event_response( + update_event_response( None, instance.status, dict(instance.headers.items()), partial_response ) yield partial_response diff --git a/src/test/component/test_component.py b/src/test/component/test_component.py index 063258be..119dc943 100644 --- a/src/test/component/test_component.py +++ b/src/test/component/test_component.py @@ -80,7 +80,7 @@ def lambda_test_function(): ) lambda_test_function() - events = SpansContainer.get_span().http_spans + events = SpansContainer.get_span().spans assert len(events) == 1 assert events[0]["info"]["httpInfo"]["host"] == f"dynamodb.{region}.amazonaws.com" assert events[0]["info"]["resourceName"] == ddb_resource @@ -95,7 +95,7 @@ def lambda_test_function(): boto3.resource("sns").Topic(sns_resource).publish(Message=json.dumps({"test": "test"})) lambda_test_function() - events = SpansContainer.get_span().http_spans + events = SpansContainer.get_span().spans assert len(events) == 1 assert events[0]["info"]["httpInfo"]["host"] == f"sns.{region}.amazonaws.com" assert events[0]["info"]["resourceName"] == sns_resource @@ -111,7 +111,7 @@ def lambda_test_function(): ) lambda_test_function() - events = SpansContainer.get_span().http_spans + events = SpansContainer.get_span().spans assert len(events) == 1 assert events[0]["info"]["httpInfo"]["host"] == f"lambda.{region}.amazonaws.com" assert events[0]["info"]["resourceName"] == lambda_resource @@ -138,7 +138,7 @@ def lambda_test_function(): ) lambda_test_function() - events = SpansContainer.get_span().http_spans + events = SpansContainer.get_span().spans assert len(events) == 2 # Single message. assert events[0]["info"]["httpInfo"]["host"] == f"kinesis.{region}.amazonaws.com" @@ -165,7 +165,7 @@ def lambda_test_function(): ) lambda_test_function() - events = SpansContainer.get_span().http_spans + events = SpansContainer.get_span().spans assert len(events) == 2 # Single message. assert events[0]["info"]["httpInfo"]["host"] == f"{region}.queue.amazonaws.com" @@ -188,7 +188,7 @@ def lambda_test_function(): s3_client.upload_file(os.path.abspath(__file__), s3_bucket_resource, "test.txt") lambda_test_function() - events = SpansContainer.get_span().http_spans + events = SpansContainer.get_span().spans assert len(events) == 2 assert events[0]["info"]["messageId"] assert events[0]["info"]["resourceName"] == s3_bucket_resource @@ -202,7 +202,7 @@ def lambda_test_function(): boto3.client("sqs").send_message(QueueUrl=sqs_resource, MessageBody="myMessage") lambda_test_function() - events = SpansContainer.get_span().http_spans + events = SpansContainer.get_span().spans # making sure there is any data in the body. body = events[0]["info"]["httpInfo"]["response"]["body"] assert body and body != "b''" diff --git a/src/test/conftest.py b/src/test/conftest.py index 05ece2c2..aaaa43b3 100644 --- a/src/test/conftest.py +++ b/src/test/conftest.py @@ -8,6 +8,7 @@ import pytest from lumigo_tracer.lumigo_utils import Configuration, get_omitting_regex, get_logger +from lumigo_tracer.wrappers.http.http_data_classes import HttpState @pytest.fixture(autouse=True) @@ -36,6 +37,7 @@ def restart_global_span(): """ yield SpansContainer._span = None + HttpState.clear() @pytest.yield_fixture(autouse=True) diff --git a/src/test/unit/test_lumigo_utils.py b/src/test/unit/test_lumigo_utils.py index 5f7371c9..cf5a7125 100644 --- a/src/test/unit/test_lumigo_utils.py +++ b/src/test/unit/test_lumigo_utils.py @@ -384,7 +384,7 @@ def test_report_json_retry(monkeypatch, reporter_mock, caplog, errors, final_log assert caplog.records[-1].levelname == final_log -@pytest.mark.parametrize("env, expected", [("True", True), ("other", False), (123, False)]) +@pytest.mark.parametrize("env, expected", [("True", True), ("other", False), ("123", False)]) def test_is_kill_switch_on(monkeypatch, env, expected): monkeypatch.setenv(KILL_SWITCH, env) assert is_kill_switch_on() == expected diff --git a/src/test/unit/test_spans_container.py b/src/test/unit/test_spans_container.py index fb87c3b7..6e79c24f 100644 --- a/src/test/unit/test_spans_container.py +++ b/src/test/unit/test_spans_container.py @@ -2,17 +2,14 @@ import pytest -from lumigo_tracer.parsers.http_data_classes import HttpRequest -from lumigo_tracer.parsers.http_parser import HTTP_TYPE +from lumigo_tracer.wrappers.http.http_parser import HTTP_TYPE from lumigo_tracer.spans_container import SpansContainer, TimeoutMechanism, FUNCTION_TYPE -from lumigo_tracer.lumigo_utils import Configuration, EXECUTION_TAGS_KEY, DEFAULT_MAX_ENTRY_SIZE +from lumigo_tracer.lumigo_utils import Configuration, EXECUTION_TAGS_KEY -@pytest.fixture() -def dummy_http_request(): - return HttpRequest( - host="dummy", method="dummy", uri="dummy", headers={"dummy": "dummy"}, body="dummy" - ) +@pytest.fixture +def dummy_span(): + return {"id": "span1", "type": "http", "info": {"hello": "world"}} def _is_start_span_sent(): @@ -41,28 +38,28 @@ def test_spans_container_end_function_got_none_return_value(monkeypatch): def test_spans_container_end_function_not_send_spans_on_send_only_on_errors_mode( - monkeypatch, dummy_http_request + monkeypatch, dummy_span ): Configuration.send_only_if_error = True SpansContainer.create_span() SpansContainer.get_span().start() - SpansContainer.get_span().add_request_event(dummy_http_request) + SpansContainer.get_span().add_span(dummy_span) reported_ttl = SpansContainer.get_span().end({}) assert reported_ttl is None def test_spans_container_end_function_send_spans_on_send_only_on_errors_mode( - monkeypatch, dummy_http_request + monkeypatch, dummy_span ): Configuration.send_only_if_error = True SpansContainer.create_span() SpansContainer.get_span().start() - SpansContainer.get_span().add_request_event(dummy_http_request) + SpansContainer.get_span().add_span(dummy_span) try: 1 / 0 except Exception: @@ -73,26 +70,26 @@ def test_spans_container_end_function_send_spans_on_send_only_on_errors_mode( def test_spans_container_end_function_send_only_on_errors_mode_false_not_effecting( - monkeypatch, dummy_http_request + monkeypatch, dummy_span ): SpansContainer.create_span() SpansContainer.get_span().start() - SpansContainer.get_span().add_request_event(dummy_http_request) + SpansContainer.get_span().add_span(dummy_span) reported_ttl = SpansContainer.get_span().end({}) assert reported_ttl is not None def test_spans_container_timeout_mechanism_send_only_on_errors_mode( - monkeypatch, context, reporter_mock, dummy_http_request + monkeypatch, context, reporter_mock, dummy_span ): monkeypatch.setattr(Configuration, "send_only_if_error", True) SpansContainer.create_span() SpansContainer.get_span().start() - SpansContainer.get_span().add_request_event(dummy_http_request) + SpansContainer.get_span().add_span(dummy_span) SpansContainer.get_span().handle_timeout() @@ -119,44 +116,26 @@ def test_timeout_mechanism_too_short_time(monkeypatch, context): assert not TimeoutMechanism.is_activated() -def test_timeout_mechanism_timeout_occurred_doesnt_send_span_twice(monkeypatch, context): +def test_timeout_mechanism_timeout_occurred_doesnt_send_span_twice( + monkeypatch, context, dummy_span +): SpansContainer.create_span() SpansContainer.get_span().start(context=context) - SpansContainer.get_span().add_request_event( - HttpRequest(host="google.com", method="", uri="", headers=None, body="") - ) - - assert SpansContainer.get_span().http_span_ids_to_send - SpansContainer.get_span().handle_timeout() - assert not SpansContainer.get_span().http_span_ids_to_send + SpansContainer.get_span().add_span(dummy_span) - -def test_timeout_mechanism_timeout_occurred_send_new_spans(monkeypatch, context): - SpansContainer.create_span() - SpansContainer.get_span().start(context=context) - SpansContainer.get_span().add_request_event( - HttpRequest(host="google.com", method="", uri="", headers=None, body="") - ) + assert SpansContainer.get_span().span_ids_to_send SpansContainer.get_span().handle_timeout() - - SpansContainer.get_span().add_request_event( - HttpRequest(host="google.com", method="", uri="", headers=None, body="2") - ) - assert SpansContainer.get_span().http_span_ids_to_send + assert not SpansContainer.get_span().span_ids_to_send -def test_timeout_mechanism_timeout_occurred_send_updated_spans(monkeypatch, context): +def test_timeout_mechanism_timeout_occurred_send_new_spans(monkeypatch, context, dummy_span): SpansContainer.create_span() SpansContainer.get_span().start(context=context) - SpansContainer.get_span().add_request_event( - HttpRequest(host="google.com", method="", uri="", headers=None, body="") - ) + SpansContainer.get_span().add_span(dummy_span) SpansContainer.get_span().handle_timeout() - SpansContainer.get_span().update_event_response( - host="google.com", status_code=200, headers=None, body=b"2" - ) - assert SpansContainer.get_span().http_span_ids_to_send + SpansContainer.get_span().add_span(dummy_span) + assert SpansContainer.get_span().span_ids_to_send def test_add_tag(): @@ -173,19 +152,3 @@ def test_get_tags_len(): SpansContainer.get_span().add_tag("k0", "v0") SpansContainer.get_span().add_tag("k1", "v1") assert SpansContainer.get_span().get_tags_len() == 2 - - -def test_aggregating_response_body(dummy_http_request): - """ - This test is here to validate that we're not leaking memory on aggregating response body. - Unfortunately python doesn't give us better tools, so we must check the problematic member itself. - """ - SpansContainer.create_span() - SpansContainer.get_span().add_request_event(dummy_http_request) - - big_response_chunk = b"leak" * DEFAULT_MAX_ENTRY_SIZE - for _ in range(10): - SpansContainer.get_span().update_event_response( - host=None, status_code=200, headers=None, body=big_response_chunk - ) - assert len(SpansContainer.get_span().previous_response_body) <= len(big_response_chunk) diff --git a/src/test/unit/test_tracer.py b/src/test/unit/test_tracer.py index 876c4ae2..5d149be4 100644 --- a/src/test/unit/test_tracer.py +++ b/src/test/unit/test_tracer.py @@ -33,7 +33,7 @@ def lambda_test_function(event, context): lambda_test_function({}, context) function_span = SpansContainer.get_span().function_span - assert not SpansContainer.get_span().http_spans + assert not SpansContainer.get_span().spans assert "started" in function_span assert "ended" in function_span assert reporter_mock.call_count == 2 @@ -58,7 +58,7 @@ def lambda_test_function(event, context): assert False function_span = SpansContainer.get_span().function_span - assert not SpansContainer.get_span().http_spans + assert not SpansContainer.get_span().spans assert function_span.get("error", {}).get("type") == "ValueError" # Make sure no lumigo_tracer assert len(function_span["error"]["frames"]) == 1 @@ -297,13 +297,13 @@ def lambda_test_function(event, context): lambda_test_function(event, context) span = SpansContainer.get_span() - assert len(span.http_spans) == 1 + assert len(span.spans) == 1 assert span.function_span["info"]["triggeredBy"] == expected_triggered_by assert span.function_span["info"].get("messageId") == expected_message_id return_value = json.loads(span.function_span["return_value"]) assert return_value["result"] == 1 assert return_value[LUMIGO_EVENT_KEY][STEP_FUNCTION_UID_KEY] - assert span.http_spans[0]["info"]["httpInfo"]["host"] == "StepFunction" + assert span.spans[0]["info"]["httpInfo"]["host"] == "StepFunction" def test_omitting_keys(context): @@ -318,7 +318,7 @@ def lambda_test_function(event, context): span = SpansContainer.get_span() assert span.function_span["return_value"] == '{"secret_password": "****"}' assert span.function_span["event"] == '{"key": "****"}' - spans = json.loads(_create_request_body(SpansContainer.get_span().http_spans, True)) + spans = json.loads(_create_request_body(SpansContainer.get_span().spans, True)) assert spans[0]["info"]["httpInfo"]["request"]["body"] == json.dumps( {"a": "b", "myPassword": "****"} ) diff --git a/src/test/unit/parsers/test_http_parser.py b/src/test/unit/wrappers/http/test_http_parser.py similarity index 97% rename from src/test/unit/parsers/test_http_parser.py rename to src/test/unit/wrappers/http/test_http_parser.py index cf6c4d70..f703ac36 100644 --- a/src/test/unit/parsers/test_http_parser.py +++ b/src/test/unit/wrappers/http/test_http_parser.py @@ -2,8 +2,8 @@ import pytest -from lumigo_tracer.parsers.http_data_classes import HttpRequest -from lumigo_tracer.parsers.http_parser import ( +from lumigo_tracer.wrappers.http.http_data_classes import HttpRequest +from lumigo_tracer.wrappers.http.http_parser import ( ServerlessAWSParser, Parser, get_parser, diff --git a/src/test/unit/sync_http/test_sync_hook.py b/src/test/unit/wrappers/http/test_sync_http_wrappers.py similarity index 89% rename from src/test/unit/sync_http/test_sync_hook.py rename to src/test/unit/wrappers/http/test_sync_http_wrappers.py index 5e36b637..6ba9daf1 100644 --- a/src/test/unit/sync_http/test_sync_hook.py +++ b/src/test/unit/wrappers/http/test_sync_http_wrappers.py @@ -14,9 +14,11 @@ from lumigo_tracer import lumigo_tracer, lumigo_utils from lumigo_tracer.auto_tag import auto_tag_event -from lumigo_tracer.lumigo_utils import EXECUTION_TAGS_KEY -from lumigo_tracer.parsers.http_parser import Parser +from lumigo_tracer.lumigo_utils import EXECUTION_TAGS_KEY, DEFAULT_MAX_ENTRY_SIZE +from lumigo_tracer.wrappers.http.http_parser import Parser from lumigo_tracer.spans_container import SpansContainer +from lumigo_tracer.wrappers.http.http_data_classes import HttpState, HttpRequest +from lumigo_tracer.wrappers.http.sync_http_wrappers import add_request_event, update_event_response def test_lambda_wrapper_http(context): @@ -26,7 +28,7 @@ def lambda_test_function(event, context): http.client.HTTPConnection("www.google.com").request("POST", "/") lambda_test_function({}, context) - http_spans = SpansContainer.get_span().http_spans + http_spans = SpansContainer.get_span().spans assert http_spans assert http_spans[0].get("info", {}).get("httpInfo", {}).get("host") == "www.google.com" assert "started" in http_spans[0] @@ -41,7 +43,7 @@ def lambda_test_function(event, context): http.client.HTTPConnection("www.google.com").request("GET", "/?q=123") lambda_test_function({}, context) - http_spans = SpansContainer.get_span().http_spans + http_spans = SpansContainer.get_span().spans assert http_spans assert http_spans[0]["info"]["httpInfo"]["request"]["uri"] == "www.google.com/?q=123" @@ -55,7 +57,7 @@ def lambda_test_function(event, context): conn.send(BytesIO(b"456")) lambda_test_function({}, context) - http_spans = SpansContainer.get_span().http_spans + http_spans = SpansContainer.get_span().spans assert http_spans assert http_spans[0]["info"]["httpInfo"]["request"]["uri"] == "www.google.com/?q=123" @@ -69,7 +71,7 @@ def lambda_test_function(event, context): conn.getresponse() lambda_test_function({}, context) - http_spans = SpansContainer.get_span().http_spans + http_spans = SpansContainer.get_span().spans assert http_spans assert http_spans[0]["info"]["httpInfo"]["response"]["statusCode"] == 200 @@ -88,7 +90,7 @@ def lambda_test_function(event, context): conn.send(BytesIO(b"456")) lambda_test_function({}, context) - http_spans = SpansContainer.get_span().http_spans + http_spans = SpansContainer.get_span().spans assert http_spans assert http_spans[0]["info"]["httpInfo"]["request"]["body"] == '"123456"' assert "content-length" in http_spans[0]["info"]["httpInfo"]["request"]["headers"] @@ -100,7 +102,7 @@ def lambda_test_function(event, context): http.client.HTTPConnection("www.google.com").send(BytesIO(b"123")) lambda_test_function({}, context) - http_events = SpansContainer.get_span().http_spans + http_events = SpansContainer.get_span().spans assert len(http_events) == 1 assert http_events[0].get("info", {}).get("httpInfo", {}).get("host") == "www.google.com" assert "started" in http_events[0] @@ -114,7 +116,7 @@ def lambda_test_function(event, context): http.client.HTTPConnection("www.github.com").send(BytesIO(b"123")) lambda_test_function({}, context) - http_events = SpansContainer.get_span().http_spans + http_events = SpansContainer.get_span().spans assert len(http_events) == 2 @@ -138,9 +140,9 @@ def lambda_test_function(event, context): pass lambda_test_function({}, context) - http_events = SpansContainer.get_span().http_spans + http_events = SpansContainer.get_span().spans assert len(http_events) == 1 - span = SpansContainer.get_span().http_spans[0] + span = SpansContainer.get_span().spans[0] assert span["info"]["httpInfo"]["request"]["body"] == '"body"' @@ -161,7 +163,7 @@ def lambda_test_function(event, context): return http.client.HTTPConnection(host="www.google.com").send(b"\r\n") lambda_test_function({}, context) - http_events = SpansContainer.get_span().http_spans + http_events = SpansContainer.get_span().spans assert len(http_events) == 1 assert http_events[0].get("info", {}).get("httpInfo", {}).get("host") == "www.google.com" assert "headers" not in http_events[0]["info"]["httpInfo"]["request"] @@ -179,7 +181,7 @@ def lambda_test_function(event, context): return lambda_test_function({}, context) - http_events = SpansContainer.get_span().http_spans + http_events = SpansContainer.get_span().spans assert len(http_events) == 1 assert http_events[0].get("info", {}).get("httpInfo", {}).get("host") == ssm_url assert http_events[0]["info"]["httpInfo"]["request"]["headers"] @@ -196,7 +198,7 @@ def lambda_test_function(event, context): return 1 assert lambda_test_function({}, context) == 1 - http_events = SpansContainer.get_span().http_spans + http_events = SpansContainer.get_span().spans assert any( '"content-type": "application/json"' in event.get("info", {}).get("httpInfo", {}).get("request", {}).get("headers", "") @@ -226,8 +228,8 @@ def lambda_test_function(event, context): return b"".join(r.stream(32)) lambda_test_function({}, context) - assert len(SpansContainer.get_span().http_spans) == 1 - event = SpansContainer.get_span().http_spans[0] + assert len(SpansContainer.get_span().spans) == 1 + event = SpansContainer.get_span().spans[0] assert event["info"]["httpInfo"]["response"]["body"] assert event["info"]["httpInfo"]["response"]["statusCode"] == 200 assert event["info"]["httpInfo"]["host"] == "www.google.com" @@ -251,7 +253,7 @@ def delayed_getaddrinfo(*args, **kwargs): # validate that the added delay didn't affect the start time start_time = lambda_test_function({}, context) - span = SpansContainer.get_span().http_spans[0] + span = SpansContainer.get_span().spans[0] assert span["started"] - start_time < 100 @@ -280,7 +282,7 @@ def lambda_test_function(event, context): return {"lumigo": "rulz"} lambda_test_function({"key": "24"}, context) - spans = SpansContainer.get_span().http_spans + spans = SpansContainer.get_span().spans assert spans[0]["info"]["httpInfo"]["request"]["headers"] == json.dumps({"a": "b"}) assert spans[1]["info"]["httpInfo"]["request"]["headers"] == json.dumps({"c": "d"}) @@ -414,3 +416,21 @@ def lambda_test_function(event, context): return "ret_value" lambda_test_function(given_event, SimpleNamespace(aws_request_id="1234")) + + +def test_aggregating_response_body(): + """ + This test is here to validate that we're not leaking memory on aggregating response body. + Unfortunately python doesn't give us better tools, so we must check the problematic member itself. + """ + SpansContainer.create_span() + add_request_event( + HttpRequest( + host="dummy", method="dummy", uri="dummy", headers={"dummy": "dummy"}, body="dummy" + ) + ) + + big_response_chunk = b"leak" * DEFAULT_MAX_ENTRY_SIZE + for _ in range(10): + update_event_response(host=None, status_code=200, headers=None, body=big_response_chunk) + assert len(HttpState.previous_response_body) <= len(big_response_chunk) From 613ef65ae757b7e1c65fe07d2fc6e553d5ba30c0 Mon Sep 17 00:00:00 2001 From: saartochner Date: Mon, 5 Oct 2020 16:29:45 +0300 Subject: [PATCH 05/24] refactor http hooks to be a wrapper --- src/lumigo_tracer/lumigo_utils.py | 14 ++++++++++++++ src/lumigo_tracer/spans_container.py | 6 ++---- 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/src/lumigo_tracer/lumigo_utils.py b/src/lumigo_tracer/lumigo_utils.py index 50daeb50..19bdebb0 100644 --- a/src/lumigo_tracer/lumigo_utils.py +++ b/src/lumigo_tracer/lumigo_utils.py @@ -4,6 +4,7 @@ import logging import os import re +import uuid from functools import reduce, lru_cache import time import http.client @@ -371,6 +372,19 @@ def is_api_gw_event(event: dict) -> bool: ) +def create_step_function_span(message_id: str): + return { + "id": str(uuid.uuid4()), + "type": "http", + "info": { + "resourceName": "StepFunction", + "messageId": message_id, + "httpInfo": {"host": "StepFunction", "request": {"method": "", "body": ""}}, + }, + "started": int(time.time() * 1000), + } + + def get_timeout_buffer(remaining_time: float): buffer = Configuration.timeout_timer_buffer if not buffer: diff --git a/src/lumigo_tracer/spans_container.py b/src/lumigo_tracer/spans_container.py index 6602f1e1..8b7c4e3f 100644 --- a/src/lumigo_tracer/spans_container.py +++ b/src/lumigo_tracer/spans_container.py @@ -20,6 +20,7 @@ get_timeout_buffer, get_logger, _is_span_has_error, + create_step_function_span, ) from lumigo_tracer import lumigo_utils from lumigo_tracer.parsing_utils import parse_trace_id, safe_split_get, recursive_json_join @@ -190,11 +191,8 @@ def add_exception_event( ) def add_step_end_event(self, ret_val): - # TODO - from lumigo_tracer.wrappers.http.http_parser import StepFunctionParser - message_id = str(uuid.uuid4()) - step_function_span = StepFunctionParser().create_span(message_id) + step_function_span = create_step_function_span(message_id) self.spans.append(recursive_json_join(step_function_span, self.base_msg)) self.span_ids_to_send.add(step_function_span["id"]) if isinstance(ret_val, dict): From 077f1be8e53d4a46d5ba48d45c11ce4eb43299ab Mon Sep 17 00:00:00 2001 From: saartochner Date: Mon, 5 Oct 2020 16:31:34 +0300 Subject: [PATCH 06/24] refactor http hooks to be a wrapper --- src/lumigo_tracer/wrappers/http/http_parser.py | 14 -------------- 1 file changed, 14 deletions(-) diff --git a/src/lumigo_tracer/wrappers/http/http_parser.py b/src/lumigo_tracer/wrappers/http/http_parser.py index 06e9534e..3a781e4d 100644 --- a/src/lumigo_tracer/wrappers/http/http_parser.py +++ b/src/lumigo_tracer/wrappers/http/http_parser.py @@ -229,20 +229,6 @@ def parse_response(self, url: str, status_code: int, headers, body: bytes) -> di ) -class StepFunctionParser(ServerlessAWSParser): - def create_span(self, message_id: str) -> dict: - return recursive_json_join( - { - "info": { - "resourceName": "StepFunction", - "httpInfo": {"host": "StepFunction"}, - "messageId": message_id, - } - }, - super().parse_request(None), # type: ignore - ) - - class ApiGatewayV2Parser(ServerlessAWSParser): # API-GW V1 covered by ServerlessAWSParser From 1c0d3ec96f5440009b687f3b0f648ea77fd58648 Mon Sep 17 00:00:00 2001 From: saartochner Date: Mon, 5 Oct 2020 16:51:33 +0300 Subject: [PATCH 07/24] increase coverage --- src/test/unit/test_parsing_utils.py | 21 +++++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/src/test/unit/test_parsing_utils.py b/src/test/unit/test_parsing_utils.py index 00e8b894..18ca59e2 100644 --- a/src/test/unit/test_parsing_utils.py +++ b/src/test/unit/test_parsing_utils.py @@ -13,6 +13,7 @@ parse_trace_id, safe_key_from_query, safe_key_from_xml, + safe_get_list, ) from lumigo_tracer.lumigo_utils import config, Configuration @@ -23,12 +24,27 @@ (("a.b.c", ".", 0), "a"), # happy flow (("a.b.c", ".", 1), "b"), (("a.b.c", ".", 5, "d"), "d"), # return the default + ((123, ".", 5, "d"), "d"), # no string ], ) def test_safe_split_get(input_params, expected_output): assert safe_split_get(*input_params) == expected_output +@pytest.mark.parametrize( + ("input_params", "expected_output"), + [ + ((["a", "b"], 1), "b"), # happy flow + (([], 2, 1), 1), # out of bound index + ((["a", "b"], "1", 1), "b"), # string index + ((["a", "b"], "bla", "default"), "default"), # string index - cant cast + ((123, 2, "default"), "default"), # Non iterable + ], +) +def test_safe_get_list(input_params, expected_output): + assert safe_get_list(*input_params) == expected_output + + @pytest.mark.parametrize( ("input_params", "expected_output"), [ @@ -77,6 +93,7 @@ def test_key_from_query(input_params, expected_output): ("Root=1-2;", ("1-2", "", ";")), ("a;1", ("", "", ";1")), ("123", ("", "", "123")), + (None, ("", "", "")), ], ) def test_parse_trace_id(trace_id, result): @@ -143,7 +160,7 @@ def test_str_to_list(): def test_str_to_list_exception(): - assert str_to_list("") is None + assert str_to_list([1]) is None def test_str_to_tuple(): @@ -151,4 +168,4 @@ def test_str_to_tuple(): def test_str_to_tuple_exception(): - assert str_to_tuple([]) is None + assert str_to_tuple([1]) is None From 1764e732899338b8978f099b6a7c61c6a3878830 Mon Sep 17 00:00:00 2001 From: saartochner Date: Tue, 6 Oct 2020 15:19:20 +0300 Subject: [PATCH 08/24] refactor wrap once --- src/lumigo_tracer/wrappers/__init__.py | 8 ++++- .../wrappers/http/sync_http_wrappers.py | 32 +++++++------------ 2 files changed, 19 insertions(+), 21 deletions(-) diff --git a/src/lumigo_tracer/wrappers/__init__.py b/src/lumigo_tracer/wrappers/__init__.py index 3617073d..c8d9130a 100644 --- a/src/lumigo_tracer/wrappers/__init__.py +++ b/src/lumigo_tracer/wrappers/__init__.py @@ -1,5 +1,11 @@ from .http.sync_http_wrappers import wrap_http_calls +already_wrapped = False + + def wrap(): - wrap_http_calls() + global already_wrapped + if not already_wrapped: + wrap_http_calls() + already_wrapped = True diff --git a/src/lumigo_tracer/wrappers/http/sync_http_wrappers.py b/src/lumigo_tracer/wrappers/http/sync_http_wrappers.py index 8d1a3ac6..d10fe4e0 100644 --- a/src/lumigo_tracer/wrappers/http/sync_http_wrappers.py +++ b/src/lumigo_tracer/wrappers/http/sync_http_wrappers.py @@ -16,7 +16,6 @@ _BODY_HEADER_SPLITTER = b"\r\n\r\n" _FLAGS_HEADER_SPLITTER = b"\r\n" MAX_READ_SIZE = 1024 -already_wrapped = False LUMIGO_HEADERS_HOOK_KEY = "_lumigo_headers_hook" @@ -218,24 +217,17 @@ def _putheader_wrapper(func, instance, args, kwargs): def wrap_http_calls(): - global already_wrapped - if not already_wrapped: - with lumigo_safe_execute("wrap http calls"): - get_logger().debug("wrapping the http request") - wrap_function_wrapper("http.client", "HTTPConnection.send", _http_send_wrapper) + with lumigo_safe_execute("wrap http calls"): + get_logger().debug("wrapping the http request") + wrap_function_wrapper("http.client", "HTTPConnection.send", _http_send_wrapper) + wrap_function_wrapper("http.client", "HTTPConnection.request", _headers_reminder_wrapper) + if importlib.util.find_spec("botocore"): + wrap_function_wrapper("botocore.awsrequest", "AWSRequest.__init__", _putheader_wrapper) + wrap_function_wrapper("http.client", "HTTPConnection.getresponse", _response_wrapper) + wrap_function_wrapper("http.client", "HTTPResponse.read", _read_wrapper) + if importlib.util.find_spec("urllib3"): wrap_function_wrapper( - "http.client", "HTTPConnection.request", _headers_reminder_wrapper + "urllib3.response", "HTTPResponse.read_chunked", _read_stream_wrapper ) - if importlib.util.find_spec("botocore"): - wrap_function_wrapper( - "botocore.awsrequest", "AWSRequest.__init__", _putheader_wrapper - ) - wrap_function_wrapper("http.client", "HTTPConnection.getresponse", _response_wrapper) - wrap_function_wrapper("http.client", "HTTPResponse.read", _read_wrapper) - if importlib.util.find_spec("urllib3"): - wrap_function_wrapper( - "urllib3.response", "HTTPResponse.read_chunked", _read_stream_wrapper - ) - if importlib.util.find_spec("requests"): - wrap_function_wrapper("requests.api", "request", _requests_wrapper) - already_wrapped = True + if importlib.util.find_spec("requests"): + wrap_function_wrapper("requests.api", "request", _requests_wrapper) From 838e077625f367603d2d125251ddd9494b5f122e Mon Sep 17 00:00:00 2001 From: saartochner Date: Thu, 8 Oct 2020 12:53:12 +0300 Subject: [PATCH 09/24] support pymongo --- requirements.txt | 1 + src/lumigo_tracer/wrappers/__init__.py | 2 + .../wrappers/pymongo/__init__.py | 0 .../wrappers/pymongo/pymongo_wrapper.py | 63 ++++++++++++++++++ .../wrappers/pymongo/test_pymongo_wrapper.py | 64 +++++++++++++++++++ 5 files changed, 130 insertions(+) create mode 100644 src/lumigo_tracer/wrappers/pymongo/__init__.py create mode 100644 src/lumigo_tracer/wrappers/pymongo/pymongo_wrapper.py create mode 100644 src/test/unit/wrappers/pymongo/test_pymongo_wrapper.py diff --git a/requirements.txt b/requirements.txt index a2b07bfa..628ac38d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,3 +6,4 @@ pytest-cov==2.6.1 capturer==2.4 attrs==19.1.0 requests==2.24.0 +pymongo==3.11.0 \ No newline at end of file diff --git a/src/lumigo_tracer/wrappers/__init__.py b/src/lumigo_tracer/wrappers/__init__.py index c8d9130a..cb3ef5c3 100644 --- a/src/lumigo_tracer/wrappers/__init__.py +++ b/src/lumigo_tracer/wrappers/__init__.py @@ -1,4 +1,5 @@ from .http.sync_http_wrappers import wrap_http_calls +from .pymongo.pymongo_wrapper import wrap_pymongo already_wrapped = False @@ -8,4 +9,5 @@ def wrap(): global already_wrapped if not already_wrapped: wrap_http_calls() + wrap_pymongo() already_wrapped = True diff --git a/src/lumigo_tracer/wrappers/pymongo/__init__.py b/src/lumigo_tracer/wrappers/pymongo/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/lumigo_tracer/wrappers/pymongo/pymongo_wrapper.py b/src/lumigo_tracer/wrappers/pymongo/pymongo_wrapper.py new file mode 100644 index 00000000..876520ef --- /dev/null +++ b/src/lumigo_tracer/wrappers/pymongo/pymongo_wrapper.py @@ -0,0 +1,63 @@ +import time +import uuid + +from lumigo_tracer.lumigo_utils import lumigo_safe_execute, get_logger, lumigo_dumps +from lumigo_tracer.spans_container import SpansContainer + +try: + from pymongo import monitoring +except Exception: + monitoring = None + +if not monitoring: + LumigoMongoMonitoring = None +else: + + class LumigoMongoMonitoring(monitoring.CommandListener): # type: ignore + def started(self, event): + with lumigo_safe_execute("pymongo started"): + SpansContainer.get_span().add_span( + { + "id": str(uuid.uuid4()), + "started": int(time.time() / 1000), + "databaseName": event.database_name, + "commandName": event.command_name, + "request": lumigo_dumps(event.command), + "mongoRequestId": event.request_id, + "mongoOperationId": event.operation_id, + "mongoConnectionId": event.connection_id, + } + ) + + def succeeded(self, event): + with lumigo_safe_execute("pymongo succeed"): + last_span = SpansContainer.get_span().get_last_span() + if not last_span or last_span.get("mongoRequestId") != event.request_id: + get_logger().warning("Mongo span ended without a record on its start") + return + last_span.update( + { + "ended": last_span["started"] + (event.duration_micros / 1000), + "response": lumigo_dumps(event.reply), + } + ) + + def failed(self, event): + with lumigo_safe_execute("pymongo succeed"): + last_span = SpansContainer.get_span().get_last_span() + if not last_span or last_span.get("mongoRequestId") != event.request_id: + get_logger().warning("Mongo span ended without a record on its start") + return + last_span.update( + { + "ended": last_span["started"] + (event.duration_micros / 1000), + "error": lumigo_dumps(event.failure), + } + ) + + +def wrap_pymongo(): + with lumigo_safe_execute("wrap pymogno"): + if monitoring: + get_logger().debug("wrapping pymongo") + monitoring.register(LumigoMongoMonitoring()) diff --git a/src/test/unit/wrappers/pymongo/test_pymongo_wrapper.py b/src/test/unit/wrappers/pymongo/test_pymongo_wrapper.py new file mode 100644 index 00000000..79c972fd --- /dev/null +++ b/src/test/unit/wrappers/pymongo/test_pymongo_wrapper.py @@ -0,0 +1,64 @@ +from types import SimpleNamespace + +import pytest + +from lumigo_tracer.spans_container import SpansContainer +from lumigo_tracer.wrappers.pymongo.pymongo_wrapper import LumigoMongoMonitoring + + +@pytest.fixture +def start_event(): + return SimpleNamespace( + database_name="dname", + command_name="cname", + command="cmd", + request_id="rid", + operation_id="oid", + connection_id="cid", + ) + + +@pytest.fixture +def success_event(): + return SimpleNamespace(duration_micros=5000, reply={"code": 200}, request_id="rid") + + +@pytest.fixture +def fail_event(): + return SimpleNamespace(duration_micros=5000, failure={"code": 500}, request_id="rid") + + +def test_pymongo_happy_flow(monkeypatch, start_event, success_event): + monitor = LumigoMongoMonitoring() + monitor.started(start_event) + monitor.succeeded(success_event) + + spans = SpansContainer.get_span().spans + assert len(spans) == 1 + assert spans[0]["request"] == '"cmd"' + assert spans[0]["ended"] > spans[0]["started"] + assert spans[0]["response"] == '{"code": 200}' + assert "error" not in spans[0] + + +def test_pymongo_only_start(monkeypatch, start_event): + monitor = LumigoMongoMonitoring() + monitor.started(start_event) + + spans = SpansContainer.get_span().spans + assert len(spans) == 1 + assert spans[0]["request"] == '"cmd"' + assert "duration" not in spans[0] + + +def test_pymongo_error(monkeypatch, start_event, fail_event): + monitor = LumigoMongoMonitoring() + monitor.started(start_event) + monitor.failed(fail_event) + + spans = SpansContainer.get_span().spans + assert len(spans) == 1 + assert spans[0]["request"] == '"cmd"' + assert spans[0]["ended"] > spans[0]["started"] + assert spans[0]["error"] == '{"code": 500}' + assert "response" not in spans[0] From fd8a27d9741bf26b7993d3839dde7bbb37ffc02a Mon Sep 17 00:00:00 2001 From: saartochner Date: Sun, 11 Oct 2020 10:18:11 +0300 Subject: [PATCH 10/24] wip --- src/test/unit/event/test_event_dumper.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/test/unit/event/test_event_dumper.py b/src/test/unit/event/test_event_dumper.py index 449fa0d4..3487b677 100644 --- a/src/test/unit/event/test_event_dumper.py +++ b/src/test/unit/event/test_event_dumper.py @@ -1,5 +1,4 @@ from collections import OrderedDict -import json import pytest from lumigo_tracer.event.event_dumper import ( @@ -418,7 +417,7 @@ def test_is_s3_event(s3_event): def test_parse_s3_event(s3_event): ordered_s3_event = EventDumper.dump_event(event=s3_event) - assert json.dumps(ordered_s3_event) == json.dumps( + assert ordered_s3_event == lumigo_dumps( OrderedDict( { "Records": [ @@ -448,7 +447,7 @@ def test_is_cloudfront_event(cloudfront_event): def test_parse_cloudfront_event(cloudfront_event): ordered_cloudfront_event = EventDumper.dump_event(event=cloudfront_event) - assert json.dumps(ordered_cloudfront_event) == json.dumps( + assert ordered_cloudfront_event == lumigo_dumps( OrderedDict( { "Records": [ From c5caf0952077450cea1eb64cd4f20169590c2e01 Mon Sep 17 00:00:00 2001 From: saartochner Date: Mon, 12 Oct 2020 13:45:32 +0300 Subject: [PATCH 11/24] fix Dori's CR --- src/lumigo_tracer/spans_container.py | 6 ++++ .../wrappers/pymongo/pymongo_wrapper.py | 28 ++++++++++++------- src/test/unit/test_spans_container.py | 9 ++++++ .../wrappers/pymongo/test_pymongo_wrapper.py | 25 +++++++++++++++++ .../unit/wrappers/test_no_wrapping_library.py | 12 ++++++++ 5 files changed, 70 insertions(+), 10 deletions(-) create mode 100644 src/test/unit/wrappers/test_no_wrapping_library.py diff --git a/src/lumigo_tracer/spans_container.py b/src/lumigo_tracer/spans_container.py index 8b7c4e3f..d25b9fcc 100644 --- a/src/lumigo_tracer/spans_container.py +++ b/src/lumigo_tracer/spans_container.py @@ -144,6 +144,12 @@ def get_last_span(self) -> Optional[dict]: return None return self.spans[-1] + def get_span_by_id(self, span_id: str) -> Optional[dict]: + for span in self.spans: + if span.get("id") == span_id: + return span + return None + def remove_last_span(self) -> Optional[dict]: return self.spans.pop() if self.spans else None diff --git a/src/lumigo_tracer/wrappers/pymongo/pymongo_wrapper.py b/src/lumigo_tracer/wrappers/pymongo/pymongo_wrapper.py index 876520ef..06345d0b 100644 --- a/src/lumigo_tracer/wrappers/pymongo/pymongo_wrapper.py +++ b/src/lumigo_tracer/wrappers/pymongo/pymongo_wrapper.py @@ -1,3 +1,5 @@ +from typing import Dict + import time import uuid @@ -14,11 +16,15 @@ else: class LumigoMongoMonitoring(monitoring.CommandListener): # type: ignore + request_to_span_id: Dict[str, str] = {} + def started(self, event): with lumigo_safe_execute("pymongo started"): + span_id = str(uuid.uuid4()) + LumigoMongoMonitoring.request_to_span_id[event.request_id] = span_id SpansContainer.get_span().add_span( { - "id": str(uuid.uuid4()), + "id": span_id, "started": int(time.time() / 1000), "databaseName": event.database_name, "commandName": event.command_name, @@ -31,26 +37,28 @@ def started(self, event): def succeeded(self, event): with lumigo_safe_execute("pymongo succeed"): - last_span = SpansContainer.get_span().get_last_span() - if not last_span or last_span.get("mongoRequestId") != event.request_id: + if event.request_id not in LumigoMongoMonitoring.request_to_span_id: get_logger().warning("Mongo span ended without a record on its start") return - last_span.update( + span_id = LumigoMongoMonitoring.request_to_span_id.pop(event.request_id) + span = SpansContainer.get_span().get_span_by_id(span_id) + span.update( { - "ended": last_span["started"] + (event.duration_micros / 1000), + "ended": span["started"] + (event.duration_micros / 1000), "response": lumigo_dumps(event.reply), } ) def failed(self, event): - with lumigo_safe_execute("pymongo succeed"): - last_span = SpansContainer.get_span().get_last_span() - if not last_span or last_span.get("mongoRequestId") != event.request_id: + with lumigo_safe_execute("pymongo failed"): + if event.request_id not in LumigoMongoMonitoring.request_to_span_id: get_logger().warning("Mongo span ended without a record on its start") return - last_span.update( + span_id = LumigoMongoMonitoring.request_to_span_id.pop(event.request_id) + span = SpansContainer.get_span().get_span_by_id(span_id) + span.update( { - "ended": last_span["started"] + (event.duration_micros / 1000), + "ended": span["started"] + (event.duration_micros / 1000), "error": lumigo_dumps(event.failure), } ) diff --git a/src/test/unit/test_spans_container.py b/src/test/unit/test_spans_container.py index 6e79c24f..c529842c 100644 --- a/src/test/unit/test_spans_container.py +++ b/src/test/unit/test_spans_container.py @@ -152,3 +152,12 @@ def test_get_tags_len(): SpansContainer.get_span().add_tag("k0", "v0") SpansContainer.get_span().add_tag("k1", "v1") assert SpansContainer.get_span().get_tags_len() == 2 + + +def test_get_span_by_id(): + container = SpansContainer.get_span() + container.add_span({"id": 1, "extra": "a"}) + container.add_span({"id": 2, "extra": "b"}) + container.add_span({"id": 3, "extra": "c"}) + assert SpansContainer.get_span().get_span_by_id(2)["extra"] == "b" + assert SpansContainer.get_span().get_span_by_id(5) is None diff --git a/src/test/unit/wrappers/pymongo/test_pymongo_wrapper.py b/src/test/unit/wrappers/pymongo/test_pymongo_wrapper.py index 79c972fd..243cf85a 100644 --- a/src/test/unit/wrappers/pymongo/test_pymongo_wrapper.py +++ b/src/test/unit/wrappers/pymongo/test_pymongo_wrapper.py @@ -62,3 +62,28 @@ def test_pymongo_error(monkeypatch, start_event, fail_event): assert spans[0]["ended"] > spans[0]["started"] assert spans[0]["error"] == '{"code": 500}' assert "response" not in spans[0] + + +def test_pymongo_concurrent_events(monkeypatch, start_event, success_event): + monitor = LumigoMongoMonitoring() + monitor.started(start_event) + monitor.started( + SimpleNamespace( + database_name="dname", + command_name="cname", + command="cmd", + request_id="rid2", + operation_id="oid", + connection_id="cid", + ) + ) + monitor.succeeded(success_event) + + spans = SpansContainer.get_span().spans + assert len(spans) == 2 + assert spans[0]["mongoRequestId"] == "rid" + assert spans[0]["ended"] > spans[0]["started"] + assert spans[0]["response"] == '{"code": 200}' + + assert spans[1]["mongoRequestId"] == "rid2" + assert "ended" not in spans[1] diff --git a/src/test/unit/wrappers/test_no_wrapping_library.py b/src/test/unit/wrappers/test_no_wrapping_library.py new file mode 100644 index 00000000..4253885e --- /dev/null +++ b/src/test/unit/wrappers/test_no_wrapping_library.py @@ -0,0 +1,12 @@ +import importlib + +import sys + + +def test_wrapping_without_libraries(monkeypatch): + monkeypatch.setitem(sys.modules, "pymongo", None) + + wrapper = importlib.import_module("src.lumigo_tracer.wrappers") + wrapper.wrap() # should succeed + + assert wrapper.pymongo.pymongo_wrapper.LumigoMongoMonitoring is None From df61e658ba676157ea08b8e5ab8f32e42519e3ad Mon Sep 17 00:00:00 2001 From: saartochner Date: Mon, 12 Oct 2020 13:57:54 +0300 Subject: [PATCH 12/24] fix Dori's CR --- src/test/unit/wrappers/test_no_wrapping_library.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/src/test/unit/wrappers/test_no_wrapping_library.py b/src/test/unit/wrappers/test_no_wrapping_library.py index 4253885e..07d900e9 100644 --- a/src/test/unit/wrappers/test_no_wrapping_library.py +++ b/src/test/unit/wrappers/test_no_wrapping_library.py @@ -1,12 +1,13 @@ import importlib - import sys +import lumigo_tracer + def test_wrapping_without_libraries(monkeypatch): + # remove pymongo from path and reload module monkeypatch.setitem(sys.modules, "pymongo", None) + wrapper = importlib.reload(lumigo_tracer.wrappers.pymongo.pymongo_wrapper) + assert wrapper.LumigoMongoMonitoring is None - wrapper = importlib.import_module("src.lumigo_tracer.wrappers") - wrapper.wrap() # should succeed - - assert wrapper.pymongo.pymongo_wrapper.LumigoMongoMonitoring is None + lumigo_tracer.wrappers.wrap() # should succeed From 2e83944ecd8e49f5f61b5d169e28da904ac682d1 Mon Sep 17 00:00:00 2001 From: saartochner Date: Mon, 12 Oct 2020 14:42:52 +0300 Subject: [PATCH 13/24] wip --- src/lumigo_tracer/wrappers/pymongo/pymongo_wrapper.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/lumigo_tracer/wrappers/pymongo/pymongo_wrapper.py b/src/lumigo_tracer/wrappers/pymongo/pymongo_wrapper.py index 06345d0b..7cbf68c6 100644 --- a/src/lumigo_tracer/wrappers/pymongo/pymongo_wrapper.py +++ b/src/lumigo_tracer/wrappers/pymongo/pymongo_wrapper.py @@ -17,6 +17,7 @@ class LumigoMongoMonitoring(monitoring.CommandListener): # type: ignore request_to_span_id: Dict[str, str] = {} + MONGO_SPAN = "mongoDb" def started(self, event): with lumigo_safe_execute("pymongo started"): @@ -25,6 +26,7 @@ def started(self, event): SpansContainer.get_span().add_span( { "id": span_id, + "type": self.MONGO_SPAN, "started": int(time.time() / 1000), "databaseName": event.database_name, "commandName": event.command_name, From f06a1c3f2bbe71be45eb663f6eb6d12b589578f4 Mon Sep 17 00:00:00 2001 From: saartochner Date: Mon, 12 Oct 2020 17:19:17 +0300 Subject: [PATCH 14/24] support redis --- requirements.txt | 3 +- src/lumigo_tracer/lumigo_utils.py | 2 +- src/lumigo_tracer/wrappers/__init__.py | 2 + src/lumigo_tracer/wrappers/redis/__init__.py | 0 .../wrappers/redis/redis_wrapper.py | 83 +++++++++++++++++++ .../unit/wrappers/redis/test_redis_wrapper.py | 83 +++++++++++++++++++ .../unit/wrappers/test_no_wrapping_library.py | 9 +- 7 files changed, 179 insertions(+), 3 deletions(-) create mode 100644 src/lumigo_tracer/wrappers/redis/__init__.py create mode 100644 src/lumigo_tracer/wrappers/redis/redis_wrapper.py create mode 100644 src/test/unit/wrappers/redis/test_redis_wrapper.py diff --git a/requirements.txt b/requirements.txt index 628ac38d..48b65bdc 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,4 +6,5 @@ pytest-cov==2.6.1 capturer==2.4 attrs==19.1.0 requests==2.24.0 -pymongo==3.11.0 \ No newline at end of file +pymongo==3.11.0 +redispy==3.0.0 diff --git a/src/lumigo_tracer/lumigo_utils.py b/src/lumigo_tracer/lumigo_utils.py index 19bdebb0..a125c62a 100644 --- a/src/lumigo_tracer/lumigo_utils.py +++ b/src/lumigo_tracer/lumigo_utils.py @@ -480,7 +480,7 @@ def aws_dump(d: Any, decimal_safe=False, **kwargs) -> str: def lumigo_dumps( - d: Union[bytes, str, dict, OrderedDict, list], + d: Optional[Union[bytes, str, dict, OrderedDict, list]], max_size: Optional[int] = None, regexes: Optional[Pattern[str]] = None, enforce_jsonify: bool = False, diff --git a/src/lumigo_tracer/wrappers/__init__.py b/src/lumigo_tracer/wrappers/__init__.py index cb3ef5c3..5fe934eb 100644 --- a/src/lumigo_tracer/wrappers/__init__.py +++ b/src/lumigo_tracer/wrappers/__init__.py @@ -1,5 +1,6 @@ from .http.sync_http_wrappers import wrap_http_calls from .pymongo.pymongo_wrapper import wrap_pymongo +from .redis.redis_wrapper import wrap_redis already_wrapped = False @@ -10,4 +11,5 @@ def wrap(): if not already_wrapped: wrap_http_calls() wrap_pymongo() + wrap_redis() already_wrapped = True diff --git a/src/lumigo_tracer/wrappers/redis/__init__.py b/src/lumigo_tracer/wrappers/redis/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/lumigo_tracer/wrappers/redis/redis_wrapper.py b/src/lumigo_tracer/wrappers/redis/redis_wrapper.py new file mode 100644 index 00000000..fa7c08ab --- /dev/null +++ b/src/lumigo_tracer/wrappers/redis/redis_wrapper.py @@ -0,0 +1,83 @@ +import importlib +from typing import Optional, Dict, List, Union + +import time +import uuid + +from lumigo_tracer.libs.wrapt import wrap_function_wrapper +from lumigo_tracer.lumigo_utils import lumigo_safe_execute, get_logger, lumigo_dumps +from lumigo_tracer.spans_container import SpansContainer + + +REDIS_SPAN = "redis" + + +def command_started( + command: str, request_args: Union[Dict, List[Dict]], connection_options: Optional[Dict] +): + span_id = str(uuid.uuid4()) + SpansContainer.get_span().add_span( + { + "id": span_id, + "type": REDIS_SPAN, + "started": int(time.time() / 1000), + "requestCommand": command, + "requestArgs": lumigo_dumps(request_args), + "connectionOptions": lumigo_dumps(connection_options), + } + ) + + +def command_finished(ret_val): + span = SpansContainer.get_span().get_last_span() + if not span: + get_logger().warning("Redis span ended without a record on its start") + return + span.update({"ended": int(time.time() / 1000), "response": lumigo_dumps(ret_val)}) + + +def command_failed(exception: Exception): + span = SpansContainer.get_span().get_last_span() + if not span: + get_logger().warning("Redis span ended without a record on its start") + return + span.update( + {"ended": int(time.time() / 1000), "error": exception.args[0] if exception.args else None} + ) + + +def execute_command_wrapper(func, instance, args, kwargs): + with lumigo_safe_execute("redis start"): + command = args[0] if args else None + request_args = args[1] if args and len(args) > 1 else None + connection_options = instance.connection_pool.connection_kwargs + command_started(command, request_args, connection_options) + try: + ret_val = func(*args, **kwargs) + command_finished(ret_val) + except Exception as e: + command_failed(e) + raise + + +def execute_wrapper(func, instance, args, kwargs): + with lumigo_safe_execute("redis start"): + commands = instance.command_stack + command = [cmd[0] for cmd in commands if cmd] or None + request_args = [cmd[1] for cmd in commands if cmd and len(cmd) > 1] or None + connection_options = instance.connection_pool.connection_kwargs + command_started(lumigo_dumps(command), request_args, connection_options) + try: + ret_val = func(*args, **kwargs) + command_finished(ret_val) + except Exception as e: + command_failed(e) + raise + + +def wrap_redis(): + with lumigo_safe_execute("wrap redis"): + if importlib.util.find_spec("redis"): + get_logger().debug("wrapping redis") + wrap_function_wrapper("redis.client", "execute_command", execute_command_wrapper) + wrap_function_wrapper("redis.client", "execute", execute_wrapper) diff --git a/src/test/unit/wrappers/redis/test_redis_wrapper.py b/src/test/unit/wrappers/redis/test_redis_wrapper.py new file mode 100644 index 00000000..147a8d23 --- /dev/null +++ b/src/test/unit/wrappers/redis/test_redis_wrapper.py @@ -0,0 +1,83 @@ +from types import SimpleNamespace + +import pytest + +from lumigo_tracer.spans_container import SpansContainer +from lumigo_tracer.wrappers.redis.redis_wrapper import execute_command_wrapper, execute_wrapper + + +@pytest.fixture +def instance(): + return SimpleNamespace( + connection_pool=SimpleNamespace(connection_kwargs={"Host": "bla", "Port": 5000}), + command_stack=None, + ) + + +def func(*args, **kwargs): + return True + + +def test_execute_command_wrapper_happy_flow(instance): + execute_command_wrapper(func, instance, ["SET", {"a": 1}], {}) + + spans = SpansContainer.get_span().spans + assert len(spans) == 1 + assert spans[0]["requestCommand"] == "SET" + assert spans[0]["requestArgs"] == '{"a": 1}' + assert spans[0]["ended"] >= spans[0]["started"] + assert spans[0]["response"] == "true" + assert "error" not in spans[0] + + +def test_execute_command_wrapper_failing_command(instance): + with pytest.raises(ZeroDivisionError): + execute_command_wrapper(lambda *args, **kwargs: 1 / 0, instance, ["SET", {"a": 1}], {}) + + spans = SpansContainer.get_span().spans + assert len(spans) == 1 + assert spans[0]["requestCommand"] == "SET" + assert spans[0]["ended"] >= spans[0]["started"] + assert spans[0]["error"] == "division by zero" + assert "response" not in spans[0] + + +def test_execute_command_wrapper_unexpected_params(instance): + execute_command_wrapper(func, instance, {"not": "list"}, {}) + + spans = SpansContainer.get_span().spans + assert len(spans) == 0 + + +def test_execute_wrapper_happy_flow(instance, monkeypatch): + monkeypatch.setattr(instance, "command_stack", [["SET", {"a": 1}], ["GET", "a"]]) + execute_wrapper(func, instance, [], {}) + + spans = SpansContainer.get_span().spans + assert len(spans) == 1 + assert spans[0]["requestCommand"] == '["SET", "GET"]' + assert spans[0]["requestArgs"] == '[{"a": 1}, "a"]' + assert spans[0]["ended"] >= spans[0]["started"] + assert spans[0]["response"] == "true" + assert "error" not in spans[0] + + +def test_execute_wrapper_failing_command(instance, monkeypatch): + monkeypatch.setattr(instance, "command_stack", [["SET", {"a": 1}], ["GET", "a"]]) + with pytest.raises(ZeroDivisionError): + execute_wrapper(lambda *args, **kwargs: 1 / 0, instance, [], {}) + + spans = SpansContainer.get_span().spans + assert len(spans) == 1 + assert spans[0]["requestCommand"] == '["SET", "GET"]' + assert spans[0]["ended"] >= spans[0]["started"] + assert spans[0]["error"] == "division by zero" + assert "response" not in spans[0] + + +def test_execute_wrapper_unexpected_params(instance, monkeypatch): + monkeypatch.setattr(instance, "command_stack", [{"not": "list"}]) + execute_wrapper(func, instance, [], {}) + + spans = SpansContainer.get_span().spans + assert len(spans) == 0 diff --git a/src/test/unit/wrappers/test_no_wrapping_library.py b/src/test/unit/wrappers/test_no_wrapping_library.py index 07d900e9..364a76a2 100644 --- a/src/test/unit/wrappers/test_no_wrapping_library.py +++ b/src/test/unit/wrappers/test_no_wrapping_library.py @@ -5,9 +5,16 @@ def test_wrapping_without_libraries(monkeypatch): - # remove pymongo from path and reload module + # remove library from path and reload module monkeypatch.setitem(sys.modules, "pymongo", None) wrapper = importlib.reload(lumigo_tracer.wrappers.pymongo.pymongo_wrapper) assert wrapper.LumigoMongoMonitoring is None + monkeypatch.setitem(sys.modules, "redis", None) + importlib.reload(lumigo_tracer.wrappers.redis.redis_wrapper) + lumigo_tracer.wrappers.wrap() # should succeed + + monkeypatch.undo() + importlib.reload(lumigo_tracer.wrappers.pymongo.pymongo_wrapper) + importlib.reload(lumigo_tracer.wrappers.redis.redis_wrapper) From 188bf41a9951e03a662c8a6b80d481ef59b0fc4a Mon Sep 17 00:00:00 2001 From: saartochner Date: Mon, 12 Oct 2020 19:04:53 +0300 Subject: [PATCH 15/24] wip --- src/test/unit/wrappers/test_no_wrapping_library.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/test/unit/wrappers/test_no_wrapping_library.py b/src/test/unit/wrappers/test_no_wrapping_library.py index 07d900e9..2e89c5c4 100644 --- a/src/test/unit/wrappers/test_no_wrapping_library.py +++ b/src/test/unit/wrappers/test_no_wrapping_library.py @@ -5,9 +5,12 @@ def test_wrapping_without_libraries(monkeypatch): - # remove pymongo from path and reload module + # remove library from path and reload module monkeypatch.setitem(sys.modules, "pymongo", None) wrapper = importlib.reload(lumigo_tracer.wrappers.pymongo.pymongo_wrapper) assert wrapper.LumigoMongoMonitoring is None lumigo_tracer.wrappers.wrap() # should succeed + + monkeypatch.undo() + importlib.reload(lumigo_tracer.wrappers.pymongo.pymongo_wrapper) From 03713d88b7af7fbf6f49bea8daccd5d1ccc444b1 Mon Sep 17 00:00:00 2001 From: saartochner Date: Tue, 13 Oct 2020 11:40:12 +0300 Subject: [PATCH 16/24] times --- src/lumigo_tracer/wrappers/pymongo/pymongo_wrapper.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/lumigo_tracer/wrappers/pymongo/pymongo_wrapper.py b/src/lumigo_tracer/wrappers/pymongo/pymongo_wrapper.py index 7cbf68c6..4fcbbd9d 100644 --- a/src/lumigo_tracer/wrappers/pymongo/pymongo_wrapper.py +++ b/src/lumigo_tracer/wrappers/pymongo/pymongo_wrapper.py @@ -27,7 +27,7 @@ def started(self, event): { "id": span_id, "type": self.MONGO_SPAN, - "started": int(time.time() / 1000), + "started": int(time.time() * 1000), "databaseName": event.database_name, "commandName": event.command_name, "request": lumigo_dumps(event.command), From f837d3854e02e4536fb41f761a305457f4ad9406 Mon Sep 17 00:00:00 2001 From: saartochner Date: Tue, 13 Oct 2020 13:48:58 +0300 Subject: [PATCH 17/24] IT --- .../wrappers/http/sync_http_wrappers.py | 2 +- .../wrappers/redis/redis_wrapper.py | 18 ++++++++++-------- .../unit/wrappers/redis/test_redis_wrapper.py | 10 +++++----- 3 files changed, 16 insertions(+), 14 deletions(-) diff --git a/src/lumigo_tracer/wrappers/http/sync_http_wrappers.py b/src/lumigo_tracer/wrappers/http/sync_http_wrappers.py index d10fe4e0..5a97a86d 100644 --- a/src/lumigo_tracer/wrappers/http/sync_http_wrappers.py +++ b/src/lumigo_tracer/wrappers/http/sync_http_wrappers.py @@ -218,7 +218,7 @@ def _putheader_wrapper(func, instance, args, kwargs): def wrap_http_calls(): with lumigo_safe_execute("wrap http calls"): - get_logger().debug("wrapping the http request") + get_logger().debug("wrapping http requests") wrap_function_wrapper("http.client", "HTTPConnection.send", _http_send_wrapper) wrap_function_wrapper("http.client", "HTTPConnection.request", _headers_reminder_wrapper) if importlib.util.find_spec("botocore"): diff --git a/src/lumigo_tracer/wrappers/redis/redis_wrapper.py b/src/lumigo_tracer/wrappers/redis/redis_wrapper.py index fa7c08ab..26fd762a 100644 --- a/src/lumigo_tracer/wrappers/redis/redis_wrapper.py +++ b/src/lumigo_tracer/wrappers/redis/redis_wrapper.py @@ -16,14 +16,16 @@ def command_started( command: str, request_args: Union[Dict, List[Dict]], connection_options: Optional[Dict] ): span_id = str(uuid.uuid4()) + host = (connection_options or {}).get("host") + port = (connection_options or {}).get("port") SpansContainer.get_span().add_span( { "id": span_id, "type": REDIS_SPAN, - "started": int(time.time() / 1000), + "started": int(time.time() * 1000), "requestCommand": command, "requestArgs": lumigo_dumps(request_args), - "connectionOptions": lumigo_dumps(connection_options), + "connectionOptions": {"host": host, "port": port}, } ) @@ -33,7 +35,7 @@ def command_finished(ret_val): if not span: get_logger().warning("Redis span ended without a record on its start") return - span.update({"ended": int(time.time() / 1000), "response": lumigo_dumps(ret_val)}) + span.update({"ended": int(time.time() * 1000), "response": lumigo_dumps(ret_val)}) def command_failed(exception: Exception): @@ -42,14 +44,14 @@ def command_failed(exception: Exception): get_logger().warning("Redis span ended without a record on its start") return span.update( - {"ended": int(time.time() / 1000), "error": exception.args[0] if exception.args else None} + {"ended": int(time.time() * 1000), "error": exception.args[0] if exception.args else None} ) def execute_command_wrapper(func, instance, args, kwargs): with lumigo_safe_execute("redis start"): command = args[0] if args else None - request_args = args[1] if args and len(args) > 1 else None + request_args = args[1:] if args and len(args) > 1 else None connection_options = instance.connection_pool.connection_kwargs command_started(command, request_args, connection_options) try: @@ -64,7 +66,7 @@ def execute_wrapper(func, instance, args, kwargs): with lumigo_safe_execute("redis start"): commands = instance.command_stack command = [cmd[0] for cmd in commands if cmd] or None - request_args = [cmd[1] for cmd in commands if cmd and len(cmd) > 1] or None + request_args = [cmd[1:] for cmd in commands if cmd and len(cmd) > 1] or None connection_options = instance.connection_pool.connection_kwargs command_started(lumigo_dumps(command), request_args, connection_options) try: @@ -79,5 +81,5 @@ def wrap_redis(): with lumigo_safe_execute("wrap redis"): if importlib.util.find_spec("redis"): get_logger().debug("wrapping redis") - wrap_function_wrapper("redis.client", "execute_command", execute_command_wrapper) - wrap_function_wrapper("redis.client", "execute", execute_wrapper) + wrap_function_wrapper("redis.client", "Redis.execute_command", execute_command_wrapper) + wrap_function_wrapper("redis.client", "Pipeline.execute", execute_wrapper) diff --git a/src/test/unit/wrappers/redis/test_redis_wrapper.py b/src/test/unit/wrappers/redis/test_redis_wrapper.py index 147a8d23..fde8b6f5 100644 --- a/src/test/unit/wrappers/redis/test_redis_wrapper.py +++ b/src/test/unit/wrappers/redis/test_redis_wrapper.py @@ -9,8 +9,7 @@ @pytest.fixture def instance(): return SimpleNamespace( - connection_pool=SimpleNamespace(connection_kwargs={"Host": "bla", "Port": 5000}), - command_stack=None, + connection_pool=SimpleNamespace(connection_kwargs={"host": "lumigo"}), command_stack=None ) @@ -19,12 +18,13 @@ def func(*args, **kwargs): def test_execute_command_wrapper_happy_flow(instance): - execute_command_wrapper(func, instance, ["SET", {"a": 1}], {}) + execute_command_wrapper(func, instance, ["SET", {"a": 1}, "b"], {}) spans = SpansContainer.get_span().spans assert len(spans) == 1 assert spans[0]["requestCommand"] == "SET" - assert spans[0]["requestArgs"] == '{"a": 1}' + assert spans[0]["requestArgs"] == '[{"a": 1}, "b"]' + assert spans[0]["connectionOptions"] == {"host": "lumigo", "port": None} assert spans[0]["ended"] >= spans[0]["started"] assert spans[0]["response"] == "true" assert "error" not in spans[0] @@ -56,7 +56,7 @@ def test_execute_wrapper_happy_flow(instance, monkeypatch): spans = SpansContainer.get_span().spans assert len(spans) == 1 assert spans[0]["requestCommand"] == '["SET", "GET"]' - assert spans[0]["requestArgs"] == '[{"a": 1}, "a"]' + assert spans[0]["requestArgs"] == '[[{"a": 1}], ["a"]]' assert spans[0]["ended"] >= spans[0]["started"] assert spans[0]["response"] == "true" assert "error" not in spans[0] From 64f8b1d4ca1671538f2cb5070a5d43cf2800c091 Mon Sep 17 00:00:00 2001 From: saartochner Date: Sun, 18 Oct 2020 16:02:40 +0300 Subject: [PATCH 18/24] support sqlalchemy --- requirements.txt | 3 +- src/lumigo_tracer/wrappers/__init__.py | 2 + src/lumigo_tracer/wrappers/sql/__init__.py | 0 .../wrappers/sql/sqlalchemy_wrapper.py | 78 +++++++++++++++++++ .../wrappers/sql/test_sqlalchemt_wrapper.py | 63 +++++++++++++++ .../unit/wrappers/test_no_wrapping_library.py | 4 + 6 files changed, 149 insertions(+), 1 deletion(-) create mode 100644 src/lumigo_tracer/wrappers/sql/__init__.py create mode 100644 src/lumigo_tracer/wrappers/sql/sqlalchemy_wrapper.py create mode 100644 src/test/unit/wrappers/sql/test_sqlalchemt_wrapper.py diff --git a/requirements.txt b/requirements.txt index 628ac38d..5d70ed16 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,4 +6,5 @@ pytest-cov==2.6.1 capturer==2.4 attrs==19.1.0 requests==2.24.0 -pymongo==3.11.0 \ No newline at end of file +pymongo==3.11.0 +sqlalchemy==1.3.20 diff --git a/src/lumigo_tracer/wrappers/__init__.py b/src/lumigo_tracer/wrappers/__init__.py index cb3ef5c3..9f51c8ba 100644 --- a/src/lumigo_tracer/wrappers/__init__.py +++ b/src/lumigo_tracer/wrappers/__init__.py @@ -1,5 +1,6 @@ from .http.sync_http_wrappers import wrap_http_calls from .pymongo.pymongo_wrapper import wrap_pymongo +from .sql.sqlalchemy_wrapper import wrap_sqlalchemy already_wrapped = False @@ -10,4 +11,5 @@ def wrap(): if not already_wrapped: wrap_http_calls() wrap_pymongo() + wrap_sqlalchemy() already_wrapped = True diff --git a/src/lumigo_tracer/wrappers/sql/__init__.py b/src/lumigo_tracer/wrappers/sql/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/lumigo_tracer/wrappers/sql/sqlalchemy_wrapper.py b/src/lumigo_tracer/wrappers/sql/sqlalchemy_wrapper.py new file mode 100644 index 00000000..4eb9c29e --- /dev/null +++ b/src/lumigo_tracer/wrappers/sql/sqlalchemy_wrapper.py @@ -0,0 +1,78 @@ +import importlib + +import time +import uuid + +from lumigo_tracer.libs.wrapt import wrap_function_wrapper +from lumigo_tracer.lumigo_utils import lumigo_safe_execute, get_logger, lumigo_dumps +from lumigo_tracer.spans_container import SpansContainer + +try: + from sqlalchemy.event import listen +except Exception: + listen = None + + +SQL_SPAN = "mySql" + + +def _before_cursor_execute(conn, cursor, statement, parameters, context, executemany): + SpansContainer.get_span().add_span( + { + "id": str(uuid.uuid4()), + "type": SQL_SPAN, + "started": int(time.time() * 1000), + "connectionParameters": { + "host": conn.engine.url.host, + "port": conn.engine.url.port, + "database": conn.engine.url.database, + "user": conn.engine.url.username, + }, + "query": statement, + "values": lumigo_dumps(parameters), + } + ) + + +def _after_cursor_execute(conn, cursor, statement, parameters, context, executemany): + span = SpansContainer.get_span().get_last_span() + if not span: + get_logger().warning("Redis span ended without a record on its start") + return + span.update({"ended": int(time.time() * 1000), "response": ""}) + + +def _handle_error(context): + span = SpansContainer.get_span().get_last_span() + if not span: + get_logger().warning("Redis span ended without a record on its start") + return + span.update( + { + "ended": int(time.time() * 1000), + "error": lumigo_dumps( + { + "type": context.original_exception.__class__.__name__, + "args": context.original_exception.args, + } + ), + } + ) + + +def execute_wrapper(func, instance, args, kwargs): + result = func(*args, **kwargs) + with lumigo_safe_execute("sqlalchemy: listen to engine"): + listen(result, "before_cursor_execute", _before_cursor_execute) + listen(result, "after_cursor_execute", _after_cursor_execute) + listen(result, "handle_error", _handle_error) + return result + + +def wrap_sqlalchemy(): + with lumigo_safe_execute("wrap sqlalchemy"): + if importlib.util.find_spec("sqlalchemy"): + get_logger().debug("wrapping sqlalchemy") + wrap_function_wrapper( + "sqlalchemy.engine.strategies", "DefaultEngineStrategy.create", execute_wrapper + ) diff --git a/src/test/unit/wrappers/sql/test_sqlalchemt_wrapper.py b/src/test/unit/wrappers/sql/test_sqlalchemt_wrapper.py new file mode 100644 index 00000000..d04f7447 --- /dev/null +++ b/src/test/unit/wrappers/sql/test_sqlalchemt_wrapper.py @@ -0,0 +1,63 @@ +import pytest +from sqlalchemy.exc import OperationalError + +from lumigo_tracer.spans_container import SpansContainer +from lumigo_tracer.tracer import lumigo_tracer +from sqlalchemy import create_engine, Table, Column, Integer, String, MetaData +from sqlalchemy.sql import select + + +md = MetaData() +Users = Table("users", md, Column("id", Integer, primary_key=True), Column("name", String)) + + +@pytest.fixture +def db(tmp_path): + path = tmp_path / "file.db" + engine = create_engine(fr"sqlite:///{path}") + md.create_all(engine) + yield f"sqlite:///{path}" + + +def test_happy_flow(context, db): + @lumigo_tracer(token="123") + def lambda_test_function(event, context): + engine = create_engine(db) + conn = engine.connect() + conn.execute(Users.insert().values(name="saart")) + result = conn.execute(select([Users])) + return result.fetchone() + + assert lambda_test_function({}, context) == (1, "saart") + http_spans = SpansContainer.get_span().spans + + assert len(http_spans) == 2 + assert http_spans[0]["query"] == "INSERT INTO users (name) VALUES (?)" + assert http_spans[0]["values"] == '["saart"]' + assert http_spans[0]["ended"] >= http_spans[0]["started"] + + assert http_spans[1]["query"] == "SELECT users.id, users.name \nFROM users" + assert http_spans[1]["values"] == "[]" + assert http_spans[0]["ended"] >= http_spans[0]["started"] + + +def test_non_existing_table(context, db): + @lumigo_tracer(token="123") + def lambda_test_function(event, context): + others = Table("others", md, Column("id", Integer, primary_key=True)) + engine = create_engine(db) + conn = engine.connect() + result = conn.execute(select([others])) + return result.fetchone() + + with pytest.raises(OperationalError): + lambda_test_function({}, context) + + http_spans = SpansContainer.get_span().spans + + assert len(http_spans) == 1 + assert http_spans[0]["query"] == "SELECT others.id \nFROM others" + assert ( + http_spans[0]["error"] == '{"type": "OperationalError", "args": ["no such table: others"]}' + ) + assert http_spans[0]["ended"] >= http_spans[0]["started"] diff --git a/src/test/unit/wrappers/test_no_wrapping_library.py b/src/test/unit/wrappers/test_no_wrapping_library.py index 2e89c5c4..8b985955 100644 --- a/src/test/unit/wrappers/test_no_wrapping_library.py +++ b/src/test/unit/wrappers/test_no_wrapping_library.py @@ -10,7 +10,11 @@ def test_wrapping_without_libraries(monkeypatch): wrapper = importlib.reload(lumigo_tracer.wrappers.pymongo.pymongo_wrapper) assert wrapper.LumigoMongoMonitoring is None + monkeypatch.setitem(sys.modules, "sqlalchemy", None) + importlib.reload(lumigo_tracer.wrappers.sql.sqlalchemy_wrapper) + lumigo_tracer.wrappers.wrap() # should succeed monkeypatch.undo() importlib.reload(lumigo_tracer.wrappers.pymongo.pymongo_wrapper) + importlib.reload(lumigo_tracer.wrappers.sql.sqlalchemy_wrapper) From 578eabf906aa37419e88d5f99887843d74e22ca7 Mon Sep 17 00:00:00 2001 From: saartochner Date: Sun, 18 Oct 2020 17:06:58 +0300 Subject: [PATCH 19/24] wip --- .../wrappers/sql/sqlalchemy_wrapper.py | 4 ++-- .../wrappers/sql/test_sqlalchemt_wrapper.py | 21 ++++++++++++++++--- 2 files changed, 20 insertions(+), 5 deletions(-) diff --git a/src/lumigo_tracer/wrappers/sql/sqlalchemy_wrapper.py b/src/lumigo_tracer/wrappers/sql/sqlalchemy_wrapper.py index 4eb9c29e..0b0c83f8 100644 --- a/src/lumigo_tracer/wrappers/sql/sqlalchemy_wrapper.py +++ b/src/lumigo_tracer/wrappers/sql/sqlalchemy_wrapper.py @@ -23,12 +23,12 @@ def _before_cursor_execute(conn, cursor, statement, parameters, context, execute "type": SQL_SPAN, "started": int(time.time() * 1000), "connectionParameters": { - "host": conn.engine.url.host, + "host": conn.engine.url.host or conn.engine.url.database, "port": conn.engine.url.port, "database": conn.engine.url.database, "user": conn.engine.url.username, }, - "query": statement, + "query": lumigo_dumps(statement), "values": lumigo_dumps(parameters), } ) diff --git a/src/test/unit/wrappers/sql/test_sqlalchemt_wrapper.py b/src/test/unit/wrappers/sql/test_sqlalchemt_wrapper.py index d04f7447..62c88d54 100644 --- a/src/test/unit/wrappers/sql/test_sqlalchemt_wrapper.py +++ b/src/test/unit/wrappers/sql/test_sqlalchemt_wrapper.py @@ -1,4 +1,5 @@ import pytest +from lumigo_tracer.lumigo_utils import DEFAULT_MAX_ENTRY_SIZE from sqlalchemy.exc import OperationalError from lumigo_tracer.spans_container import SpansContainer @@ -32,11 +33,11 @@ def lambda_test_function(event, context): http_spans = SpansContainer.get_span().spans assert len(http_spans) == 2 - assert http_spans[0]["query"] == "INSERT INTO users (name) VALUES (?)" + assert http_spans[0]["query"] == '"INSERT INTO users (name) VALUES (?)"' assert http_spans[0]["values"] == '["saart"]' assert http_spans[0]["ended"] >= http_spans[0]["started"] - assert http_spans[1]["query"] == "SELECT users.id, users.name \nFROM users" + assert http_spans[1]["query"] == '"SELECT users.id, users.name \\nFROM users"' assert http_spans[1]["values"] == "[]" assert http_spans[0]["ended"] >= http_spans[0]["started"] @@ -56,8 +57,22 @@ def lambda_test_function(event, context): http_spans = SpansContainer.get_span().spans assert len(http_spans) == 1 - assert http_spans[0]["query"] == "SELECT others.id \nFROM others" + assert http_spans[0]["query"] == '"SELECT others.id \\nFROM others"' assert ( http_spans[0]["error"] == '{"type": "OperationalError", "args": ["no such table: others"]}' ) assert http_spans[0]["ended"] >= http_spans[0]["started"] + + +def test_pruning_long_strings(context, db): + @lumigo_tracer(token="123") + def lambda_test_function(event, context): + engine = create_engine(db) + conn = engine.connect() + conn.execute(Users.insert().values(name="a" * (DEFAULT_MAX_ENTRY_SIZE * 5))) + + lambda_test_function({}, context) + http_spans = SpansContainer.get_span().spans + + assert len(http_spans) == 1 + assert len(http_spans[0]["values"]) <= DEFAULT_MAX_ENTRY_SIZE * 2 From 1a81a51551bd2df4866b1cec7b34a6f26ba904e5 Mon Sep 17 00:00:00 2001 From: saartochner Date: Mon, 19 Oct 2020 18:22:03 +0300 Subject: [PATCH 20/24] wip --- src/lumigo_tracer/wrappers/__init__.py | 4 ++-- src/test/unit/wrappers/test_no_wrapping_library.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/lumigo_tracer/wrappers/__init__.py b/src/lumigo_tracer/wrappers/__init__.py index 5fe934eb..4719318e 100644 --- a/src/lumigo_tracer/wrappers/__init__.py +++ b/src/lumigo_tracer/wrappers/__init__.py @@ -6,9 +6,9 @@ already_wrapped = False -def wrap(): +def wrap(force: bool = False): global already_wrapped - if not already_wrapped: + if force or not already_wrapped: wrap_http_calls() wrap_pymongo() wrap_redis() diff --git a/src/test/unit/wrappers/test_no_wrapping_library.py b/src/test/unit/wrappers/test_no_wrapping_library.py index 364a76a2..b0e6fd57 100644 --- a/src/test/unit/wrappers/test_no_wrapping_library.py +++ b/src/test/unit/wrappers/test_no_wrapping_library.py @@ -13,7 +13,7 @@ def test_wrapping_without_libraries(monkeypatch): monkeypatch.setitem(sys.modules, "redis", None) importlib.reload(lumigo_tracer.wrappers.redis.redis_wrapper) - lumigo_tracer.wrappers.wrap() # should succeed + lumigo_tracer.wrappers.wrap(force=True) # should succeed monkeypatch.undo() importlib.reload(lumigo_tracer.wrappers.pymongo.pymongo_wrapper) From dc34b79c3a128a865ca36374680d72842ada309b Mon Sep 17 00:00:00 2001 From: saartochner Date: Mon, 19 Oct 2020 18:37:05 +0300 Subject: [PATCH 21/24] wip --- src/lumigo_tracer/wrappers/__init__.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/lumigo_tracer/wrappers/__init__.py b/src/lumigo_tracer/wrappers/__init__.py index 4719318e..0d37658b 100644 --- a/src/lumigo_tracer/wrappers/__init__.py +++ b/src/lumigo_tracer/wrappers/__init__.py @@ -8,8 +8,9 @@ def wrap(force: bool = False): global already_wrapped - if force or not already_wrapped: + if not already_wrapped: wrap_http_calls() + if force or not already_wrapped: wrap_pymongo() wrap_redis() already_wrapped = True From 8cc3c649e60506b6bd9472997734ba4519f79ae0 Mon Sep 17 00:00:00 2001 From: saartochner Date: Tue, 20 Oct 2020 13:55:58 +0300 Subject: [PATCH 22/24] wip --- src/lumigo_tracer/lumigo_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/lumigo_tracer/lumigo_utils.py b/src/lumigo_tracer/lumigo_utils.py index c13e4bf9..b958366a 100644 --- a/src/lumigo_tracer/lumigo_utils.py +++ b/src/lumigo_tracer/lumigo_utils.py @@ -494,7 +494,7 @@ def aws_dump(d: Any, decimal_safe=False, **kwargs) -> str: def lumigo_dumps( - d: Optional[Union[bytes, str, dict, OrderedDict, list, None]], + d: Union[bytes, str, dict, OrderedDict, list, None], max_size: Optional[int] = None, regexes: Optional[Pattern[str]] = None, enforce_jsonify: bool = False, From 8b6430f4d1a9d62a48b46f48a915794fcc51da45 Mon Sep 17 00:00:00 2001 From: saartochner Date: Tue, 20 Oct 2020 13:57:43 +0300 Subject: [PATCH 23/24] wip --- .../wrappers/sql/sqlalchemy_wrapper.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/src/lumigo_tracer/wrappers/sql/sqlalchemy_wrapper.py b/src/lumigo_tracer/wrappers/sql/sqlalchemy_wrapper.py index 0b0c83f8..d4e16d72 100644 --- a/src/lumigo_tracer/wrappers/sql/sqlalchemy_wrapper.py +++ b/src/lumigo_tracer/wrappers/sql/sqlalchemy_wrapper.py @@ -1,10 +1,13 @@ import importlib - -import time import uuid from lumigo_tracer.libs.wrapt import wrap_function_wrapper -from lumigo_tracer.lumigo_utils import lumigo_safe_execute, get_logger, lumigo_dumps +from lumigo_tracer.lumigo_utils import ( + lumigo_safe_execute, + get_logger, + lumigo_dumps, + get_current_ms_time, +) from lumigo_tracer.spans_container import SpansContainer try: @@ -21,7 +24,7 @@ def _before_cursor_execute(conn, cursor, statement, parameters, context, execute { "id": str(uuid.uuid4()), "type": SQL_SPAN, - "started": int(time.time() * 1000), + "started": get_current_ms_time(), "connectionParameters": { "host": conn.engine.url.host or conn.engine.url.database, "port": conn.engine.url.port, @@ -39,7 +42,7 @@ def _after_cursor_execute(conn, cursor, statement, parameters, context, executem if not span: get_logger().warning("Redis span ended without a record on its start") return - span.update({"ended": int(time.time() * 1000), "response": ""}) + span.update({"ended": get_current_ms_time(), "response": ""}) def _handle_error(context): @@ -49,7 +52,7 @@ def _handle_error(context): return span.update( { - "ended": int(time.time() * 1000), + "ended": get_current_ms_time(), "error": lumigo_dumps( { "type": context.original_exception.__class__.__name__, From f773f262fec6fa8af29760e34724013f0bb25e3d Mon Sep 17 00:00:00 2001 From: saartochner Date: Tue, 20 Oct 2020 15:39:04 +0300 Subject: [PATCH 24/24] fix Nirhod's CR --- .../wrappers/sql/sqlalchemy_wrapper.py | 75 ++++++++++--------- ..._wrapper.py => test_sqlalchemy_wrapper.py} | 23 +++++- 2 files changed, 59 insertions(+), 39 deletions(-) rename src/test/unit/wrappers/sql/{test_sqlalchemt_wrapper.py => test_sqlalchemy_wrapper.py} (85%) diff --git a/src/lumigo_tracer/wrappers/sql/sqlalchemy_wrapper.py b/src/lumigo_tracer/wrappers/sql/sqlalchemy_wrapper.py index d4e16d72..d4aff570 100644 --- a/src/lumigo_tracer/wrappers/sql/sqlalchemy_wrapper.py +++ b/src/lumigo_tracer/wrappers/sql/sqlalchemy_wrapper.py @@ -20,47 +20,50 @@ def _before_cursor_execute(conn, cursor, statement, parameters, context, executemany): - SpansContainer.get_span().add_span( - { - "id": str(uuid.uuid4()), - "type": SQL_SPAN, - "started": get_current_ms_time(), - "connectionParameters": { - "host": conn.engine.url.host or conn.engine.url.database, - "port": conn.engine.url.port, - "database": conn.engine.url.database, - "user": conn.engine.url.username, - }, - "query": lumigo_dumps(statement), - "values": lumigo_dumps(parameters), - } - ) + with lumigo_safe_execute("handle sqlalchemy before execute"): + SpansContainer.get_span().add_span( + { + "id": str(uuid.uuid4()), + "type": SQL_SPAN, + "started": get_current_ms_time(), + "connectionParameters": { + "host": conn.engine.url.host or conn.engine.url.database, + "port": conn.engine.url.port, + "database": conn.engine.url.database, + "user": conn.engine.url.username, + }, + "query": lumigo_dumps(statement), + "values": lumigo_dumps(parameters), + } + ) def _after_cursor_execute(conn, cursor, statement, parameters, context, executemany): - span = SpansContainer.get_span().get_last_span() - if not span: - get_logger().warning("Redis span ended without a record on its start") - return - span.update({"ended": get_current_ms_time(), "response": ""}) + with lumigo_safe_execute("handle sqlalchemy after execute"): + span = SpansContainer.get_span().get_last_span() + if not span: + get_logger().warning("Redis span ended without a record on its start") + return + span.update({"ended": get_current_ms_time(), "response": ""}) def _handle_error(context): - span = SpansContainer.get_span().get_last_span() - if not span: - get_logger().warning("Redis span ended without a record on its start") - return - span.update( - { - "ended": get_current_ms_time(), - "error": lumigo_dumps( - { - "type": context.original_exception.__class__.__name__, - "args": context.original_exception.args, - } - ), - } - ) + with lumigo_safe_execute("handle sqlalchemy error"): + span = SpansContainer.get_span().get_last_span() + if not span: + get_logger().warning("Redis span ended without a record on its start") + return + span.update( + { + "ended": get_current_ms_time(), + "error": lumigo_dumps( + { + "type": context.original_exception.__class__.__name__, + "args": context.original_exception.args, + } + ), + } + ) def execute_wrapper(func, instance, args, kwargs): @@ -74,7 +77,7 @@ def execute_wrapper(func, instance, args, kwargs): def wrap_sqlalchemy(): with lumigo_safe_execute("wrap sqlalchemy"): - if importlib.util.find_spec("sqlalchemy"): + if importlib.util.find_spec("sqlalchemy") and listen: get_logger().debug("wrapping sqlalchemy") wrap_function_wrapper( "sqlalchemy.engine.strategies", "DefaultEngineStrategy.create", execute_wrapper diff --git a/src/test/unit/wrappers/sql/test_sqlalchemt_wrapper.py b/src/test/unit/wrappers/sql/test_sqlalchemy_wrapper.py similarity index 85% rename from src/test/unit/wrappers/sql/test_sqlalchemt_wrapper.py rename to src/test/unit/wrappers/sql/test_sqlalchemy_wrapper.py index 62c88d54..b3e0816c 100644 --- a/src/test/unit/wrappers/sql/test_sqlalchemt_wrapper.py +++ b/src/test/unit/wrappers/sql/test_sqlalchemy_wrapper.py @@ -1,11 +1,14 @@ +import uuid + import pytest -from lumigo_tracer.lumigo_utils import DEFAULT_MAX_ENTRY_SIZE + from sqlalchemy.exc import OperationalError +from sqlalchemy import create_engine, Table, Column, Integer, String, MetaData +from sqlalchemy.sql import select +from lumigo_tracer.lumigo_utils import DEFAULT_MAX_ENTRY_SIZE from lumigo_tracer.spans_container import SpansContainer from lumigo_tracer.tracer import lumigo_tracer -from sqlalchemy import create_engine, Table, Column, Integer, String, MetaData -from sqlalchemy.sql import select md = MetaData() @@ -76,3 +79,17 @@ def lambda_test_function(event, context): assert len(http_spans) == 1 assert len(http_spans[0]["values"]) <= DEFAULT_MAX_ENTRY_SIZE * 2 + + +def test_exception_in_wrapper(context, db, monkeypatch): + @lumigo_tracer(token="123") + def lambda_test_function(event, context): + engine = create_engine(db) + conn = engine.connect() + conn.execute(Users.insert().values(name="a" * (DEFAULT_MAX_ENTRY_SIZE * 5))) + + monkeypatch.setattr(uuid, "uuid4", lambda: 1 / 0) + + lambda_test_function({}, context) # No exception + + assert not SpansContainer.get_span().spans