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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: async-lambda-unstable
-Version: 0.6.6
+Version: 0.6.8
 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.8}/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.8"

{async_lambda_unstable-0.6.6 → async_lambda_unstable-0.6.8}/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.8}/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)
@@ -423,6 +425,12 @@ class AsyncLambdaController:
                                 "Ref": "AsyncLambdaPayloadBucket"
                             },
                             "ASYNC_LAMBDA_ACCOUNT_ID": {"Ref": "AWS::AccountId"},
+                            "ASYNC_LAMBDA_DELAY_SCHEDULER_ROLE_ARN": {
+                                "Fn::GetAtt": ["AsyncLambdaDelaySchedulerRole", "Arn"]
+                            },
+                            "ASYNC_LAMBDA_DELAY_SCHEDULE_GROUP": {
+                                "Ref": "AsyncLambdaDelayScheduleGroup"
+                            },
                             **build_config.environment_variables,
                         },
                     },
@@ -451,6 +459,63 @@ class AsyncLambdaController:
                         ),
                     },
                 },
+                "AsyncLambdaDelayDLQ": {
+                    "Type": "AWS::SQS::Queue",
+                    "Properties": {
+                        "QueueName": f"{config.name}-delay-dlq",
+                        "MessageRetentionPeriod": 1_209_600,  # 14 days
+                        "Tags": make_cf_tags(
+                            {
+                                **build_config.tags,
+                                "async-lambda-queue-type": "dlq",
+                            }
+                        ),
+                    },
+                },
+                "AsyncLambdaDelayScheduleGroup": {
+                    "Type": "AWS::Scheduler::ScheduleGroup",
+                    "Properties": {
+                        "Name": f"{config.name}-delay",
+                        "Tags": make_cf_tags(build_config.tags),
+                    },
+                },
+                "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"
+                                        }
+                                    },
+                                },
+                            ],
+                        },
+                    },
+                },
             },
         }
         _task_list = list(self.tasks.values())
@@ -466,13 +531,61 @@ class AsyncLambdaController:
                 external_async_task_id
             )
         ]
-        task_ref_policies = {}
-        if len(managed_tasks_resources) > 0:
-            task_ref_policies = self._build_send_to_all_async_lambda_queues_policies(
+
+        send_all_queues_task_ref_policies = (
+            self._build_send_to_all_async_lambda_queues_policies(
                 managed_tasks_resources
             )
-            for key in task_ref_policies.keys():
-                template["Resources"][key] = task_ref_policies[key]
+        )
+        for key in send_all_queues_task_ref_policies.keys():
+            template["Resources"][key] = send_all_queues_task_ref_policies[key]
+
+        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_ref}
+                    for policy_ref in send_all_queues_task_ref_policies.keys()
+                ],
+                "Policies": [
+                    {
+                        "PolicyName": "SendToDelayDLQ",
+                        "PolicyDocument": {
+                            "Version": "2012-10-17",
+                            "Statement": [
+                                {
+                                    "Sid": "SendToDelayDLQ",
+                                    "Effect": "Allow",
+                                    "Action": "sqs:SendMessage",
+                                    "Resource": {
+                                        "Fn::GetAtt": [
+                                            "AsyncLambdaDelayDLQ",
+                                            "Arn",
+                                        ]
+                                    },
+                                },
+                            ],
+                        },
+                    },
+                ],
+            },
+        }
+        task_ref_policies = {
+            **send_all_queues_task_ref_policies,
+            "AsyncLambdaCreateDelaySchedulesPolicy": template["Resources"][
+                "AsyncLambdaCreateDelaySchedulesPolicy"
+            ],
+        }
 
         has_api_tasks = False
         for task in _task_list:
@@ -768,6 +881,7 @@ class AsyncLambdaController:
         force_sync: bool = False,
         lane: Optional[int] = None,
         message_group_id: Optional[str] = None,
