async-lambda-unstable 0.4.18__tar.gz → 0.5.2__tar.gz

{async-lambda-unstable-0.4.18 → async-lambda-unstable-0.5.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: async-lambda-unstable
-Version: 0.4.18
+Version: 0.5.2
 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.18 → async-lambda-unstable-0.5.2}/async_lambda/__init__.py

@@ -19,4 +19,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.18"
+__version__ = "0.5.2"

{async-lambda-unstable-0.4.18 → async-lambda-unstable-0.5.2}/async_lambda/client.py

@@ -7,10 +7,12 @@ import boto3
 class Clients:
     s3_client: Optional[Any] = None
     sqs_client: Optional[Any] = None
+    sts_client: Optional[Any] = None
 
     def reset(self):
         self.s3_client = None
         self.sqs_client = None
+        self.sts_client = None
 
 
 clients = Clients()
@@ -33,3 +35,9 @@ def get_sqs_client():
     if clients.sqs_client is None:
         clients.sqs_client = boto3.client("sqs", **get_client_kwargs())
     return clients.sqs_client
+
+
+def get_sts_client():
+    if clients.sts_client is None:
+        clients.sts_client = boto3.client("sts", **get_client_kwargs())
+    return clients.sts_client

{async-lambda-unstable-0.4.18 → async-lambda-unstable-0.5.2}/async_lambda/controller.py

@@ -4,7 +4,18 @@ import logging
 import re
 import time
 from datetime import datetime, timezone
-from typing import Any, Callable, Dict, List, Optional, Sequence, Type, TypeVar, Union
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Sequence,
+    Set,
+    Type,
+    TypeVar,
+    Union,
+)
 from uuid import uuid4
 
 from . import env
@@ -41,6 +52,7 @@ class AsyncLambdaController:
     lane_count: Optional[int] = None
     propagate_lane_assignment: Optional[bool] = None
     tasks: Dict[str, AsyncLambdaTask]
+    external_async_tasks: Set[str]
     current_task_id: Optional[str] = None
     current_lane: Optional[int] = None
     current_invocation_id: Optional[str] = None
@@ -61,6 +73,7 @@ class AsyncLambdaController:
         controller_name: Optional[str] = None,
     ):
         self.tasks = dict()
+        self.external_async_tasks = set()
         self.is_sub = is_sub
         self.lane_count = lane_count
         self.propagate_lane_assignment = propagate_lane_assignment
@@ -97,6 +110,16 @@ class AsyncLambdaController:
             )
         self.tasks[task.task_id] = task
 
+    def add_external_task(self, external_task_id: str):
+        """
+        Adds an external task to the controller, which will then be enabled to async_invoke.
+
+        `external_task_id` is `{config.name}-{task_id}` of the external task (the queue name)
+
+        Sending payloads via S3 is not currently supported with external tasks.
+        """
+        self.external_async_tasks.add(external_task_id)
+
     def get_lane_count(self) -> int:
         if self.lane_count is not None:
             return self.lane_count
@@ -193,7 +216,11 @@ class AsyncLambdaController:
             if task.trigger_type == TaskTriggerType.API_EVENT:
                 has_api_tasks = True
             for logical_id, resource in task.get_sam_template(
-                module, _task_list, config_dict, stage
+                module=module,
+                tasks=_task_list,
+                external_async_tasks=self.external_async_tasks,
+                config_dict=config_dict,
+                stage=stage,
             ).items():
                 template["Resources"][logical_id] = resource
 
@@ -354,28 +381,44 @@ class AsyncLambdaController:
                 force_sync=force_sync,
                 lane=lane,
             )
+        is_external_task = False
         if destination_task_id not in self.tasks:
-            raise Exception(
-                f"No such task exists with the task_id {destination_task_id}"
-            )
+            if destination_task_id in self.external_async_tasks:
+                is_external_task = True
+            else:
+                raise Exception(
+                    f"No such task exists with the task_id {destination_task_id}"
+                )
+        if not is_external_task:
+            destination_task = self.tasks[destination_task_id]
+            if destination_task.trigger_type not in MANAGED_SQS_TASK_TYPES:
+                raise Exception(
+                    f"Unable to invoke task '{destination_task_id}' because it is a {destination_task.trigger_type} task"
+                )
 
