Annual Limit daily delivered & failed notification counts and seeding implementation

.github/actions/waffles/requirements.txt

@@ -1,4 +1,4 @@
 docopt==0.6.2
 Flask==2.3.3
 markupsafe==2.1.5
-git+https://github.com/cds-snc/[email protected].6#egg=notifications-utils
+git+https://github.com/cds-snc/[email protected].7#egg=notifications-utils

notifications_utils/clients/redis/annual_limit.py

@@ -1,4 +1,27 @@
-"""This module stores daily notification counts and annual limit statuses for a service in Redis."""
+"""
+This module stores daily notification counts and annual limit statuses for a service in Redis using a hash structure:
+
+
+annual-limit: {
+    {service_id}: {
+        notifications: {
+            sms_delivered: int,
+            email_delivered: int,
+            sms_failed: int,
+            email_failed: int
+        },
+        status: {
+            near_sms_limit: Datetime,
+            near_email_limit: Datetime,
+            over_sms_limit: Datetime,
+            over_email_limit: Datetime
+            seeded_at: Datetime
+        }
+    }
+}
+
+
+"""
 
 from datetime import datetime
 
@@ -9,13 +32,18 @@
 SMS_FAILED = "sms_failed"
 EMAIL_FAILED = "email_failed"
 
+NOTIFICATIONS = [SMS_DELIVERED, EMAIL_DELIVERED, SMS_FAILED, EMAIL_FAILED]
+
 NEAR_SMS_LIMIT = "near_sms_limit"
 NEAR_EMAIL_LIMIT = "near_email_limit"
 OVER_SMS_LIMIT = "over_sms_limit"
 OVER_EMAIL_LIMIT = "over_email_limit"
+SEEDED_AT = "seeded_at"
+
+STATUSES = [NEAR_SMS_LIMIT, NEAR_EMAIL_LIMIT, OVER_SMS_LIMIT, OVER_EMAIL_LIMIT]
 
 
-def notifications_key(service_id):
+def annual_limit_notifications_key(service_id):
     """
     Generates the Redis hash key for storing daily metrics of a service.
     """
@@ -44,25 +72,48 @@ def init_app(self, app, *args, **kwargs):
         pass
 
     def increment_notification_count(self, service_id: str, field: str):
-        self._redis_client.increment_hash_value(notifications_key(service_id), field)
+        self._redis_client.increment_hash_value(annual_limit_notifications_key(service_id), field)
 
     def get_notification_count(self, service_id: str, field: str):
         """
         Retrieves the specified daily notification count for a service. (e.g. SMS_DELIVERED, EMAIL_FAILED, etc.)
         """
-        return int(self._redis_client.get_hash_field(notifications_key(service_id), field))
+        return int(self._redis_client.get_hash_field(annual_limit_notifications_key(service_id), field))
 
     def get_all_notification_counts(self, service_id: str):
         """
         Retrieves all daily notification metrics for a service.
         """
-        return decode_byte_dict(self._redis_client.get_all_from_hash(notifications_key(service_id)))
+        return decode_byte_dict(self._redis_client.get_all_from_hash(annual_limit_notifications_key(service_id)))
+
+    def reset_all_notification_counts(self, service_ids=None):
+        """
+        Resets all daily notification metrics.
+        :param: service_ids: list of service_ids to reset, if None, resets all services
+        """
+        hashes = (
+            annual_limit_notifications_key("*")
+            if not service_ids
+            else [annual_limit_notifications_key(service_id) for service_id in service_ids]
+        )
+
+        self._redis_client.delete_hash_fields(hashes=hashes)
+
+    def was_seeded_today(self, service_id):
+        last_seeded_time = self.get_seeded_at(service_id)
+        return last_seeded_time == datetime.utcnow().strftime("%Y-%m-%d") if last_seeded_time else False
+
+    def get_seeded_at(self, service_id: str):
+        return self._redis_client.get_hash_field(annual_limit_status_key(service_id), SEEDED_AT).decode("utf-8")
+
+    def set_seeded_at(self, service_id):
+        self._redis_client.set_hash_value(annual_limit_status_key(service_id), SEEDED_AT, datetime.utcnow().strftime("%Y-%m-%d"))
 
     def clear_notification_counts(self, service_id: str):
         """
         Clears all daily notification metrics for a service.
         """
