diff --git a/.gitignore b/.gitignore index 6031323..3227424 100644 --- a/.gitignore +++ b/.gitignore @@ -8,4 +8,5 @@ samconfig.toml .aws-sam # Other ignores +envs.json hello_world_function diff --git a/README.md b/README.md index d3e42fb..e28958f 100644 --- a/README.md +++ b/README.md @@ -67,6 +67,16 @@ We have one EventBridge rule which must trigger at 12:00 in the local time of Li * `scheduled-event-adjuster:local-timezone` = `Europe/Lisbon` * `scheduled-event-adjuster:local-time` = `12:00` +### Using custom tag prefixes + +As explained above, the solution expects all tags to be prefixed with `scheduled-event-adjuster` by default. However, this behavior can be customized by providing a value as the `TagPrefix` SAM parameter. If the solution were deployed like this: + +``` +sam deploy --parameter-overrides "TagPrefix=foo" +``` + +Then it would act on resources tagged with `foo:enabled` instead `scheduled-event-adjuster:enabled`. + ## Roadmap These are the features we're working on: @@ -90,6 +100,12 @@ The function can be locally invoked using `sam local invoke`. The `events` direc sam local invoke AdjustScheduleFunction --event events/event.json ``` +If you want to define environment variables for the function, copy and modify `envs.dist.json` (for example, to `envs.json`) and then use it when invoking the function: + +```bash +sam local invoke AdjustScheduleFunction --event events/event.json -n envs.json +``` + ### Unit tests Tests are defined in the `tests` folder in this project, and can be run using the [pytest](https://docs.pytest.org/en/latest/): diff --git a/adjust_schedule_function/adjust_schedule/app.py b/adjust_schedule_function/adjust_schedule/app.py index 2f7d06e..7a3149b 100644 --- a/adjust_schedule_function/adjust_schedule/app.py +++ b/adjust_schedule_function/adjust_schedule/app.py @@ -9,11 +9,15 @@ import os +tag_prefix = '' +if 'TAG_PREFIX' in os.environ and os.environ['TAG_PREFIX'].strip(): + tag_prefix = os.environ['TAG_PREFIX'].strip() + asg_service = AutoScalingService(boto3.client('autoscaling')) eventbridge_service = EventBridgeService(boto3.client('events')) processors = [ - AutoScalingGroupProcessor(asg_service, RecurrenceCalculator()), - EventBridgeProcessor(eventbridge_service, RecurrenceCalculator()), + AutoScalingGroupProcessor(tag_prefix, asg_service, RecurrenceCalculator()), + EventBridgeProcessor(tag_prefix, eventbridge_service, RecurrenceCalculator()), ] bus = EventBus(boto3.client('events')) diff --git a/adjust_schedule_function/lib/processors/autoscaling.py b/adjust_schedule_function/lib/processors/autoscaling.py index 47ebde9..b09b556 100644 --- a/adjust_schedule_function/lib/processors/autoscaling.py +++ b/adjust_schedule_function/lib/processors/autoscaling.py @@ -2,7 +2,8 @@ from lib.processors.base import ResourceProcessor class AutoScalingGroupProcessor(ResourceProcessor): - def __init__(self, asg_service, recurrence_calculator): + def __init__(self, tag_prefix, asg_service, recurrence_calculator): + super().__init__(tag_prefix) self._asg_service = asg_service self._recurrence_calculator = recurrence_calculator @@ -21,15 +22,15 @@ def _process_asg(self, asg): print("Processing ASG '{}'".format(asg_name)) - if utils.get_tag_by_key(asg['Tags'], self.ENABLED_TAG) == None: + if utils.get_tag_by_key(asg['Tags'], self._get_enabled_tag()) == None: print("Skipping: ASG '{}' is not enabled (missing tag '{}')".format(asg_name, - self.ENABLED_TAG)) + self._get_enabled_tag())) return result - local_timezone = utils.get_tag_by_key(asg['Tags'], self.LOCAL_TIMEZONE_TAG) + local_timezone = utils.get_tag_by_key(asg['Tags'], self._get_local_timezone_tag()) if not local_timezone: print("Skipping: ASG '{}' has no timezone defined (missing tag '{}')".format(asg_name, - self.LOCAL_TIMEZONE_TAG)) + self._get_local_timezone_tag())) return result scheduled_actions = self._asg_service.get_asg_scheduled_actions(asg_name) @@ -39,7 +40,7 @@ def _process_asg(self, asg): action_name = action['ScheduledActionName'] current_recurrence = action['Recurrence'] - local_time_tag_key = self.LOCAL_TIME_TAG + ':' + action_name + local_time_tag_key = self._get_local_time_tag() + ':' + action_name local_time = utils.get_tag_by_key(asg['Tags'], local_time_tag_key) if not local_time: print("Skipping: action '{}' does not have local time tag (missing tag '{}')".format(action_name, local_time_tag_key)) diff --git a/adjust_schedule_function/lib/processors/base.py b/adjust_schedule_function/lib/processors/base.py index 007b3bd..5cdb5b2 100644 --- a/adjust_schedule_function/lib/processors/base.py +++ b/adjust_schedule_function/lib/processors/base.py @@ -1,10 +1,20 @@ class ResourceProcessor: - # The tag that determines whether the resource should be processed. - ENABLED_TAG = 'scheduled-event-adjuster:enabled' + def __init__(self, tag_prefix): + self._tag_prefix = tag_prefix - # The tag that determines the timezone of the local time. - LOCAL_TIMEZONE_TAG = 'scheduled-event-adjuster:local-timezone' + def _get_enabled_tag(self): + """Returns the tag that, when present, determines whether a resource + must be processed. + """ + return '%s:%s' % (self._tag_prefix, 'enabled') - # The tag that determines the local time at which the scheduled event is - # expected to run. - LOCAL_TIME_TAG = 'scheduled-event-adjuster:local-time' + def _get_local_timezone_tag(self): + """Returns the tag that specifies the timezone of the local time. + """ + return '%s:%s' % (self._tag_prefix, 'local-timezone') + + def _get_local_time_tag(self): + """Returns the tag that specifies the local time at which scheduled + events must run. + """ + return '%s:%s' % (self._tag_prefix, 'local-time') diff --git a/adjust_schedule_function/lib/processors/eventbridge.py b/adjust_schedule_function/lib/processors/eventbridge.py index bf5aa5d..e19fa8e 100644 --- a/adjust_schedule_function/lib/processors/eventbridge.py +++ b/adjust_schedule_function/lib/processors/eventbridge.py @@ -2,7 +2,8 @@ from lib.processors.base import ResourceProcessor class EventBridgeProcessor(ResourceProcessor): - def __init__(self, eventbridge_service, recurrence_calculator): + def __init__(self, tag_prefix, eventbridge_service, recurrence_calculator): + super().__init__(tag_prefix) self._eventbridge_service = eventbridge_service self._recurrence_calculator = recurrence_calculator @@ -17,22 +18,22 @@ def process_resources(self): tags = self._eventbridge_service.get_rule_tags(rule['Arn']) - if utils.get_tag_by_key(tags, self.ENABLED_TAG) == None: + if utils.get_tag_by_key(tags, self._get_enabled_tag()) == None: print("Skipping: EventBridge rule '{}' is not enabled (missing tag '{}')".format(rule['Name'], - self.ENABLED_TAG)) + self._get_enabled_tag())) continue - local_timezone = utils.get_tag_by_key(tags, self.LOCAL_TIMEZONE_TAG) - local_time = utils.get_tag_by_key(tags, self.LOCAL_TIME_TAG) + local_timezone = utils.get_tag_by_key(tags, self._get_local_timezone_tag()) + local_time = utils.get_tag_by_key(tags, self._get_local_time_tag()) if not local_timezone: print("Skipping: EventBridge rule '{}' has no timezone defined (missing tag '{}')".format(rule['Name'], - self.LOCAL_TIMEZONE_TAG)) + self._get_local_timezone_tag())) continue if not local_time: print("Skipping: EventBridge rule '{}' does not have local time tag (missing tag '{}')".format(rule['Name'], - self.LOCAL_TIME_TAG)) + self._get_local_time_tag())) continue # Remove the 'cron()' surrounding the cron expression itself, diff --git a/envs.dist.json b/envs.dist.json new file mode 100644 index 0000000..eb6e80d --- /dev/null +++ b/envs.dist.json @@ -0,0 +1,5 @@ +{ + "AdjustScheduleFunction": { + "TAG_PREFIX": "" + } +} diff --git a/template.yaml b/template.yaml index e394f2e..c16fd35 100644 --- a/template.yaml +++ b/template.yaml @@ -13,6 +13,11 @@ Parameters: Type: String Description: The email address that should be notified when relevant events occur. + TagPrefix: + Type: String + Description: (Optional) The tag prefix to use when checking resources' + tags. Leave empty to use the default prefix. + Default: 'scheduled-event-adjuster' Globals: Function: @@ -25,6 +30,9 @@ Resources: CodeUri: adjust_schedule_function Handler: adjust_schedule/app.lambda_handler Runtime: python3.7 + Environment: + Variables: + TAG_PREFIX: !Ref TagPrefix Policies: - Version: '2012-10-17' Statement: diff --git a/tests/unit/processors/test_autoscaling.py b/tests/unit/processors/test_autoscaling.py index 8cfda08..314e46e 100644 --- a/tests/unit/processors/test_autoscaling.py +++ b/tests/unit/processors/test_autoscaling.py @@ -11,15 +11,15 @@ def test_process_resources_with_different_recurrence(mocker): 'AutoScalingGroupARN': 'MyAsgARN', 'Tags': [ { - 'Key': 'scheduled-event-adjuster:enabled', + 'Key': 'foo:bar:enabled', 'Value': '' }, { - 'Key': 'scheduled-event-adjuster:local-timezone', + 'Key': 'foo:bar:local-timezone', 'Value': 'Europe/Madrid' }, { - 'Key': 'scheduled-event-adjuster:local-time:ActionOne', + 'Key': 'foo:bar:local-time:ActionOne', 'Value': '10:00' } ] @@ -35,7 +35,73 @@ def test_process_resources_with_different_recurrence(mocker): ] asg_svc = AutoScalingService() rec_calc = RecurrenceCalculator() - processor = AutoScalingGroupProcessor(asg_svc, rec_calc) + processor = AutoScalingGroupProcessor('foo:bar', asg_svc, rec_calc) + mocker.patch.object(asg_svc, 'get_asgs', return_value=asgs) + mocker.patch.object(asg_svc, 'get_asg_scheduled_actions', return_value=scheduled_actions) + mocker.patch.object(asg_svc, 'update_asg_scheduled_actions') + mocker.patch.object(rec_calc, 'calculate_recurrence', return_value='NewRecurrence') + + result = processor.process_resources() + + rec_calc.calculate_recurrence.assert_called_once_with('OriginalRecurrence', + '10:00', + 'Europe/Madrid') + asg_svc.update_asg_scheduled_actions.assert_called_once_with( + 'MyAsg', + [ + { + 'ScheduledActionName': 'ActionOne', + 'Recurrence': 'NewRecurrence', + 'DesiredCapacity': 123 + } + ] + ) + assert len(result) == 1 + assert result[0] == { + 'Type': 'AutoScalingGroupScalingPolicy', + 'ResourceName': 'MyAsg', + 'ResourceArn': 'MyAsgARN', + 'OriginalRecurrence': 'OriginalRecurrence', + 'NewRecurrence': 'NewRecurrence', + 'LocalTime': '10:00', + 'LocalTimezone': 'Europe/Madrid', + 'AdditionalDetails': { + 'ActionName': 'ActionOne' + } + } + +def test_process_resources_with_different_recurrence_and_custom_tag_prefix(mocker): + asgs = [ + { + 'AutoScalingGroupName': 'MyAsg', + 'AutoScalingGroupARN': 'MyAsgARN', + 'Tags': [ + { + 'Key': 'foo:bar:enabled', + 'Value': '' + }, + { + 'Key': 'foo:bar:local-timezone', + 'Value': 'Europe/Madrid' + }, + { + 'Key': 'foo:bar:local-time:ActionOne', + 'Value': '10:00' + } + ] + } + ] + scheduled_actions = [ + { + 'ScheduledActionName': 'ActionOne', + 'ScheduledActionARN': 'ActionOneARN', + 'Recurrence': 'OriginalRecurrence', + 'DesiredCapacity': 123 + } + ] + asg_svc = AutoScalingService() + rec_calc = RecurrenceCalculator() + processor = AutoScalingGroupProcessor('foo:bar', asg_svc, rec_calc) mocker.patch.object(asg_svc, 'get_asgs', return_value=asgs) mocker.patch.object(asg_svc, 'get_asg_scheduled_actions', return_value=scheduled_actions) mocker.patch.object(asg_svc, 'update_asg_scheduled_actions') @@ -77,15 +143,15 @@ def test_process_resources_with_same_recurrence(mocker): 'AutoScalingGroupARN': 'MyAsgARN', 'Tags': [ { - 'Key': 'scheduled-event-adjuster:enabled', + 'Key': 'foo:bar:enabled', 'Value': '' }, { - 'Key': 'scheduled-event-adjuster:local-timezone', + 'Key': 'foo:bar:local-timezone', 'Value': 'Europe/Madrid' }, { - 'Key': 'scheduled-event-adjuster:local-time:ActionOne', + 'Key': 'foo:bar:local-time:ActionOne', 'Value': '10:00' } ] @@ -101,7 +167,7 @@ def test_process_resources_with_same_recurrence(mocker): ] asg_svc = AutoScalingService() rec_calc = RecurrenceCalculator() - processor = AutoScalingGroupProcessor(asg_svc, rec_calc) + processor = AutoScalingGroupProcessor('foo:bar', asg_svc, rec_calc) mocker.patch.object(asg_svc, 'get_asgs', return_value=asgs) mocker.patch.object(asg_svc, 'get_asg_scheduled_actions', return_value=scheduled_actions) mocker.patch.object(asg_svc, 'update_asg_scheduled_actions') @@ -122,7 +188,7 @@ def test_process_resources_without_enabled_tag(mocker): ] asg_svc = AutoScalingService() rec_calc = RecurrenceCalculator() - processor = AutoScalingGroupProcessor(asg_svc, rec_calc) + processor = AutoScalingGroupProcessor('foo:bar', asg_svc, rec_calc) mocker.patch.object(asg_svc, 'get_asgs', return_value=asgs) result = processor.process_resources() @@ -134,12 +200,12 @@ def test_process_resources_without_timezone_tag(mocker): { 'AutoScalingGroupName': 'MyAsg', 'AutoScalingGroupARN': 'MyAsgARN', - 'Tags': [{'Key': 'scheduled-event-adjuster:enabled', 'Value': ''}] + 'Tags': [{'Key': 'foo:bar:enabled', 'Value': ''}] } ] asg_svc = AutoScalingService() rec_calc = RecurrenceCalculator() - processor = AutoScalingGroupProcessor(asg_svc, rec_calc) + processor = AutoScalingGroupProcessor('foo:bar', asg_svc, rec_calc) mocker.patch.object(asg_svc, 'get_asgs', return_value=asgs) result = processor.process_resources() @@ -152,8 +218,8 @@ def test_process_asg_without_action_time_tag(mocker): 'AutoScalingGroupName': 'MyAsg', 'AutoScalingGroupARN': 'MyAsgARN', 'Tags': [ - {'Key': 'scheduled-event-adjuster:enabled', 'Value': ''}, - {'Key': 'scheduled-event-adjuster:local-timezone', 'Value': 'Europe/Madrid'} + {'Key': 'foo:bar:enabled', 'Value': ''}, + {'Key': 'foo:bar:local-timezone', 'Value': 'Europe/Madrid'} ] } ] @@ -167,7 +233,7 @@ def test_process_asg_without_action_time_tag(mocker): ] asg_svc = AutoScalingService() rec_calc = RecurrenceCalculator() - processor = AutoScalingGroupProcessor(asg_svc, rec_calc) + processor = AutoScalingGroupProcessor('foo:bar', asg_svc, rec_calc) mocker.patch.object(asg_svc, 'get_asgs', return_value=asgs) mocker.patch.object(asg_svc, 'get_asg_scheduled_actions', return_value=scheduled_actions) diff --git a/tests/unit/processors/test_eventbridge.py b/tests/unit/processors/test_eventbridge.py index 3538166..5790f50 100644 --- a/tests/unit/processors/test_eventbridge.py +++ b/tests/unit/processors/test_eventbridge.py @@ -7,13 +7,13 @@ def test_process_resources_updates_schedule_when_recurrences_are_different(mocker): rules = [{'Name': 'ruleName', 'Arn': 'ruleArn', 'ScheduleExpression': 'cron(foo)'}] tags = [ - {'Key': 'scheduled-event-adjuster:enabled', 'Value': ''}, - {'Key': 'scheduled-event-adjuster:local-timezone', 'Value': 'Europe/Madrid'}, - {'Key': 'scheduled-event-adjuster:local-time', 'Value': '10:00'} + {'Key': 'foo:bar:enabled', 'Value': ''}, + {'Key': 'foo:bar:local-timezone', 'Value': 'Europe/Madrid'}, + {'Key': 'foo:bar:local-time', 'Value': '10:00'} ] eb_svc = EventBridgeService() rec_calc = RecurrenceCalculator() - processor = EventBridgeProcessor(eb_svc, rec_calc) + processor = EventBridgeProcessor('foo:bar', eb_svc, rec_calc) mocker.patch.object(eb_svc, 'get_scheduled_rules', return_value=rules) mocker.patch.object(eb_svc, 'get_rule_tags', return_value=tags) mocker.patch.object(eb_svc, 'update_rule_schedule', return_value=None) @@ -38,13 +38,13 @@ def test_process_resources_updates_schedule_when_recurrences_are_different(mocke def test_process_resources_skips_update_when_an_exception_is_thrown_when_updating_schedule(mocker): rules = [{'Name': 'ruleName', 'Arn': 'ruleArn', 'ScheduleExpression': 'cron(foo)'}] tags = [ - {'Key': 'scheduled-event-adjuster:enabled', 'Value': ''}, - {'Key': 'scheduled-event-adjuster:local-timezone', 'Value': 'Europe/Madrid'}, - {'Key': 'scheduled-event-adjuster:local-time', 'Value': '10:00'} + {'Key': 'foo:bar:enabled', 'Value': ''}, + {'Key': 'foo:bar:local-timezone', 'Value': 'Europe/Madrid'}, + {'Key': 'foo:bar:local-time', 'Value': '10:00'} ] eb_svc = EventBridgeService() rec_calc = RecurrenceCalculator() - processor = EventBridgeProcessor(eb_svc, rec_calc) + processor = EventBridgeProcessor('foo:bar', eb_svc, rec_calc) mocker.patch.object(eb_svc, 'get_scheduled_rules', return_value=rules) mocker.patch.object(eb_svc, 'get_rule_tags', return_value=tags) mocker.patch.object(eb_svc, 'update_rule_schedule', side_effect=Exception('mocked exception')) @@ -59,12 +59,12 @@ def test_process_resources_skips_update_when_an_exception_is_thrown_when_updatin def test_process_resources_with_same_recurrence(mocker): rules = [{'Name': 'ruleName', 'Arn': 'ruleArn', 'ScheduleExpression': 'cron(foo)'}] tags = [ - {'Key': 'scheduled-event-adjuster:enabled', 'Value': ''}, - {'Key': 'scheduled-event-adjuster:local-timezone', 'Value': 'Europe/Madrid'} + {'Key': 'foo:bar:enabled', 'Value': ''}, + {'Key': 'foo:bar:local-timezone', 'Value': 'Europe/Madrid'} ] eb_svc = EventBridgeService() rec_calc = RecurrenceCalculator() - processor = EventBridgeProcessor(eb_svc, rec_calc) + processor = EventBridgeProcessor('foo:bar', eb_svc, rec_calc) mocker.patch.object(eb_svc, 'get_scheduled_rules', return_value=rules) mocker.patch.object(eb_svc, 'get_rule_tags', return_value=tags) mocker.patch.object(rec_calc, 'calculate_recurrence', return_value='foo') @@ -79,7 +79,7 @@ def test_process_resources_without_enabled_tag(mocker): tags = [{'Key': 'nope', 'Value': 'nope'}] eb_svc = EventBridgeService() rec_calc = RecurrenceCalculator() - processor = EventBridgeProcessor(eb_svc, rec_calc) + processor = EventBridgeProcessor('foo:bar', eb_svc, rec_calc) mocker.patch.object(eb_svc, 'get_scheduled_rules', return_value=rules) mocker.patch.object(eb_svc, 'get_rule_tags', return_value=tags) @@ -90,10 +90,10 @@ def test_process_resources_without_enabled_tag(mocker): def test_process_resources_without_timezone_tag(mocker): rules = [{'Name': 'ruleName', 'Arn': 'ruleArn'}] - tags = [{'Key': 'scheduled-event-adjuster:enabled', 'Value': ''}] + tags = [{'Key': 'foo:bar:enabled', 'Value': ''}] eb_svc = EventBridgeService() rec_calc = RecurrenceCalculator() - processor = EventBridgeProcessor(eb_svc, rec_calc) + processor = EventBridgeProcessor('foo:bar', eb_svc, rec_calc) mocker.patch.object(eb_svc, 'get_scheduled_rules', return_value=rules) mocker.patch.object(eb_svc, 'get_rule_tags', return_value=tags) @@ -105,12 +105,12 @@ def test_process_resources_without_timezone_tag(mocker): def test_process_resources_without_local_time_tag(mocker): rules = [{'Name': 'ruleName', 'Arn': 'ruleArn'}] tags = [ - {'Key': 'scheduled-event-adjuster:enabled', 'Value': ''}, - {'Key': 'scheduled-event-adjuster:local-timezone', 'Value': 'Europe/Madrid'}, + {'Key': 'foo:bar:enabled', 'Value': ''}, + {'Key': 'foo:bar:local-timezone', 'Value': 'Europe/Madrid'}, ] eb_svc = EventBridgeService() rec_calc = RecurrenceCalculator() - processor = EventBridgeProcessor(eb_svc, rec_calc) + processor = EventBridgeProcessor('foo:bar', eb_svc, rec_calc) mocker.patch.object(eb_svc, 'get_scheduled_rules', return_value=rules) mocker.patch.object(eb_svc, 'get_rule_tags', return_value=tags)