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 markdown,
the preserved `.md` sidecar. 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."""


def artifact_disk_size(relative_path, kind):
    """Bytes an artifact occupies on disk, including the markdown sidecar."""
    root = settings.ARTIFACTS_ROOT
    total = 0
    primary = root / relative_path
    if primary.exists():
        total += primary.stat().st_size
    if kind == Artifact.MARKDOWN:
        sidecar = primary.with_suffix(".md")
        if sidecar.exists():
            total += sidecar.stat().st_size
    return total


def delete_artifact_files(artifact):
    """Remove an artifact's files from disk (primary, markdown 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)
    if artifact.kind == Artifact.MARKDOWN:
        primary.with_suffix(".md").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 list of evicted artifacts (already deleted). Raises QuotaExceeded
    if even removing every unkept artifact leaves no 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 []

    evicted = []
    unkept = (
        Artifact.objects.filter(owner=user, is_kept=False).order_by("created_at")
    )
    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)
        delete_artifact_files(artifact)
        artifact.delete()
        evicted.append(artifact)

    if used + incoming_bytes > quota:
        raise QuotaExceeded("Quota is full of kept artifacts; free space to post.")
    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