pigway-frame-flow/pigway_frame_flow/extensions/storage.py

"""Opt-in local retention, protecting stars, current content and durable references."""

import re
import time
from datetime import UTC, datetime, timedelta

from pydantic import Field

from pigway_frame_flow.models import StrictModel
from pigway_frame_flow.storage.database import encode, now_iso


class StoragePolicy(StrictModel):
    enabled: bool = False
    generated_days: int = Field(default=7, ge=1, le=3650)
    cache_days: int = Field(default=2, ge=1, le=3650)
    max_bytes: int = Field(default=1024 * 1024 * 1024, ge=1024 * 1024, le=1024**4)


class LocalRetention:
    def __init__(self, service):
        self.service = service
        self.last = 0.0

    def config(self):
        row = self.service.db.one(
            "SELECT payload FROM app_config WHERE key='local_retention'"
        )
        return (
            StoragePolicy.model_validate_json(row["payload"])
            if row
            else StoragePolicy()
        )

    def protected(self):
        db = self.service.db
        ids = {
            row["asset_id"]
            for row in db.all("SELECT asset_id FROM asset_metadata WHERE starred=1")
        }
        payloads = []
        for query in [
            "SELECT payload FROM jobs",
            "SELECT payload FROM template_config",
            "SELECT payload FROM operations WHERE status IN ('QUEUED','RENDERING','UPLOADING','SYNCING','UNKNOWN')",
        ]:
            payloads += [row["payload"] for row in db.all(query)]
        for payload in payloads:
            ids.update(re.findall(r"[a-f0-9]{32}", payload))
        images: set[tuple[str, int]] = set()
        for row in db.all("SELECT * FROM device_state"):
            images.update(
                (row["device_id"], row[key])
                for key in ["previous_image_id", "base_image_id", "managed_image_id"]
                if row[key]
            )
        images.update(
            (row["device_id"], row["image_id"])
            for row in db.all("SELECT device_id,image_id FROM overrides")
        )
        for device_id, image_id in images:
            row = db.one(
                "SELECT asset_id FROM image_sources WHERE device_id=? AND image_id=?",
                (device_id, image_id),
            )
            if row:
                ids.add(row["asset_id"])
        return ids

    def inspect(self):
        rows = []
        for row in self.service.db.all("SELECT * FROM assets ORDER BY created_at"):
            path = self.service.assets.path(row["id"])
            if path.exists():
                rows.append((row, path, path.stat().st_size))
        return rows

    def stats(self):
        rows = self.inspect()
        protected = self.protected()
        return {
            "policy": self.config().model_dump(),
            "total_bytes": sum(x[2] for x in rows),
            "asset_count": len(rows),
            "protected_count": sum(x[0]["id"] in protected for x in rows),
        }

    def sweep(self, force=False, dry_run=False):
        if not force and self.last and time.monotonic() - self.last < 3600:
            return {}
        if not dry_run:
            self.last = time.monotonic()
        config = self.config()
        if not config.enabled:
            return {"deleted": 0, "bytes_freed": 0, "disabled": True}
        with self.service.lock:
            protected = self.protected()
            rows = self.inspect()
            total = sum(x[2] for x in rows)
            removed = []
            freed = 0
            for row, path, size in rows:
                # Manual uploaded photos are never removed by automatic retention.
                if row["id"] in protected or row["kind"] not in ("generated", "cache"):
                    continue
                days = (
                    config.cache_days
                    if row["kind"] == "cache"
                    else config.generated_days
                )
                expired = datetime.fromisoformat(row["created_at"]) < datetime.now(
                    UTC
                ) - timedelta(days=days)
                if not expired and total <= config.max_bytes:
                    continue
                removed.append(row["id"])
                total -= size
                freed += size
                if not dry_run:
                    path.unlink(missing_ok=True)
                    for table, key in [
                        ("asset_metadata", "asset_id"),
                        ("asset_sources", "asset_id"),
                        ("assets", "id"),
                    ]:
                        self.service.db.execute(
                            f"DELETE FROM {table} WHERE {key}=?", (row["id"],)
                        )
            result = {"deleted": len(removed), "bytes_freed": freed, "dry_run": dry_run}
            if not dry_run:
                self.service.db.execute(
                    "INSERT OR REPLACE INTO app_config VALUES('local_retention_last',?)",
                    (encode({**result, "timestamp": now_iso()}),),
                )
            return result