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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (30) hide show
  1. {async_lambda_unstable-0.6.7 → async_lambda_unstable-0.6.8}/PKG-INFO +1 -1
  2. {async_lambda_unstable-0.6.7 → async_lambda_unstable-0.6.8}/async_lambda/__init__.py +1 -1
  3. {async_lambda_unstable-0.6.7 → async_lambda_unstable-0.6.8}/async_lambda/controller.py +260 -256
  4. {async_lambda_unstable-0.6.7 → async_lambda_unstable-0.6.8}/async_lambda/env.py +16 -0
  5. {async_lambda_unstable-0.6.7 → async_lambda_unstable-0.6.8}/.gitignore +0 -0
  6. {async_lambda_unstable-0.6.7 → async_lambda_unstable-0.6.8}/README.md +0 -0
  7. {async_lambda_unstable-0.6.7 → async_lambda_unstable-0.6.8}/async_lambda/build_config.py +0 -0
  8. {async_lambda_unstable-0.6.7 → async_lambda_unstable-0.6.8}/async_lambda/cli.py +0 -0
  9. {async_lambda_unstable-0.6.7 → async_lambda_unstable-0.6.8}/async_lambda/client.py +0 -0
  10. {async_lambda_unstable-0.6.7 → async_lambda_unstable-0.6.8}/async_lambda/config.py +0 -0
  11. {async_lambda_unstable-0.6.7 → async_lambda_unstable-0.6.8}/async_lambda/defer.py +0 -0
  12. {async_lambda_unstable-0.6.7 → async_lambda_unstable-0.6.8}/async_lambda/middleware.py +0 -0
  13. {async_lambda_unstable-0.6.7 → async_lambda_unstable-0.6.8}/async_lambda/models/__init__.py +0 -0
  14. {async_lambda_unstable-0.6.7 → async_lambda_unstable-0.6.8}/async_lambda/models/api_response.py +0 -0
  15. {async_lambda_unstable-0.6.7 → async_lambda_unstable-0.6.8}/async_lambda/models/case_insensitive_dict.py +0 -0
  16. {async_lambda_unstable-0.6.7 → async_lambda_unstable-0.6.8}/async_lambda/models/events/__init__.py +0 -0
  17. {async_lambda_unstable-0.6.7 → async_lambda_unstable-0.6.8}/async_lambda/models/events/api_event.py +0 -0
  18. {async_lambda_unstable-0.6.7 → async_lambda_unstable-0.6.8}/async_lambda/models/events/base_event.py +0 -0
  19. {async_lambda_unstable-0.6.7 → async_lambda_unstable-0.6.8}/async_lambda/models/events/dynamodb_event.py +0 -0
  20. {async_lambda_unstable-0.6.7 → async_lambda_unstable-0.6.8}/async_lambda/models/events/managed_sqs_batch_event.py +0 -0
  21. {async_lambda_unstable-0.6.7 → async_lambda_unstable-0.6.8}/async_lambda/models/events/managed_sqs_event.py +0 -0
  22. {async_lambda_unstable-0.6.7 → async_lambda_unstable-0.6.8}/async_lambda/models/events/scheduled_event.py +0 -0
  23. {async_lambda_unstable-0.6.7 → async_lambda_unstable-0.6.8}/async_lambda/models/events/unmanaged_sqs_event.py +0 -0
  24. {async_lambda_unstable-0.6.7 → async_lambda_unstable-0.6.8}/async_lambda/models/mock/mock_context.py +0 -0
  25. {async_lambda_unstable-0.6.7 → async_lambda_unstable-0.6.8}/async_lambda/models/mock/mock_event.py +0 -0
  26. {async_lambda_unstable-0.6.7 → async_lambda_unstable-0.6.8}/async_lambda/models/task.py +0 -0
  27. {async_lambda_unstable-0.6.7 → async_lambda_unstable-0.6.8}/async_lambda/payload_encoder.py +0 -0
  28. {async_lambda_unstable-0.6.7 → async_lambda_unstable-0.6.8}/async_lambda/py.typed +0 -0
  29. {async_lambda_unstable-0.6.7 → async_lambda_unstable-0.6.8}/async_lambda/util.py +0 -0
  30. {async_lambda_unstable-0.6.7 → async_lambda_unstable-0.6.8}/pyproject.toml +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: async-lambda-unstable
