perf(billing): compute usage balances once per limit check DEV-926

kobo/apps/openrosa/apps/api/tests/viewsets/test_xform_submission_api.py

@@ -10,6 +10,7 @@
 from constance.test import override_config
 from django.conf import settings
 from django.contrib.auth.models import AnonymousUser
+from django.core.cache import cache
 from django.core.exceptions import PermissionDenied
 from django.core.files.base import ContentFile
 from django.core.files.uploadedfile import InMemoryUploadedFile
@@ -144,7 +145,7 @@ def test_over_limit_submission_rejection_anonymous(self):
     def test_over_limit_submission_rejection_authenticated(self, mock_usage):
         """
         Ensure submissions by an authenticated user are rejected if asset owner
-        is over their storage or submission limit and that check_exceeded_limit
+        is over their storage or submission limit and that check_exceeded_limits
         is run.
         """
         path = os.path.join(
@@ -167,15 +168,16 @@ def test_over_limit_submission_rejection_authenticated(self, mock_usage):
             }
             mock_usage.return_value = mock_balances
             with patch(
-                'kobo.apps.openrosa.libs.utils.logger_tools.check_exceeded_limit',
-                return_value=None,
+                'kobo.apps.openrosa.libs.utils.logger_tools.check_exceeded_limits',
+                return_value={},
             ) as patched:
                 request = self.factory.post('/submission', data, format='json')
                 auth = DigestAuth('bob', 'bobbob')
                 request.META.update(auth(request.META, response))
                 response = self.view(request, username=self.user.username)
-                patched.assert_any_call(self.user, UsageType.SUBMISSION)
-                patched.assert_any_call(self.user, UsageType.STORAGE_BYTES)
+                patched.assert_called_once_with(
+                    self.user, [UsageType.SUBMISSION, UsageType.STORAGE_BYTES]
+                )
                 self.assertEqual(response.status_code, status.HTTP_402_PAYMENT_REQUIRED)
 
             mock_balances = {
@@ -186,8 +188,8 @@ def test_over_limit_submission_rejection_authenticated(self, mock_usage):
             }
             mock_usage.return_value = mock_balances
             with patch(
-                'kobo.apps.openrosa.libs.utils.logger_tools.check_exceeded_limit',
-                return_value=None,
+                'kobo.apps.openrosa.libs.utils.logger_tools.check_exceeded_limits',
+                return_value={},
             ) as patched:
                 request = self.factory.post('/submission', data, format='json')
                 response = self.view(request)
@@ -196,10 +198,59 @@ def test_over_limit_submission_rejection_authenticated(self, mock_usage):
                 auth = DigestAuth('bob', 'bobbob')
                 request.META.update(auth(request.META, response))
                 response = self.view(request, username=self.user.username)
-                patched.assert_any_call(self.user, UsageType.SUBMISSION)
-                patched.assert_any_call(self.user, UsageType.STORAGE_BYTES)
+                patched.assert_called_once_with(
+                    self.user, [UsageType.SUBMISSION, UsageType.STORAGE_BYTES]
+                )
                 self.assertEqual(response.status_code, status.HTTP_402_PAYMENT_REQUIRED)
 
+    @pytest.mark.skipif(
+        not settings.STRIPE_ENABLED, reason='Requires stripe functionality'
+    )
+    @patch(
+        'kobo.apps.openrosa.libs.utils.logger_tools.ServiceUsageCalculator.get_usage_balances'  # noqa: E501
+    )
+    @patch('kobo.apps.stripe.utils.limit_enforcement._get_usage_balances')
+    def test_over_limit_rejection_counts_on_fresh_usage(
+        self, mock_fresh_usage, mock_cached_usage
+    ):
+        """
+        Ensure a rejected submission records exceeded-limit counters from fresh
+        usage, not from the cached balances that triggered the rejection
+        """
+        from kobo.apps.stripe.models import ExceededLimitCounter
+
+        mock_cached_usage.return_value = {
+            UsageType.STORAGE_BYTES: {'exceeded': True},
+            UsageType.SUBMISSION: None,
+        }
+        # e.g. the owner deleted attachments after the balances were cached
+        mock_fresh_usage.return_value = {
+            UsageType.STORAGE_BYTES: {'exceeded': False},
+            UsageType.SUBMISSION: None,
+        }
+        for usage_type in (UsageType.SUBMISSION, UsageType.STORAGE_BYTES):
+            cache.delete(f'{self.user.id}_checked_exceeded_{usage_type}_limit')
+        path = os.path.join(
+            os.path.dirname(os.path.abspath(__file__)),
+            '..',
+            'fixtures',
+            'transport_submission.json',
+        )
+        with open(path, 'rb') as f:
+            data = json.loads(f.read())
+
+        request = self.factory.post('/submission', data, format='json')
+        response = self.view(request)
+        self.assertEqual(response.status_code, 401)
+        request = self.factory.post('/submission', data, format='json')
+        auth = DigestAuth('bob', 'bobbob')
+        request.META.update(auth(request.META, response))
+        response = self.view(request, username=self.user.username)
+
+        self.assertEqual(response.status_code, status.HTTP_402_PAYMENT_REQUIRED)
+        mock_fresh_usage.assert_called_once()
+        assert not ExceededLimitCounter.objects.exists()
+
     @pytest.mark.skipif(
         not settings.STRIPE_ENABLED, reason='Requires stripe functionality'
     )

