async-lambda-unstable 0.6.1__tar.gz → 0.6.2__tar.gz

{async_lambda_unstable-0.6.1 → async_lambda_unstable-0.6.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: async-lambda-unstable
-Version: 0.6.1
+Version: 0.6.2
 Summary: A framework for creating AWS Lambda Async Workflows. - Unstable Branch
 Author-email: "Nuclei, Inc" <engineering@nuclei.ai>
 Requires-Dist: click>=8.0.0

{async_lambda_unstable-0.6.1 → async_lambda_unstable-0.6.2}/async_lambda/__init__.py

@@ -21,4 +21,4 @@ from .models.events.managed_sqs_event import ManagedSQSEvent as ManagedSQSEvent
 from .models.events.scheduled_event import ScheduledEvent as ScheduledEvent
 from .models.events.unmanaged_sqs_event import UnmanagedSQSEvent as UnmanagedSQSEvent
 
-__version__ = "0.6.1"
+__version__ = "0.6.2"

{async_lambda_unstable-0.6.1 → async_lambda_unstable-0.6.2}/async_lambda/controller.py

@@ -814,6 +814,7 @@ class AsyncLambdaController:
         destination_task_id: str,
         sqs_payloads: Sequence[dict],
         delay: Union[int, Sequence[int]] = 0,
+        message_group_id: Union[Optional[str], Sequence[Optional[str]]] = None,
         force_sync: bool = False,
         lane: Optional[int] = None,
         index: int = 0,
@@ -899,8 +900,28 @@ class AsyncLambdaController:
                         time.sleep(delay[i])
                     else:
                         time.sleep(delay)
+
+                _message_group_id = (
+                    message_group_id[i]
+                    if isinstance(message_group_id, Sequence)
+                    and not isinstance(message_group_id, str)
+                    else message_group_id
+                )
+                if (
+                    _message_group_id is None
+                    and (
+                        current_message_group_id := self.get_current_message_group_id()
+                    )
+                    is not None
+                    and self.should_propagate_message_group_id()
+                ):
+                    _message_group_id = current_message_group_id
+
+                current_message_group_id = self.get_current_message_group_id()
                 mock_event = MockSQSLambdaEvent(
-                    json.dumps(sqs_payload), source_queue_arn=queue_arn
+                    json.dumps(sqs_payload),
+                    source_queue_arn=queue_arn,
+                    message_group_id=_message_group_id,
                 )
                 mock_context = MockLambdaContext(destination_task.task_id)
                 self.handle_invocation(
@@ -908,6 +929,7 @@ class AsyncLambdaController:
                 )
                 self.set_current_lane(current_lane)
                 self.set_current_task_id(current_task_id)
+                self.set_current_message_group_id(current_message_group_id)
         else:
             entries: List[dict] = []
             for i, sqs_payload in enumerate(sqs_payloads):
@@ -915,13 +937,31 @@ class AsyncLambdaController:
                     _delay = delay[i]
                 else:
                     _delay = delay
-                entries.append(
-                    {
-                        "MessageBody": json.dumps(sqs_payload),
-                        "DelaySeconds": _delay,
-                        "Id": f"index_{index + i}",
-                    }
+
+                _message_group_id = (
+                    message_group_id[i]
+                    if isinstance(message_group_id, Sequence)
+                    and not isinstance(message_group_id, str)
+                    else message_group_id
                 )
+                if (
+                    _message_group_id is None
+                    and (
+                        current_message_group_id := self.get_current_message_group_id()
+                    )
+                    is not None
+                    and self.should_propagate_message_group_id()
+                ):
+                    _message_group_id = current_message_group_id
+
+                entry = {
+                    "MessageBody": json.dumps(sqs_payload),
+                    "DelaySeconds": _delay,
+                    "Id": f"index_{index + i}",
+                }
+                if _message_group_id:
+                    entry["MessageGroupId"] = _message_group_id
+                entries.append(entry)
             if is_external_task:
                 url = f"https://sqs.{env.get_aws_region()}.amazonaws.com/{env.get_aws_account_id()}/{destination_task_id}"
             else:
@@ -1088,6 +1128,7 @@ class AsyncLambdaController:
         task_id: str,
         payloads: Sequence[Any],
         delay: Union[int, Sequence[int]] = 0,
+        message_group_id: Union[Optional[str], Sequence[Optional[str]]] = None,
         force_sync: bool = False,
         lane: Optional[int] = None,
     ):
@@ -1120,6 +1161,7 @@ class AsyncLambdaController:
                 destination_task_id=task_id,
                 sqs_payloads=sqs_payloads,
                 delay=delay,
+                message_group_id=message_group_id,
                 force_sync=force_sync,
                 lane=lane,
                 index=i,