-        destination_task = self.tasks[destination_task_id]
-        if destination_task.trigger_type not in MANAGED_SQS_TASK_TYPES:
-            raise Exception(
-                f"Unable to invoke task '{destination_task_id}' because it is a {destination_task.trigger_type} task"
-            )
+            if lane is None and destination_task.should_propagate_lane_assignment():
+                lane = self.get_current_lane()
+            if lane is None:
+                lane = 0
 
-        if lane is None and destination_task.should_propagate_lane_assignment():
-            lane = self.get_current_lane()
-        if lane is None:
+            if lane < 0 or lane >= destination_task.get_lane_count():
+                raise Exception(
+                    f"Unable to invoke task {destination_task_id} in lane {lane} because it is not a valid lane for the task."
+                )
+        else:
+            if lane is None:
+                lane = self.get_current_lane()
+            if lane != 0:
+                logger.warning(
+                    f"Selected lane is {lane}. External tasks only support lane 0."
+                )
             lane = 0
 
-        if lane < 0 or lane >= destination_task.get_lane_count():
-            raise Exception(
-                f"Unable to invoke task {destination_task_id} in lane {lane} because it is not a valid lane for the task."
-            )
-
         if force_sync or env.get_force_sync_mode():
+            if is_external_task:
+                raise NotImplementedError(
+                    f"Unable to run external task {destination_task_id} locally."
+                )
             if delay:
                 time.sleep(delay)
             # Sync invocation with mock event/context
@@ -393,8 +436,12 @@ class AsyncLambdaController:
             self.set_current_lane(current_lane)
             return result
         else:
+            if is_external_task:
+                url = f"https://sqs.{env.get_aws_region()}.amazonaws.com/{env.get_aws_account_id()}/{destination_task_id}"
+            else:
+                url = destination_task.get_managed_queue_url(lane=lane)
             get_sqs_client().send_message(
-                QueueUrl=destination_task.get_managed_queue_url(lane=lane),
+                QueueUrl=url,
                 MessageBody=json.dumps(sqs_payload),
                 DelaySeconds=delay,
             )
@@ -419,29 +466,45 @@ class AsyncLambdaController:
                 force_sync=force_sync,
                 lane=lane,
             )
+
+        is_external_task = False
         if destination_task_id not in self.tasks:
-            raise Exception(
-                f"No such task exists with the task_id {destination_task_id}"
-            )
+            if destination_task_id in self.external_async_tasks:
+                is_external_task = True
+            else:
+                raise Exception(
+                    f"No such task exists with the task_id {destination_task_id}"
+                )
+        if not is_external_task:
+            destination_task = self.tasks[destination_task_id]
+            if destination_task.trigger_type not in MANAGED_SQS_TASK_TYPES:
+                raise Exception(
+                    f"Unable to invoke task '{destination_task_id}' because it is a {destination_task.trigger_type} task"
+                )
 
-        destination_task = self.tasks[destination_task_id]
-        if destination_task.trigger_type not in MANAGED_SQS_TASK_TYPES:
-            raise Exception(
-                f"Unable to invoke task '{destination_task_id}' because it is a {destination_task.trigger_type} task"
-            )
+            if lane is None and destination_task.should_propagate_lane_assignment():
+                lane = self.get_current_lane()
+            if lane is None:
+                lane = 0
 
-        if lane is None and destination_task.should_propagate_lane_assignment():
-            lane = self.get_current_lane()
-        if lane is None:
+            if lane < 0 or lane >= destination_task.get_lane_count():
+                raise Exception(
+                    f"Unable to invoke task {destination_task_id} in lane {lane} because it is not a valid lane for the task."
+                )
+        else:
+            if lane is None:
+                lane = self.get_current_lane()
+            if lane != 0:
+                logger.warning(
+                    f"Selected lane is {lane}. External tasks only support lane 0."
+                )
             lane = 0
 
-        if lane < 0 or lane >= destination_task.get_lane_count():
-            raise Exception(
-                f"Unable to invoke task {destination_task_id} in lane {lane} because it is not a valid lane for the task."
-            )
-
         if force_sync or env.get_force_sync_mode():
-
+            if is_external_task:
+                raise NotImplementedError(
+                    f"Unable to run external task {destination_task_id} locally."
+                )
             # Sync invocation with mock event/context
             current_task_id = self.current_task_id
             current_lane = self.get_current_lane()