kobo/apps/openrosa/apps/logger/tests/test_simple_submission.py

@@ -153,12 +153,13 @@ def test_corrupted_submission(self):
     )
     def test_check_exceeded_limit_on_submission(self):
         with patch(
-            'kobo.apps.openrosa.libs.utils.logger_tools.check_exceeded_limit',
-            return_value=None,
+            'kobo.apps.openrosa.libs.utils.logger_tools.check_exceeded_limits',
+            return_value={},
         ) as patched:
             self._submit_simple_yes()
-            patched.assert_any_call(self.user, UsageType.SUBMISSION)
-            patched.assert_any_call(self.user, UsageType.STORAGE_BYTES)
+            patched.assert_called_once_with(
+                self.user, [UsageType.SUBMISSION, UsageType.STORAGE_BYTES]
+            )
 
     def test_parsed_instance_submitted_by_value(self):
         class MockRequest:

kobo/apps/openrosa/libs/utils/logger_tools.py

@@ -93,7 +93,7 @@
 from kobo.apps.openrosa.libs.utils.model_tools import queryset_iterator, set_uuid
 from kobo.apps.openrosa.libs.utils.viewer_tools import get_mongo_userform_id
 from kobo.apps.organizations.constants import UsageType
-from kobo.apps.stripe.utils.limit_enforcement import check_exceeded_limit
+from kobo.apps.stripe.utils.limit_enforcement import check_exceeded_limits
 from kpi.constants import PERM_ADD_SUBMISSIONS, PERM_CHANGE_SUBMISSIONS
 from kpi.deployment_backends.kc_access.storage import (
     default_kobocat_storage as default_storage,
@@ -320,13 +320,17 @@ def create_instance(
     ):
         calculator = ServiceUsageCalculator(xform.user)
         balances = calculator.get_usage_balances()
-        for usage_type in [UsageType.STORAGE_BYTES, UsageType.SUBMISSION]:
-            balance = balances[usage_type]
-            if balance and balance['exceeded']:
-                check_exceeded_limit(xform.user, UsageType.SUBMISSION)
-                check_exceeded_limit(xform.user, UsageType.STORAGE_BYTES)
+        exceeded_types = [
+            usage_type
+            for usage_type in [UsageType.STORAGE_BYTES, UsageType.SUBMISSION]
+            if balances[usage_type] and balances[usage_type]['exceeded']
+        ]
+        if exceeded_types:
+            check_exceeded_limits(
+                xform.user, [UsageType.SUBMISSION, UsageType.STORAGE_BYTES]
+            )
 
-                raise ExceededUsageLimitError({'type': usage_type})
+            raise ExceededUsageLimitError({'type': exceeded_types[0]})
 
     # get root uuid
     root_uuid, fallback_on_uuid = get_root_uuid_from_xml(xml)
@@ -433,8 +437,10 @@ def create_instance(
             )
 
             if settings.STRIPE_ENABLED:
-                check_exceeded_limit(xform.user, UsageType.SUBMISSION)
-                check_exceeded_limit(xform.user, UsageType.STORAGE_BYTES)
+                check_exceeded_limits(
+                    xform.user,
+                    [UsageType.SUBMISSION, UsageType.STORAGE_BYTES],
+                )
 
             return instance
 

kobo/apps/stripe/tests/test_stripe_utils.py

@@ -37,6 +37,7 @@
 )
 from kobo.apps.stripe.utils.limit_enforcement import (
     check_exceeded_limit,
+    check_exceeded_limits,
     update_or_remove_limit_counter,
 )
 from kobo.apps.stripe.utils.manual_subscription import (
@@ -762,8 +763,8 @@ def test_check_exceeded_limit_adds_counters(self):
         # We want to test this function directly here, so we patch it out when
         # it is called on submission to avoid cache restrictions
         with patch(
-            'kobo.apps.openrosa.libs.utils.logger_tools.check_exceeded_limit',
-            return_value=None,
+            'kobo.apps.openrosa.libs.utils.logger_tools.check_exceeded_limits',
+            return_value={},
         ):
             self.add_submissions(count=2, asset=self.asset, username='someuser')
         self.add_nlp_trackers()
@@ -790,8 +791,8 @@ def test_check_exceeded_limit_updates_counters(self):
         # We want to test this function directly here, so we patch it out when
         # is called on submission to avoid cache restrictions
         with patch(
-            'kobo.apps.openrosa.libs.utils.logger_tools.check_exceeded_limit',
-            return_value=None,
+            'kobo.apps.openrosa.libs.utils.logger_tools.check_exceeded_limits',
+            return_value={},
         ):
             self.add_submissions(count=2, asset=self.asset, username='someuser')
         self.add_nlp_trackers()
@@ -866,3 +867,28 @@ def test_update_or_remove_limit_counter(self):
         ):
             update_or_remove_limit_counter(counter)
             assert ExceededLimitCounter.objects.count() == 0
