async-lambda-unstable 0.4.7__tar.gz → 0.4.9__tar.gz

{async-lambda-unstable-0.4.7 → async-lambda-unstable-0.4.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: async-lambda-unstable
-Version: 0.4.7
+Version: 0.4.9
 Summary: A framework for creating AWS Lambda Async Workflows. - Unstable Branch
 Author-email: "Nuclei, Inc" <engineering@nuclei.ai>
 Description-Content-Type: text/markdown

{async-lambda-unstable-0.4.7 → async-lambda-unstable-0.4.9}/async_lambda/__init__.py

@@ -14,4 +14,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.4.7"
+__version__ = "0.4.9"

{async-lambda-unstable-0.4.7 → async-lambda-unstable-0.4.9}/async_lambda/cli.py

@@ -221,7 +221,10 @@ def invoke(
     mock_context = MockLambdaContext(task.get_function_name())
     if task.trigger_type == TaskTriggerType.UNMANAGED_SQS:
         mock_event = MockSQSLambdaEvent(body or "")
-    elif task.trigger_type == TaskTriggerType.MANAGED_SQS:
+    elif task.trigger_type in [
+        TaskTriggerType.MANAGED_SQS,
+        TaskTriggerType.MANAGED_SQS_BATCH,
+    ]:
         sqs_payload = json.dumps(
             {
                 "source_task_id": None,

{async-lambda-unstable-0.4.7 → async-lambda-unstable-0.4.9}/async_lambda/controller.py

@@ -4,7 +4,14 @@ import logging
 import re
 import time
 from datetime import datetime, timezone
-from typing import Any, Callable, Dict, List, Optional, Type
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Type,
+)
 from uuid import uuid4
 
 from . import env
@@ -15,6 +22,7 @@ from .middleware import MET, RT, MiddlewareFunction, MiddlewareRegistration
 from .models.events.api_event import APIEvent
 from .models.events.base_event import BaseEvent
 from .models.events.dynamodb_event import DynamoDBEvent
+from .models.events.managed_sqs_batch_event import ManagedSQSBatchEvent
 from .models.events.managed_sqs_event import ManagedSQSEvent
 from .models.events.scheduled_event import ScheduledEvent
 from .models.events.unmanaged_sqs_event import UnmanagedSQSEvent
@@ -280,6 +288,22 @@ class AsyncLambdaController:
                         self.set_current_lane(lane=lane_index)
                         break
             self.set_current_invocation_id(_event.invocation_id)
+        elif task.trigger_type == TaskTriggerType.MANAGED_SQS_BATCH:
+            _event = ManagedSQSBatchEvent(*args)
+            lane_count = task.get_lane_count()
+            if lane_count == 1:
+                self.set_current_lane(lane=0)
+            else:
+                assert all(
+                    event.event_source_arn == _event.events[0].event_source_arn
+                    for event in _event.events
+                )
+                for lane_index in range(lane_count):
+                    if _event.events[0].event_source_arn == task.get_managed_queue_arn(
+                        lane=lane_index
+                    ):
+                        self.set_current_lane(lane=lane_index)
+                        break
         elif task.trigger_type == TaskTriggerType.UNMANAGED_SQS:
             _event = UnmanagedSQSEvent(*args)
         elif task.trigger_type == TaskTriggerType.SCHEDULED_EVENT:
@@ -307,7 +331,7 @@ class AsyncLambdaController:
         self,
         destination_task_id: str,
         sqs_payload: dict,
-        delay: Optional[int] = None,
+        delay: int = 0,
         force_sync: bool = False,
         lane: Optional[int] = None,
     ):
@@ -422,7 +446,7 @@ class AsyncLambdaController:
         self,
         task_id: str,
         payload: Any,
-        delay: Optional[int] = 0,
+        delay: int = 0,
         force_sync: bool = False,
         lane: Optional[int] = None,
     ):
@@ -481,6 +505,52 @@ class AsyncLambdaController:
                         "is_dlq_task": is_dlq_task,
                         "lane_count": lane_count,
                         "propagate_lane_assignment": propagate_lane_assignment,
+                        "batch_size": 1,
+                    },
+                    **kwargs,
+                )
+            )
+            return inner
+
+        return _task
+
+    def async_batch_task(
+        self,
+        task_id: str,
+        max_receive_count: int = 1,
+        dlq_task_id: Optional[str] = None,
+        is_dlq_task: bool = False,
+        lane_count: Optional[int] = None,
+        propagate_lane_assignment: Optional[bool] = None,
+        batch_size: int = 20,
+        **kwargs,
+    ):
+        """
+        Decorate a function to register it as an async batch task.
+        These functions can be asynchronously invoked with the `async_invoke` function
+        via their `task_id`.
+        """
+        logger.debug(f"Registering async batch task '{task_id}' with the controller.")
+
+        def _task(func: Callable[[ManagedSQSBatchEvent], Any]):
+            @functools.wraps(func)
+            def inner(*args, **kwargs):
+                self.set_current_task_id(task_id)
+                return func(*args, **kwargs)
+
+            self.add_task(
+                AsyncLambdaTask(
+                    controller=self,
+                    executable=inner,
+                    task_id=task_id,
+                    trigger_type=TaskTriggerType.MANAGED_SQS_BATCH,
+                    trigger_config={
+                        "max_receive_count": max_receive_count,
+                        "dlq_task_id": dlq_task_id,
+                        "is_dlq_task": is_dlq_task,
+                        "lane_count": lane_count,
+                        "propagate_lane_assignment": propagate_lane_assignment,
+                        "batch_size": batch_size,
                     },
                     **kwargs,
                 )

async-lambda-unstable-0.4.9/async_lambda/models/events/managed_sqs_batch_event.py

@@ -0,0 +1,16 @@
+from typing import List
+
+from .base_event import BaseEvent
+from .managed_sqs_event import ManagedSQSEvent
+
+
+class ManagedSQSBatchEvent(BaseEvent):
+    events: List[ManagedSQSEvent]
+
+    def _hydrate_event(self):
+        self.events = []
+        for i in range(len(self._event["Records"])):
+            managed_sqs_event = ManagedSQSEvent(
+                {"Records": [self._event["Records"][i]]}, self._context, self._task
+            )
+            self.events.append(managed_sqs_event)

{async-lambda-unstable-0.4.7 → async-lambda-unstable-0.4.9}/async_lambda/models/mock/mock_event.py

@@ -29,6 +29,36 @@ def MockSQSLambdaEvent(body: str, source_queue_arn: Optional[str] = None) -> dic
     }
 
 