@@ -462,7 +525,7 @@ class AsyncLambdaController:
                 self.set_current_lane(current_lane)
                 self.set_current_task_id(current_task_id)
         else:
-            entries = []
+            entries: List[dict] = []
             for i, sqs_payload in enumerate(sqs_payloads):
                 if isinstance(delay, Sequence):
                     _delay = delay[i]
@@ -475,13 +538,27 @@ class AsyncLambdaController:
                         "Id": f"index_{index + i}",
                     }
                 )
-            response = get_sqs_client().send_message_batch(
-                QueueUrl=destination_task.get_managed_queue_url(lane=lane),
-                Entries=entries,
-            )
-            if len(response.get("Failed", [])) > 0:
-                logger.error(response["Failed"])
-                raise Exception(f"Failed to send {len(response['Failed'])} messages.")
+            if is_external_task:
+                url = f"https://sqs.{env.get_aws_region()}.amazonaws.com/{env.get_aws_account_id()}/{destination_task_id}"
+            else:
+                url = destination_task.get_managed_queue_url(lane=lane)
+            failed_messages = []
+            for _ in range(5):
+                response = get_sqs_client().send_message_batch(
+                    QueueUrl=url,
+                    Entries=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
+                ]
+            logger.error(failed_messages)
+            raise Exception(f"Failed to send {len(failed_messages)} messages.")
 
     def new_payload(
         self,
@@ -504,6 +581,14 @@ class AsyncLambdaController:
             "destination_task_id": destination_task_id,
             "invocation_id": invocation_id,
         }
+        is_external_task = False
+        if (
+            self.get_task(destination_task_id) is None
+            and destination_task_id in self.external_async_tasks
+        ):
+            is_external_task = True
+            raw_sqs_body["source_app"] = config.name
+
         serialized_payload = json.dumps(payload, cls=PayloadEncoder)
         payload_size = len(serialized_payload.encode())
         if payload_size < 250_000:  # we need to double encode to be sure
@@ -512,6 +597,10 @@ class AsyncLambdaController:
         if force_sync or env.get_force_sync_mode():
             raw_sqs_body["payload"] = serialized_payload
         elif payload_size >= 250_000:  # payload is bigger than max SQS size
+            if is_external_task:
+                raise NotImplementedError(
+                    "Payload is too large for SQS and S3 payloads are not supported with external invocations"
+                )
             date_part = datetime.now(tz=timezone.utc).strftime("%Y/%m/%d")
             key = f"{date_part}/{uuid4().hex}.json"
             logger.info(f"Utilizing S3 Payload because of payload size. Key: {key}")
@@ -529,6 +618,9 @@ class AsyncLambdaController:
         for task in controller.tasks.values():
             self.add_task(task)
 