+        unique_delay_id: Optional[str] = None,
     ):
         """
         Sends an asynchronous invocation payload to a managed or external task via SQS.
@@ -783,6 +897,8 @@ class AsyncLambdaController:
             delay (int, optional): Delay in seconds before sending the message. Defaults to 0.
             force_sync (bool, optional): If True, invokes the task synchronously. Defaults to False.
             lane (Optional[int], optional): The lane to use for invocation. If None, lane assignment is determined automatically.
+            message_group_id (Optional[str], optional): Optional message group ID for enabling SQS Fair Queues. Defaults to None.
+            unique_delay_id (Optional[str], optional): Unique name for EventBridge scheduled delays when delay is longer than 900 seconds. Random if not provided.
 
         Returns:
             Any: The result of the synchronous invocation if `force_sync` is True; otherwise, None.
@@ -808,6 +924,7 @@ class AsyncLambdaController:
                 raise Exception(
                     f"No such task exists with the task_id {destination_task_id}"
                 )
+
         destination_task = None
         if not is_external_task:
             destination_task = self.tasks[destination_task_id]
@@ -878,16 +995,72 @@ 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:
+                _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,
-            )
+                get_sqs_client().send_message(
+                    QueueUrl=url,
+                    MessageBody=json.dumps(sqs_payload),
+                    DelaySeconds=delay,
+                    **_kwargs,
+                )
+            else:
+                self._send_via_scheduler(
+                    queue_url=url,
+                    message_body=json.dumps(sqs_payload),
+                    delay=delay,
+                    schedule_name=unique_delay_id,
+                    message_group_id=_message_group_id,
+                )
+
+            return None
+
+    @staticmethod
+    def _send_via_scheduler(
+        queue_url: str,
+        message_body: str,
+        delay: int,
+        schedule_name: Optional[str] = None,
+        message_group_id: Optional[str] = None,
+    ):
+        """
+        Schedules a one-time EventBridge scheduler SQS sendMessage at some point in the future.
+
+        Args:
+            schedule_name (str): Unique name for the schedule.
+            queue_url (str): URL 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')})"
+
+        sqs_input: dict = {
+            "QueueUrl": queue_url,
+            "MessageBody": message_body,
+        }
+        if message_group_id:
+            sqs_input["MessageGroupId"] = message_group_id
+
+        sqs_universal_target: dict = {
+            "Arn": "arn:aws:scheduler:::aws-sdk:sqs:sendMessage",
+            "RoleArn": env.get_delay_scheduler_role_arn(),
+            "Input": json.dumps(sqs_input),
+            "DeadLetterConfig": {"Arn": env.get_delay_schedule_dlq_arn()},
+        }
+
+        get_scheduler_client().create_schedule(
+            Name=schedule_name or uuid4(),
+            GroupName=env.get_delay_schedule_group_name(),
+            ScheduleExpression=schedule_expression,
+            ScheduleExpressionTimezone="UTC",
+            FlexibleTimeWindow={"Mode": "OFF"},
+            Target=sqs_universal_target,
+            ActionAfterCompletion="DELETE",
+        )
 
     def send_async_invoke_payload_batch(
         self,
@@ -1172,6 +1345,7 @@ class AsyncLambdaController:
         force_sync: bool = False,
         lane: Optional[int] = None,
         message_group_id: Optional[str] = None,
+        unique_delay_id: Optional[str] = None,
     ):
         """
         Asynchronously invokes a task by sending a payload to the specified destination.
@@ -1186,6 +1360,8 @@ class AsyncLambdaController:
             delay (int, optional): Delay in seconds before invoking the task. Defaults to 0.
             force_sync (bool, optional): If True, forces synchronous invocation. Defaults to False.
             lane (Optional[int], optional): Optional lane identifier for routing. Defaults to None.
+            message_group_id (Optional[str], optional): Optional message group ID. Defaults to None.
+            unique_delay_id (Optional[str], optional): Unique name for EventBridge scheduled delays when delay is longer than 900 seconds. Random if not provided.
 
         Returns:
             Any: The result of sending the asynchronous invocation payload.
@@ -1201,6 +1377,7 @@ class AsyncLambdaController:
             force_sync=force_sync,
             lane=lane,
             message_group_id=message_group_id,
+            unique_delay_id=unique_delay_id,
         )
 
     def async_invoke_batch(

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

@@ -206,3 +206,50 @@ 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"]
+
+
+def get_delay_schedule_dlq_arn() -> str:
+    """
+    Retrieves the EventBridge Scheduler delay DLQ for this stack.
+
+    One SQS DLQ is provisioned per CloudFormation stack. EventBridge Scheduler performs a very large number of retries,
+    so ending up here would most likely be due to permissions or target deletion.
+
+    Returns:
+        str: The value of the 'ASYNC_LAMBDA_DELAY_SCHEDULE_DLQ_ARN' environment variable.
+
+    Raises:
+        KeyError: If 'ASYNC_LAMBDA_DELAY_SCHEDULE_DLQ_ARN' is not set in the environment.
+    """
+    return os.environ["ASYNC_LAMBDA_DELAY_SCHEDULE_DLQ_ARN"]