Allow this library to be used as a general parsing library (#268)

* move the tracer-specific logic into a well-marked module and don't run code in global time if not lambda

* support ReceiveMessage trigger in sqs parser

Author: saartochner-lumigo

.secrets.baseline

@@ -115,14 +115,14 @@
         "filename": "src/lumigo_tracer/lumigo_utils.py",
         "hashed_secret": "116b0dcbb1dfb3f19c902ebfc29f35d908cec6a4",
         "is_verified": false,
-        "line_number": 69
+        "line_number": 48
       },
       {
         "type": "Secret Keyword",
         "filename": "src/lumigo_tracer/lumigo_utils.py",
         "hashed_secret": "2a93870f6e6a2158cb9631b349e75d8f65fd81af",
         "is_verified": false,
-        "line_number": 70
+        "line_number": 49
       }
     ],
     "src/test/unit/event/test_event_dumper.py": [
@@ -168,14 +168,14 @@
         "filename": "src/test/unit/event/test_event_trigger.py",
         "hashed_secret": "885bb9903f72e004ff2974807b70e7c970d3e6d5",
         "is_verified": false,
-        "line_number": 567
+        "line_number": 612
       },
       {
         "type": "Hex High Entropy String",
         "filename": "src/test/unit/event/test_event_trigger.py",
         "hashed_secret": "3fae06dc55a618caed1d794dcd512bfe7e76c9f1",
         "is_verified": false,
-        "line_number": 612
+        "line_number": 657
       }
     ],
     "src/test/unit/test_lumigo_utils.py": [
@@ -184,7 +184,7 @@
         "filename": "src/test/unit/test_lumigo_utils.py",
         "hashed_secret": "f32b67c7e26342af42efabc674d441dca0a281c5",
         "is_verified": false,
-        "line_number": 263
+        "line_number": 199
       }
     ],
     "src/test/unit/test_tracer.py": [
@@ -220,5 +220,5 @@
       }
     ]
   },
-  "generated_at": "2022-11-08T16:26:37Z"
+  "generated_at": "2023-01-22T10:15:06Z"
 }

src/lumigo_tracer/__init__.py

@@ -1,4 +1,4 @@
-from .tracer import lumigo_tracer, LumigoChalice  # noqa
+from lumigo_tracer.lambda_tracer.tracer import lumigo_tracer, LumigoChalice  # noqa
 from .user_utils import (  # noqa
     report_error,
     add_execution_tag,
@@ -9,3 +9,6 @@
     error,
 )
 from .auto_instrument_handler import _handler  # noqa
+from .lambda_tracer.global_scope_exec import global_scope_exec
+
+global_scope_exec()

src/lumigo_tracer/auto_instrument_handler.py

@@ -1,5 +1,7 @@
 import os
 
+from lumigo_tracer.lumigo_utils import is_aws_environment
+
 try:
     # Try to import AWS's current _get_handler logic
     from bootstrap import _get_handler as aws_get_handler
@@ -27,8 +29,15 @@ def _handler(*args, **kwargs):  # type: ignore[no-untyped-def]
     return original_handler(*args, **kwargs)
 
 
-try:
-    # import handler during runtime initialization, as usual.
-    get_original_handler()
-except Exception:
-    pass
+def prefetch_handler_import() -> None:
+    """
+    This function imports the handler.
+    When we call it in the global scope, it will be executed during the lambda initialization,
+        thus will mimic the usual behavior.
+    """
+    if not is_aws_environment():
+        return
+    try:
+        get_original_handler()
+    except Exception:
+        pass

src/lumigo_tracer/event/event_trigger.py

@@ -6,8 +6,8 @@
 from lumigo_tracer.lumigo_utils import Configuration, get_logger
 
 
