extensions-website/stats/models.py

from django.db import models, transaction
from django.http import HttpRequest

from utils import clean_ip_address
import extensions.models


class _FromRequestMixin:
    @classmethod
    def create_from_request(cls, request: HttpRequest, object_id: int):
        """Create a new record for the given Extension ID based on the given request."""
        if object_id is None:
            return
        ip_address = clean_ip_address(request) if request.user.is_anonymous else None
        user_id = request.user.pk if request.user.is_authenticated else None
        cls.objects.bulk_create(
            [
                cls(
                    ip_address=ip_address,
                    user_id=user_id,
                    **{cls.target_object_id_field: object_id},
                )
            ],
            ignore_conflicts=True,
        )

    @classmethod
    @transaction.atomic
    def update_counters(cls, to_field: str):
        last_seen = cls.counter_index.objects.filter(field=to_field).first()
        last_seen_id = last_seen.last_seen_id if last_seen else 0
        object_id_count = (
            cls.objects.filter(id__gt=last_seen_id)
            .values(cls.target_object_id_field)
            .annotate(count=models.Count(cls.target_object_id_field))
        )
        if not object_id_count:  # nothing to to
            return

        all_object_ids = {_[cls.target_object_id_field] for _ in object_id_count}
        affected_objects = {_.pk: _ for _ in cls.target_model.objects.filter(id__in=all_object_ids)}
        to_update = []
        for row in object_id_count:
            object_id = row[cls.target_object_id_field]
            count = row['count']
            obj = affected_objects[object_id]
            setattr(obj, to_field, getattr(obj, to_field) + count)
            to_update.append(obj)
        cls.target_model.objects.bulk_update(
            to_update, batch_size=500, fields={to_field, 'date_modified'}
        )

        # Update the last seen ID to make sure next count starts from it
        max_id = cls.objects.aggregate(max_id=models.Max('pk')).get('max_id')
        cls.counter_index.objects.update_or_create(
            field=to_field, defaults={'last_seen_id': max_id}
        )

    def __str__(self) -> str:
        ip_address_f = f', IP {self.ip_address}' if self.ip_address else ''
        user_id_f = f', user ID {self.user_id}' if self.user_id else ''
        return (
            f'{self.__class__.__name__}'
            f' #{getattr(self, self.target_object_id_field)}{ip_address_f}{user_id_f}'
        )


class _ExtensionStatMixin(_FromRequestMixin, models.Model):
    class Meta:
        indexes = [
            models.Index(fields=('extension_id', 'user_id', 'ip_address')),
        ]
        constraints = [
            models.UniqueConstraint(
                fields=('extension_id', 'ip_address'),
                name='%(app_label)s_%(class)s_ip_address_key',
                # Unique constraint can only be enforced on non-null values
                condition=models.Q(ip_address__isnull=False),
            ),
            models.UniqueConstraint(
                fields=('extension_id', 'user_id'),
                name='%(app_label)s_%(class)s_user_id_key',
                # Unique constraint can only be enforced on non-null values
                condition=models.Q(user_id__isnull=False),
            ),
        ]
        abstract = True

    # While IP address certainly doesn't represent a unique visitor,
    # it is a good enough compromise between being able to count unique anonymous visits and
    # having to sacrifice a lot of storage space for it.
    ip_address = models.GenericIPAddressField(protocol='both', null=True)
    user_id = models.PositiveIntegerField(null=True)
    extension = models.ForeignKey('extensions.Extension', null=False, on_delete=models.CASCADE)


class ExtensionCountedStat(models.Model):
    """Store last counted ID of unique visits/downloads."""

    class _Field(models.TextChoices):
        view_count = 'view_count'
        download_count = 'download_count'

    field = models.CharField(
        null=False, blank=False, max_length=20, choices=_Field.choices, primary_key=True
    )
    last_seen_id = models.PositiveIntegerField(null=False, blank=False)


class VersionCountedStat(models.Model):
    """Store last counted ID of unique visits/downloads."""

    class _Field(models.TextChoices):
        download_count = 'download_count'

    field = models.CharField(
        null=False, blank=False, max_length=20, choices=_Field.choices, primary_key=True
    )
    last_seen_id = models.PositiveIntegerField(null=False, blank=False)


class ExtensionView(_ExtensionStatMixin, models.Model):
    counter_index = ExtensionCountedStat
    target_model = extensions.models.Extension
    target_object_id_field = 'extension_id'


class ExtensionDownload(_ExtensionStatMixin, models.Model):
    counter_index = ExtensionCountedStat
    target_model = extensions.models.Extension
    target_object_id_field = 'extension_id'


class VersionDownload(_FromRequestMixin, models.Model):
    counter_index = VersionCountedStat
    target_model = extensions.models.Version
    target_object_id_field = 'version_id'

    class Meta:
        indexes = [
            models.Index(fields=('version_id', 'user_id', 'ip_address')),
        ]
        constraints = [
            models.UniqueConstraint(
                fields=('version_id', 'ip_address'),
                name='%(app_label)s_%(class)s_ip_address_key',
                # Unique constraint can only be enforced on non-null values
                condition=models.Q(ip_address__isnull=False),
            ),
            models.UniqueConstraint(
                fields=('version_id', 'user_id'),
                name='%(app_label)s_%(class)s_user_id_key',
                # Unique constraint can only be enforced on non-null values
                condition=models.Q(user_id__isnull=False),
            ),
        ]

    # While IP address certainly doesn't represent a unique visitor,
    # it is a good enough compromise between being able to count unique anonymous visits and
    # having to sacrifice a lot of storage space for it.
    ip_address = models.GenericIPAddressField(protocol='both', null=True)
    user_id = models.PositiveIntegerField(null=True)
    version = models.ForeignKey('extensions.Version', null=False, on_delete=models.CASCADE)