RD 3761 - non-http libraries preparations (#150)

* split event and parsers
* split sync_hook to tracer and http hooks
* refactor http hooks to be a wrapper
* refactor wrap once
* increase coverage

Author: saartochner-lumigo

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

src/lumigo_tracer/sync_http/handler.py renamed to src/lumigo_tracer/auto_instrument_handler.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 []

src/lumigo_tracer/auto_tag/auto_tag_event.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",
@@ -77,7 +77,7 @@ def is_supported(event) -> bool:
 
     @staticmethod
     @abstractmethod
-    def parse(event) -> Dict:
+    def parse(event) -> OrderedDict:
         raise NotImplementedError()
 
 
@@ -87,7 +87,7 @@ def is_supported(event) -> bool:
         return safe_get(event, ["Records", 0, "eventSource"]) == "aws:s3"
 
     @staticmethod
-    def parse(event) -> Dict:
+    def parse(event) -> OrderedDict:
         new_event: OrderedDict = OrderedDict()
         new_event["Records"] = []
 
@@ -116,7 +116,7 @@ def is_supported(event) -> bool:
         return bool(safe_get(event, ["Records", 0, "cf", "config", "distributionId"], {}))
 
     @staticmethod
-    def parse(event) -> Dict:
+    def parse(event) -> OrderedDict:
         new_event: OrderedDict = OrderedDict()
         new_event["Records"] = []
 
@@ -144,7 +144,7 @@ def is_supported(event) -> bool:
         return is_api_gw_event(event=event)
 
     @staticmethod
-    def parse(event) -> Dict:
+    def parse(event) -> OrderedDict:
         new_event: OrderedDict = OrderedDict()
         # Add order keys
         for order_key in API_GW_KEYS_ORDER:
@@ -173,7 +173,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
@@ -192,7 +192,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
@@ -205,9 +205,9 @@ 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(),
@@ -218,10 +218,10 @@ def parse_event(event: Dict, handlers: List[EventParseHandler] = None):
         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)

src/lumigo_tracer/parsers/__init__.py renamed to src/lumigo_tracer/event/__init__.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: <host>, 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]}

src/lumigo_tracer/parsers/event_parser.py renamed to src/lumigo_tracer/event/event_dumper.py

@@ -5,10 +5,10 @@
 from datetime import datetime
 
 from lumigo_tracer.extension.extension_utils import get_current_bandwidth, get_extension_logger
-from lumigo_tracer import utils
+from lumigo_tracer import lumigo_utils
 from lumigo_tracer.extension.lambda_service import LambdaService
 from lumigo_tracer.extension.sampler import Sampler
-from lumigo_tracer.utils import lumigo_safe_execute
+from lumigo_tracer.lumigo_utils import lumigo_safe_execute
 
 SPAN_TYPE = "extensionExecutionEnd"
 
@@ -49,7 +49,7 @@ def shutdown(self):
 
     def _finish_previous_invocation(self, current_bandwidth: Optional[int]):
         self.sampler.stop_sampling()
-        token = os.environ.get(utils.LUMIGO_TOKEN_KEY)
+        token = os.environ.get(lumigo_utils.LUMIGO_TOKEN_KEY)
         if not token:
             get_extension_logger().warning(
                 f"Skip sending data: No token was found. Request id: {self.request_id}"
@@ -71,4 +71,4 @@ def _finish_previous_invocation(self, current_bandwidth: Optional[int]):
             networkBytesUsed=current_bandwidth - self.bandwidth,
             cpuUsageTime=[s.dump() for s in self.sampler.get_samples()],
         )
-        utils.report_json(os.environ.get("AWS_REGION", "us-east-1"), msgs=[asdict(span)])
+        lumigo_utils.report_json(os.environ.get("AWS_REGION", "us-east-1"), msgs=[asdict(span)])

src/lumigo_tracer/event/event_trigger.py

@@ -3,7 +3,7 @@
 import urllib.request
 from typing import Optional, Dict
 
-from lumigo_tracer.utils import get_logger, lumigo_safe_execute
+from lumigo_tracer.lumigo_utils import get_logger, lumigo_safe_execute
 
 
 def get_current_cpu_time() -> Optional[int]:

src/lumigo_tracer/extension/extension.py

@@ -5,7 +5,7 @@
 from lumigo_tracer.extension.extension_utils import get_extension_logger
 from lumigo_tracer.extension.extension import LumigoExtension
 from lumigo_tracer.extension.lambda_service import LambdaService
-from lumigo_tracer.utils import is_kill_switch_on, lumigo_safe_execute, config
+from lumigo_tracer.lumigo_utils import is_kill_switch_on, lumigo_safe_execute, config
 
 STOP_EXTENSION_KEY = "LUMIGO_EXTENSION_STOP"
 LUMIGO_EXTENSION_NAME = "lumigo"

src/lumigo_tracer/extension/extension_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:
@@ -471,7 +485,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