Stoken annotation: move it all to one place to reduce duplication.

This commit is contained in:
Tom Hacohen 2020-12-04 19:00:14 +02:00
parent 4ce96e043e
commit 2d0bcbdc20
2 changed files with 17 additions and 13 deletions

View File

@ -31,6 +31,11 @@ from .exceptions import EtebaseValidationError
UidValidator = RegexValidator(regex=r"^[a-zA-Z0-9\-_]{20,}$", message="Not a valid UID")
def stoken_annotation_builder(stoken_id_fields):
aggr_fields = [Coalesce(Max(field), V(0)) for field in stoken_id_fields]
return Greatest(*aggr_fields) if len(aggr_fields) > 1 else aggr_fields[0]
class CollectionType(models.Model):
owner = models.ForeignKey(settings.AUTH_USER_MODEL, on_delete=models.CASCADE)
uid = models.BinaryField(editable=True, blank=False, null=False, db_index=True, unique=True)
@ -40,7 +45,7 @@ class Collection(models.Model):
main_item = models.OneToOneField("CollectionItem", related_name="parent", null=True, on_delete=models.SET_NULL)
owner = models.ForeignKey(settings.AUTH_USER_MODEL, on_delete=models.CASCADE)
stoken_id_fields = ["items__revisions__stoken", "members__stoken"]
stoken_annotation = stoken_annotation_builder(["items__revisions__stoken", "members__stoken"])
def __str__(self):
return self.uid
@ -59,11 +64,9 @@ class Collection(models.Model):
@cached_property
def stoken(self):
aggr_fields = [Coalesce(Max(field), V(0)) for field in self.stoken_id_fields]
max_stoken = Greatest(*aggr_fields) if len(aggr_fields) > 1 else aggr_fields[0]
stoken_id = (
self.__class__.objects.filter(main_item=self.main_item)
.annotate(max_stoken=max_stoken)
.annotate(max_stoken=self.stoken_annotation)
.values("max_stoken")
.first()["max_stoken"]
)
@ -94,6 +97,8 @@ class CollectionItem(models.Model):
version = models.PositiveSmallIntegerField()
encryptionKey = models.BinaryField(editable=True, blank=False, null=True)
stoken_annotation = stoken_annotation_builder(["revisions__stoken"])
class Meta:
unique_together = ("uid", "collection")
@ -191,6 +196,8 @@ class CollectionMember(models.Model):
default=AccessLevels.READ_ONLY,
)
stoken_annotation = stoken_annotation_builder(["stoken"])
class Meta:
unique_together = ("user", "collection")

View File

@ -18,8 +18,7 @@ from django.conf import settings
from django.contrib.auth import get_user_model, user_logged_in, user_logged_out
from django.core.exceptions import PermissionDenied
from django.db import transaction, IntegrityError
from django.db.models import Max, Value as V, Q
from django.db.models.functions import Coalesce, Greatest
from django.db.models import Q
from django.http import HttpResponseBadRequest, HttpResponse, Http404
from django.shortcuts import get_object_or_404
@ -94,7 +93,7 @@ class BaseViewSet(viewsets.ModelViewSet):
permission_classes = tuple(app_settings.API_PERMISSIONS)
renderer_classes = [JSONRenderer, MessagePackRenderer] + ([BrowsableAPIRenderer] if settings.DEBUG else [])
parser_classes = [JSONParser, MessagePackParser, FormParser, MultiPartParser]
stoken_id_fields = None
stoken_annotation = None
def get_serializer_class(self):
serializer_class = self.serializer_class
@ -125,9 +124,7 @@ class BaseViewSet(viewsets.ModelViewSet):
def filter_by_stoken(self, request, queryset):
stoken_rev = self.get_stoken_obj(request)
aggr_fields = [Coalesce(Max(field), V(0)) for field in self.stoken_id_fields]
max_stoken = Greatest(*aggr_fields) if len(aggr_fields) > 1 else aggr_fields[0]
queryset = queryset.annotate(max_stoken=max_stoken).order_by("max_stoken")
queryset = queryset.annotate(max_stoken=self.stoken_annotation).order_by("max_stoken")
if stoken_rev is not None:
queryset = queryset.filter(max_stoken__gt=stoken_rev.id)
@ -179,7 +176,7 @@ class CollectionViewSet(BaseViewSet):
serializer_class = CollectionSerializer
lookup_field = "main_item__uid"
lookup_url_kwarg = "uid"
stoken_id_fields = ["items__revisions__stoken__id", "members__stoken__id"]
stoken_annotation = Collection.stoken_annotation
def get_queryset(self, queryset=None):
if queryset is None:
@ -262,7 +259,7 @@ class CollectionItemViewSet(BaseViewSet):
queryset = CollectionItem.objects.all()
serializer_class = CollectionItemSerializer
lookup_field = "uid"
stoken_id_fields = ["revisions__stoken__id"]
stoken_annotation = CollectionItem.stoken_annotation
def get_queryset(self):
collection_uid = self.kwargs["collection_uid"]
@ -482,7 +479,7 @@ class CollectionMemberViewSet(BaseViewSet):
serializer_class = CollectionMemberSerializer
lookup_field = f"user__{User.USERNAME_FIELD}__iexact"
lookup_url_kwarg = "username"
stoken_id_fields = ["stoken__id"]
stoken_annotation = CollectionMember.stoken_annotation
# FIXME: need to make sure that there's always an admin, and maybe also don't let an owner remove adm access
# (if we want to transfer, we need to do that specifically)