mirror of
https://github.com/alexta69/metube.git
synced 2026-07-18 00:58:18 +00:00
237 lines
8.4 KiB
Python
237 lines
8.4 KiB
Python
from __future__ import annotations
|
|
|
|
import base64
|
|
import collections.abc
|
|
import errno
|
|
import json
|
|
import logging
|
|
import os
|
|
import shelve
|
|
import tempfile
|
|
import time
|
|
from datetime import datetime
|
|
from typing import Any, Optional
|
|
|
|
log = logging.getLogger("state_store")
|
|
|
|
STATE_SCHEMA_VERSION = 2
|
|
_BYTES_MARKER = "__metube_bytes__"
|
|
_DATETIME_MARKER = "__metube_datetime__"
|
|
|
|
# Errnos that signal the filesystem cannot support the temp-file + rename
|
|
# atomic-write strategy (for example an NFS-backed state dir returning EPERM on
|
|
# mkstemp). These are safe to fall back on because they mean the atomic
|
|
# mechanism is unavailable, not that the data write itself failed. Errors like
|
|
# ENOSPC/EIO are deliberately excluded so a genuine storage failure surfaces
|
|
# instead of silently truncating an existing good state file.
|
|
_ATOMIC_UNSUPPORTED_ERRNOS = frozenset(
|
|
e
|
|
for e in (
|
|
errno.EPERM,
|
|
errno.EACCES,
|
|
errno.ENOSYS,
|
|
errno.EINVAL,
|
|
getattr(errno, "EOPNOTSUPP", None),
|
|
getattr(errno, "ENOTSUP", None),
|
|
)
|
|
if e is not None
|
|
)
|
|
|
|
|
|
def to_json_compatible(value: Any) -> Any:
|
|
if value is None or isinstance(value, (bool, int, float, str)):
|
|
return value
|
|
if isinstance(value, bytes):
|
|
return {_BYTES_MARKER: base64.b64encode(value).decode("ascii")}
|
|
if isinstance(value, datetime):
|
|
return {_DATETIME_MARKER: value.isoformat()}
|
|
if isinstance(value, collections.abc.Mapping):
|
|
return {str(k): to_json_compatible(v) for k, v in value.items()}
|
|
if isinstance(value, (list, tuple, set, frozenset)):
|
|
return [to_json_compatible(v) for v in value]
|
|
if isinstance(value, collections.abc.Iterable):
|
|
return [to_json_compatible(v) for v in value]
|
|
raise TypeError(f"Value of type {type(value).__name__} is not JSON serializable")
|
|
|
|
|
|
def from_json_compatible(value: Any) -> Any:
|
|
if isinstance(value, list):
|
|
return [from_json_compatible(v) for v in value]
|
|
if isinstance(value, dict):
|
|
if set(value.keys()) == {_BYTES_MARKER}:
|
|
return base64.b64decode(value[_BYTES_MARKER].encode("ascii"))
|
|
if set(value.keys()) == {_DATETIME_MARKER}:
|
|
return datetime.fromisoformat(value[_DATETIME_MARKER])
|
|
return {k: from_json_compatible(v) for k, v in value.items()}
|
|
return value
|
|
|
|
|
|
def read_legacy_shelf(path: str) -> Optional[list[tuple[Any, Any]]]:
|
|
if not os.path.exists(path):
|
|
return None
|
|
try:
|
|
with shelve.open(path, "r") as shelf:
|
|
return list(shelf.items())
|
|
except Exception as exc:
|
|
log.warning("Could not read legacy shelf at %s: %s", path, exc)
|
|
return None
|
|
|
|
|
|
class AtomicJsonStore:
|
|
def __init__(self, path: str, *, kind: str, schema_version: int = STATE_SCHEMA_VERSION):
|
|
self.path = path
|
|
self.kind = kind
|
|
self.schema_version = schema_version
|
|
self._direct_write_fallback_warned = False
|
|
|
|
def _ensure_parent(self) -> None:
|
|
parent = os.path.dirname(self.path)
|
|
if parent and not os.path.isdir(parent):
|
|
os.makedirs(parent, exist_ok=True)
|
|
|
|
def _build_payload(self, data: dict[str, Any]) -> dict[str, Any]:
|
|
payload = {
|
|
"schema_version": self.schema_version,
|
|
"kind": self.kind,
|
|
}
|
|
payload.update(data)
|
|
return payload
|
|
|
|
def load(self) -> Optional[dict[str, Any]]:
|
|
if not os.path.exists(self.path):
|
|
return None
|
|
try:
|
|
with open(self.path, encoding="utf-8") as f:
|
|
payload = json.load(f)
|
|
if not isinstance(payload, dict):
|
|
raise ValueError("State file must contain a JSON object")
|
|
if payload.get("kind") != self.kind:
|
|
raise ValueError(
|
|
f"State file kind mismatch: expected {self.kind}, got {payload.get('kind')}"
|
|
)
|
|
return payload
|
|
except Exception as exc:
|
|
self.quarantine_invalid_file(exc)
|
|
return None
|
|
|
|
def save(self, data: dict[str, Any]) -> None:
|
|
self._ensure_parent()
|
|
payload = self._build_payload(data)
|
|
try:
|
|
self._atomic_write(payload)
|
|
except OSError as exc:
|
|
if exc.errno not in _ATOMIC_UNSUPPORTED_ERRNOS:
|
|
raise
|
|
self._warn_direct_write_fallback(exc)
|
|
self._direct_write(payload)
|
|
|
|
def _atomic_write(self, payload: dict[str, Any]) -> None:
|
|
text = self._serialize(payload)
|
|
parent = os.path.dirname(self.path) or "."
|
|
fd, tmp_path = tempfile.mkstemp(
|
|
prefix=f".{os.path.basename(self.path)}.",
|
|
suffix=".tmp",
|
|
dir=parent,
|
|
text=True,
|
|
)
|
|
try:
|
|
with os.fdopen(fd, "w", encoding="utf-8") as f:
|
|
f.write(text)
|
|
f.flush()
|
|
self._best_effort_fsync(f.fileno())
|
|
os.replace(tmp_path, self.path)
|
|
self._fsync_directory(parent)
|
|
except Exception:
|
|
try:
|
|
os.remove(tmp_path)
|
|
except OSError:
|
|
pass
|
|
raise
|
|
|
|
def _direct_write(self, payload: dict[str, Any]) -> None:
|
|
# Serialize before truncating so a serialization failure never destroys
|
|
# the existing state file (the atomic path gets this for free via its
|
|
# temp file).
|
|
text = self._serialize(payload)
|
|
# Create with 0o600 so the fallback keeps the owner-only permissions the
|
|
# atomic path gets from mkstemp; state files can contain URLs and
|
|
# per-download option overrides that must not leak on shared mounts.
|
|
fd = os.open(self.path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
|
|
with os.fdopen(fd, "w", encoding="utf-8") as f:
|
|
# The 0o600 mode above only applies when the file is created; force
|
|
# it on rewrites too so an existing, broadly-permissioned state file
|
|
# is tightened to match the atomic path. Best-effort because some
|
|
# network filesystems reject chmod, and that must not re-crash save.
|
|
try:
|
|
os.fchmod(f.fileno(), 0o600)
|
|
except OSError:
|
|
pass
|
|
f.write(text)
|
|
f.flush()
|
|
self._best_effort_fsync(f.fileno())
|
|
# Make the new directory entry durable too, matching the atomic path.
|
|
self._fsync_directory(os.path.dirname(self.path) or ".")
|
|
|
|
@staticmethod
|
|
def _best_effort_fsync(fileno: int) -> None:
|
|
# Tolerate fsync being unsupported on the underlying filesystem (for
|
|
# example a network mount that returns EINVAL/ENOSYS), but let genuine
|
|
# storage failures such as ENOSPC/EIO surface so a non-durable write is
|
|
# never reported as success. An unsupported fsync must not by itself
|
|
# abandon the atomic rename path.
|
|
try:
|
|
os.fsync(fileno)
|
|
except OSError as exc:
|
|
if exc.errno not in _ATOMIC_UNSUPPORTED_ERRNOS:
|
|
raise
|
|
|
|
@staticmethod
|
|
def _serialize(payload: dict[str, Any]) -> str:
|
|
return json.dumps(payload, ensure_ascii=False, separators=(",", ":")) + "\n"
|
|
|
|
def _warn_direct_write_fallback(self, exc: OSError) -> None:
|
|
if self._direct_write_fallback_warned:
|
|
return
|
|
self._direct_write_fallback_warned = True
|
|
log.warning(
|
|
"Atomic state write failed for %s (%s); falling back to direct write",
|
|
self.path,
|
|
exc,
|
|
)
|
|
|
|
def quarantine_invalid_file(self, exc: Exception) -> None:
|
|
if not os.path.exists(self.path):
|
|
return
|
|
ts = time.strftime("%Y%m%d%H%M%S")
|
|
backup_path = f"{self.path}.invalid.{ts}"
|
|
try:
|
|
os.replace(self.path, backup_path)
|
|
log.warning(
|
|
"State file at %s was invalid (%s); moved it to %s",
|
|
self.path,
|
|
exc,
|
|
backup_path,
|
|
)
|
|
except OSError as move_exc:
|
|
log.warning(
|
|
"State file at %s was invalid (%s) and could not be moved aside: %s",
|
|
self.path,
|
|
exc,
|
|
move_exc,
|
|
)
|
|
|
|
@staticmethod
|
|
def _fsync_directory(path: str) -> None:
|
|
try:
|
|
flags = os.O_RDONLY
|
|
if hasattr(os, "O_DIRECTORY"):
|
|
flags |= os.O_DIRECTORY
|
|
fd = os.open(path, flags)
|
|
except OSError:
|
|
return
|
|
try:
|
|
os.fsync(fd)
|
|
except OSError:
|
|
pass
|
|
finally:
|
|
os.close(fd)
|