+def MockSQSBatchLambdaEvent(
+    bodies: List[str], source_queue_arn: Optional[str] = None
+) -> dict:
+    now_timestamp = int(time())
+    if not source_queue_arn:
+        source_queue_arn = "arn:aws:sqs:us-east-1:123456789012:my-queue"
+
+    records = []
+    for body in bodies:
+        records.append(
+            {
+                "messageId": "059f36b4-87a3-44ab-83d2-661975830a7e",
+                "receiptHandle": "AQEBwJnKyrHigUMZj6rYigCgxlaS3SLy0b",
+                "body": body,
+                "attributes": {
+                    "ApproximateReceiveCount": "1",
+                    "SentTimestamp": str(now_timestamp),
+                    "SenderId": "AIDAIENQZJOLO23YVJ4VO",
+                    "ApproximateFirstReceiveTimestamp": str(now_timestamp + 5),
+                },
+                "messageAttributes": {},
+                "md5OfBody": "e4e68fb7bd0e697a0ae8f1bb342846b4",
+                "eventSource": "aws:sqs",
+                "eventSourceARN": source_queue_arn,
+                "awsRegion": "us-east-1",
+            }
+        )
+    return {"Records": records}
+
+
 def MockAPILambdaEvent(
     path: str,
     method: str,

{async-lambda-unstable-0.4.7 → async-lambda-unstable-0.4.9}/async_lambda/models/task.py

@@ -15,6 +15,7 @@ if TYPE_CHECKING:
 
 from ..middleware import RT, MiddlewareStackExecutor
 from .events.dynamodb_event import DynamoDBEvent
+from .events.managed_sqs_batch_event import ManagedSQSBatchEvent
 from .events.managed_sqs_event import ManagedSQSEvent
 from .events.scheduled_event import ScheduledEvent
 from .events.unmanaged_sqs_event import UnmanagedSQSEvent
@@ -22,15 +23,22 @@ from .events.unmanaged_sqs_event import UnmanagedSQSEvent
 
 class TaskTriggerType(Enum):
     MANAGED_SQS = 1
-    UNMANAGED_SQS = 2
-    SCHEDULED_EVENT = 3
-    API_EVENT = 4
-    DYNAMODB_EVENT = 5
+    MANAGED_SQS_BATCH = 2
+    UNMANAGED_SQS = 3
+    SCHEDULED_EVENT = 4
+    API_EVENT = 5
+    DYNAMODB_EVENT = 6
 
 
 EventType = TypeVar(
     "EventType",
-    bound=Union[ManagedSQSEvent, ScheduledEvent, UnmanagedSQSEvent, DynamoDBEvent],
+    bound=Union[
+        ManagedSQSEvent,
+        ManagedSQSBatchEvent,
+        ScheduledEvent,
+        UnmanagedSQSEvent,
+        DynamoDBEvent,
+    ],
 )
 
 
@@ -81,7 +89,8 @@ class AsyncLambdaTask(Generic[EventType, RT]):
         self._has_run_init_tasks = False
 
         if (
-            self.trigger_type == TaskTriggerType.MANAGED_SQS
+            self.trigger_type
+            in [TaskTriggerType.MANAGED_SQS, TaskTriggerType.MANAGED_SQS_BATCH]
             and "dlq_task_id" in self.trigger_config
         ):
             dlq_task_id = self.trigger_config["dlq_task_id"]
@@ -91,7 +100,10 @@ class AsyncLambdaTask(Generic[EventType, RT]):
                     raise Exception(
                         f"Error setting DLQ Task ID: No task with the task_id {dlq_task_id} exists."
                     )
-                if dlq_task.trigger_type != TaskTriggerType.MANAGED_SQS:
+                if dlq_task.trigger_type not in [
+                    TaskTriggerType.MANAGED_SQS,
+                    TaskTriggerType.MANAGED_SQS_BATCH,
+                ]:
                     raise Exception(
                         f"Error setting DLQ Task ID: Task {dlq_task_id} is not an async-task."
                     )
@@ -119,7 +131,10 @@ class AsyncLambdaTask(Generic[EventType, RT]):
             raise ValueError("Task ID must be less than 32 characters long.")
 
     def get_lane_count(self) -> int:
-        if self.trigger_type != TaskTriggerType.MANAGED_SQS:
+        if self.trigger_type not in [
+            TaskTriggerType.MANAGED_SQS,
+            TaskTriggerType.MANAGED_SQS_BATCH,
+        ]:
             raise Exception(f"The task {self.task_id} is not a managed queue task.")
         if "lane_count" in self.trigger_config and isinstance(
             self.trigger_config["lane_count"], int
@@ -128,7 +143,10 @@ class AsyncLambdaTask(Generic[EventType, RT]):
         return self.controller.get_lane_count()
 
     def should_propagate_lane_assignment(self) -> bool:
-        if self.trigger_type != TaskTriggerType.MANAGED_SQS:
+        if self.trigger_type not in [
+            TaskTriggerType.MANAGED_SQS,
+            TaskTriggerType.MANAGED_SQS_BATCH,
+        ]:
             raise Exception(f"The task {self.task_id} is not a managed queue task.")
         if "propagate_lane_assignment" in self.trigger_config and isinstance(
             self.trigger_config["propagate_lane_assignment"], bool
@@ -140,7 +158,10 @@ class AsyncLambdaTask(Generic[EventType, RT]):
         """
         Returns the managed queue's name for this task.
         """
-        if self.trigger_type != TaskTriggerType.MANAGED_SQS:
+        if self.trigger_type not in [
+            TaskTriggerType.MANAGED_SQS,
+            TaskTriggerType.MANAGED_SQS_BATCH,
+        ]:
             raise Exception(f"The task {self.task_id} is not a managed queue task.")
         if lane == 0:
             return f"{config.name}-{self.task_id}"
@@ -150,12 +171,18 @@ class AsyncLambdaTask(Generic[EventType, RT]):
         return f"{config.name}-{self.task_id}"
 
     def get_managed_queue_arn(self, lane: int = 0):
-        if self.trigger_type != TaskTriggerType.MANAGED_SQS:
+        if self.trigger_type not in [
+            TaskTriggerType.MANAGED_SQS,
+            TaskTriggerType.MANAGED_SQS_BATCH,
+        ]:
             raise Exception(f"The task {self.task_id} is not a managed queue task.")
         return f"arn:aws:sqs:{env.get_aws_region()}:{env.get_aws_account_id()}:{self.get_managed_queue_name(lane=lane)}"
 
     def get_managed_queue_url(self, lane: int = 0):
-        if self.trigger_type != TaskTriggerType.MANAGED_SQS:
+        if self.trigger_type not in [
+            TaskTriggerType.MANAGED_SQS,
+            TaskTriggerType.MANAGED_SQS_BATCH,
+        ]:
             raise Exception(f"The task {self.task_id} is not a managed queue task.")
         return f"https://sqs.{env.get_aws_region()}.amazonaws.com/{env.get_aws_account_id()}/{self.get_managed_queue_name(lane=lane)}"
 
@@ -163,21 +190,30 @@ class AsyncLambdaTask(Generic[EventType, RT]):
         return f"{self.task_id}ALFunc"
 
     def get_managed_queue_logical_id(self, lane: int = 0):
-        if self.trigger_type != TaskTriggerType.MANAGED_SQS:
+        if self.trigger_type not in [
+            TaskTriggerType.MANAGED_SQS,
+            TaskTriggerType.MANAGED_SQS_BATCH,
+        ]:
             raise Exception(f"The task {self.task_id} is not a managed queue task.")
         if lane == 0:
             return f"{self.task_id}ALQueue"
         return f"{self.task_id}ALQueueL{lane}"
 
     def get_managed_queue_extra_logical_id(self, index: int, lane: int = 0):
-        if self.trigger_type != TaskTriggerType.MANAGED_SQS:
+        if self.trigger_type not in [
+            TaskTriggerType.MANAGED_SQS,
+            TaskTriggerType.MANAGED_SQS_BATCH,
+        ]:
             raise Exception(f"The task {self.task_id} is not a managed queue task.")
         if lane == 0:
             return f"{self.get_function_logical_id()}Extra{index}"
         return f"{self.get_function_logical_id()}Extra{index}L{lane}"
 
     def get_managed_queue_event_logical_id(self, lane: int = 0):
-        if self.trigger_type != TaskTriggerType.MANAGED_SQS:
+        if self.trigger_type not in [
+            TaskTriggerType.MANAGED_SQS,
+            TaskTriggerType.MANAGED_SQS_BATCH,
+        ]:
             raise Exception(f"The task {self.task_id} is not a managed queue task.")
         if lane == 0:
             return "ManagedSQS"
@@ -185,16 +221,17 @@ class AsyncLambdaTask(Generic[EventType, RT]):
 
     def get_template_events(self) -> dict:
         sqs_properties = {}
-        if (
-            isinstance(self.maximum_concurrency, list)
-            and self.trigger_type != TaskTriggerType.MANAGED_SQS
-        ):
+        if isinstance(self.maximum_concurrency, list) and self.trigger_type not in [
+            TaskTriggerType.MANAGED_SQS,
+            TaskTriggerType.MANAGED_SQS_BATCH,
+        ]:
             raise Exception(
                 f"Invalid maximum concurrency configuration for task {self.task_id}. Must be an int, not a list of ints. Lanes are only supported for ManagedSQS tasks."
             )
         if (
             isinstance(self.maximum_concurrency, list)
-            and self.trigger_type == TaskTriggerType.MANAGED_SQS
+            and self.trigger_type
+            in [TaskTriggerType.MANAGED_SQS, TaskTriggerType.MANAGED_SQS_BATCH]
             and len(self.maximum_concurrency) != self.get_lane_count()
         ):
             raise Exception(
@@ -204,7 +241,10 @@ class AsyncLambdaTask(Generic[EventType, RT]):
             sqs_properties["ScalingConfig"] = {
                 "MaximumConcurrency": self.maximum_concurrency
             }
-        if self.trigger_type == TaskTriggerType.MANAGED_SQS:
+        if self.trigger_type in [
+            TaskTriggerType.MANAGED_SQS,
+            TaskTriggerType.MANAGED_SQS_BATCH,
+        ]:
             events = {}
             for lane_index in range(self.get_lane_count()):
                 sqs_properties = {}
@@ -219,7 +259,7 @@ class AsyncLambdaTask(Generic[EventType, RT]):
                 events[self.get_managed_queue_event_logical_id(lane=lane_index)] = {
                     "Type": "SQS",
                     "Properties": {
-                        "BatchSize": 1,
+                        "BatchSize": self.trigger_config["batch_size"],
                         "Enabled": True,
                         "Queue": {
                             "Fn::GetAtt": [
@@ -284,7 +324,10 @@ class AsyncLambdaTask(Generic[EventType, RT]):
         raise NotImplementedError()
 
     def get_policy_sqs_resources(self) -> List[dict]:
-        if self.trigger_type == TaskTriggerType.MANAGED_SQS:
+        if self.trigger_type in [
+            TaskTriggerType.MANAGED_SQS,
+            TaskTriggerType.MANAGED_SQS_BATCH,
+        ]:
             return [
                 {
                     "Fn::GetAtt": [
@@ -332,7 +375,8 @@ class AsyncLambdaTask(Generic[EventType, RT]):
         managed_tasks_resources = [
             resource
             for task in tasks
-            if task.trigger_type == TaskTriggerType.MANAGED_SQS
+            if task.trigger_type
+            in [TaskTriggerType.MANAGED_SQS, TaskTriggerType.MANAGED_SQS_BATCH]
             for resource in task.get_policy_sqs_resources()
         ]
         if len(managed_tasks_resources) > 0:
@@ -408,7 +452,10 @@ class AsyncLambdaTask(Generic[EventType, RT]):
             }
         }
 
-        if self.trigger_type == TaskTriggerType.MANAGED_SQS:
+        if self.trigger_type in [
+            TaskTriggerType.MANAGED_SQS,
+            TaskTriggerType.MANAGED_SQS_BATCH,
+        ]:
             dlq_task = self.get_dlq_task()
             if dlq_task is None:
                 dead_letter_target_arn = {
@@ -429,10 +476,7 @@ class AsyncLambdaTask(Generic[EventType, RT]):
                     "async-lambda-lane": str(lane_index),
                     "async-lambda-queue-type": "managed",
                 }
-                if (
-                    self.trigger_type == TaskTriggerType.MANAGED_SQS
-                    and self.trigger_config["is_dlq_task"]
-                ):
+                if self.trigger_config["is_dlq_task"]:
                     _extra_tags["async-lambda-queue-type"] = "dlq-task"
                 template[self.get_managed_queue_logical_id(lane=lane_index)] = {
                     "Type": "AWS::SQS::Queue",
@@ -474,7 +518,10 @@ class AsyncLambdaTask(Generic[EventType, RT]):
         return json.loads(stringified_extra)
 
     def get_dlq_task(self) -> Optional["AsyncLambdaTask"]:
-        if self.trigger_type != TaskTriggerType.MANAGED_SQS:
+        if self.trigger_type not in [
+            TaskTriggerType.MANAGED_SQS,
+            TaskTriggerType.MANAGED_SQS,
+        ]:
             return None
         if self.trigger_config.get("is_dlq_task"):
             return None

{async-lambda-unstable-0.4.7 → async-lambda-unstable-0.4.9}/async_lambda_unstable.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: async-lambda-unstable
-Version: 0.4.7
+Version: 0.4.9
 Summary: A framework for creating AWS Lambda Async Workflows. - Unstable Branch
 Author-email: "Nuclei, Inc" <engineering@nuclei.ai>
 Description-Content-Type: text/markdown

{async-lambda-unstable-0.4.7 → async-lambda-unstable-0.4.9}/async_lambda_unstable.egg-info/SOURCES.txt

@@ -19,6 +19,7 @@ async_lambda/models/events/__init__.py
 async_lambda/models/events/api_event.py
 async_lambda/models/events/base_event.py
 async_lambda/models/events/dynamodb_event.py
+async_lambda/models/events/managed_sqs_batch_event.py
 async_lambda/models/events/managed_sqs_event.py
 async_lambda/models/events/scheduled_event.py
 async_lambda/models/events/unmanaged_sqs_event.py