Skip to content

Commit a3371d4

Browse files
aahungaahung
authored
feat: Add ScheduleV2 for AWS::Scheduler::Schedule (#308)
1 parent e742fca commit a3371d4

27 files changed

+2186
-3
lines changed

‎samtranslator/model/eventsources/scheduler.py

Lines changed: 234 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,234 @@
1+
from enum import Enum, auto
2+
from typing import Any, Dict, List, Optional, Tuple, Union, cast
3+
4+
from samtranslator.metrics.method_decorator import cw_timer
5+
from samtranslator.model import PropertyType, Resource, ResourceMacro
6+
from samtranslator.model.iam import IAMRole
7+
from samtranslator.model.sqs import SQSQueue
8+
from samtranslator.model.types import is_str, is_type
9+
from samtranslator.model.eventsources import FUNCTION_EVETSOURCE_METRIC_PREFIX
10+
from samtranslator.model.eventbridge_utils import EventBridgeRuleUtils
11+
from samtranslator.model.exceptions import InvalidEventException
12+
from samtranslator.model.iam import IAMRolePolicies
13+
from samtranslator.model.scheduler import SchedulerSchedule
14+
from samtranslator.translator.logical_id_generator import LogicalIdGenerator
15+
16+
17+
class _SchedulerScheduleTargetType(Enum):
18+
FUNCTION = auto()
19+
STATE_MACHINE = auto()
20+
21+
22+
class SchedulerEventSource(ResourceMacro):
23+
"""
24+
Scheduler event source for SAM Functions and SAM State Machine.
25+
26+
It will translate into an "AWS::Scheduler::Schedule."
27+
Because a Scheduler Schedule resource requires an execution role,
28+
this macro will also create an IAM role with permissions to invoke
29+
the function/state machine.
30+
"""
31+
32+
resource_type = "ScheduleV2"
33+
34+
# As the first version, the properties of Scheduler schedule event will be the
35+
# same as the original "Schedule" event.
36+
# See class "Schedule" in samtranslator.model.eventsources.push and samtranslator.model.stepfunctions.events.
37+
property_types = {
38+
"PermissionsBoundary": PropertyType(False, is_str()),
39+
"ScheduleExpression": PropertyType(True, is_str()),
40+
"FlexibleTimeWindow": PropertyType(False, is_type(dict)),
41+
"Name": PropertyType(False, is_str()),
42+
"State": PropertyType(False, is_str()),
43+
"Description": PropertyType(False, is_str()),
44+
"StartDate": PropertyType(False, is_str()),
45+
"EndDate": PropertyType(False, is_str()),
46+
"ScheduleExpressionTimezone": PropertyType(False, is_str()),
47+
"GroupName": PropertyType(False, is_str()),
48+
"KmsKeyArn": PropertyType(False, is_str()),
49+
"Input": PropertyType(False, is_str()),
50+
"RoleArn": PropertyType(False, is_str()),
51+
"DeadLetterConfig": PropertyType(False, is_type(dict)),
52+
"RetryPolicy": PropertyType(False, is_type(dict)),
53+
}
54+
55+
# Below are type hints, must maintain consistent with properties_types
56+
# - pass-through to generated IAM role
57+
PermissionsBoundary: Optional[str]
58+
# - pass-through to AWS::Scheduler::Schedule
59+
ScheduleExpression: str
60+
FlexibleTimeWindow: Optional[Dict[str, Any]]
61+
Name: Optional[str]
62+
State: Optional[str]
63+
Description: Optional[str]
64+
StartDate: Optional[str]
65+
EndDate: Optional[str]
66+
ScheduleExpressionTimezone: Optional[str]
67+
GroupName: Optional[str]
68+
KmsKeyArn: Optional[str]
69+
# - pass-through to AWS::Scheduler::Schedule's Target
70+
Input: Optional[str]
71+
RoleArn: Optional[str]
72+
DeadLetterConfig: Optional[Dict[str, Any]]
73+
RetryPolicy: Optional[Dict[str, Any]]
74+
75+
DEFAULT_FLEXIBLE_TIME_WINDOW = {"Mode": "OFF"}
76+
77+
@cw_timer(prefix=FUNCTION_EVETSOURCE_METRIC_PREFIX) # type: ignore
78+
def to_cloudformation(self, **kwargs: Dict[str, Any]) -> List[Resource]:
79+
"""Returns the Scheduler Schedule and an IAM role.
80+
81+
:param dict kwargs: no existing resources need to be modified
82+
:returns: a list of vanilla CloudFormation Resources, to which this push event expands
83+
:rtype: list
84+
"""
85+
86+
target: Resource
87+
88+
# For SAM statemachine, the resource object is passed using kwargs["resource"],
89+
# https://github.com/aws/serverless-application-model/blob/a25933379e1cad3d0df4b35729ee2ec335402fdf/samtranslator/model/stepfunctions/generators.py#L266
90+
if kwargs.get("resource"):
91+
target_type = _SchedulerScheduleTargetType.STATE_MACHINE
92+
target = cast(Resource, kwargs["resource"])
93+
# for SAM function, the resource object is passed using kwargs["function"],
94+
# unlike SFN using "resource" keyword argument:
95+
# https://github.com/aws/serverless-application-model/blob/a25933379e1cad3d0df4b35729ee2ec335402fdf/samtranslator/model/sam_resources.py#L681
96+
elif kwargs.get("function"):
97+
target_type = _SchedulerScheduleTargetType.FUNCTION
98+
target = cast(Resource, kwargs["function"])
99+
else:
100+
raise TypeError("Missing required keyword argument: function/resource")
101+
102+
passthrough_resource_attributes = target.get_passthrough_resource_attributes() # type: ignore[no-untyped-call]
103+
104+
resources: List[Resource] = []
105+
106+
scheduler_schedule = self._construct_scheduler_schedule_without_target(passthrough_resource_attributes)
107+
resources.append(scheduler_schedule)
108+
109+
dlq_queue_arn: Optional[str] = None
110+
if self.DeadLetterConfig is not None:
111+
# The dql config spec is the same as normal "Schedule" event,
112+
# so continue to use EventBridgeRuleUtils for validation.
113+
# However, Scheduler doesn't use AWS::SQS::QueuePolicy to grant permissions.
114+
# so we cannot use EventBridgeRuleUtils.get_dlq_queue_arn_and_resources() here.
115+
EventBridgeRuleUtils.validate_dlq_config(self.logical_id, self.DeadLetterConfig) # type: ignore[no-untyped-call]
116+
dlq_queue_arn, dlq_resources = self._get_dlq_queue_arn_and_resources(
117+
self.DeadLetterConfig, passthrough_resource_attributes
118+
)
119+
resources.extend(dlq_resources)
120+
121+
execution_role_arn: Union[str, Dict[str, Any]] = self.RoleArn # type: ignore[assignment]
122+
if not execution_role_arn:
123+
execution_role = self._construct_execution_role(
124+
target, target_type, passthrough_resource_attributes, dlq_queue_arn, self.PermissionsBoundary
125+
)
126+
resources.append(execution_role)
127+
execution_role_arn = execution_role.get_runtime_attr("arn") # type: ignore[no-untyped-call]
128+
129+
scheduler_schedule.Target = self._construct_scheduler_schedule_target(target, execution_role_arn, dlq_queue_arn)
130+
131+
return resources
132+
133+
def _construct_scheduler_schedule_without_target(
134+
self, passthrough_resource_attributes: Dict[str, Any]
135+
) -> SchedulerSchedule:
136+
scheduler_schedule = SchedulerSchedule(self.logical_id, attributes=passthrough_resource_attributes)
137+
scheduler_schedule.ScheduleExpression = self.ScheduleExpression
138+
139+
if self.State:
140+
scheduler_schedule.State = self.State
141+
142+
# Scheduler schedule's Name is a required property
143+
scheduler_schedule.Name = self.Name or self.logical_id
144+
145+
# pass-through other properties
146+
scheduler_schedule.Description = self.Description
147+
scheduler_schedule.FlexibleTimeWindow = self.FlexibleTimeWindow or self.DEFAULT_FLEXIBLE_TIME_WINDOW
148+
scheduler_schedule.StartDate = self.StartDate
149+
scheduler_schedule.EndDate = self.EndDate
150+
scheduler_schedule.ScheduleExpressionTimezone = self.ScheduleExpressionTimezone
151+
scheduler_schedule.GroupName = self.GroupName
152+
scheduler_schedule.KmsKeyArn = self.KmsKeyArn
153+
154+
return scheduler_schedule
155+
156+
def _construct_execution_role(
157+
self,
158+
target: Resource,
159+
target_type: _SchedulerScheduleTargetType,
160+
passthrough_resource_attributes: Dict[str, Any],
161+
dlq_queue_arn: Optional[str],
162+
permissions_boundary: Optional[str],
163+
) -> IAMRole:
164+
"""Constructs the execution role for Scheduler Schedule."""
165+
if target_type == _SchedulerScheduleTargetType.FUNCTION:
166+
policy = IAMRolePolicies.lambda_invoke_function_role_policy(target.get_runtime_attr("arn"), self.logical_id) # type: ignore[no-untyped-call, no-untyped-call]
167+
elif target_type == _SchedulerScheduleTargetType.STATE_MACHINE:
168+
policy = IAMRolePolicies.step_functions_start_execution_role_policy( # type: ignore[no-untyped-call]
169+
target.get_runtime_attr("arn"), self.logical_id # type: ignore[no-untyped-call]
170+
)
171+
else:
172+
raise RuntimeError(f"Unexpected target type {target_type.name}")
173+
174+
role_logical_id = LogicalIdGenerator(self.logical_id + "Role").gen() # type: ignore[no-untyped-call, no-untyped-call]
175+
execution_role = IAMRole(role_logical_id, attributes=passthrough_resource_attributes)
176+
execution_role.AssumeRolePolicyDocument = IAMRolePolicies.scheduler_assume_role_policy()
177+
178+
policies = [policy]
179+
if dlq_queue_arn:
180+
policies.append(IAMRolePolicies.sqs_send_message_role_policy(dlq_queue_arn, self.logical_id))
181+
execution_role.Policies = policies
182+
183+
if permissions_boundary:
184+
execution_role.PermissionsBoundary = permissions_boundary
185+
return execution_role
186+
187+
def _construct_scheduler_schedule_target(
188+
self, target: Resource, execution_role_arn: Union[str, Dict[str, Any]], dead_letter_queue_arn: Optional[Any]
189+
) -> Dict[str, Any]:
190+
"""Constructs the Target property for the Scheduler Schedule.
191+
192+
:returns: the Target property
193+
:rtype: dict
194+
195+
Inspired by https://github.com/aws/serverless-application-model/blob/a25933379e1cad3d0df4b35729ee2ec335402fdf/samtranslator/model/eventsources/push.py#L157
196+
"""
197+
target_dict: Dict[str, Any] = {
198+
"Arn": target.get_runtime_attr("arn"), # type: ignore[no-untyped-call]
199+
"RoleArn": execution_role_arn,
200+
}
201+
if self.Input is not None:
202+
target_dict["Input"] = self.Input
203+
204+
if self.DeadLetterConfig is not None:
205+
target_dict["DeadLetterConfig"] = {"Arn": dead_letter_queue_arn}
206+
207+
if self.RetryPolicy is not None:
208+
target_dict["RetryPolicy"] = self.RetryPolicy
209+
210+
return target_dict
211+
212+
def _get_dlq_queue_arn_and_resources(
213+
self, dlq_config: Dict[str, Any], passthrough_resource_attributes: Optional[Dict[str, Any]]
214+
) -> Tuple[Any, List[Resource]]:
215+
"""
216+
Returns dlq queue arn and dlq_resources, assuming self.DeadLetterConfig has been validated.
217+
218+
Inspired by https://github.com/aws/serverless-application-model/blob/a25933379e1cad3d0df4b35729ee2ec335402fdf/samtranslator/model/eventbridge_utils.py#L44
219+
"""
220+
dlq_queue_arn = dlq_config.get("Arn")
221+
if dlq_queue_arn is not None:
222+
return dlq_queue_arn, []
223+
queue_logical_id = dlq_config.get("QueueLogicalId")
224+
if queue_logical_id is not None and not isinstance(queue_logical_id, str):
225+
raise InvalidEventException(
226+
self.logical_id,
227+
"QueueLogicalId must be a string",
228+
)
229+
dlq_resources: List[Resource] = []
230+
queue = SQSQueue(queue_logical_id or self.logical_id + "Queue", attributes=passthrough_resource_attributes)
231+
dlq_resources.append(queue)
232+
233+
dlq_queue_arn = queue.get_runtime_attr("arn") # type: ignore[no-untyped-call]
234+
return dlq_queue_arn, dlq_resources

‎samtranslator/model/iam.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,16 @@ def cloud_watch_log_assume_role_policy(cls): # type: ignore[no-untyped-def]
8585
}
8686
return document
8787

