Skip to content

Commit

Permalink
perf(organization): reduce DB queries even further for owner label in…
Browse files Browse the repository at this point in the history
… asset list (#5324)

### 📣 Summary
Reduced query count to minimize load and enhance scalability.
  • Loading branch information
noliveleger authored Dec 4, 2024
1 parent 2ae607e commit c595290
Show file tree
Hide file tree
Showing 3 changed files with 64 additions and 24 deletions.
9 changes: 7 additions & 2 deletions kobo/apps/organizations/views.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,10 +57,14 @@ def get_queryset(self, *args, **kwargs):
# validated within `OrganizationViewSet.assets()`.
raise AttributeError('`permissions_checked` is missing')

organization = getattr(
self.request, 'organization', self.request.user.organization
)

queryset = super().get_queryset(*args, **kwargs)
if self.action == 'list':
return queryset.filter(
owner=self.request.user.organization.owner_user_object
owner=organization.owner_user_object
)
else:
raise NotImplementedError
Expand Down Expand Up @@ -97,12 +101,13 @@ def assets(self, request: Request, *args, **kwargs):
### Additional Information
For more details, please refer to `/api/v2/assets/`.
"""
self.get_object() # Call check permissions
organization = self.get_object() # Call check permissions

# Permissions check is done by `OrganizationAssetViewSet` permission classes
asset_view = OrganizationAssetViewSet.as_view({'get': 'list'})
django_http_request = request._request
django_http_request.permissions_checked = True
django_http_request.organization = organization
return asset_view(request=django_http_request)

def get_queryset(self) -> QuerySet:
Expand Down
36 changes: 25 additions & 11 deletions kpi/serializers/v2/asset.py
Original file line number Diff line number Diff line change
Expand Up @@ -745,14 +745,14 @@ def get_export_settings(self, obj: Asset) -> ReturnList:
context=self.context,
).data

def get_access_types(self, obj):
def get_access_types(self, asset):
"""
Handles the detail endpoint but also takes advantage of the
`AssetViewSet.get_serializer_context()` "cache" for the list endpoint,
if it is present
"""
# Avoid extra queries if obj is not a collection
if obj.asset_type != ASSET_TYPE_COLLECTION:
if asset.asset_type != ASSET_TYPE_COLLECTION:
return None

# User is the owner
Expand All @@ -762,17 +762,17 @@ def get_access_types(self, obj):
return None

access_types = []
if request.user == obj.owner:
if request.user == asset.owner:
access_types.append('owned')

# User can view the collection.
try:
# The list view should provide a cache
asset_permission_assignments = self.context[
'object_permissions_per_asset'
].get(obj.pk)
].get(asset.pk)
except KeyError:
asset_permission_assignments = obj.permissions.all()
asset_permission_assignments = asset.permissions.all()

# We test at the same time whether the collection is public or not
for obj_permission in asset_permission_assignments:
Expand All @@ -784,13 +784,13 @@ def get_access_types(self, obj):
):
access_types.append('public')

if request.user == obj.owner:
if request.user == asset.owner:
# Do not go further, `access_type` cannot be `shared`
# and `owned`
break

if (
request.user != obj.owner
request.user != asset.owner
and not obj_permission.deny
and obj_permission.user == request.user
):
Expand All @@ -805,10 +805,10 @@ def get_access_types(self, obj):
try:
# The list view should provide a cache
subscriptions = self.context['user_subscriptions_per_asset'].get(
obj.pk, []
asset.pk, []
)
except KeyError:
subscribed = obj.has_subscribed_user(request.user.pk)
subscribed = asset.has_subscribed_user(request.user.pk)
else:
subscribed = request.user.pk in subscriptions
if subscribed:
Expand All @@ -818,7 +818,16 @@ def get_access_types(self, obj):
if request.user.is_superuser:
access_types.append('superuser')

if obj.owner.organization.get_user_role(request.user) == ORG_ADMIN_ROLE:
try:
organization = self.context['organization_by_asset'].get(asset.id)
except KeyError:
# Fallback on context if it exists (i.e.: asset lists of an organization).
# Otherwise, retrieve from the asset owner.
organization = self.context.get(
'organization', asset.owner.organization
)

if organization.get_user_role(request.user) == ORG_ADMIN_ROLE:
access_types.extend(['shared', 'org-admin'])
access_types = list(set(access_types))

Expand All @@ -833,7 +842,12 @@ def get_owner_label(self, asset):
try:
organization = self.context['organization_by_asset'].get(asset.id)
except KeyError:
organization = asset.owner.organization
# Fallback on context if it exists (i.e.: asset lists of an organization).
# Otherwise, retrieve from the asset owner.
organization = self.context.get(
'organization', asset.owner.organization
)

if (
organization
and organization.is_owner(asset.owner)
Expand Down
43 changes: 32 additions & 11 deletions kpi/views/v2/asset.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
from collections import OrderedDict, defaultdict
from operator import itemgetter

from django.db.models import Count
from django.db.models import Count, Prefetch
from django.http import Http404
from django.shortcuts import get_object_or_404
from rest_framework import exceptions, renderers, status
Expand All @@ -13,6 +13,7 @@

from kobo.apps.audit_log.base_views import AuditLoggedModelViewSet
from kobo.apps.audit_log.models import AuditType
from kobo.apps.organizations.models import Organization
from kpi.constants import (
ASSET_TYPE_ARG_NAME,
ASSET_TYPE_SURVEY,
Expand Down Expand Up @@ -722,16 +723,36 @@ def get_serializer_context(self):

context_['children_count_per_asset'] = children_count_per_asset

# 5) Get organization per asset
assets = (
Asset.objects.filter(id__in=asset_ids)
.prefetch_related('owner__organizations_organization')
)
organization_by_asset = defaultdict(dict)
for asset in assets:
organization = getattr(asset.owner, 'organization', None)
organization_by_asset[asset.id] = organization
context_['organization_by_asset'] = organization_by_asset
# 5) Get organization…
if organization := getattr(self.request, 'organization', None):
# …from request.
# e.g.: /api/v2/organizations/<organization_id>/assets/`
context_['organization'] = organization
else:
# …per asset
# e.g.: /api/v2/organizations/assets/`
assets = (
Asset.objects.only('owner', 'uid', 'name')
.filter(id__in=asset_ids)
.select_related('owner')
.prefetch_related(
Prefetch(
'owner__organizations_organization',
queryset=Organization.objects.all().order_by(
'-organization_users__created'
),
to_attr='organizations',
)
)
)
organization_by_asset = defaultdict(dict)
for asset in assets:
try:
organization_by_asset[asset.id] = asset.owner.organizations[0]
except IndexError:
pass

context_['organization_by_asset'] = organization_by_asset

return context_

Expand Down

0 comments on commit c595290

Please sign in to comment.