-def _recursive_parse_trigger_by(
-    message: Dict[Any, Any], parent_id: Optional[str], level: int
+def recursive_parse_trigger(
+    message: Dict[Any, Any], parent_id: Optional[str] = None, level: int = 0
 ) -> List[TriggerType]:
     triggers = []
     if level >= Configuration.chained_services_max_depth:
@@ -27,7 +27,7 @@ def _recursive_parse_trigger_by(
                 if INNER_MESSAGES_MAGIC_PATTERN.search(sub_message):
                     # We want to load only relevant messages, so first run a quick scan
                     triggers.extend(
-                        _recursive_parse_trigger_by(
+                        recursive_parse_trigger(
                             json.loads(sub_message), parent_id=current_trigger_id, level=level + 1
                         )
                     )
@@ -36,4 +36,4 @@ def _recursive_parse_trigger_by(
 
 
 def parse_triggers(event: Dict[Any, Any]) -> List[Dict[Any, Any]]:
-    return _recursive_parse_trigger_by(event, parent_id=None, level=0)
+    return recursive_parse_trigger(event, parent_id=None, level=0)

src/lumigo_tracer/event/trigger_parsing/sqs_parser.py

@@ -10,32 +10,40 @@
 class SqsEventTriggerParser(EventTriggerParser):
     @staticmethod
     def _should_handle(event: Dict[Any, Any]) -> bool:
-        return bool(event.get("Records", [{}])[0].get("eventSource") == "aws:sqs")
+        return bool(event.get("Records", [{}])[0].get("eventSource") == "aws:sqs") or bool(
+            event.get("service_name") == "sqs" and event.get("operation_name") == "ReceiveMessage"
+        )
+
+    @staticmethod
+    def _get_messages(event: Dict[Any, Any]) -> List[Dict[Any, Any]]:
+        return event.get("Records", []) + event.get("Messages", [])  # type: ignore
 
     @staticmethod
     def handle(event: Dict[Any, Any], target_id: Optional[str]) -> TriggerType:
+        messages = SqsEventTriggerParser._get_messages(event)
         message_ids = []
-        for record in event.get("Records", []):
-            record_message_id = record.get("messageId")
+        for record in messages:
+            record_message_id = record.get("messageId") or record.get("MessageId")
             if not record_message_id:
                 continue
             message_ids.append(record_message_id)
 
+        arn = event.get("Records", [{}])[0].get("eventSourceARN") or "Unknown"
         return EventTriggerParser.build_trigger(
             target_id=target_id,
             resource_type="sqs",
             from_message_ids=message_ids,
             extra={
-                ExtraKeys.ARN: event["Records"][0]["eventSourceARN"],
-                ExtraKeys.RECORDS_NUM: len(event["Records"]),
+                ExtraKeys.ARN: arn,
+                ExtraKeys.RECORDS_NUM: len(messages),
             },
         )
 
     @staticmethod
     def extract_inner(event: Dict[Any, Any]) -> List[str]:
         inner_messages = []
-        for record in event.get("Records", []):
-            body = record.get("body")
+        for record in SqsEventTriggerParser._get_messages(event):
+            body = record.get("body") or record.get("Body")
             if isinstance(body, str):
                 inner_messages.append(body)
         return inner_messages

src/lumigo_tracer/extension/extension.py

@@ -8,6 +8,7 @@
 from lumigo_tracer import lumigo_utils
 from lumigo_tracer.extension.lambda_service import LambdaService
 from lumigo_tracer.extension.sampler import Sampler
+from lumigo_tracer.lambda_tracer import lambda_reporter
 from lumigo_tracer.lumigo_utils import lumigo_safe_execute
 
 SPAN_TYPE = "extensionExecutionEnd"
@@ -74,4 +75,4 @@ def _finish_previous_invocation(self, current_bandwidth: Optional[int]):  # type
             cpuUsageTime=[s.dump() for s in self.sampler.get_cpu_samples()],
             memoryUsage=[s.dump() for s in self.sampler.get_memory_samples()],
         )
-        lumigo_utils.report_json(os.environ.get("AWS_REGION", "us-east-1"), msgs=[asdict(span)])
+        lambda_reporter.report_json(os.environ.get("AWS_REGION", "us-east-1"), msgs=[asdict(span)])

src/lumigo_tracer/lambda_tracer/__init__.py

@@ -0,0 +1,14 @@
+from lumigo_tracer import auto_instrument_handler
+from lumigo_tracer.lambda_tracer import lambda_reporter
+from lumigo_tracer.lumigo_utils import is_aws_environment
+from lumigo_tracer import wrappers
+
+
+def global_scope_exec() -> None:
+    if is_aws_environment():
+        # Connection to edge: build the session
+        lambda_reporter.establish_connection_global()
+        # auto_instrument: import handler during runtime initialization, as usual.
+        auto_instrument_handler.prefetch_handler_import()
+        # follow requests to third party services
+        wrappers.wrap()