diff --git a/kobo/apps/organizations/exceptions.py b/kobo/apps/organizations/exceptions.py deleted file mode 100644 index bf5b33ad5f..0000000000 --- a/kobo/apps/organizations/exceptions.py +++ /dev/null @@ -1,2 +0,0 @@ -class UsageLimitExceeded(Exception): - pass diff --git a/kobo/apps/organizations/types.py b/kobo/apps/organizations/types.py index f6e2054635..e9570c319d 100644 --- a/kobo/apps/organizations/types.py +++ b/kobo/apps/organizations/types.py @@ -1,3 +1,3 @@ from typing import Literal -UsageType = Literal['character', 'seconds', 'submission', 'storage'] +UsageType = Literal['characters', 'seconds', 'submission', 'storage'] diff --git a/kobo/apps/stripe/constants.py b/kobo/apps/stripe/constants.py index fc8e573a0c..e93a869920 100644 --- a/kobo/apps/stripe/constants.py +++ b/kobo/apps/stripe/constants.py @@ -22,14 +22,14 @@ ORGANIZATION_USAGE_MAX_CACHE_AGE = timedelta(minutes=15) USAGE_LIMIT_MAP = { - 'character': 'mt_characters', + 'characters': 'mt_characters', 'seconds': 'asr_seconds', 'storage': 'storage_bytes', 'submission': 'submission', } USAGE_LIMIT_MAP_STRIPE = { - 'character': 'nlp_character', + 'characters': 'nlp_character', 'seconds': 'nlp_seconds', 'storage': 'storage_bytes', 'submission': 'submission', diff --git a/kobo/apps/stripe/models.py b/kobo/apps/stripe/models.py index 12bf64b423..09623ca076 100644 --- a/kobo/apps/stripe/models.py +++ b/kobo/apps/stripe/models.py @@ -16,6 +16,7 @@ from kobo.apps.stripe.constants import ACTIVE_STRIPE_STATUSES, USAGE_LIMIT_MAP from kobo.apps.stripe.utils import get_default_add_on_limits from kpi.fields import KpiUidField +from kpi.utils.django_orm_helper import DeductUsageValue class PlanAddOn(models.Model): @@ -49,56 +50,12 @@ class Meta: models.Index(fields=['organization', 'limits_remaining', 'charge']), ] - @property - def is_expended(self): - """ - Whether the addon is at/over its usage limits. - """ - for limit_type, limit_value in self.limits_remaining.items(): - if limit_value > 0: - return False - return True - - @property - def total_usage_limits(self): - """ - The total usage limits for this add-on, based on the usage_limits for a single add-on and the quantity. - """ - return {key: value * self.quantity for key, value in self.usage_limits.items()} - - @property - def valid_tags(self) -> List: - """ - The tag metadata (on the subscription product/price) needed to view/purchase this add-on. - If the org that purchased this add-on no longer has that a plan with those tags, the add-on will be inactive. - If the add-on doesn't require a tag, this property will return an empty list. - """ - return self.product.metadata.get('valid_tags', '').split(',') - - @admin.display(boolean=True, description='available') - def is_available(self): - return not ( - self.is_expended or self.charge.refunded - ) and bool(self.organization) - - def increment(self, limit_type, amount_used): - """ - Increments the usage counter for limit_type by amount_used. - Returns the amount of this add-on that was used (up to its limit). - Will return 0 if limit_type does not apply to this add-on. - """ - if limit_type in self.usage_limits.keys(): - limit_available = self.limits_remaining.get(limit_type) - amount_to_use = min(amount_used, limit_available) - self.limits_remaining[limit_type] -= amount_to_use - self.save() - return amount_to_use - return 0 - @staticmethod def create_or_update_one_time_add_on(charge: Charge): """ - Create a PlanAddOn object from a Charge object, if the Charge is for a one-time add-on. + Create a PlanAddOn object from a Charge object, if the Charge is for a + one-time add-on. + Returns True if a PlanAddOn was created, false otherwise. """ if ( @@ -154,26 +111,12 @@ def create_or_update_one_time_add_on(charge: Charge): add_on.save() return add_on_created - @staticmethod - def make_add_ons_from_existing_charges(): - """ - Create a PlanAddOn object for each eligible Charge object in the database. - Does not refresh Charge data from Stripe. - Returns the number of PlanAddOns created. - """ - created_count = 0 - # TODO: This should filter out charges that are already matched to an add on - for charge in Charge.objects.all().iterator(chunk_size=500): - if PlanAddOn.create_or_update_one_time_add_on(charge): - created_count += 1 - return created_count - @staticmethod def get_organization_totals( organization: 'Organization', usage_type: UsageType ) -> (int, int): """ - Returns the total limit and the total remaining usage for a given org. + Returns the total limit and the total remaining usage for a given organization and usage type. """ usage_mapped = USAGE_LIMIT_MAP[usage_type] @@ -199,13 +142,63 @@ def get_organization_totals( return totals['total_usage_limit'], totals['total_remaining'] + @property + def is_expended(self): + """ + Whether the addon is at/over its usage limits. + """ + for limit_type, limit_value in self.limits_remaining.items(): + if limit_value > 0: + return False + return True + + @admin.display(boolean=True, description='available') + def is_available(self): + return ( + self.charge.payment_intent.status == PaymentIntentStatus.succeeded and not ( + self.is_expended or self.charge.refunded + ) and bool(self.organization) + ) + + def deduct(self, limit_type, amount_used): + """ + Deducts the add on usage counter for limit_type by amount_used. + Returns the amount of this add-on that was used (up to its limit). + Will return 0 if limit_type does not apply to this add-on. + """ + if limit_type in self.usage_limits.keys(): + limit_available = self.limits_remaining.get(limit_type) + amount_to_use = min(amount_used, limit_available) + PlanAddOn.objects.filter(pk=self.id).update( + limits_remaining=DeductUsageValue( + 'limits_remaining', keyname=limit_type, amount=amount_used) + ) + return amount_to_use + return 0 + + @staticmethod + def make_add_ons_from_existing_charges(): + """ + Create a PlanAddOn object for each eligible Charge object in the database. + Does not refresh Charge data from Stripe. + Returns the number of PlanAddOns created. + """ + created_count = 0 + # TODO: This should filter out charges that are already matched to an add on + for charge in Charge.objects.all().iterator(chunk_size=500): + if PlanAddOn.create_or_update_one_time_add_on(charge): + created_count += 1 + return created_count + @staticmethod - def increment_add_ons_for_organization( + def deduct_add_ons_for_organization( organization: 'Organization', usage_type: UsageType, amount: int ): """ - Increments the usage counter for limit_type by amount_used for a given user. - Will always increment the add-on with the most used first, so that add-ons are used up in FIFO order. + Deducts the usage counter for limit_type by amount_used for a given user. + Will always spend the add-on with the most used first, so that add-ons + are used up in FIFO order. + Returns the amount of usage that was not applied to an add-on. """ usage_mapped = USAGE_LIMIT_MAP[usage_type] @@ -220,11 +213,29 @@ def increment_add_ons_for_organization( remaining = amount for add_on in add_ons.iterator(): if add_on.is_available(): - remaining -= add_on.increment( + remaining -= add_on.deduct( limit_type=limit_key, amount_used=remaining ) return remaining + @property + def total_usage_limits(self): + """ + The total usage limits for this add-on, based on the usage_limits for a single + add-on and the quantity. + """ + return {key: value * self.quantity for key, value in self.usage_limits.items()} + + @property + def valid_tags(self) -> List: + """ + The tag metadata (on the subscription product/price) needed to view/purchase + this add-on. If the org that purchased this add-on no longer has that a plan + with those tags, the add-on will be inactive. If the add-on doesn't require a + tag, this property will return an empty list. + """ + return self.product.metadata.get('valid_tags', '').split(',') + @receiver(post_save, sender=Charge) def make_add_on_for_charge(sender, instance, created, **kwargs): diff --git a/kobo/apps/stripe/tests/test_one_time_addons_api.py b/kobo/apps/stripe/tests/test_one_time_addons_api.py index 2f060e9544..643bb8798c 100644 --- a/kobo/apps/stripe/tests/test_one_time_addons_api.py +++ b/kobo/apps/stripe/tests/test_one_time_addons_api.py @@ -180,7 +180,7 @@ def test_not_own_addon(self): assert response_get_list.status_code == status.HTTP_200_OK assert response_get_list.data['results'] == [] - @data('character', 'seconds') + @data('characters', 'seconds') def test_get_user_totals(self, usage_type): limit = 2000 quantity = 5 @@ -202,7 +202,7 @@ def test_get_user_totals(self, usage_type): assert total_limit == limit * (quantity + 2) assert remaining == limit * (quantity + 2) - PlanAddOn.increment_add_ons_for_organization( + PlanAddOn.deduct_add_ons_for_organization( self.organization, usage_type, limit * quantity ) total_limit, remaining = PlanAddOn.get_organization_totals( diff --git a/kobo/apps/stripe/tests/test_organization_usage.py b/kobo/apps/stripe/tests/test_organization_usage.py index 286cb494b6..602e949d72 100644 --- a/kobo/apps/stripe/tests/test_organization_usage.py +++ b/kobo/apps/stripe/tests/test_organization_usage.py @@ -9,6 +9,7 @@ import pytest from dateutil.relativedelta import relativedelta +from ddt import ddt, data from django.core.cache import cache from django.test import override_settings from django.urls import reverse @@ -463,6 +464,7 @@ def test_users_without_enterprise_see_only_their_usage(self): assert response.data['count'] == 1 +@ddt class OrganizationsUtilsTestCase(BaseTestCase): fixtures = ['test_data'] @@ -479,16 +481,11 @@ def test_get_plan_community_limit(self): generate_enterprise_subscription(self.organization) limit = get_organization_plan_limit(self.organization, 'seconds') assert limit == 2000 # TODO get the limits from the community plan, overrides - limit = get_organization_plan_limit(self.organization, 'character') + limit = get_organization_plan_limit(self.organization, 'characters') assert limit == 2000 # TODO get the limits from the community plan, overrides - def test_get_subscription_limits_characters(self): - self._test_get_suscription_limit('character') - - def test_get_subscription_limits_seconds(self): - self._test_get_suscription_limit('seconds') - - def _test_get_suscription_limit(self, usage_type): + @data('characters', 'seconds') + def test_get_suscription_limit(self, usage_type): stripe_key = f'{USAGE_LIMIT_MAP_STRIPE[usage_type]}_limit' product_metadata = { stripe_key: 1234, diff --git a/kobo/apps/trackers/tests/test_utils.py b/kobo/apps/trackers/tests/test_utils.py index f0a640a7c3..d7e5153e58 100644 --- a/kobo/apps/trackers/tests/test_utils.py +++ b/kobo/apps/trackers/tests/test_utils.py @@ -84,7 +84,7 @@ def _make_payment( charge.save() return charge - @data('character', 'seconds') + @data('characters', 'seconds') def test_organization_usage_utils(self, usage_type): stripe_key = f'{USAGE_LIMIT_MAP_STRIPE[usage_type]}_limit' usage_key = f'{USAGE_LIMIT_MAP[usage_type]}_limit' diff --git a/kobo/apps/trackers/utils.py b/kobo/apps/trackers/utils.py index f59a17812c..2115324358 100644 --- a/kobo/apps/trackers/utils.py +++ b/kobo/apps/trackers/utils.py @@ -53,11 +53,11 @@ def update_nlp_counter( if service.endswith('asr_seconds'): kwargs['total_asr_seconds'] = F('total_asr_seconds') + amount if asset_id is not None and organization is not None: - handle_usage_increment(organization, 'seconds', amount) + handle_usage_deduction(organization, 'seconds', amount) if service.endswith('mt_characters'): kwargs['total_mt_characters'] = F('total_mt_characters') + amount if asset_id is not None and organization is not None: - handle_usage_increment(organization, 'character', amount) + handle_usage_deduction(organization, 'characters', amount) NLPUsageCounter.objects.filter(pk=counter_id).update( counters=IncrementValue('counters', keyname=service, increment=amount), @@ -71,9 +71,8 @@ def get_organization_usage(organization: Organization, usage_type: UsageType) -> Get the used amount for a given organization and usage type """ usage_calc = ServiceUsageCalculator( - organization.owner.organization_user.user, organization + organization.owner.organization_user.user, organization, disable_cache=True ) - usage_calc._clear_cache() # Do not use cached values usage = usage_calc.get_nlp_usage_by_type(USAGE_LIMIT_MAP[usage_type]) return usage @@ -101,11 +100,11 @@ def get_organization_remaining_usage( return total_remaining -def handle_usage_increment( +def handle_usage_deduction( organization: Organization, usage_type: UsageType, amount: int ): """ - Increment the given usage type for this organization by the given amount + Deducts the specified usage type for this organization by the given amount """ PlanAddOn = apps.get_model('stripe', 'PlanAddOn') @@ -115,9 +114,9 @@ def handle_usage_increment( current_usage = 0 new_total_usage = current_usage + amount if new_total_usage > plan_limit: - increment = ( + deduction = ( amount if current_usage >= plan_limit else new_total_usage - plan_limit ) - PlanAddOn.increment_add_ons_for_organization( - organization, usage_type, increment + PlanAddOn.deduct_add_ons_for_organization( + organization, usage_type, deduction ) diff --git a/kpi/tests/test_usage_calculator.py b/kpi/tests/test_usage_calculator.py index f121a4500e..96e918d8da 100644 --- a/kpi/tests/test_usage_calculator.py +++ b/kpi/tests/test_usage_calculator.py @@ -161,6 +161,20 @@ def setUp(self): self.add_nlp_trackers() self.add_submissions(count=5) + def test_disable_cache(self): + calculator = ServiceUsageCalculator(self.anotheruser, None, disable_cache=True) + nlp_usage_A = calculator.get_nlp_usage_counters() + self.add_nlp_trackers() + nlp_usage_B = calculator.get_nlp_usage_counters() + assert ( + 2*nlp_usage_A['asr_seconds_current_month'] == + nlp_usage_B['asr_seconds_current_month'] + ) + assert ( + 2*nlp_usage_A['mt_characters_current_month'] == + nlp_usage_B['mt_characters_current_month'] + ) + def test_nlp_usage_counters(self): calculator = ServiceUsageCalculator(self.anotheruser, None) nlp_usage = calculator.get_nlp_usage_counters() @@ -169,16 +183,6 @@ def test_nlp_usage_counters(self): assert nlp_usage['mt_characters_current_month'] == 5473 assert nlp_usage['mt_characters_all_time'] == 6726 - def test_storage_usage(self): - calculator = ServiceUsageCalculator(self.anotheruser, None) - assert calculator.get_storage_usage() == 5 * self.expected_file_size() - - def test_submission_counters(self): - calculator = ServiceUsageCalculator(self.anotheruser, None) - submission_counters = calculator.get_submission_counters() - assert submission_counters['current_month'] == 5 - assert submission_counters['all_time'] == 5 - def test_no_data(self): calculator = ServiceUsageCalculator(self.someuser, None) nlp_usage = calculator.get_nlp_usage_counters() @@ -212,5 +216,15 @@ def test_organization_setup(self): assert calculator.get_storage_usage() == 5 * self.expected_file_size() - assert calculator.get_nlp_usage_by_type(USAGE_LIMIT_MAP['character']) == 5473 + assert calculator.get_nlp_usage_by_type(USAGE_LIMIT_MAP['characters']) == 5473 assert calculator.get_nlp_usage_by_type(USAGE_LIMIT_MAP['seconds']) == 4586 + + def test_storage_usage(self): + calculator = ServiceUsageCalculator(self.anotheruser, None) + assert calculator.get_storage_usage() == 5 * self.expected_file_size() + + def test_submission_counters(self): + calculator = ServiceUsageCalculator(self.anotheruser, None) + submission_counters = calculator.get_submission_counters() + assert submission_counters['current_month'] == 5 + assert submission_counters['all_time'] == 5 diff --git a/kpi/utils/cache.py b/kpi/utils/cache.py index 5c1a0a1a2d..d2cb5f5189 100644 --- a/kpi/utils/cache.py +++ b/kpi/utils/cache.py @@ -86,6 +86,9 @@ def _setup_cache(self): """ Sets up the cache client and the cache hash name for the hset """ + if getattr(self, '_cache_available', None) is False: + return + self._redis_client = None self._cache_available = True self._cached_hset = {} diff --git a/kpi/utils/django_orm_helper.py b/kpi/utils/django_orm_helper.py index 57e2dd55ee..92aac2ce89 100644 --- a/kpi/utils/django_orm_helper.py +++ b/kpi/utils/django_orm_helper.py @@ -40,6 +40,31 @@ def __init__(self, expression: str, keyname: str, increment: int, **extra): ) +class DeductUsageValue(Func): + + function = 'jsonb_set' + usage_value = "COALESCE(%(expressions)s ->> '%(keyname)s', '0')::int" + template = ( + '%(function)s(%(expressions)s,' + '\'{"%(keyname)s"}\',' + '(' + f'CASE WHEN {usage_value} > %(amount)s ' + f'THEN {usage_value} - %(amount)s ' + 'ELSE 0 ' + 'END ' + ')::text::jsonb)' + ) + arity = 1 + + def __init__(self, expression: str, keyname: str, amount: int, **extra): + super().__init__( + expression, + keyname=keyname, + amount=amount, + **extra, + ) + + class OrderCustomCharField(Func): """ DO NOT use on fields other than CharField while the application maintains diff --git a/kpi/utils/usage_calculator.py b/kpi/utils/usage_calculator.py index 70eca3fa51..f190c02574 100644 --- a/kpi/utils/usage_calculator.py +++ b/kpi/utils/usage_calculator.py @@ -20,10 +20,14 @@ class ServiceUsageCalculator(CachedClass): CACHE_TTL = settings.ENDPOINT_CACHE_DURATION - def __init__(self, user: User, organization: Optional['Organization']): + def __init__( + self, user: User, + organization: Optional['Organization'], + disable_cache: bool = False + ): self.user = user self.organization = organization - + self._cache_available = not disable_cache self._user_ids = [user.pk] self._user_id_query = self._filter_by_user([user.pk]) if organization and settings.STRIPE_ENABLED: