75 lines
2.5 KiB
Python
75 lines
2.5 KiB
Python
from __future__ import annotations
|
|
|
|
from collections.abc import Iterable
|
|
from dataclasses import dataclass
|
|
from datetime import UTC, datetime, timedelta
|
|
from typing import Protocol
|
|
|
|
|
|
class RetentionError(ValueError):
|
|
pass
|
|
|
|
|
|
class BackupLike(Protocol):
|
|
id: str
|
|
created_at: datetime
|
|
pinned: bool
|
|
tombstoned_at: datetime | None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RetentionPolicy:
|
|
keep_last: int = 1
|
|
keep_days: int = 0
|
|
keep_daily: int = 0
|
|
keep_weekly: int = 0
|
|
keep_monthly: int = 0
|
|
|
|
@classmethod
|
|
def from_dict(cls, value: dict[str, object]) -> RetentionPolicy:
|
|
allowed = {"keep_last", "keep_days", "keep_daily", "keep_weekly", "keep_monthly"}
|
|
if set(value) - allowed:
|
|
raise RetentionError("retention policy contains unknown keys")
|
|
kwargs: dict[str, int] = {}
|
|
for key in allowed:
|
|
raw = value.get(key, 0 if key != "keep_last" else 1)
|
|
if not isinstance(raw, int) or isinstance(raw, bool) or raw < 0:
|
|
raise RetentionError("retention values must be non-negative integers")
|
|
kwargs[key] = raw
|
|
return cls(**kwargs)
|
|
|
|
|
|
def retained_ids(
|
|
backups: Iterable[BackupLike], policy: RetentionPolicy, now: datetime | None = None
|
|
) -> set[str]:
|
|
"""Return union retention set; the newest non-tombstoned backup is always protected."""
|
|
reference = (now or datetime.now(UTC)).astimezone(UTC)
|
|
items = sorted(
|
|
(item for item in backups if item.tombstoned_at is None),
|
|
key=lambda item: item.created_at.astimezone(UTC),
|
|
reverse=True,
|
|
)
|
|
if not items:
|
|
return set()
|
|
kept = {items[0].id}
|
|
kept.update(item.id for item in items[: policy.keep_last])
|
|
kept.update(item.id for item in items if item.pinned)
|
|
if policy.keep_days:
|
|
cutoff = reference - timedelta(days=policy.keep_days)
|
|
kept.update(item.id for item in items if item.created_at.astimezone(UTC) >= cutoff)
|
|
for count, key in (
|
|
(policy.keep_daily, lambda stamp: stamp.date()),
|
|
(policy.keep_weekly, lambda stamp: stamp.isocalendar()[:2]),
|
|
(policy.keep_monthly, lambda stamp: (stamp.year, stamp.month)),
|
|
):
|
|
buckets: set[object] = set()
|
|
for item in items:
|
|
stamp = item.created_at.astimezone(UTC)
|
|
bucket = key(stamp)
|
|
if len(buckets) >= count and bucket not in buckets:
|
|
continue
|
|
buckets.add(bucket)
|
|
if bucket in buckets:
|
|
kept.add(item.id)
|
|
return kept
|