+        for task_id in controller.external_async_tasks:
+            self.external_async_tasks.add(task_id)
+
     def async_invoke(
         self,
         task_id: str,

async-lambda-unstable-0.5.2/async_lambda/env.py

@@ -0,0 +1,81 @@
+import logging
+import os
+from typing import Optional
+
+from .client import get_sts_client
+
+
+class AWSConfig:
+    _aws_region: Optional[str] = None
+    _account_id: Optional[str] = None
+
+    @property
+    def aws_region(self):
+        if self._aws_region:
+            return self._aws_region
+        self._aws_region = os.environ.get(
+            "AWS_REGION", os.environ.get("AWS_DEFAULT_REGION")
+        )
+        if self._aws_region is None:
+            raise ValueError("Unable to find AWS_REGION for constructing ARN or URL")
+        return self._aws_region
+
+    @property
+    def account_id(self):
+        if self._account_id:
+            return self._account_id
+        self._account_id = os.environ.get("ASYNC_LAMBDA_ACCOUNT_ID")
+        if self._account_id is not None:
+            return self._account_id
+        self._account_id = get_sts_client().get_caller_identity().get("Account")
+        logging.info(f"Fetched account_id from sts: {self._account_id}")
+        if self._account_id is None:
+            raise ValueError("Unable to get ACCOUNT_ID from env or STS.")
+        return self._account_id
+
+    def reset(self):
+        self._aws_region = None
+        self._account_id = None
+
+
+aws_config = AWSConfig()
+
+
+def reset():
+    aws_config.reset()
+
+
+def is_build_mode() -> bool:
+    return bool(os.environ.get("ASYNC_LAMBDA_BUILD_MODE", False))
+
+
+def get_aws_region() -> str:
+    return aws_config.aws_region
+
+
+def get_aws_account_id() -> str:
+    return aws_config.account_id
+
+
+def get_payload_bucket() -> str:
+    return os.environ["ASYNC_LAMBDA_PAYLOAD_S3_BUCKET"]
+
+
+def get_current_task_id() -> str:
+    return os.environ["ASYNC_LAMBDA_TASK_ID"]
+
+
+def is_cloud() -> bool:
+    return bool(os.environ.get("AWS_LAMBDA_FUNCTION_NAME"))
+
+
+def enable_force_sync_mode():
+    os.environ["ASYNC_LAMBDA_FORCE_SYNC"] = "1"
+
+
+def disable_force_sync_mode():
+    del os.environ["ASYNC_LAMBDA_FORCE_SYNC"]
+
+
+def get_force_sync_mode() -> bool:
+    return bool(os.environ.get("ASYNC_LAMBDA_FORCE_SYNC", ""))

{async-lambda-unstable-0.4.18 → async-lambda-unstable-0.5.2}/async_lambda/models/task.py

@@ -2,7 +2,17 @@ import inspect
 import json
 import re
 from enum import Enum
-from typing import TYPE_CHECKING, Any, Callable, Generic, List, Optional, TypeVar, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Generic,
+    List,
+    Optional,
+    Set,
+    TypeVar,
+    Union,
+)
 
 from async_lambda.util import make_cf_tags
 
@@ -322,10 +332,21 @@ class AsyncLambdaTask(Generic[EventType, RT]):
             return [self.trigger_config["queue_arn"]]
         return []
 
+    @classmethod
+    def get_policy_external_task_resources(cls, external_task_id: str) -> list:
+        return [
+            {
+                "Fn::Sub": "arn:aws:sqs:${AWS::Region}:${AWS::AccountId}:external_task_id".replace(
+                    "external_task_id", external_task_id
+                )
+            }
+        ]
+
     def get_sam_template(
         self,
         module: str,
         tasks: List["AsyncLambdaTask"],
+        external_async_tasks: Set[str],
         config_dict: dict,
         stage: Optional[str] = None,
     ) -> dict:
@@ -358,6 +379,12 @@ class AsyncLambdaTask(Generic[EventType, RT]):
             for task in tasks
             if task.trigger_type in MANAGED_SQS_TASK_TYPES
             for resource in task.get_policy_sqs_resources()
+        ] + [
+            resource
+            for external_async_task_id in external_async_tasks
+            for resource in self.get_policy_external_task_resources(
+                external_async_task_id
+            )
         ]
         if len(managed_tasks_resources) > 0:
             policy_statements.append(

{async-lambda-unstable-0.4.18 → async-lambda-unstable-0.5.2}/async_lambda_unstable.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: async-lambda-unstable
-Version: 0.4.18
+Version: 0.5.2
 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.18/async_lambda/env.py

@@ -1,37 +0,0 @@
-import os
-
-
-def is_build_mode() -> bool:
-    return bool(os.environ.get("ASYNC_LAMBDA_BUILD_MODE", False))
-
-
-def get_aws_region() -> str:
-    return os.environ.get("AWS_REGION", "local")
-
-
-def get_aws_account_id() -> str:
-    return os.environ.get("ASYNC_LAMBDA_ACCOUNT_ID", "localaccount")
-
-
-def get_payload_bucket() -> str:
-    return os.environ["ASYNC_LAMBDA_PAYLOAD_S3_BUCKET"]
-
-
-def get_current_task_id() -> str:
-    return os.environ["ASYNC_LAMBDA_TASK_ID"]
-
-
-def is_cloud() -> bool:
-    return bool(os.environ.get("AWS_LAMBDA_FUNCTION_NAME"))
-
-
-def enable_force_sync_mode():
-    os.environ["ASYNC_LAMBDA_FORCE_SYNC"] = "1"
-
-
-def disable_force_sync_mode():
-    del os.environ["ASYNC_LAMBDA_FORCE_SYNC"]
-
-
-def get_force_sync_mode() -> bool:
-    return bool(os.environ.get("ASYNC_LAMBDA_FORCE_SYNC", ""))