async-lambda-unstable 0.6.6__tar.gz → 0.6.7__tar.gz

{async_lambda_unstable-0.6.6 → async_lambda_unstable-0.6.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: async-lambda-unstable
-Version: 0.6.6
+Version: 0.6.7
 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.6 → async_lambda_unstable-0.6.7}/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.6"
+__version__ = "0.6.7"

{async_lambda_unstable-0.6.6 → async_lambda_unstable-0.6.7}/async_lambda/client.py

@@ -12,6 +12,7 @@ class Clients:
         s3_client (Optional[Any]): The AWS S3 client instance.
         sqs_client (Optional[Any]): The AWS SQS client instance.
         sts_client (Optional[Any]): The AWS STS client instance.
+        scheduler_client (Optional[Any]): The AWS EventBridge Scheduler client instance.
 
     Methods:
         reset():
@@ -21,14 +22,16 @@ class Clients:
     s3_client: Optional[Any] = None
     sqs_client: Optional[Any] = None
     sts_client: Optional[Any] = None
+    scheduler_client: Optional[Any] = None
 
     def reset(self):
         """
-        Resets the AWS service clients (S3, SQS, STS) by setting them to None.
+        Resets the AWS service clients (S3, SQS, STS, Scheduler) by setting them to None.
         """
         self.s3_client = None
         self.sqs_client = None
         self.sts_client = None
+        self.scheduler_client = None
 
 
 clients = Clients()
@@ -92,3 +95,18 @@ def get_sts_client():
     if clients.sts_client is None:
         clients.sts_client = boto3.client("sts", **get_client_kwargs())
     return clients.sts_client
+
+
+def get_scheduler_client():
+    """
+    Returns a cached AWS EventBridge Scheduler client instance.
+
+    If the client does not exist, it creates a new one using boto3 with the provided keyword arguments.
+    Subsequent calls will return the cached client.
+
+    Returns:
+        boto3.client: An AWS EventBridge Scheduler client instance.
+    """
+    if clients.scheduler_client is None:
+        clients.scheduler_client = boto3.client("scheduler", **get_client_kwargs())
+    return clients.scheduler_client

{async_lambda_unstable-0.6.6 → async_lambda_unstable-0.6.7}/async_lambda/controller.py

@@ -4,7 +4,7 @@ import logging
 import random
 import re
 import time
