pigway-frame-flow/pigway_frame_flow/display/retention.py

"""Opt-in retention for known uploads; never sweep the vendor's entire gallery."""

import json
import time
from datetime import UTC, datetime, timedelta

from pigway_frame_flow.errors import ErrorCode as E
from pigway_frame_flow.errors import FlowError
from pigway_frame_flow.storage.database import now_iso


class Retention:
    def __init__(self, service):
        self.service = service
        self.db = service.db
        self.checked = 0.0

    def info(self, device_id, image_id):
        row = self.db.one(
            "SELECT * FROM image_retention WHERE device_id=? AND image_id=?",
            (device_id, image_id),
        )
        return {
            "starred": bool(row and row["starred"]),
            "expires_at": row["expires_at"] if row else None,
        }

    def star(self, device_id, image_id, starred):
        with self.service.lock:
            row = self.db.one(
                "SELECT cleanup_operation FROM image_retention WHERE device_id=? AND image_id=?",
                (device_id, image_id),
            )
            if row and row["cleanup_operation"]:
                op = self.db.one(
                    "SELECT write_started FROM operations WHERE id=?",
                    (row["cleanup_operation"],),
                )
                if op and op["write_started"]:
                    raise FlowError(
                        E.CONFLICT, "Deletion has begun; reconcile before starring", 409
                    )
            if image_id not in {
                x.image_id for x in self.service.registry.get(device_id).images()
            }:
                raise FlowError(E.NOT_FOUND, "Gallery image not found", 404)
            self.db.execute(
                "INSERT INTO image_retention(device_id,image_id,starred) VALUES(?,?,?) ON CONFLICT(device_id,image_id) DO UPDATE SET starred=excluded.starred",
                (device_id, image_id, int(starred)),
            )
            return self.info(device_id, image_id)

    def register(self, device_id, image_id, operation_id):
        expires = (datetime.now(UTC) + timedelta(days=7)).isoformat()
        self.db.execute(
            "INSERT INTO image_retention(device_id,image_id,starred,expires_at,upload_operation) VALUES(?,?,0,?,?) ON CONFLICT(device_id,image_id) DO NOTHING",
            (device_id, image_id, expires, operation_id),
        )

    def protected(self, device_id, image_id):
        state = self.service.state(device_id)
        ids = {state["previous_image_id"], state["base_image_id"]}
        ids.update(
            x["image_id"]
            for x in self.db.all(
                "SELECT image_id FROM overrides WHERE device_id=?", (device_id,)
            )
        )
        for row in self.db.all("SELECT payload FROM playlists"):
            p = json.loads(row["payload"])
            if p.get("device_id") == device_id:
                ids.update(p.get("image_ids", []))
        for row in self.db.all("SELECT payload FROM jobs"):
            job = json.loads(row["payload"])
            if job.get("device_id") == device_id:
                ids.add(job.get("source", {}).get("image_id"))
        return image_id in ids

    def eligible(self, device_id, image_id):
        row = self.db.one(
            "SELECT * FROM image_retention WHERE device_id=? AND image_id=?",
            (device_id, image_id),
        )
        return bool(
            row
            and row["upload_operation"]
            and not row["starred"]
            and row["expires_at"] <= now_iso()
            and not self.protected(device_id, image_id)
        )

    def sweep(self):
        if self.checked and time.monotonic() - self.checked < 3600:
            return
        self.checked = time.monotonic()
        with self.service.lock:
            if (
                not self.service.settings.mock
                and not self.service.settings.write_enabled
            ):
                return
            rows = self.db.all(
                "SELECT * FROM image_retention WHERE starred=0 AND expires_at<=? AND cleanup_operation IS NULL",
                (now_iso(),),
            )
            # Bound work per cycle; device reads occur only when there is due work.
            for row in rows[:20]:
                device_id, image_id = row["device_id"], row["image_id"]
                if self.service.blocked(device_id) or not self.eligible(
                    device_id, image_id
                ):
                    continue
                driver = self.service.registry.get(device_id)
                device = driver.device(refresh=True)
                if device.status != "online" or device.current_content.id == image_id:
                    continue
                id = self.service.submit(
                    "delete",
                    {"image_id": image_id, "confirm_current": False, "retention": True},
                    device_id,
                    "retention",
                    "retention-" + row["upload_operation"],
                    "retention-"
                    + row["upload_operation"]
                    + "-"
                    + datetime.now(UTC).strftime("%Y%m%d%H"),
                )
                self.db.execute(
                    "UPDATE image_retention SET cleanup_operation=? WHERE device_id=? AND image_id=?",
                    (id, device_id, image_id),
                )