3
- Version: 0.6.7
3
+ Version: 0.6.8
4
4
  Summary: A framework for creating AWS Lambda Async Workflows. - Unstable Branch
5
5
  Author-email: "Nuclei, Inc" <engineering@nuclei.ai>
6
6
  Requires-Dist: click>=8.0.0
@@ -21,4 +21,4 @@ from .models.events.managed_sqs_event import ManagedSQSEvent as ManagedSQSEvent
21
21
  from .models.events.scheduled_event import ScheduledEvent as ScheduledEvent
22
22
  from .models.events.unmanaged_sqs_event import UnmanagedSQSEvent as UnmanagedSQSEvent
23
23
 
24
- __version__ = "0.6.7"
24
+ __version__ = "0.6.8"
@@ -69,30 +69,6 @@ class BatchInvokeException(Exception):
69
69
  super().__init__(msg)
70
70
 
71
71
 
72
- # Exception raised when async_invoke_* is called to a non-internal task with a delay greater than _SQS_MAX_DELAY_SECONDS
73
- class AsyncInvokeInvalidDelay(Exception):
74
- pass
75
-
76
-
77
- class _BatchEntry:
78
- enqueue_id: str
79
- payload: dict
80
- delay: int
81
- message_group_id: Optional[str]
82
-
83
- def __init__(
84
- self,
85
- enqueue_id: str,
86
- payload: dict,
87
- delay: int,
88
- message_group_id: Optional[str],
89
- ):
90
- self.enqueue_id = enqueue_id
91
- self.payload = payload
92
- self.delay = delay
93
- self.message_group_id = message_group_id
94
-
95
-
96
72
  class AsyncLambdaController:
97
73
  """
98
74
  AsyncLambdaController manages async tasks, middleware, and invocation logic for the async-lambda framework.
@@ -449,6 +425,12 @@ class AsyncLambdaController:
449
425
  "Ref": "AsyncLambdaPayloadBucket"
450
426
  },
451
427
  "ASYNC_LAMBDA_ACCOUNT_ID": {"Ref": "AWS::AccountId"},
428
+ "ASYNC_LAMBDA_DELAY_SCHEDULER_ROLE_ARN": {
429
+ "Fn::GetAtt": ["AsyncLambdaDelaySchedulerRole", "Arn"]
430
+ },
431
+ "ASYNC_LAMBDA_DELAY_SCHEDULE_GROUP": {
432
+ "Ref": "AsyncLambdaDelayScheduleGroup"
433
+ },
452
434
  **build_config.environment_variables,
453
435
  },
454
436
  },
@@ -477,102 +459,133 @@ class AsyncLambdaController:
477
459
  ),
478
460
  },
479
461
  },
462
+ "AsyncLambdaDelayDLQ": {
463
+ "Type": "AWS::SQS::Queue",
464
+ "Properties": {
465
+ "QueueName": f"{config.name}-delay-dlq",
466
+ "MessageRetentionPeriod": 1_209_600, # 14 days
467
+ "Tags": make_cf_tags(
468
+ {
469
+ **build_config.tags,
470
+ "async-lambda-queue-type": "dlq",
471
+ }
472
+ ),
473
+ },
474
+ },
475
+ "AsyncLambdaDelayScheduleGroup": {
476
+ "Type": "AWS::Scheduler::ScheduleGroup",
477
+ "Properties": {
478
+ "Name": f"{config.name}-delay",
479
+ "Tags": make_cf_tags(build_config.tags),
480
+ },
481
+ },
482
+ "AsyncLambdaCreateDelaySchedulesPolicy": {
483
+ "Type": "AWS::IAM::ManagedPolicy",
484
+ "Properties": {
485
+ "ManagedPolicyName": {
486
+ "Fn::Sub": "${AWS::StackName}-create-delay-schedules"
487
+ },
488
+ "PolicyDocument": {
489
+ "Version": "2012-10-17",
490
+ "Statement": [
491
+ {
492
+ "Sid": "CreateSchedules",
493
+ "Effect": "Allow",
494
+ "Action": "scheduler:CreateSchedule",
495
+ "Resource": {
496
+ "Fn::Sub": "arn:aws:scheduler:${AWS::Region}:${AWS::AccountId}:schedule/${AsyncLambdaDelayScheduleGroup}/*"
497
+ },
498
+ },
499
+ {
500
+ "Sid": "PassSchedulerExecutionRole",
501
+ "Effect": "Allow",
502
+ "Action": "iam:PassRole",
503
+ "Resource": {
504
+ "Fn::GetAtt": [
505
+ "AsyncLambdaDelaySchedulerRole",
506
+ "Arn",
507
+ ]
508
+ },
509
+ "Condition": {
510
+ "StringEquals": {
511
+ "iam:PassedToService": "scheduler.amazonaws.com"
512
+ }
513
+ },
514
+ },
515
+ ],
516
+ },
517
+ },
518
+ },
480
519
  },
481
520
  }
482
521
  _task_list = list(self.tasks.values())
483
- internal_tasks_resources = [
522
+ managed_tasks_resources = [
484
523
  resource
485
524
  for task in _task_list
486
525
  if task.trigger_type in MANAGED_SQS_TASK_TYPES
487
526
  for resource in task.get_policy_sqs_resources()
488
- ]
489
- external_tasks_resources = [
527
+ ] + [
490
528
  resource
491
529
  for external_async_task_id in self.external_async_tasks
492
530
  for resource in AsyncLambdaTask.get_policy_external_task_resources(
493
531
  external_async_task_id
494
532
  )
495
533
  ]
496
- managed_tasks_resources = internal_tasks_resources + external_tasks_resources
497
- task_ref_policies = {}
498
- if len(managed_tasks_resources) > 0:
499
- task_ref_policies = self._build_send_to_all_async_lambda_queues_policies(
534
+
535
+ send_all_queues_task_ref_policies = (
536
+ self._build_send_to_all_async_lambda_queues_policies(
500
537
  managed_tasks_resources
501
538
  )
502
- for key in task_ref_policies.keys():
503
- template["Resources"][key] = task_ref_policies[key]
504
-
505
- if len(internal_tasks_resources) > 0:
506
- template["Resources"]["AsyncLambdaDelayScheduleGroup"] = {
507
- "Type": "AWS::Scheduler::ScheduleGroup",
508
- "Properties": {
509
- "Name": f"{config.name}-delay",
510
- "Tags": make_cf_tags(build_config.tags),
511
- },
512
- }
513
- template["Resources"]["AsyncLambdaDelaySchedulerRole"] = {
514
- "Type": "AWS::IAM::Role",
515
- "Properties": {
516
- "AssumeRolePolicyDocument": {
517
- "Version": "2012-10-17",
518
- "Statement": [
519
- {
520
- "Effect": "Allow",
521
- "Principal": {"Service": "scheduler.amazonaws.com"},
522
- "Action": "sts:AssumeRole",
523
- }
524
- ],
525
- },
526
- "ManagedPolicyArns": [
527
- {"Ref": policy_id} for policy_id in task_ref_policies.keys()
539
+ )
540
+ for key in send_all_queues_task_ref_policies.keys():
541
+ template["Resources"][key] = send_all_queues_task_ref_policies[key]
542
+
543
+ template["Resources"]["AsyncLambdaDelaySchedulerRole"] = {
544
+ "Type": "AWS::IAM::Role",
545
+ "Properties": {
546
+ "AssumeRolePolicyDocument": {
547
+ "Version": "2012-10-17",
548
+ "Statement": [
549
+ {
550
+ "Effect": "Allow",
551
+ "Principal": {"Service": "scheduler.amazonaws.com"},
552
+ "Action": "sts:AssumeRole",
553
+ }
528
554
  ],
529
555
  },
530
- }
531
- template["Resources"][
532
- "AsyncLambdaCreateDelaySchedulesPolicy"
533
- ] = task_ref_policies["AsyncLambdaCreateDelaySchedulesPolicy"] = {
534
- "Type": "AWS::IAM::ManagedPolicy",
535
- "Properties": {
536
- "ManagedPolicyName": {
537
- "Fn::Sub": "${AWS::StackName}-create-delay-schedules"
538
- },
539
- "PolicyDocument": {
540
- "Version": "2012-10-17",
541
- "Statement": [
542
- {
543
- "Sid": "CreateSchedules",
544
- "Effect": "Allow",
545
- "Action": "scheduler:CreateSchedule",
546
- "Resource": {
547
- "Fn::Sub": "arn:aws:scheduler:${AWS::Region}:${AWS::AccountId}:schedule/${AsyncLambdaDelayScheduleGroup}/*"
548
- },
549
- },
550
- {
551
- "Sid": "PassSchedulerExecutionRole",
552
- "Effect": "Allow",
553
- "Action": "iam:PassRole",
554
- "Resource": {
555
- "Fn::GetAtt": [
556
- "AsyncLambdaDelaySchedulerRole",
557
- "Arn",
558
- ]
559
- },
560
- "Condition": {
561
- "StringEquals": {
562
- "iam:PassedToService": "scheduler.amazonaws.com"
563
- }
556
+ "ManagedPolicyArns": [
557
+ {"Ref": policy_ref}
558
+ for policy_ref in send_all_queues_task_ref_policies.keys()
559
+ ],
560
+ "Policies": [
561
+ {
562
+ "PolicyName": "SendToDelayDLQ",
563
+ "PolicyDocument": {
564
+ "Version": "2012-10-17",
565
+ "Statement": [
566
+ {
567
+ "Sid": "SendToDelayDLQ",
568
+ "Effect": "Allow",
569
+ "Action": "sqs:SendMessage",
570
+ "Resource": {
571
+ "Fn::GetAtt": [
572
+ "AsyncLambdaDelayDLQ",
573
+ "Arn",
574
+ ]
575
+ },
564
576
  },
565
- },
566
- ],
577
+ ],
578
+ },
567
579
  },
568
- },
569
- }
570
- template["Globals"]["Function"]["Environment"]["Variables"][
571
- "ASYNC_LAMBDA_DELAY_SCHEDULE_GROUP"
572
- ] = {"Ref": "AsyncLambdaDelayScheduleGroup"}
573
- template["Globals"]["Function"]["Environment"]["Variables"][
574
- "ASYNC_LAMBDA_DELAY_SCHEDULER_ROLE_ARN"
575
- ] = {"Fn::GetAtt": ["AsyncLambdaDelaySchedulerRole", "Arn"]}
580
+ ],
581
+ },
582
+ }
583
+ task_ref_policies = {
584
+ **send_all_queues_task_ref_policies,
585
+ "AsyncLambdaCreateDelaySchedulesPolicy": template["Resources"][
586
+ "AsyncLambdaCreateDelaySchedulesPolicy"
587
+ ],
588
+ }
576
589
 
577
590
  has_api_tasks = False
578
591
  for task in _task_list:
@@ -868,6 +881,7 @@ class AsyncLambdaController:
868
881
  force_sync: bool = False,
869
882
  lane: Optional[int] = None,
870
883
  message_group_id: Optional[str] = None,
884
+ unique_delay_id: Optional[str] = None,
871
885
  ):
872
886
  """