-from datetime import datetime, timezone
+from datetime import datetime, timedelta, timezone
 from typing import (
     Any,
     Callable,
@@ -21,7 +21,7 @@ from uuid import uuid4
 
 from . import env
 from .build_config import get_build_config_for_stage
-from .client import get_s3_client, get_sqs_client
+from .client import get_s3_client, get_scheduler_client, get_sqs_client
 from .config import config
 from .middleware import MET, RT, MiddlewareFunction, MiddlewareRegistration
 from .models.events.api_event import APIEvent
@@ -39,6 +39,8 @@ from .util import make_cf_tags
 
 logger = logging.getLogger(__name__)
 
+_SQS_MAX_DELAY_SECONDS = 900
+
 BaseEventT = TypeVar("BaseEventT", bound=BaseEvent)
 APIEventT = TypeVar("APIEventT", bound=APIEvent)
 ManagedSQSEventT = TypeVar("ManagedSQSEventT", bound=ManagedSQSEvent)
@@ -67,6 +69,30 @@ class BatchInvokeException(Exception):
         super().__init__(msg)
 
 
+# Exception raised when async_invoke_* is called to a non-internal task with a delay greater than _SQS_MAX_DELAY_SECONDS
+class AsyncInvokeInvalidDelay(Exception):
+    pass
+
+
+class _BatchEntry:
+    enqueue_id: str
+    payload: dict
+    delay: int
+    message_group_id: Optional[str]
+
+    def __init__(
+        self,
+        enqueue_id: str,
+        payload: dict,
+        delay: int,
+        message_group_id: Optional[str],
+    ):
+        self.enqueue_id = enqueue_id
+        self.payload = payload
+        self.delay = delay
+        self.message_group_id = message_group_id
+
+
 class AsyncLambdaController:
     """
     AsyncLambdaController manages async tasks, middleware, and invocation logic for the async-lambda framework.
@@ -454,18 +480,20 @@ class AsyncLambdaController:
             },
         }
         _task_list = list(self.tasks.values())
-        managed_tasks_resources = [
+        internal_tasks_resources = [
             resource
             for task in _task_list
             if task.trigger_type in MANAGED_SQS_TASK_TYPES
             for resource in task.get_policy_sqs_resources()
-        ] + [
+        ]
+        external_tasks_resources = [
             resource
             for external_async_task_id in self.external_async_tasks
             for resource in AsyncLambdaTask.get_policy_external_task_resources(
                 external_async_task_id
             )
         ]
+        managed_tasks_resources = internal_tasks_resources + external_tasks_resources
         task_ref_policies = {}
         if len(managed_tasks_resources) > 0:
             task_ref_policies = self._build_send_to_all_async_lambda_queues_policies(
@@ -474,6 +502,78 @@ class AsyncLambdaController:
             for key in task_ref_policies.keys():
                 template["Resources"][key] = task_ref_policies[key]
 
+        if len(internal_tasks_resources) > 0:
+            template["Resources"]["AsyncLambdaDelayScheduleGroup"] = {
+                "Type": "AWS::Scheduler::ScheduleGroup",
+                "Properties": {
+                    "Name": f"{config.name}-delay",
+                    "Tags": make_cf_tags(build_config.tags),
+                },
+            }
+            template["Resources"]["AsyncLambdaDelaySchedulerRole"] = {
+                "Type": "AWS::IAM::Role",
+                "Properties": {
+                    "AssumeRolePolicyDocument": {
+                        "Version": "2012-10-17",
+                        "Statement": [
+                            {
+                                "Effect": "Allow",
+                                "Principal": {"Service": "scheduler.amazonaws.com"},
+                                "Action": "sts:AssumeRole",
+                            }
+                        ],
+                    },
+                    "ManagedPolicyArns": [
+                        {"Ref": policy_id} for policy_id in task_ref_policies.keys()
+                    ],
+                },
+            }
+            template["Resources"][
+                "AsyncLambdaCreateDelaySchedulesPolicy"
+            ] = task_ref_policies["AsyncLambdaCreateDelaySchedulesPolicy"] = {
+                "Type": "AWS::IAM::ManagedPolicy",
+                "Properties": {
+                    "ManagedPolicyName": {
+                        "Fn::Sub": "${AWS::StackName}-create-delay-schedules"
+                    },
+                    "PolicyDocument": {
+                        "Version": "2012-10-17",
+                        "Statement": [
+                            {
+                                "Sid": "CreateSchedules",
+                                "Effect": "Allow",
+                                "Action": "scheduler:CreateSchedule",
+                                "Resource": {
+                                    "Fn::Sub": "arn:aws:scheduler:${AWS::Region}:${AWS::AccountId}:schedule/${AsyncLambdaDelayScheduleGroup}/*"
+                                },
+                            },
+                            {
+                                "Sid": "PassSchedulerExecutionRole",
+                                "Effect": "Allow",
+                                "Action": "iam:PassRole",
+                                "Resource": {
+                                    "Fn::GetAtt": [
+                                        "AsyncLambdaDelaySchedulerRole",
+                                        "Arn",
+                                    ]
+                                },
+                                "Condition": {
+                                    "StringEquals": {
+                                        "iam:PassedToService": "scheduler.amazonaws.com"
+                                    }
+                                },
+                            },
+                        ],
+                    },
+                },
+            }
+            template["Globals"]["Function"]["Environment"]["Variables"][
+                "ASYNC_LAMBDA_DELAY_SCHEDULE_GROUP"
+            ] = {"Ref": "AsyncLambdaDelayScheduleGroup"}
+            template["Globals"]["Function"]["Environment"]["Variables"][
+                "ASYNC_LAMBDA_DELAY_SCHEDULER_ROLE_ARN"
+            ] = {"Fn::GetAtt": ["AsyncLambdaDelaySchedulerRole", "Arn"]}
+
         has_api_tasks = False
         for task in _task_list:
             if task.trigger_type == TaskTriggerType.API_EVENT:
@@ -808,6 +908,12 @@ class AsyncLambdaController:
                 raise Exception(
                     f"No such task exists with the task_id {destination_task_id}"
                 )
+
+        if is_external_task and delay > _SQS_MAX_DELAY_SECONDS:
+            raise AsyncInvokeInvalidDelay(
+                f"Unable to invoke task {destination_task_id} with delay {delay} seconds. External tasks only support delays up to {_SQS_MAX_DELAY_SECONDS} seconds."
+            )
+
         destination_task = None
         if not is_external_task:
             destination_task = self.tasks[destination_task_id]
@@ -878,16 +984,29 @@ class AsyncLambdaController:
                 assert destination_task is not None
                 url = destination_task.get_managed_queue_url(lane=lane)
 
-            _kwargs = {}
-            if _message_group_id:
-                _kwargs["MessageGroupId"] = _message_group_id
+            if delay > _SQS_MAX_DELAY_SECONDS:
+                assert not is_external_task
 
-            get_sqs_client().send_message(
-                QueueUrl=url,
-                MessageBody=json.dumps(sqs_payload),
-                DelaySeconds=delay,
-                **_kwargs,
-            )
+                queue_arn = destination_task.get_managed_queue_arn(lane=lane)
+                # TODO: Use _message_group_id when AWS tells me how
+                self._send_via_scheduler(
+                    queue_arn=queue_arn,
+                    message_body=json.dumps(sqs_payload),
+                    delay=delay,
+                    message_group_id=_message_group_id,
+                )
+            else:
+                _kwargs = {}
+                if _message_group_id:
+                    _kwargs["MessageGroupId"] = _message_group_id
+
+                get_sqs_client().send_message(
+                    QueueUrl=url,
+                    MessageBody=json.dumps(sqs_payload),
+                    DelaySeconds=delay,
+                    **_kwargs,
+                )
+            return None
 
     def send_async_invoke_payload_batch(
         self,
@@ -964,6 +1083,35 @@ class AsyncLambdaController:
                 )
             lane = 0
 
+        batch_entries: List[_BatchEntry] = []
+        for i, sqs_payload in enumerate(sqs_payloads):
+            if isinstance(delay, Sequence):
+                _delay = delay[i]
+            else:
+                _delay = 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
+            batch_entries.append(
+                _BatchEntry(
+                    enqueue_id=f"index_{index + i}",
+                    payload=sqs_payload,
+                    delay=_delay,
+                    message_group_id=_message_group_id,
+                )
+            )
+
         if force_sync or env.get_force_sync_mode():
             if is_external_task:
                 raise NotImplementedError(
@@ -974,34 +1122,16 @@ class AsyncLambdaController:
             current_lane = self.get_current_lane()
             assert destination_task is not None
             queue_arn = destination_task.get_managed_queue_arn(lane=lane)
-            for i, sqs_payload in enumerate(sqs_payloads):
-                if delay:
-                    if isinstance(delay, Sequence):
-                        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
+
+            for batch in batch_entries:
+                if batch.delay:
+                    time.sleep(batch.delay)
 
                 current_message_group_id = self.get_current_message_group_id()
                 mock_event = MockSQSLambdaEvent(
-                    json.dumps(sqs_payload),
+                    json.dumps(batch.payload),
                     source_queue_arn=queue_arn,
-                    message_group_id=_message_group_id,
+                    message_group_id=batch.message_group_id,
                 )
                 mock_context = MockLambdaContext(destination_task.task_id)
                 self.handle_invocation(
@@ -1011,69 +1141,112 @@ class AsyncLambdaController:
                 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):
-                if isinstance(delay, Sequence):
-                    _delay = delay[i]
+            sqs_entries: List[dict] = []
+            for batch in batch_entries:
+                if batch.delay > _SQS_MAX_DELAY_SECONDS:
+                    if is_external_task:
+                        raise AsyncInvokeInvalidDelay(
+                            f"send_async_invoke_payload_batch does not support a delay longer than {_SQS_MAX_DELAY_SECONDS} seconds for external tasks"
+                        )
+                    assert destination_task is not None
+                    queue_arn = destination_task.get_managed_queue_arn(lane=lane)
+                    self._send_via_scheduler(
+                        queue_arn=queue_arn,
+                        message_body=json.dumps(batch.payload),
+                        delay=batch.delay,
+                    )
                 else:
-                    _delay = 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()
+                    entry: dict = {
+                        "MessageBody": json.dumps(batch.payload),
+                        "DelaySeconds": batch.delay,
+                        "Id": batch.enqueue_id,
+                    }
+                    if batch.message_group_id:
+                        entry["MessageGroupId"] = batch.message_group_id
+                    sqs_entries.append(entry)
+            if sqs_entries:
+                if is_external_task:
+                    url = f"https://sqs.{env.get_aws_region()}.amazonaws.com/{env.get_aws_account_id()}/{destination_task_id}"
+                else:
+                    assert destination_task is not None
+                    url = destination_task.get_managed_queue_url(lane=lane)
+                failed_messages: List[dict] = []
+                batch_retry_count = env.get_batch_failure_retry_count() + 1
+                entries = sqs_entries
+                for i in range(batch_retry_count):
+                    response = get_sqs_client().send_message_batch(
+                        QueueUrl=url,
+                        Entries=entries,
                     )
-                    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:
-                assert destination_task is not None
-                url = destination_task.get_managed_queue_url(lane=lane)
-            failed_messages = []
-            batch_retry_count = env.get_batch_failure_retry_count() + 1
-            for i in range(batch_retry_count):
-                response = get_sqs_client().send_message_batch(
-                    QueueUrl=url,
-                    Entries=entries,
+                    failed_messages = response.get("Failed", [])
+                    if len(failed_messages) == 0:
+                        return
+                    logger.warning(failed_messages)
+                    logger.warning(f"{len(failed_messages)} messages failed to send. ")
+                    failed_message_ids = {message["Id"] for message in failed_messages}
+                    entries = [
+                        entry for entry in entries if entry["Id"] in failed_message_ids
+                    ]
+                    if i < batch_retry_count:
+                        send_delay = 0.5 + random.random()
+                        logger.info(
+                            f"Waiting {send_delay:.3f} before attempting batch failures again."
+                        )
+                        time.sleep(send_delay)
+                logger.error(failed_messages)
+                raise BatchInvokeException(
+                    f"Failed to send {len(failed_messages)} messages.",
+                    failed_payloads=[
+                        int(entry["Id"].split("_")[-1]) for entry in entries
+                    ],
                 )
-                failed_messages: List[dict] = response.get("Failed", [])
-                if len(failed_messages) == 0:
-                    return
-                logger.warning(failed_messages)
-                logger.warning(f"{len(failed_messages)} messages failed to send. ")
-                failed_message_ids = {message["Id"] for message in failed_messages}
-                entries = [
-                    entry for entry in entries if entry["Id"] in failed_message_ids
-                ]
-                if i < batch_retry_count:
-                    send_delay = 0.5 + random.random()
-                    logger.info(
-                        f"Waiting {send_delay:.3f} before attempting batch failures again."
-                    )
-                    time.sleep(send_delay)
-            logger.error(failed_messages)
-            raise BatchInvokeException(
-                f"Failed to send {len(failed_messages)} messages.",
-                failed_payloads=[int(entry["Id"].split("_")[-1]) for entry in entries],
-            )
+
+    @staticmethod
+    def _send_via_scheduler(
+        queue_arn: str,
+        message_body: str,
+        delay: int,
+        message_group_id: Optional[str] = None,
+    ):
+        """
+        Schedules a message for future delivery to an SQS queue via EventBridge Scheduler.
+
+        Used when the requested delay exceeds SQS's maximum of 900 seconds (15 minutes).
+        Creates a one-time schedule in the stack-level schedule group (ASYNC_LAMBDA_DELAY_SCHEDULE_GROUP),
+        which is provisioned by CloudFormation. The schedule auto-deletes after firing.
+
+        Requires the following environment variables:
+        - ASYNC_LAMBDA_DELAY_SCHEDULE_GROUP: name of the CloudFormation-managed schedule group.
+        - ASYNC_LAMBDA_DELAY_SCHEDULER_ROLE_ARN: ARN of the IAM role EventBridge Scheduler uses
+          to deliver to the target SQS queue (must have sqs:SendMessage on all queues in the group).
+
+        Args:
+            queue_arn (str): ARN of the destination SQS queue.
+            message_body (str): JSON-serialized message body to deliver to the queue.
+            delay (int): Delay in seconds before delivering the message.
+            message_group_id (Optional[str]): MessageGroupId for FIFO queues. Defaults to None.
+        """
+        schedule_time = datetime.now(tz=timezone.utc) + timedelta(seconds=delay)
+        schedule_expression = f"at({schedule_time.strftime('%Y-%m-%dT%H:%M:%S')})"
+
+        target: dict = {
+            "Arn": queue_arn,
+            "RoleArn": env.get_delay_scheduler_role_arn(),
+            "Input": message_body,
+        }
+        # TODO: Set message_group_id based on feedback from AWS
+        # if message_group_id:
+        #     target["SqsParameters"] = {"MessageGroupId": message_group_id}
+
+        get_scheduler_client().create_schedule(
+            Name=uuid4().hex,
+            GroupName=env.get_delay_schedule_group_name(),
+            ScheduleExpression=schedule_expression,
+            ScheduleExpressionTimezone="UTC",
+            FlexibleTimeWindow={"Mode": "OFF"},
+            Target=target,
+            ActionAfterCompletion="DELETE",
+        )
 
     def new_payload(
         self,

{async_lambda_unstable-0.6.6 → async_lambda_unstable-0.6.7}/async_lambda/env.py

@@ -206,3 +206,34 @@ def get_batch_failure_retry_count() -> int:
         int: The number of times to retry batch failures. Defaults to 20 if the environment variable is not set.
     """
     return int(os.environ.get("ASYNC_LAMBDA_BATCH_FAILURE_RETRY_COUNT", 20))
+
+
+def get_delay_scheduler_role_arn() -> str:
+    """
+    Retrieves the IAM role ARN used by EventBridge Scheduler to send messages to SQS.
+
+    The role must have sqs:SendMessage permission on the target queues.
+
+    Returns:
+        str: The value of the 'ASYNC_LAMBDA_DELAY_SCHEDULER_ROLE_ARN' environment variable.
+
+    Raises:
+        KeyError: If 'ASYNC_LAMBDA_DELAY_SCHEDULER_ROLE_ARN' is not set in the environment.
+    """
+    return os.environ["ASYNC_LAMBDA_DELAY_SCHEDULER_ROLE_ARN"]
+
+
+def get_delay_schedule_group_name() -> str:
+    """
+    Retrieves the EventBridge Scheduler delay schedule group name for this stack.
+
+    One group is provisioned per CloudFormation stack (set via AWS::Scheduler::ScheduleGroup
+    in the SAM template) and injected as ASYNC_LAMBDA_DELAY_SCHEDULE_GROUP when there internally managed queues.
+
+    Returns:
+        str: The value of the 'ASYNC_LAMBDA_DELAY_SCHEDULE_GROUP' environment variable.
+
+    Raises:
+        KeyError: If 'ASYNC_LAMBDA_DELAY_SCHEDULE_GROUP' is not set in the environment.
+    """
+    return os.environ["ASYNC_LAMBDA_DELAY_SCHEDULE_GROUP"]