Add support for custom tag prefixes

.gitignore

@@ -8,4 +8,5 @@ samconfig.toml
 .aws-sam
 
 # Other ignores
+envs.json
 hello_world_function

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/):

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'))
 

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))

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')

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,

envs.dist.json

@@ -0,0 +1,5 @@
+{
+	"AdjustScheduleFunction": {
+		"TAG_PREFIX": ""
+	}
+}

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:

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)