873
887
  Sends an asynchronous invocation payload to a managed or external task via SQS.
@@ -883,6 +897,8 @@ class AsyncLambdaController:
883
897
  delay (int, optional): Delay in seconds before sending the message. Defaults to 0.
884
898
  force_sync (bool, optional): If True, invokes the task synchronously. Defaults to False.
885
899
  lane (Optional[int], optional): The lane to use for invocation. If None, lane assignment is determined automatically.
900
+ message_group_id (Optional[str], optional): Optional message group ID for enabling SQS Fair Queues. Defaults to None.
901
+ unique_delay_id (Optional[str], optional): Unique name for EventBridge scheduled delays when delay is longer than 900 seconds. Random if not provided.
886
902
 
887
903
  Returns:
888
904
  Any: The result of the synchronous invocation if `force_sync` is True; otherwise, None.
@@ -909,11 +925,6 @@ class AsyncLambdaController:
909
925
  f"No such task exists with the task_id {destination_task_id}"
910
926
  )
911
927
 
912
- if is_external_task and delay > _SQS_MAX_DELAY_SECONDS:
913
- raise AsyncInvokeInvalidDelay(
914
- f"Unable to invoke task {destination_task_id} with delay {delay} seconds. External tasks only support delays up to {_SQS_MAX_DELAY_SECONDS} seconds."
915
- )
916
-
917
928
  destination_task = None
918
929
  if not is_external_task:
919
930
  destination_task = self.tasks[destination_task_id]
@@ -984,18 +995,7 @@ class AsyncLambdaController:
984
995
  assert destination_task is not None
985
996
  url = destination_task.get_managed_queue_url(lane=lane)
986
997
 
