Add Modulo Range Condition under Feature Flags

Author: ajwad-shaikh

aws_lambda_powertools/utilities/feature_flags/time_conditions.py renamed to aws_lambda_powertools/utilities/feature_flags/comparators.py

@@ -1,9 +1,9 @@
 from datetime import datetime, tzinfo
-from typing import Dict, Optional
+from typing import Any, Dict, Optional
 
 from dateutil.tz import gettz
 
-from .schema import HOUR_MIN_SEPARATOR, TimeValues
+from .schema import HOUR_MIN_SEPARATOR, ModuloRangeValues, TimeValues
 
 
 def _get_now_from_timezone(timezone: Optional[tzinfo]) -> datetime:
@@ -16,23 +16,23 @@ def _get_now_from_timezone(timezone: Optional[tzinfo]) -> datetime:
     return datetime.now(timezone)
 
 
-def compare_days_of_week(action: str, values: Dict) -> bool:
-    timezone_name = values.get(TimeValues.TIMEZONE.value, "UTC")
+def compare_days_of_week(context_value: Any, condition_value: Dict) -> bool:
+    timezone_name = condition_value.get(TimeValues.TIMEZONE.value, "UTC")
 
     # %A = Weekday as locale’s full name.
     current_day = _get_now_from_timezone(gettz(timezone_name)).strftime("%A").upper()
 
-    days = values.get(TimeValues.DAYS.value, [])
+    days = condition_value.get(TimeValues.DAYS.value, [])
     return current_day in days
 
 
-def compare_datetime_range(action: str, values: Dict) -> bool:
-    timezone_name = values.get(TimeValues.TIMEZONE.value, "UTC")
+def compare_datetime_range(context_value: Any, condition_value: Dict) -> bool:
+    timezone_name = condition_value.get(TimeValues.TIMEZONE.value, "UTC")
     timezone = gettz(timezone_name)
     current_time: datetime = _get_now_from_timezone(timezone)
 
-    start_date_str = values.get(TimeValues.START.value, "")
-    end_date_str = values.get(TimeValues.END.value, "")
+    start_date_str = condition_value.get(TimeValues.START.value, "")
+    end_date_str = condition_value.get(TimeValues.END.value, "")
 
     # Since start_date and end_date doesn't include timezone information, we mark the timestamp
     # with the same timezone as the current_time. This way all the 3 timestamps will be on
@@ -42,12 +42,12 @@ def compare_datetime_range(action: str, values: Dict) -> bool:
     return start_date <= current_time <= end_date
 
 
-def compare_time_range(action: str, values: Dict) -> bool:
-    timezone_name = values.get(TimeValues.TIMEZONE.value, "UTC")
+def compare_time_range(context_value: Any, condition_value: Dict) -> bool:
+    timezone_name = condition_value.get(TimeValues.TIMEZONE.value, "UTC")
     current_time: datetime = _get_now_from_timezone(gettz(timezone_name))
 
-    start_hour, start_min = values.get(TimeValues.START.value, "").split(HOUR_MIN_SEPARATOR)
-    end_hour, end_min = values.get(TimeValues.END.value, "").split(HOUR_MIN_SEPARATOR)
+    start_hour, start_min = condition_value.get(TimeValues.START.value, "").split(HOUR_MIN_SEPARATOR)
+    end_hour, end_min = condition_value.get(TimeValues.END.value, "").split(HOUR_MIN_SEPARATOR)
 
     start_time = current_time.replace(hour=int(start_hour), minute=int(start_min))
     end_time = current_time.replace(hour=int(end_hour), minute=int(end_min))
@@ -71,3 +71,14 @@ def compare_time_range(action: str, values: Dict) -> bool:
     else:
         # In normal circumstances, we need to assert **both** conditions
         return start_time <= current_time <= end_time
+
+
+def compare_modulo_range(context_value: int, condition_value: Dict) -> bool:
+    """
+    Returns for a given context 'a' and modulo condition 'b' -> b.start <= a % b.base <= b.end
+    """
+    base = condition_value.get(ModuloRangeValues.BASE.value, 1)
+    start = condition_value.get(ModuloRangeValues.START.value, 1)
+    end = condition_value.get(ModuloRangeValues.END.value, 1)
+
+    return start <= context_value % base <= end

aws_lambda_powertools/utilities/feature_flags/feature_flags.py

@@ -5,12 +5,13 @@
 from ...shared.types import JSONType
 from . import schema
 from .base import StoreProvider
