diff --git a/kobo/apps/project_ownership/tests/api/v2/test_api.py b/kobo/apps/project_ownership/tests/api/v2/test_api.py index 8f85ab2aa4..59ec50bced 100644 --- a/kobo/apps/project_ownership/tests/api/v2/test_api.py +++ b/kobo/apps/project_ownership/tests/api/v2/test_api.py @@ -13,7 +13,6 @@ InviteStatusChoices, Transfer, ) -from kobo.apps.project_ownership.tests.utils import MockServiceUsageSerializer from kobo.apps.trackers.utils import update_nlp_counter from kpi.constants import PERM_VIEW_ASSET diff --git a/kobo/apps/project_ownership/tests/utils.py b/kobo/apps/project_ownership/tests/utils.py deleted file mode 100644 index fe46f81c9e..0000000000 --- a/kobo/apps/project_ownership/tests/utils.py +++ /dev/null @@ -1,35 +0,0 @@ -from django.db.models import F -from kpi.models.asset import Asset - - -class MockServiceUsageSerializer: - - def _get_storage_usage(self): - - assets = Asset.objects.annotate(user_id=F('owner_id')).filter( - self._user_id_query - ) - - self._total_storage_bytes = 0 - for asset in assets: - if asset.has_deployment: - for submission in asset.deployment.get_submissions(asset.owner): - self._total_storage_bytes += sum( - [att['bytes'] for att in submission['_attachments']] - ) - - def _get_submission_counters(self, month_filter, year_filter): - self._total_submission_count = { - 'all_time': 0, - 'current_year': 0, - 'current_month': 0, - } - assets = Asset.objects.annotate(user_id=F('owner_id')).filter( - self._user_id_query - ) - for asset in assets: - if asset.has_deployment: - submissions = asset.deployment.get_submissions(asset.owner) - self._total_submission_count['all_time'] += len(submissions) - self._total_submission_count['current_year'] += len(submissions) - self._total_submission_count['current_month'] += len(submissions) diff --git a/kobo/apps/stripe/tests/test_organization_usage.py b/kobo/apps/stripe/tests/test_organization_usage.py index 9ab4f2fa9b..9de13ff2a9 100644 --- a/kobo/apps/stripe/tests/test_organization_usage.py +++ b/kobo/apps/stripe/tests/test_organization_usage.py @@ -14,20 +14,20 @@ from kobo.apps.kobo_auth.shortcuts import User from kobo.apps.organizations.models import Organization, OrganizationUser -from kobo.apps.stripe.tests.utils import ( - generate_enterprise_subscription, - generate_plan_subscription, -) from kobo.apps.trackers.tests.submission_utils import ( create_mock_assets, add_mock_submissions, ) -from kpi.tests.api.v2.test_api_service_usage import ServiceUsageAPIBase +from kobo.apps.stripe.tests.utils import ( + generate_enterprise_subscription, + generate_plan_subscription, +) +from kpi.tests.test_usage_calculator import BaseServiceUsageTestCase from kpi.tests.api.v2.test_api_asset_usage import AssetUsageAPITestCase from rest_framework import status -class OrganizationServiceUsageAPIMultiUserTestCase(ServiceUsageAPIBase): +class OrganizationServiceUsageAPIMultiUserTestCase(BaseServiceUsageTestCase): """ Test organization service usage when Stripe is enabled. @@ -162,7 +162,7 @@ def test_endpoint_is_cached(self): ) -class OrganizationServiceUsageAPITestCase(ServiceUsageAPIBase): +class OrganizationServiceUsageAPITestCase(BaseServiceUsageTestCase): org_id = 'orgAKWMFskafsngf' @classmethod diff --git a/kpi/serializers/v2/service_usage.py b/kpi/serializers/v2/service_usage.py index 0d53560542..e82b160911 100644 --- a/kpi/serializers/v2/service_usage.py +++ b/kpi/serializers/v2/service_usage.py @@ -1,21 +1,14 @@ -from django.conf import settings -from django.db.models import Sum, Q -from django.db.models.functions import Coalesce -from django.utils import timezone from rest_framework import serializers from rest_framework.fields import empty -from kobo.apps.kobo_auth.shortcuts import User -from kobo.apps.openrosa.apps.logger.models import ( - DailyXFormSubmissionCounter, - XForm, -) from kobo.apps.organizations.models import Organization -from kobo.apps.organizations.utils import get_monthly_billing_dates, get_yearly_billing_dates -from kobo.apps.stripe.constants import ACTIVE_STRIPE_STATUSES -from kobo.apps.trackers.models import NLPUsageCounter +from kobo.apps.organizations.utils import ( + get_monthly_billing_dates, + get_yearly_billing_dates, +) from kpi.deployment_backends.openrosa_backend import OpenRosaDeploymentBackend from kpi.models.asset import Asset +from kpi.utils.usage_calculator import ServiceUsageCalculator class AssetUsageSerializer(serializers.HyperlinkedModelSerializer): @@ -117,158 +110,32 @@ class ServiceUsageSerializer(serializers.Serializer): def __init__(self, instance=None, data=empty, **kwargs): super().__init__(instance=instance, data=data, **kwargs) + organization = None + organization_id = self.context.get('organization_id', None) + if organization_id: + organization = Organization.objects.filter( + organization_users__user_id=instance.id, + id=organization_id, + ).first() + self.calculator = ServiceUsageCalculator(instance, organization) - self._total_nlp_usage = {} - self._total_storage_bytes = 0 - self._total_submission_count = {} - self._current_month_start = None - self._current_month_end = None - self._current_year_start = None - self._current_year_end = None - self._organization = None - self._now = timezone.now() - self._get_per_asset_usage(instance) - - def get_total_nlp_usage(self, user): - return self._total_nlp_usage - - def get_total_submission_count(self, user): - return self._total_submission_count - - def get_total_storage_bytes(self, user): - return self._total_storage_bytes - - def get_current_month_start(self, user): - return self._current_month_start.isoformat() - def get_current_month_end(self, user): - return self._current_month_end.isoformat() + return self.calculator.current_month_end.isoformat() - def get_current_year_start(self, user): - return self._current_year_start.isoformat() + def get_current_month_start(self, user): + return self.calculator.current_month_start.isoformat() def get_current_year_end(self, user): - return self._current_year_end.isoformat() - - def _filter_by_user(self, user_ids: list) -> Q: - """ - Turns a list of user ids into a query object to filter by - """ - return Q(user_id__in=user_ids) - - def _get_nlp_user_counters(self, month_filter, year_filter): - nlp_tracking = NLPUsageCounter.objects.only( - 'date', 'total_asr_seconds', 'total_mt_characters' - ).filter(self._user_id_query).aggregate( - asr_seconds_current_year=Coalesce( - Sum('total_asr_seconds', filter=year_filter), 0 - ), - mt_characters_current_year=Coalesce( - Sum('total_mt_characters', filter=year_filter), 0 - ), - asr_seconds_current_month=Coalesce( - Sum('total_asr_seconds', filter=month_filter), 0 - ), - mt_characters_current_month=Coalesce( - Sum('total_mt_characters', filter=month_filter), 0 - ), - asr_seconds_all_time=Coalesce(Sum('total_asr_seconds'), 0), - mt_characters_all_time=Coalesce(Sum('total_mt_characters'), 0), - ) - - for nlp_key, count in nlp_tracking.items(): - self._total_nlp_usage[nlp_key] = count if count is not None else 0 - - def _get_organization_details(self, user_id: int): - # Get the organization ID from the request - organization_id = self.context.get( - 'organization_id', None - ) - - if not organization_id: - return + return self.calculator.current_year_end.isoformat() - self._organization = Organization.objects.filter( - organization_users__user_id=user_id, - id=organization_id, - ).first() - - if not self._organization: - # Couldn't find organization, proceed as normal - return - - if settings.STRIPE_ENABLED: - # if the user is in an organization and has an enterprise plan, get all org users - # we evaluate this queryset instead of using it as a subquery because it's referencing - # fields from the auth_user tables on kpi *and* kobocat, making getting results in a - # single query not feasible until those tables are combined - user_ids = list( - User.objects.filter( - organizations_organization__id=organization_id, - organizations_organization__djstripe_customers__subscriptions__status__in=ACTIVE_STRIPE_STATUSES, - organizations_organization__djstripe_customers__subscriptions__items__price__product__metadata__has_key='plan_type', - organizations_organization__djstripe_customers__subscriptions__items__price__product__metadata__plan_type='enterprise', - ).values_list('pk', flat=True)[:settings.ORGANIZATION_USER_LIMIT] - ) - if user_ids: - self._user_id_query = self._filter_by_user(user_ids) - - def _get_per_asset_usage(self, user): - self._user_id = user.pk - self._user_id_query = self._filter_by_user([self._user_id]) - # get the billing data and list of organization users (if applicable) - self._get_organization_details(self._user_id) - - self._get_storage_usage() - - self._current_month_start, self._current_month_end = get_monthly_billing_dates(self._organization) - self._current_year_start, self._current_year_end = get_yearly_billing_dates(self._organization) - - current_month_filter = Q( - date__range=[self._current_month_start, self._now] - ) - current_year_filter = Q( - date__range=[self._current_year_start, self._now] - ) - - self._get_submission_counters(current_month_filter, current_year_filter) - self._get_nlp_user_counters(current_month_filter, current_year_filter) - - def _get_storage_usage(self): - """ - Get the storage used by non-(soft-)deleted projects for all users - - Users are represented by their ids with `self._user_ids` - """ - xforms = XForm.objects.only('attachment_storage_bytes', 'id').exclude( - pending_delete=True - ).filter(self._user_id_query) - - total_storage_bytes = xforms.aggregate( - bytes_sum=Coalesce(Sum('attachment_storage_bytes'), 0), - ) - - self._total_storage_bytes = total_storage_bytes['bytes_sum'] or 0 + def get_current_year_start(self, user): + return self.calculator.current_year_start.isoformat() - def _get_submission_counters(self, month_filter, year_filter): - """ - Calculate submissions for all users' projects even their deleted ones + def get_total_nlp_usage(self, user): + return self.calculator.get_nlp_usage_counters() - Users are represented by their ids with `self._user_ids` - """ - submission_count = DailyXFormSubmissionCounter.objects.only( - 'counter', 'user_id' - ).filter(self._user_id_query).aggregate( - all_time=Coalesce(Sum('counter'), 0), - current_year=Coalesce( - Sum('counter', filter=year_filter), 0 - ), - current_month=Coalesce( - Sum('counter', filter=month_filter), 0 - ), - ) + def get_total_submission_count(self, user): + return self.calculator.get_submission_counters() - for submission_key, count in submission_count.items(): - self._total_submission_count[submission_key] = ( - count if count is not None else 0 - ) + def get_total_storage_bytes(self, user): + return self.calculator.get_storage_usage() diff --git a/kpi/tests/api/v2/test_api_service_usage.py b/kpi/tests/api/v2/test_api_service_usage.py index 9d927d93c8..68af6b2bbb 100644 --- a/kpi/tests/api/v2/test_api_service_usage.py +++ b/kpi/tests/api/v2/test_api_service_usage.py @@ -1,159 +1,11 @@ -# coding: utf-8 -import os.path -import uuid - -from dateutil.relativedelta import relativedelta -from django.conf import settings from django.urls import reverse -from django.utils import timezone from rest_framework import status -from kobo.apps.kobo_auth.shortcuts import User -from kobo.apps.trackers.models import NLPUsageCounter from kpi.models import Asset -from kpi.tests.base_test_case import BaseAssetTestCase -from kpi.urls.router_api_v2 import URL_NAMESPACE as ROUTER_URL_NAMESPACE - - -class ServiceUsageAPIBase(BaseAssetTestCase): - """ - This class contains setup logic and utility functions to test submissions/usage - """ - fixtures = ['test_data'] - - URL_NAMESPACE = ROUTER_URL_NAMESPACE - - counter = None - attachment_id = 0 - - @classmethod - def setUpTestData(cls): - super().setUpTestData() - cls.anotheruser = User.objects.get(username='anotheruser') - cls.someuser = User.objects.get(username='someuser') - - def setUp(self): - super().setUp() - self.client.login(username='anotheruser', password='anotheruser') - - def _create_asset(self, user=None): - owner = user or self.anotheruser - content_source_asset = { - 'survey': [ - { - 'type': 'audio', - 'label': 'q1', - 'required': 'false', - '$kuid': 'abcd', - }, - { - 'type': 'file', - 'label': 'q2', - 'required': 'false', - '$kuid': 'efgh', - }, - ] - } - self.asset = Asset.objects.create( - content=content_source_asset, - owner=owner, - asset_type='survey', - ) - - self.asset.deploy(backend='mock', active=True) - self.asset.save() - - self.asset.deployment.set_namespace(self.URL_NAMESPACE) - self.submission_list_url = reverse( - self._get_endpoint('submission-list'), - kwargs={'format': 'json', 'parent_lookup_asset': self.asset.uid}, - ) - self._deployment = self.asset.deployment - - def add_nlp_trackers(self): - """ - Add nlp data to an asset - """ - # this month - today = timezone.now().date() - counter_1 = { - 'google_asr_seconds': 4586, - 'google_mt_characters': 5473, - } - NLPUsageCounter.objects.create( - user_id=self.anotheruser.id, - asset_id=self.asset.id, - date=today, - counters=counter_1, - total_asr_seconds=counter_1['google_asr_seconds'], - total_mt_characters=counter_1['google_mt_characters'], - ) - - # last month - last_month = today - relativedelta(months=1) - counter_2 = { - 'google_asr_seconds': 142, - 'google_mt_characters': 1253, - } - NLPUsageCounter.objects.create( - user_id=self.anotheruser.id, - asset_id=self.asset.id, - date=last_month, - counters=counter_2, - total_asr_seconds=counter_2['google_asr_seconds'], - total_mt_characters=counter_2['google_mt_characters'], - ) - - def add_submissions(self, count=2): - """ - Add one or more submissions to an asset (TWO by default) - """ - submissions = [] - v_uid = self.asset.latest_deployed_version.uid - - for x in range(count): - submission = { - '__version__': v_uid, - 'q1': 'audio_conversion_test_clip.3gp', - 'q2': 'audio_conversion_test_image.jpg', - '_uuid': str(uuid.uuid4()), - '_attachments': [ - { - 'id': self.attachment_id, - 'download_url': 'http://testserver/anotheruser/audio_conversion_test_clip.3gp', - 'filename': 'anotheruser/audio_conversion_test_clip.3gp', - 'mimetype': 'video/3gpp', - }, - { - 'id': self.attachment_id + 1, - 'download_url': 'http://testserver/anotheruser/audio_conversion_test_image.jpg', - 'filename': 'anotheruser/audio_conversion_test_image.jpg', - 'mimetype': 'image/jpeg', - }, - ], - '_submitted_by': 'anotheruser', - } - # increment the attachment ID for each attachment created - self.attachment_id = self.attachment_id + 2 - submissions.append(submission) - - self.asset.deployment.mock_submissions(submissions) - - @staticmethod - def expected_file_size(): - """ - Calculate the expected combined file size for the test audio clip and image - """ - return os.path.getsize( - settings.BASE_DIR - + '/kpi/fixtures/attachments/audio_conversion_test_clip.3gp' - ) + os.path.getsize( - settings.BASE_DIR - + '/kpi/fixtures/attachments/audio_conversion_test_image.jpg' - ) +from kpi.tests.test_usage_calculator import BaseServiceUsageTestCase -class ServiceUsageAPITestCase(ServiceUsageAPIBase): +class ServiceUsageAPITestCase(BaseServiceUsageTestCase): def test_anonymous_user(self): """ Test that the endpoint is forbidden to anonymous user diff --git a/kpi/tests/test_usage_calculator.py b/kpi/tests/test_usage_calculator.py new file mode 100644 index 0000000000..921ab99c96 --- /dev/null +++ b/kpi/tests/test_usage_calculator.py @@ -0,0 +1,214 @@ +import os.path +import uuid + +from dateutil.relativedelta import relativedelta +from django.conf import settings +from django.test import override_settings +from django.urls import reverse +from django.utils import timezone +from model_bakery import baker + +from kobo.apps.kobo_auth.shortcuts import User +from kobo.apps.openrosa.apps.logger.models import ( + XForm, + DailyXFormSubmissionCounter, +) +from kobo.apps.organizations.models import Organization +from kobo.apps.stripe.tests.utils import generate_enterprise_subscription +from kobo.apps.trackers.models import NLPUsageCounter +from kpi.models import Asset +from kpi.tests.base_test_case import BaseAssetTestCase +from kpi.utils.usage_calculator import ServiceUsageCalculator +from kpi.urls.router_api_v2 import URL_NAMESPACE as ROUTER_URL_NAMESPACE + + +class BaseServiceUsageTestCase(BaseAssetTestCase): + """ + This class contains setup logic and utility functions to test usage + calculations + """ + fixtures = ['test_data'] + + URL_NAMESPACE = ROUTER_URL_NAMESPACE + + def setUp(self): + super().setUp() + self.client.login(username='anotheruser', password='anotheruser') + + @classmethod + def setUpTestData(cls): + super().setUpTestData() + cls.anotheruser = User.objects.get(username='anotheruser') + cls.someuser = User.objects.get(username='someuser') + + def _create_asset(self, user=None): + owner = user or self.anotheruser + content_source_asset = { + 'survey': [ + { + 'type': 'audio', + 'label': 'q1', + 'required': 'false', + '$kuid': 'abcd', + }, + { + 'type': 'file', + 'label': 'q2', + 'required': 'false', + '$kuid': 'efgh', + }, + ] + } + self.asset = Asset.objects.create( + content=content_source_asset, + owner=owner, + asset_type='survey', + ) + + self.asset.deploy(backend='mock', active=True) + self.asset.save() + + self.asset.deployment.set_namespace(self.URL_NAMESPACE) + self.submission_list_url = reverse( + self._get_endpoint('submission-list'), + kwargs={'format': 'json', 'parent_lookup_asset': self.asset.uid}, + ) + self._deployment = self.asset.deployment + + + def add_nlp_trackers(self): + """ + Add nlp data to an asset + """ + # this month + today = timezone.now().date() + counter_1 = { + 'google_asr_seconds': 4586, + 'google_mt_characters': 5473, + } + NLPUsageCounter.objects.create( + user_id=self.anotheruser.id, + asset_id=self.asset.id, + date=today, + counters=counter_1, + total_asr_seconds=counter_1['google_asr_seconds'], + total_mt_characters=counter_1['google_mt_characters'], + ) + + # last month + last_month = today - relativedelta(months=1) + counter_2 = { + 'google_asr_seconds': 142, + 'google_mt_characters': 1253, + } + NLPUsageCounter.objects.create( + user_id=self.anotheruser.id, + asset_id=self.asset.id, + date=last_month, + counters=counter_2, + total_asr_seconds=counter_2['google_asr_seconds'], + total_mt_characters=counter_2['google_mt_characters'], + ) + + def add_submissions(self, count=2): + """ + Add one or more submissions to an asset (TWO by default) + """ + submissions = [] + v_uid = self.asset.latest_deployed_version.uid + + for x in range(count): + submission = { + '__version__': v_uid, + 'q1': 'audio_conversion_test_clip.3gp', + 'q2': 'audio_conversion_test_image.jpg', + '_uuid': str(uuid.uuid4()), + '_attachments': [ + { + 'download_url': 'http://testserver/anotheruser/audio_conversion_test_clip.3gp', + 'filename': 'anotheruser/audio_conversion_test_clip.3gp', + 'mimetype': 'video/3gpp', + }, + { + 'download_url': 'http://testserver/anotheruser/audio_conversion_test_image.jpg', + 'filename': 'anotheruser/audio_conversion_test_image.jpg', + 'mimetype': 'image/jpeg', + }, + ], + '_submitted_by': 'anotheruser', + } + # increment the attachment ID for each attachment created + submissions.append(submission) + + self.asset.deployment.mock_submissions(submissions) + + def expected_file_size(self): + """ + Calculate the expected combined file size for the test audio clip and image + """ + return os.path.getsize( + settings.BASE_DIR + + '/kpi/fixtures/attachments/audio_conversion_test_clip.3gp' + ) + os.path.getsize( + settings.BASE_DIR + + '/kpi/fixtures/attachments/audio_conversion_test_image.jpg' + ) + + +class ServiceUsageCalculatorTestCase(BaseServiceUsageTestCase): + def setUp(self): + super().setUp() + self._create_asset() + self.add_nlp_trackers() + self.add_submissions(count=5) + + def test_nlp_usage_counters(self): + calculator = ServiceUsageCalculator(self.anotheruser, None) + nlp_usage = calculator.get_nlp_usage_counters() + assert nlp_usage['asr_seconds_current_month'] == 4586 + assert nlp_usage['asr_seconds_all_time'] == 4728 + 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() + submission_counters = calculator.get_submission_counters() + + assert nlp_usage['asr_seconds_current_month'] == 0 + assert nlp_usage['asr_seconds_all_time'] == 0 + assert nlp_usage['mt_characters_current_month'] == 0 + assert nlp_usage['mt_characters_all_time'] == 0 + assert calculator.get_storage_usage() == 0 + assert submission_counters['current_month'] == 0 + assert submission_counters['all_time'] == 0 + + @override_settings(STRIPE_ENABLED=True) + def test_organization_setup(self): + organization = baker.make(Organization, id='org_abcd1234') + organization.add_user(user=self.anotheruser, is_admin=True) + organization.add_user(user=self.someuser, is_admin=True) + generate_enterprise_subscription(organization) + + calculator = ServiceUsageCalculator(self.someuser, organization) + submission_counters = calculator.get_submission_counters() + assert submission_counters['current_month'] == 5 + assert submission_counters['all_time'] == 5 + + nlp_usage = calculator.get_nlp_usage_counters() + assert nlp_usage['asr_seconds_current_month'] == 4586 + assert nlp_usage['asr_seconds_all_time'] == 4728 + assert nlp_usage['mt_characters_current_month'] == 5473 + assert nlp_usage['mt_characters_all_time'] == 6726 + + assert calculator.get_storage_usage() == 5 * self.expected_file_size() diff --git a/kpi/utils/usage_calculator.py b/kpi/utils/usage_calculator.py new file mode 100644 index 0000000000..71d3d24175 --- /dev/null +++ b/kpi/utils/usage_calculator.py @@ -0,0 +1,141 @@ +from typing import Optional + +from django.conf import settings +from django.db.models import Sum, Q +from django.db.models.functions import Coalesce +from django.utils import timezone + +from kobo.apps.kobo_auth.shortcuts import User +from kobo.apps.openrosa.apps.logger.models import ( + DailyXFormSubmissionCounter, + XForm, +) +from kobo.apps.organizations.models import Organization +from kobo.apps.organizations.utils import ( + get_monthly_billing_dates, + get_yearly_billing_dates, +) +from kobo.apps.stripe.constants import ACTIVE_STRIPE_STATUSES +from kobo.apps.trackers.models import NLPUsageCounter + + +class ServiceUsageCalculator: + def __init__(self, user: User, organization: Optional[Organization]): + self.user = user + self.organization = organization + + self._user_ids = [user.pk] + self._user_id_query = self._filter_by_user([user.pk]) + if organization and settings.STRIPE_ENABLED: + # if the user is in an organization and has an enterprise plan, get all org users + # we evaluate this queryset instead of using it as a subquery because it's referencing + # fields from the auth_user tables on kpi *and* kobocat, making getting results in a + # single query not feasible until those tables are combined + user_ids = list( + User.objects.filter( + organizations_organization__id=organization.id, + organizations_organization__djstripe_customers__subscriptions__status__in=ACTIVE_STRIPE_STATUSES, + organizations_organization__djstripe_customers__subscriptions__items__price__product__metadata__has_key='plan_type', + organizations_organization__djstripe_customers__subscriptions__items__price__product__metadata__plan_type='enterprise', + ).values_list('pk', flat=True)[ + : settings.ORGANIZATION_USER_LIMIT + ] + ) + if user_ids: + self._user_ids = user_ids + self._user_id_query = self._filter_by_user(user_ids) + + now = timezone.now() + self.current_month_start, self.current_month_end = ( + get_monthly_billing_dates(organization) + ) + self.current_year_start, self.current_year_end = ( + get_yearly_billing_dates(organization) + ) + self.current_month_filter = Q( + date__range=[self.current_month_start, now] + ) + self.current_year_filter = Q(date__range=[self.current_year_start, now]) + + def _filter_by_user(self, user_ids: list) -> Q: + """ + Turns a list of user ids into a query object to filter by + """ + return Q(user_id__in=user_ids) + + def get_nlp_usage_counters(self): + nlp_tracking = ( + NLPUsageCounter.objects.only( + 'date', 'total_asr_seconds', 'total_mt_characters' + ) + .filter(self._user_id_query) + .aggregate( + asr_seconds_current_year=Coalesce( + Sum('total_asr_seconds', filter=self.current_year_filter), 0 + ), + mt_characters_current_year=Coalesce( + Sum('total_mt_characters', filter=self.current_year_filter), + 0, + ), + asr_seconds_current_month=Coalesce( + Sum('total_asr_seconds', filter=self.current_month_filter), + 0, + ), + mt_characters_current_month=Coalesce( + Sum( + 'total_mt_characters', filter=self.current_month_filter + ), + 0, + ), + asr_seconds_all_time=Coalesce(Sum('total_asr_seconds'), 0), + mt_characters_all_time=Coalesce(Sum('total_mt_characters'), 0), + ) + ) + + total_nlp_usage = {} + for nlp_key, count in nlp_tracking.items(): + total_nlp_usage[nlp_key] = count if count is not None else 0 + + return total_nlp_usage + + def get_storage_usage(self): + """ + Get the storage used by non-(soft-)deleted projects for all users + + Users are represented by their ids with `self._user_ids` + """ + xforms = XForm.objects.only('attachment_storage_bytes', 'id').exclude( + pending_delete=True + ).filter(self._user_id_query) + + total_storage_bytes = xforms.aggregate( + bytes_sum=Coalesce(Sum('attachment_storage_bytes'), 0), + ) + + return total_storage_bytes['bytes_sum'] or 0 + + def get_submission_counters(self): + """ + Calculate submissions for all users' projects even their deleted ones + + Users are represented by their ids with `self._user_ids` + """ + submission_count = DailyXFormSubmissionCounter.objects.only( + 'counter', 'user_id' + ).filter(self._user_id_query).aggregate( + all_time=Coalesce(Sum('counter'), 0), + current_year=Coalesce( + Sum('counter', filter=self.current_year_filter), 0 + ), + current_month=Coalesce( + Sum('counter', filter=self.current_month_filter), 0 + ), + ) + + total_submission_count = {} + for submission_key, count in submission_count.items(): + total_submission_count[submission_key] = ( + count if count is not None else 0 + ) + + return total_submission_count