987
- if delay > _SQS_MAX_DELAY_SECONDS:
988
- assert not is_external_task
989
-
990
- queue_arn = destination_task.get_managed_queue_arn(lane=lane)
991
- # TODO: Use _message_group_id when AWS tells me how
992
- self._send_via_scheduler(
993
- queue_arn=queue_arn,
994
- message_body=json.dumps(sqs_payload),
995
- delay=delay,
996
- message_group_id=_message_group_id,
997
- )
998
- else:
998
+ if delay <= _SQS_MAX_DELAY_SECONDS:
999
999
  _kwargs = {}
1000
1000
  if _message_group_id:
1001
1001
  _kwargs["MessageGroupId"] = _message_group_id
@@ -1006,8 +1006,62 @@ class AsyncLambdaController:
1006
1006
  DelaySeconds=delay,
1007
1007
  **_kwargs,
1008
1008
  )
1009
+ else:
1010
+ self._send_via_scheduler(
1011
+ queue_url=url,
1012
+ message_body=json.dumps(sqs_payload),
1013
+ delay=delay,
1014
+ schedule_name=unique_delay_id,
1015
+ message_group_id=_message_group_id,
1016
+ )
1017
+
1009
1018
  return None
1010
1019
 
1020
+ @staticmethod
1021
+ def _send_via_scheduler(
1022
+ queue_url: str,
1023
+ message_body: str,
1024
+ delay: int,
1025
+ schedule_name: Optional[str] = None,
1026
+ message_group_id: Optional[str] = None,
1027
+ ):
1028
+ """
1029
+ Schedules a one-time EventBridge scheduler SQS sendMessage at some point in the future.
1030
+
1031
+ Args:
1032
+ schedule_name (str): Unique name for the schedule.
1033
+ queue_url (str): URL of the destination SQS queue.
1034
+ message_body (str): JSON-serialized message body to deliver to the queue.
1035
+ delay (int): Delay in seconds before delivering the message.
1036
+ message_group_id (Optional[str]): MessageGroupId for FIFO queues. Defaults to None.
1037
+ """
1038
+ schedule_time = datetime.now(tz=timezone.utc) + timedelta(seconds=delay)
1039
+ schedule_expression = f"at({schedule_time.strftime('%Y-%m-%dT%H:%M:%S')})"
1040
+
1041
+ sqs_input: dict = {
1042
+ "QueueUrl": queue_url,
1043
+ "MessageBody": message_body,
1044
+ }
1045
+ if message_group_id:
1046
+ sqs_input["MessageGroupId"] = message_group_id
1047
+
1048
+ sqs_universal_target: dict = {
1049
+ "Arn": "arn:aws:scheduler:::aws-sdk:sqs:sendMessage",
1050
+ "RoleArn": env.get_delay_scheduler_role_arn(),
1051
+ "Input": json.dumps(sqs_input),
1052
+ "DeadLetterConfig": {"Arn": env.get_delay_schedule_dlq_arn()},
1053
+ }
1054
+
1055
+ get_scheduler_client().create_schedule(
1056
+ Name=schedule_name or uuid4(),
1057
+ GroupName=env.get_delay_schedule_group_name(),
1058
+ ScheduleExpression=schedule_expression,
1059
+ ScheduleExpressionTimezone="UTC",
1060
+ FlexibleTimeWindow={"Mode": "OFF"},
1061
+ Target=sqs_universal_target,
1062
+ ActionAfterCompletion="DELETE",
1063
+ )
1064
+
1011
1065
  def send_async_invoke_payload_batch(
1012
1066
  self,
1013
1067
  destination_task_id: str,
@@ -1083,35 +1137,6 @@ class AsyncLambdaController:
1083
1137
  )
1084
1138
  lane = 0
1085
1139
 
1086
- batch_entries: List[_BatchEntry] = []
1087
- for i, sqs_payload in enumerate(sqs_payloads):
1088
- if isinstance(delay, Sequence):
1089
- _delay = delay[i]
1090
- else:
1091
- _delay = delay
1092
-
1093
- _message_group_id = (
1094
- message_group_id[i]
1095
- if isinstance(message_group_id, Sequence)
1096
- and not isinstance(message_group_id, str)
1097
- else message_group_id
1098
- )
1099
- if (
1100
- _message_group_id is None
1101
- and (current_message_group_id := self.get_current_message_group_id())
1102
- is not None
1103
- and self.should_propagate_message_group_id()
1104
- ):
1105
- _message_group_id = current_message_group_id
1106
- batch_entries.append(
1107
- _BatchEntry(
1108
- enqueue_id=f"index_{index + i}",
1109
- payload=sqs_payload,
1110
- delay=_delay,
1111
- message_group_id=_message_group_id,
1112
- )
1113
- )
1114
-
1115
1140
  if force_sync or env.get_force_sync_mode():
