Added force ttl check option to kv-table functions

global_helpers/panther_oss_helpers.py

@@ -186,12 +186,21 @@ def kv_table() -> boto3.resource:
     return _KV_TABLE
 
 
-def get_counter(key: str) -> int:
+def ttl_expired(response: dict) -> bool:
+    """Checks whether a response from the panther-kv table has passed it's TTL date"""
+    # This can be used when the TTL timing is very exacting and DDB's cleanup is too slow
+    expiration = response.get("Item", {}).get("expiresAt", 0)
+    return expiration and float(expiration) <= (datetime.now()).timestamp()
+
+
+def get_counter(key: str, force_ttl_check: bool = False) -> int:
     """Get a counter's current value (defaulting to 0 if key does not exist)."""
     response = kv_table().get_item(
         Key={"key": key},
         ProjectionExpression=_COUNT_COL,
     )
+    if force_ttl_check and ttl_expired(response):
+        return 0
     return response.get("Item", {}).get(_COUNT_COL, 0)
 
 
@@ -271,7 +280,7 @@ def put_dictionary(key: str, val: dict, epoch_seconds: int = None):
         set_key_expiration(key, epoch_seconds)
 
 
-def get_dictionary(key: str) -> dict:
+def get_dictionary(key: str, force_ttl_check: bool = False) -> dict:
     # Retrieve the item from DynamoDB
     response = kv_table().get_item(Key={"key": key})
 
@@ -281,6 +290,9 @@ def get_dictionary(key: str) -> dict:
     if not item:
         return {}
 
+    if force_ttl_check and ttl_expired(response):
+        return {}
+
     try:
         # Deserialize from JSON to a Python dictionary
         return json.loads(item)
@@ -291,12 +303,14 @@ def get_dictionary(key: str) -> dict:
         ) from exc
 
 
-def get_string_set(key: str) -> Set[str]:
+def get_string_set(key: str, force_ttl_check: bool = False) -> Set[str]:
     """Get a string set's current value (defaulting to empty set if key does not exit)."""
     response = kv_table().get_item(
         Key={"key": key},
         ProjectionExpression=_STRING_SET_COL,
     )
+    if force_ttl_check and ttl_expired(response):
+        return set()
     return response.get("Item", {}).get(_STRING_SET_COL, set())