-        self._redis_client.expire(notifications_key(service_id), -1)
+        self._redis_client.expire(annual_limit_notifications_key(service_id), -1)
 
     def set_annual_limit_status(self, service_id: str, field: str, value: datetime):
         """

notifications_utils/clients/redis/redis_client.py

@@ -81,6 +81,70 @@
             return self.scripts["delete-keys-by-pattern"](args=[pattern])
         return 0
 
+    def delete_hash_fields(self, hashes: (str | list), fields: list = None, raise_exception=False):
+        """Deletes fields from the specified hashes. if fields is `None`, then all fields from the hashes are deleted, deleting the hash entirely.
+
+        Args:
+            hashes (str|list): The hash pattern or list of hash keys to delete fields from.
+            fields (list): A list of fields to delete from the hashes. If `None`, then all fields are deleted.
+
+        Returns:
+            _type_: _description_
+        """
+        if self.active:
+            try:
+                hashes = [prepare_value(h) for h in hashes] if isinstance(hashes, list) else prepare_value(hashes)
+                # When fields are passed in, use the list as is
+                # When hashes is a list, and no fields are passed in, fetch the fields from the first hash in the list
+                # otherwise we know we're going scan iterate over a pattern so we'll fetch the fields on the first pass in the loop below
+                fields = (
+                    [prepare_value(f) for f in fields]
+                    if fields is not None
+                    else self.redis_store.hkeys(hashes[0])
+                    if isinstance(hashes, list)
+                    else None
+                )
+                # Use a pipeline to atomically delete fields from each hash.
+                pipe = self.redis_store.pipeline()
+                # if hashes is not a list, we're scan iterating over keys matching a pattern
+                for key in hashes if isinstance(hashes, list) else self.redis_store.scan_iter(hashes):
+                    if not fields:
+                        fields = self.redis_store.hkeys(key)
+                    key = prepare_value(key)
+                    pipe.hdel(key, *fields)
+                result = pipe.execute()
+                return result
+            except Exception as e:
+                self.__handle_exception(e, raise_exception, "expire_hash_fields", hashes)
+
+    def set_hash_fields_by_pattern_or_keys(self, mapping, keys: str | list = None, raise_exception=False):
+        """
+        Bulk set hash fields.
+        :param pattern: the pattern to match keys or a list of keys to set
+        :param mappting: the mapping of fields to set
+        :param raise_exception: True if we should allow the exception to bubble up
+        """
+        if self.active:
+            try:
+                for key in self.redis_store.hscan_iter(keys):
+                    self.redis_store.hmset(key, mapping)
+            except Exception as e:
+                self.__handle_exception(e, raise_exception, "bulk_set_hash_fields", keys)
+
+    def set_hash_fields_by_keys(self, keys, mapping, raise_exception=False):
+        """
+        Bulk set hash fields.
+        :param pattern: the pattern to match keys
+        :param mappting: the mapping of fields to set
+        :param raise_exception: True if we should allow the exception to bubble up
+        """
+        if self.active:
+            try:
+                for key in self.redis_store.scan_iter(keys):
+                    self.redis_store.hmset(key, mapping)
+            except Exception as e:
+                self.__handle_exception(e, raise_exception, "bulk_set_hash_fields", keys)
+
     def exceeded_rate_limit(self, cache_key, limit, interval, raise_exception=False):
         """
         Rate limiting.

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "notifications-utils"
-version = "52.3.6"
+version = "52.3.7"
 description = "Shared python code for Notification - Provides logging utils etc."
 authors = ["Canadian Digital Service"]
 license = "MIT license"

tests/test_annual_limit.py

@@ -127,6 +127,31 @@ def test_clear_notification_counts(mock_annual_limit_client, mock_notification_c
     assert len(mock_annual_limit_client.get_all_notification_counts(mocked_service_id)) == 0
 
 
+@pytest.mark.parametrize(
+    "service_ids",
+    [
+        [
+            str(uuid.uuid4()),
+            str(uuid.uuid4()),
+            str(uuid.uuid4()),
+            str(uuid.uuid4()),
+        ]
+    ],
+)
+def test_bulk_reset_notification_counts(mock_annual_limit_client, mock_notification_count_types, service_ids):
+    for service_id in service_ids:
+        for field in mock_notification_count_types:
+            mock_annual_limit_client.increment_notification_count(service_id, field)
+        assert len(mock_annual_limit_client.get_all_notification_counts(service_id)) == 4
+
+    mock_annual_limit_client.reset_all_notification_counts()
+
+    for service_id in service_ids:
+        assert len(mock_annual_limit_client.get_all_notification_counts(service_id)) == 4
+        for field in mock_notification_count_types:
+            assert mock_annual_limit_client.get_notification_count(service_id, field) == 0
+
+
 def test_set_annual_limit_status(mock_annual_limit_client, mocked_service_id):
     mock_annual_limit_client.set_annual_limit_status(mocked_service_id, NEAR_SMS_LIMIT, datetime.utcnow())
     result = mock_annual_limit_client.get_annual_limit_status(mocked_service_id, NEAR_SMS_LIMIT)

tests/test_redis_client.py

@@ -184,6 +184,51 @@ def test_set_hash_value(mocked_redis_client):
     mocked_redis_client.redis_store.hset.assert_called_with(key, field, value)
 
 
+@pytest.mark.parametrize(
+    "hash, updates, expected",
+    [
+        (
+            {
+                "key1": {
+                    "field1": "value1",
+                    "field2": 2,
+                    "field3": "value3".encode("utf-8"),
+                },
+                "key2": {
+                    "field1": "value1",
+                    "field2": 2,
+                    "field3": "value3".encode("utf-8"),
+                },
+                "key3": {
+                    "field1": "value1",
+                    "field2": 2,
+                    "field3": "value3".encode("utf-8"),
+                },
+            },
+            {
+                "field1": "value2",
+                "field2": 3,
+                "field3": "value4".encode("utf-8"),
+            },
+            {
+                b"field1": b"value2",
+                b"field2": b"3",
+                b"field3": b"value4",
+            },
+        )
+    ],
+)
+def test_bulk_set_hash_fields(better_mocked_redis_client, hash, updates, expected):
+    for key, fields in hash.items():
+        for field, value in fields.items():
+            better_mocked_redis_client.set_hash_value(key, field, value)
+
+    better_mocked_redis_client.bulk_set_hash_fields("key*", mapping=updates)
+
+    for key, _ in hash.items():
+        assert better_mocked_redis_client.redis_store.hgetall(key) == expected
+
+
 def test_decrement_hash_value_should_decrement_value_by_one_for_key(mocked_redis_client):
     key = "12345"
     value = "template-1111"