1116
1141
  if is_external_task:
1117
1142
  raise NotImplementedError(
@@ -1122,16 +1147,34 @@ class AsyncLambdaController:
1122
1147
  current_lane = self.get_current_lane()
1123
1148
  assert destination_task is not None
1124
1149
  queue_arn = destination_task.get_managed_queue_arn(lane=lane)
1125
-
1126
- for batch in batch_entries:
1127
- if batch.delay:
1128
- time.sleep(batch.delay)
1150
+ for i, sqs_payload in enumerate(sqs_payloads):
1151
+ if delay:
1152
+ if isinstance(delay, Sequence):
1153
+ time.sleep(delay[i])
1154
+ else:
1155
+ time.sleep(delay)
1156
+
1157
+ _message_group_id = (
1158
+ message_group_id[i]
1159
+ if isinstance(message_group_id, Sequence)
1160
+ and not isinstance(message_group_id, str)
1161
+ else message_group_id
1162
+ )
1163
+ if (
1164
+ _message_group_id is None
1165
+ and (
1166
+ current_message_group_id := self.get_current_message_group_id()
1167
+ )
1168
+ is not None
1169
+ and self.should_propagate_message_group_id()
1170
+ ):
1171
+ _message_group_id = current_message_group_id
1129
1172
 
1130
1173
  current_message_group_id = self.get_current_message_group_id()
1131
1174
  mock_event = MockSQSLambdaEvent(
1132
- json.dumps(batch.payload),
1175
+ json.dumps(sqs_payload),
1133
1176
  source_queue_arn=queue_arn,
1134
- message_group_id=batch.message_group_id,
1177
+ message_group_id=_message_group_id,
1135
1178
  )
1136
1179
  mock_context = MockLambdaContext(destination_task.task_id)
1137
1180
  self.handle_invocation(
@@ -1141,112 +1184,69 @@ class AsyncLambdaController:
1141
1184
  self.set_current_task_id(current_task_id)
1142
1185
  self.set_current_message_group_id(current_message_group_id)
1143
1186
  else:
1144
- sqs_entries: List[dict] = []
1145
- for batch in batch_entries:
1146
- if batch.delay > _SQS_MAX_DELAY_SECONDS:
1147
- if is_external_task:
1148
- raise AsyncInvokeInvalidDelay(
1149
- f"send_async_invoke_payload_batch does not support a delay longer than {_SQS_MAX_DELAY_SECONDS} seconds for external tasks"
1150
- )
1151
- assert destination_task is not None
1152
- queue_arn = destination_task.get_managed_queue_arn(lane=lane)
1153
- self._send_via_scheduler(
1154
- queue_arn=queue_arn,
1155
- message_body=json.dumps(batch.payload),
1156
- delay=batch.delay,
1157
- )
1187
+ entries: List[dict] = []
1188
+ for i, sqs_payload in enumerate(sqs_payloads):
1189
+ if isinstance(delay, Sequence):
1190
+ _delay = delay[i]
1158
1191
  else:
1159
- entry: dict = {
1160
- "MessageBody": json.dumps(batch.payload),
1161
- "DelaySeconds": batch.delay,
1162
- "Id": batch.enqueue_id,
1163
- }
1164
- if batch.message_group_id:
1165
- entry["MessageGroupId"] = batch.message_group_id
1166
- sqs_entries.append(entry)
1167
- if sqs_entries:
1168
- if is_external_task:
1169
- url = f"https://sqs.{env.get_aws_region()}.amazonaws.com/{env.get_aws_account_id()}/{destination_task_id}"
1170
- else:
1171
- assert destination_task is not None
1172
- url = destination_task.get_managed_queue_url(lane=lane)
1173
- failed_messages: List[dict] = []
1174
- batch_retry_count = env.get_batch_failure_retry_count() + 1
1175
- entries = sqs_entries
1176
- for i in range(batch_retry_count):
1177
- response = get_sqs_client().send_message_batch(
1178
- QueueUrl=url,
1179
- Entries=entries,
1180
- )
1181
- failed_messages = response.get("Failed", [])
1182
- if len(failed_messages) == 0:
1183
- return
1184
- logger.warning(failed_messages)
1185
- logger.warning(f"{len(failed_messages)} messages failed to send. ")
1186
- failed_message_ids = {message["Id"] for message in failed_messages}
1187
- entries = [
1188
- entry for entry in entries if entry["Id"] in failed_message_ids
1189
- ]
1190
- if i < batch_retry_count:
1191
- send_delay = 0.5 + random.random()
1192
- logger.info(
1193
- f"Waiting {send_delay:.3f} before attempting batch failures again."
1194
- )
1195
- time.sleep(send_delay)
1196
- logger.error(failed_messages)
1197
- raise BatchInvokeException(
1198
- f"Failed to send {len(failed_messages)} messages.",
1199
- failed_payloads=[
1200
- int(entry["Id"].split("_")[-1]) for entry in entries
1201
- ],
1202
- )
1203
-
1204
- @staticmethod
1205
- def _send_via_scheduler(
1206
- queue_arn: str,
1207
- message_body: str,
1208
- delay: int,
1209
- message_group_id: Optional[str] = None,
1210
- ):
1211
- """
1212
- Schedules a message for future delivery to an SQS queue via EventBridge Scheduler.
1192
+ _delay = delay
1213
1193
 
1214
- Used when the requested delay exceeds SQS's maximum of 900 seconds (15 minutes).
1215
- Creates a one-time schedule in the stack-level schedule group (ASYNC_LAMBDA_DELAY_SCHEDULE_GROUP),
1216
- which is provisioned by CloudFormation. The schedule auto-deletes after firing.
1217
-
1218
- Requires the following environment variables:
1219
- - ASYNC_LAMBDA_DELAY_SCHEDULE_GROUP: name of the CloudFormation-managed schedule group.
1220
- - ASYNC_LAMBDA_DELAY_SCHEDULER_ROLE_ARN: ARN of the IAM role EventBridge Scheduler uses
1221
- to deliver to the target SQS queue (must have sqs:SendMessage on all queues in the group).
1222
-
1223
- Args:
1224
- queue_arn (str): ARN of the destination SQS queue.
1225
- message_body (str): JSON-serialized message body to deliver to the queue.
1226
- delay (int): Delay in seconds before delivering the message.
1227
- message_group_id (Optional[str]): MessageGroupId for FIFO queues. Defaults to None.
1228
- """
1229
- schedule_time = datetime.now(tz=timezone.utc) + timedelta(seconds=delay)
1230
- schedule_expression = f"at({schedule_time.strftime('%Y-%m-%dT%H:%M:%S')})"
1231
-
1232
- target: dict = {
1233
- "Arn": queue_arn,
1234
- "RoleArn": env.get_delay_scheduler_role_arn(),
1235
- "Input": message_body,
1236
- }
1237
- # TODO: Set message_group_id based on feedback from AWS
1238
- # if message_group_id:
1239
- # target["SqsParameters"] = {"MessageGroupId": message_group_id}
1194
+ _message_group_id = (
1195
+ message_group_id[i]
1196
+ if isinstance(message_group_id, Sequence)
1197
+ and not isinstance(message_group_id, str)
1198
+ else message_group_id
1199
+ )
1200
+ if (
1201
+ _message_group_id is None
1202
+ and (
1203
+ current_message_group_id := self.get_current_message_group_id()
1204
+ )
1205
+ is not None
1206
+ and self.should_propagate_message_group_id()
1207
+ ):
1208
+ _message_group_id = current_message_group_id
1240
1209
 
1241
- get_scheduler_client().create_schedule(
1242
- Name=uuid4().hex,
1243
- GroupName=env.get_delay_schedule_group_name(),
1244
- ScheduleExpression=schedule_expression,
1245
- ScheduleExpressionTimezone="UTC",
1246
- FlexibleTimeWindow={"Mode": "OFF"},
1247
- Target=target,
1248
- ActionAfterCompletion="DELETE",
1249
- )
1210
+ entry = {
1211
+ "MessageBody": json.dumps(sqs_payload),
1212
+ "DelaySeconds": _delay,
1213
+ "Id": f"index_{index + i}",
1214
+ }
1215
+ if _message_group_id:
1216
+ entry["MessageGroupId"] = _message_group_id
1217
+ entries.append(entry)
1218
+ if is_external_task:
1219
+ url = f"https://sqs.{env.get_aws_region()}.amazonaws.com/{env.get_aws_account_id()}/{destination_task_id}"
1220
+ else:
1221
+ assert destination_task is not None
1222
+ url = destination_task.get_managed_queue_url(lane=lane)
1223
+ failed_messages = []
1224
+ batch_retry_count = env.get_batch_failure_retry_count() + 1
1225
+ for i in range(batch_retry_count):
1226
+ response = get_sqs_client().send_message_batch(
1227
+ QueueUrl=url,
1228
+ Entries=entries,
1229
+ )
1230
+ failed_messages: List[dict] = response.get("Failed", [])
1231
+ if len(failed_messages) == 0:
1232
+ return
1233
+ logger.warning(failed_messages)
1234
+ logger.warning(f"{len(failed_messages)} messages failed to send. ")
1235
+ failed_message_ids = {message["Id"] for message in failed_messages}
1236
+ entries = [
1237
+ entry for entry in entries if entry["Id"] in failed_message_ids
1238
+ ]
1239
+ if i < batch_retry_count:
1240
+ send_delay = 0.5 + random.random()
1241
+ logger.info(
1242
+ f"Waiting {send_delay:.3f} before attempting batch failures again."
1243
+ )
1244
+ time.sleep(send_delay)
1245
+ logger.error(failed_messages)
1246
+ raise BatchInvokeException(
1247
+ f"Failed to send {len(failed_messages)} messages.",
1248
+ failed_payloads=[int(entry["Id"].split("_")[-1]) for entry in entries],
1249
+ )
1250
1250
 
1251
1251
  def new_payload(
1252
1252
  self,
@@ -1345,6 +1345,7 @@ class AsyncLambdaController:
1345
1345
  force_sync: bool = False,
1346
1346
  lane: Optional[int] = None,
1347
1347
  message_group_id: Optional[str] = None,
1348
+ unique_delay_id: Optional[str] = None,
1348
1349
  ):
1349
1350
  """
1350
1351
  Asynchronously invokes a task by sending a payload to the specified destination.
@@ -1359,6 +1360,8 @@ class AsyncLambdaController:
1359
1360
  delay (int, optional): Delay in seconds before invoking the task. Defaults to 0.
1360
1361
  force_sync (bool, optional): If True, forces synchronous invocation. Defaults to False.
1361
1362
  lane (Optional[int], optional): Optional lane identifier for routing. Defaults to None.
1363
+ message_group_id (Optional[str], optional): Optional message group ID. Defaults to None.
1364
+ unique_delay_id (Optional[str], optional): Unique name for EventBridge scheduled delays when delay is longer than 900 seconds. Random if not provided.
1362
1365
 
1363
1366
  Returns:
1364
1367
  Any: The result of sending the asynchronous invocation payload.
@@ -1374,6 +1377,7 @@ class AsyncLambdaController:
1374
1377
  force_sync=force_sync,
1375
1378
  lane=lane,
1376
1379
  message_group_id=message_group_id,
1380
+ unique_delay_id=unique_delay_id,
1377
1381
  )
1378
1382
 
1379
1383
  def async_invoke_batch(
@@ -237,3 +237,19 @@ def get_delay_schedule_group_name() -> str:
237
237
  KeyError: If 'ASYNC_LAMBDA_DELAY_SCHEDULE_GROUP' is not set in the environment.
238
238
  """
239
239
  return os.environ["ASYNC_LAMBDA_DELAY_SCHEDULE_GROUP"]
240
+
241
+
242
+ def get_delay_schedule_dlq_arn() -> str:
243
+ """
244
+ Retrieves the EventBridge Scheduler delay DLQ for this stack.
245
+
246
+ One SQS DLQ is provisioned per CloudFormation stack. EventBridge Scheduler performs a very large number of retries,
247
+ so ending up here would most likely be due to permissions or target deletion.
248
+
249
+ Returns:
250
+ str: The value of the 'ASYNC_LAMBDA_DELAY_SCHEDULE_DLQ_ARN' environment variable.
251
+
252
+ Raises:
253
+ KeyError: If 'ASYNC_LAMBDA_DELAY_SCHEDULE_DLQ_ARN' is not set in the environment.
254
+ """
255
+ return os.environ["ASYNC_LAMBDA_DELAY_SCHEDULE_DLQ_ARN"]