+
+    def test_check_exceeded_limits_computes_balances_once(self):
+        mock_balances = {
+            UsageType.ASR_SECONDS: None,
+            UsageType.MT_CHARACTERS: None,
+            UsageType.STORAGE_BYTES: {'exceeded': True},
+            UsageType.SUBMISSION: {'exceeded': True},
+        }
+        with patch(
+            'kpi.utils.usage_calculator.ServiceUsageCalculator.get_usage_balances',
+            return_value=mock_balances,
+        ) as patched:
+            check_exceeded_limits(
+                self.someuser,
+                [UsageType.SUBMISSION, UsageType.STORAGE_BYTES],
+            )
+            patched.assert_called_once()
+
+        for usage_type in [UsageType.SUBMISSION, UsageType.STORAGE_BYTES]:
+            assert (
+                ExceededLimitCounter.objects.filter(
+                    user_id=self.anotheruser.id, limit_type=usage_type
+                ).count()
+                == 1
+            )

kobo/apps/stripe/utils/limit_enforcement.py

@@ -1,4 +1,5 @@
 from datetime import timedelta
+from typing import TYPE_CHECKING
 
 from django.conf import settings
 from django.core.cache import cache
@@ -9,42 +10,71 @@
 from kobo.apps.stripe.utils.import_management import requires_stripe
 from kpi.utils.usage_calculator import ServiceUsageCalculator
 
+if TYPE_CHECKING:
+    from kobo.apps.stripe.models import ExceededLimitCounter
+
 
 @requires_stripe
 def check_exceeded_limit(user, usage_type: UsageType, **kwargs):
     """
-    Checks whether user's org has exceeded its limits for a given
-    usage type and updates ExceededLimitCounters accordingly. Uses
-    cached key to avoid running checks more than once within
-    the ENDPOINT_CACHE_DURATION
+    Check a single usage type against its limit.
+
+    Thin wrapper around `check_exceeded_limits` kept for callers that work with
+    one usage type. Returns the created/updated counter, or None when the check
+    is gated or the limit is not exceeded.
+    """
+    return check_exceeded_limits(user, [usage_type]).get(usage_type)
+
+
+@requires_stripe
+def check_exceeded_limits(
+    user,
+    usage_types: list[UsageType],
+    **kwargs,
+) -> 'dict[UsageType, ExceededLimitCounter]':
+    """
+    Check whether user's org has exceeded its limits for the given usage types
+    and update ExceededLimitCounters accordingly.
+
+    Uses a cached key per usage type to avoid running checks more than once
+    within the ENDPOINT_CACHE_DURATION. Fresh usage balances are computed at most
+    once per call, and only when a usage type passes its gate.
     """
     org = user.organization
     if org.is_mmo:
         user = org.owner_user_object
 
-    cache_key = f'{user.id}_checked_exceeded_{usage_type}_limit'
+    ExceededLimitCounter = kwargs['exceeded_limit_counter_model']
+    balances = None
+    counters = {}
+
+    for usage_type in usage_types:
+        cache_key = f'{user.id}_checked_exceeded_{usage_type}_limit'
 
-    if cache.get(cache_key, None):
-        return
+        if cache.get(cache_key, None):
+            continue
 
-    ExceededLimitCounter = kwargs['exceeded_limit_counter_model']
-    balances = _get_usage_balances(user)
-    balance = balances[usage_type]
-    counter = None
-
-    if balance and balance['exceeded']:
-        counter, created = ExceededLimitCounter.objects.get_or_create(
-            user=user,
-            limit_type=usage_type,
-        )
-
-        if not created and counter.date_modified.date() < timezone.now().date():
-            delta = timezone.now().date() - counter.date_modified.date()
-            counter.days += delta.days
-            counter.save()
-
-    cache.set(cache_key, True, settings.ENDPOINT_CACHE_DURATION)
-    return counter
+        if balances is None:
+            balances = _get_usage_balances(user)
+
+        balance = balances[usage_type]
+
+        if balance and balance['exceeded']:
+            counter, created = ExceededLimitCounter.objects.get_or_create(
+                user=user,
+                limit_type=usage_type,
+            )
+
+            if not created and counter.date_modified.date() < timezone.now().date():
+                delta = timezone.now().date() - counter.date_modified.date()
+                counter.days += delta.days
+                counter.save()
+
+            counters[usage_type] = counter
+
+        cache.set(cache_key, True, settings.ENDPOINT_CACHE_DURATION)
+
+    return counters
 
 
 @requires_stripe