88+
@classmethod
89+
def scheduler_assume_role_policy(cls) -> Dict[str, Any]:
90+
document = {
91+
"Version": "2012-10-17",
92+
"Statement": [
93+
{"Action": ["sts:AssumeRole"], "Effect": "Allow", "Principal": {"Service": ["scheduler.amazonaws.com"]}}
94+
],
95+
}
96+
return document
97+
8898
@classmethod
8999
def lambda_assume_role_policy(cls): # type: ignore[no-untyped-def]
90100
document = {

‎samtranslator/model/role_utils/role_constructor.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ def construct_role_for_resource( # type: ignore[no-untyped-def]
1414
policy_documents=None,
1515
permissions_boundary=None,
1616
tags=None,
17-
):
17+
) -> IAMRole:
1818
"""
1919
Constructs an execution role for a resource.
2020
:param resource_logical_id: The logical_id of the SAM resource that the role will be associated with

‎samtranslator/model/sam_resources.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import samtranslator.model.eventsources.pull
2222
import samtranslator.model.eventsources.push
2323
import samtranslator.model.eventsources.cloudwatchlogs
24+
import samtranslator.model.eventsources.scheduler
2425
from .api.api_generator import ApiGenerator
2526
from .api.http_api_generator import HttpApiGenerator
2627
from .packagetype import ZIP, IMAGE
@@ -118,6 +119,7 @@ class SamFunction(SamResourceMacro):
118119
samtranslator.model.eventsources.pull,
119120
samtranslator.model.eventsources.push,
120121
samtranslator.model.eventsources.cloudwatchlogs,
122+
samtranslator.model.eventsources.scheduler,
121123
)
122124

123125
# DeadLetterQueue
@@ -525,7 +527,7 @@ def _construct_role(self, managed_policy_map, event_invoke_policies): # type: i
525527
if event_invoke_policies is not None:
526528
policy_documents.extend(event_invoke_policies)
527529

528-
execution_role = construct_role_for_resource( # type: ignore[no-untyped-call]
530+
execution_role = construct_role_for_resource(
529531
resource_logical_id=self.logical_id,
530532
attributes=role_attributes,
531533
managed_policy_map=managed_policy_map,
@@ -1557,6 +1559,7 @@ class SamStateMachine(SamResourceMacro):
15571559
}
15581560
event_resolver = ResourceTypeResolver( # type: ignore[no-untyped-call]
15591561
samtranslator.model.stepfunctions.events,
1562+
samtranslator.model.eventsources.scheduler,
15601563
)
15611564

15621565
@cw_timer

‎samtranslator/model/scheduler.py

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
from typing import Any, Dict, Optional
2+
3+
from samtranslator.model import PropertyType, Resource
4+
from samtranslator.model.types import is_type, is_str
5+
from samtranslator.model.intrinsics import fnGetAtt
6+
7+
8+
class SchedulerSchedule(Resource):
9+
resource_type = "AWS::Scheduler::Schedule"
10+
property_types = {
11+
"ScheduleExpression": PropertyType(True, is_str()),
12+
"FlexibleTimeWindow": PropertyType(True, is_type(dict)),
13+
"Name": PropertyType(True, is_str()),
14+
"State": PropertyType(False, is_str()),
15+
"Description": PropertyType(False, is_str()),
16+
"StartDate": PropertyType(False, is_str()),
17+
"EndDate": PropertyType(False, is_str()),
18+
"ScheduleExpressionTimezone": PropertyType(False, is_str()),
19+
"GroupName": PropertyType(False, is_str()),
20+
"KmsKeyArn": PropertyType(False, is_str()),
21+
"Target": PropertyType(True, is_type(dict)),
22+
}
23+
24+
ScheduleExpression: str
25+
FlexibleTimeWindow: Dict[str, Any]
26+
Name: str
27+
State: Optional[str]
28+
Description: Optional[str]
29+
StartDate: Optional[str]
30+
EndDate: Optional[str]
31+
ScheduleExpressionTimezone: Optional[str]
32+
GroupName: Optional[str]
33+
KmsKeyArn: Optional[str]
34+
Target: Dict[str, Any]
35+
36+
runtime_attrs = {"arn": lambda self: fnGetAtt(self.logical_id, "Arn")}

‎samtranslator/model/stepfunctions/generators.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -217,7 +217,7 @@ def _construct_role(self): # type: ignore[no-untyped-def]
217217
policy_template_processor=None,
218218
)
219219

220-
execution_role = construct_role_for_resource( # type: ignore[no-untyped-call]
220+
execution_role = construct_role_for_resource(
221221
resource_logical_id=self.logical_id,
222222
attributes=self.passthrough_resource_attributes,
223223
managed_policy_map=self.managed_policy_map,

0 commit comments

Comments
 (0)