-from .exceptions import ConfigurationStoreError
-from .time_conditions import (
+from .comparators import (
     compare_datetime_range,
     compare_days_of_week,
+    compare_modulo_range,
     compare_time_range,
 )
+from .exceptions import ConfigurationStoreError
 
 
 class FeatureFlags:
@@ -65,6 +66,7 @@ def _match_by_action(self, action: str, condition_value: Any, context_value: Any
             schema.RuleAction.SCHEDULE_BETWEEN_TIME_RANGE.value: lambda a, b: compare_time_range(a, b),
             schema.RuleAction.SCHEDULE_BETWEEN_DATETIME_RANGE.value: lambda a, b: compare_datetime_range(a, b),
             schema.RuleAction.SCHEDULE_BETWEEN_DAYS_OF_WEEK.value: lambda a, b: compare_days_of_week(a, b),
+            schema.RuleAction.MODULO_RANGE.value: lambda a, b: lambda a, b: compare_modulo_range(a, b),
         }
 
         try:

aws_lambda_powertools/utilities/feature_flags/schema.py

@@ -41,6 +41,7 @@ class RuleAction(str, Enum):
     SCHEDULE_BETWEEN_TIME_RANGE = "SCHEDULE_BETWEEN_TIME_RANGE"  # hour:min 24 hours clock
     SCHEDULE_BETWEEN_DATETIME_RANGE = "SCHEDULE_BETWEEN_DATETIME_RANGE"  # full datetime format, excluding timezone
     SCHEDULE_BETWEEN_DAYS_OF_WEEK = "SCHEDULE_BETWEEN_DAYS_OF_WEEK"  # MONDAY, TUESDAY, .... see TimeValues enum
+    MODULO_RANGE = "MODULO_RANGE"
 
 
 class TimeKeys(Enum):
@@ -71,6 +72,16 @@ class TimeValues(Enum):
     SATURDAY = "SATURDAY"
 
 
+class ModuloRangeValues(Enum):
+    """
+    Possible values when using modulo range rule
+    """
+
+    BASE = "BASE"
+    START = "START"
+    END = "END"
+
+
 class SchemaValidator(BaseValidator):
     """Validates feature flag schema configuration
 
@@ -341,6 +352,9 @@ def validate_condition_value(condition: Dict[str, Any], rule_name: str):
             )
         elif action == RuleAction.SCHEDULE_BETWEEN_DAYS_OF_WEEK.value:
             ConditionsValidator._validate_schedule_between_days_of_week(value, rule_name)
+        # modulo range condition needs validation on base, start, and end attributes
+        elif action == RuleAction.MODULO_RANGE.value:
+            ConditionsValidator._validate_modulo_range(value, rule_name)
 
     @staticmethod
     def _validate_datetime_value(datetime_str: str, rule_name: str):
@@ -434,3 +448,22 @@ def _validate_schedule_between_time_and_datetime_ranges(
         # try to see if the timezone string corresponds to any known timezone
         if not tz.gettz(timezone):
             raise SchemaValidationError(f"'TIMEZONE' value must represent a valid IANA timezone, rule={rule_name}")
+
+    @staticmethod
+    def _validate_modulo_range(value: Any, rule_name: str):
+        error_str = f"condition with a 'MODULO_RANGE' action must have a condition value type dictionary with 'BASE', 'START' and 'END' keys, rule={rule_name}"  # noqa: E501
+        if not isinstance(value, dict):
+            raise SchemaValidationError(error_str)
+
+        base = value.get(ModuloRangeValues.BASE.value)
+        start = value.get(ModuloRangeValues.START.value)
+        end = value.get(ModuloRangeValues.END.value)
+        if base is None or start is None or end is None:
+            raise SchemaValidationError(error_str)
+        if not isinstance(base, int) or not isinstance(start, int) or not isinstance(end, int):
+            raise SchemaValidationError(f"'BASE', 'START' and 'END' must be integers, rule={rule_name}")
+
+        if not 0 <= start <= end <= base - 1:
+            raise SchemaValidationError(
+                f"condition with 'MODULO_RANGE' action must satisfy 0 <= START <= END <= BASE-1, rule={rule_name}"
+            )

tests/functional/feature_flags/test_time_based_actions.py

@@ -36,7 +36,7 @@ def evaluate_mocked_schema(
 
     # Mock the current time
     year, month, day, hour, minute, second, timezone = mocked_time
-    time = mocker.patch("aws_lambda_powertools.utilities.feature_flags.time_conditions._get_now_from_timezone")
+    time = mocker.patch("aws_lambda_powertools.utilities.feature_flags.comparators._get_now_from_timezone")
     time.return_value = datetime.datetime(
         year=year, month=month, day=day, hour=hour, minute=minute, second=second, microsecond=0, tzinfo=timezone
     )