138 lines
5.1 KiB
Python
138 lines
5.1 KiB
Python
"""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')",
|
|
"SELECT payload FROM app_config WHERE key LIKE 'upload-retry:%'",
|
|
"SELECT payload FROM app_config WHERE key LIKE 'wrist-thumbnail:%'",
|
|
]:
|
|
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
|