daybook/artifacts/storage.py

"""Disk-usage accounting, per-user quota, and automatic eviction.

An artifact's on-disk footprint is the primary stored file plus, for kinds
that preserve their original source (see `SIDECAR_SUFFIXES`), the sidecar
file. We record it on the row (`size_bytes`) at ingest so usage is a cheap
SUM rather than a stat of every file.
"""

import shutil
from datetime import timedelta

from django.conf import settings
from django.db.models import Sum
from django.utils import timezone

from .models import Artifact, Asset, UserProfile


class QuotaExceeded(Exception):
    """Raised when a post cannot fit even after evicting every unkept artifact."""


# Kinds that preserve their original source alongside the rendered .html file.
SIDECAR_SUFFIXES = {
    Artifact.MARKDOWN: ".md",
    Artifact.JSON: ".json",
    Artifact.CSV: ".csv",
}


def artifact_disk_size(relative_path, kind):
    """Bytes an artifact occupies on disk, including any sidecar file."""
    root = settings.ARTIFACTS_ROOT
    total = 0
    primary = root / relative_path
    if primary.exists():
        total += primary.stat().st_size
    sidecar_suffix = SIDECAR_SUFFIXES.get(kind)
    if sidecar_suffix:
        sidecar = primary.with_suffix(sidecar_suffix)
        if sidecar.exists():
            total += sidecar.stat().st_size
    return total


def delete_artifact_files(artifact):
    """Remove an artifact's files from disk (primary, sidecar, assets).

    Assets are cascade-deleted from the DB by the FK, but their files must be
    removed here explicitly.
    """
    root = settings.ARTIFACTS_ROOT
    primary = root / artifact.path
    primary.unlink(missing_ok=True)
    sidecar_suffix = SIDECAR_SUFFIXES.get(artifact.kind)
    if sidecar_suffix:
        primary.with_suffix(sidecar_suffix).unlink(missing_ok=True)
    for asset in artifact.assets.all():
        delete_asset_file(asset)


def delete_asset_file(asset):
    (settings.ARTIFACTS_ROOT / asset.path).unlink(missing_ok=True)


def artifact_total_bytes(artifact):
    """An artifact's own bytes plus those of the assets embedded in it."""
    asset_bytes = artifact.assets.aggregate(total=Sum("size_bytes")).get("total") or 0
    return artifact.size_bytes + asset_bytes


def user_usage_bytes(user):
    artifact_bytes = (
        Artifact.objects.filter(owner=user).aggregate(total=Sum("size_bytes")).get("total") or 0
    )
    asset_bytes = (
        Asset.objects.filter(owner=user).aggregate(total=Sum("size_bytes")).get("total") or 0
    )
    return artifact_bytes + asset_bytes


def user_quota_bytes(user):
    profile, _ = UserProfile.objects.get_or_create(user=user)
    return profile.quota_bytes


def make_room_for(user, incoming_bytes):
    """Evict oldest unkept artifacts until `incoming_bytes` fits under quota.

    Returns the ids of evicted artifacts (already deleted - Django clears an
    instance's pk on .delete(), so ids are captured beforehand). Raises
    QuotaExceeded, without deleting anything, if even evicting every unkept
    artifact wouldn't free enough room.
    """
    quota = user_quota_bytes(user)
    if incoming_bytes > quota:
        raise QuotaExceeded("Artifact is larger than the entire quota.")

    used = user_usage_bytes(user)
    if used + incoming_bytes <= quota:
        return []

    unkept = list(
        Artifact.objects.filter(owner=user, is_kept=False).order_by("created_at")
    )
    # Check whether eviction can possibly succeed before deleting anything -
    # otherwise a doomed request would still destroy every unkept artifact on
    # its way to failing anyway.
    evictable_total = sum(artifact_total_bytes(artifact) for artifact in unkept)
    if used - evictable_total + incoming_bytes > quota:
        raise QuotaExceeded("Quota is full of kept artifacts; free space to post.")

    evicted = []
    for artifact in unkept:
        if used + incoming_bytes <= quota:
            break
        # Freeing an artifact also frees the assets embedded in it.
        used -= artifact_total_bytes(artifact)
        evicted_id = artifact.id
        delete_artifact_files(artifact)
        artifact.delete()
        evicted.append(evicted_id)

    return evicted


ASSET_ORPHAN_GRACE = timedelta(hours=24)


def sweep_orphan_assets(user=None):
    """Delete assets uploaded but never referenced by an artifact.

    Runs opportunistically on upload. The grace period gives an agent time to
    upload images and then post the artifact that references them.
    """
    cutoff = timezone.now() - ASSET_ORPHAN_GRACE
    orphans = Asset.objects.filter(artifact__isnull=True, created_at__lt=cutoff)
    if user is not None:
        orphans = orphans.filter(owner=user)
    count = 0
    for asset in orphans:
        delete_asset_file(asset)
        asset.delete()
        count += 1
    return count


def disk_free_bytes():
    """Free space on the volume that holds the artifact store."""
    return shutil.disk_usage(settings.ARTIFACTS_ROOT).free


def total_allocated_bytes():
    return UserProfile.objects.aggregate(total=Sum("quota_bytes")).get("total") or 0