487 lines
20 KiB
Python
487 lines
20 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import shutil
|
|
import threading
|
|
import time
|
|
import xml.etree.ElementTree as ET
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
|
|
from .config import Settings
|
|
from .database import Database, utcnow
|
|
|
|
|
|
def _path_size(path: Path) -> int:
|
|
"""Return logical file sizes without following directory symlinks."""
|
|
try:
|
|
if path.is_symlink():
|
|
return 0
|
|
if path.is_file():
|
|
return path.stat().st_size
|
|
if not path.is_dir():
|
|
return 0
|
|
except OSError:
|
|
return 0
|
|
total = 0
|
|
stack = [path]
|
|
while stack:
|
|
current = stack.pop()
|
|
try:
|
|
with os.scandir(current) as entries:
|
|
for entry in entries:
|
|
try:
|
|
if entry.is_symlink():
|
|
continue
|
|
if entry.is_dir(follow_symlinks=False):
|
|
stack.append(Path(entry.path))
|
|
elif entry.is_file(follow_symlinks=False):
|
|
total += entry.stat(follow_symlinks=False).st_size
|
|
except OSError:
|
|
continue
|
|
except OSError:
|
|
continue
|
|
return total
|
|
|
|
|
|
def _root_size(path: Path, *, missing_available: bool = False) -> tuple[int, bool, str | None]:
|
|
try:
|
|
if path.is_symlink():
|
|
return 0, False, "符号链接目录不会被统计"
|
|
if not path.exists():
|
|
if missing_available:
|
|
return 0, True, None
|
|
return 0, False, "目录不存在"
|
|
if not os.access(path, os.R_OK):
|
|
return 0, False, "目录不可读"
|
|
return _path_size(path), True, None
|
|
except OSError as exc:
|
|
return 0, False, f"无法读取目录：{exc}"
|
|
|
|
|
|
def _entry(bytes_value: int, *, available: bool = True, reason: str | None = None, **extra) -> dict:
|
|
value = {"bytes": int(bytes_value), "available": bool(available)}
|
|
if reason:
|
|
value["reason"] = reason
|
|
value.update(extra)
|
|
return value
|
|
|
|
|
|
def _dedupe_roots(paths: list[Path]) -> list[Path]:
|
|
roots: list[Path] = []
|
|
for path in sorted({str(item) for item in paths}):
|
|
candidate = Path(path)
|
|
if any(candidate == root or root in candidate.parents for root in roots):
|
|
continue
|
|
roots = [root for root in roots if root not in candidate.parents]
|
|
roots.append(candidate)
|
|
return roots
|
|
|
|
|
|
class StorageUsageService:
|
|
def __init__(
|
|
self,
|
|
db: Database,
|
|
settings: Settings,
|
|
rclone=None,
|
|
remote_cache=None,
|
|
sources=None,
|
|
cache_seconds: float = 30.0,
|
|
):
|
|
self.db = db
|
|
self.settings = settings
|
|
self.rclone = rclone
|
|
self.remote_cache = remote_cache
|
|
self.sources = sources
|
|
self.cache_seconds = cache_seconds
|
|
self._lock = threading.Lock()
|
|
self._cached_at = 0.0
|
|
self._cached: dict | None = None
|
|
self._capacity_cache: dict[str, tuple[float, dict]] = {}
|
|
|
|
@staticmethod
|
|
def _capacity_result(
|
|
source_id: str,
|
|
*,
|
|
scope: str,
|
|
status: str,
|
|
provider: str,
|
|
total: int = 0,
|
|
used: int = 0,
|
|
available: int = 0,
|
|
reason: str | None = None,
|
|
cache_seconds: float = 30,
|
|
) -> dict:
|
|
result = {
|
|
"source_id": source_id,
|
|
"scope": scope,
|
|
"status": status,
|
|
"total_bytes": max(0, int(total)),
|
|
"used_bytes": max(0, int(used)),
|
|
"available_bytes": max(0, int(available)),
|
|
"provider": provider,
|
|
"collected_at": datetime.now(UTC).isoformat(),
|
|
"cached_seconds": int(cache_seconds),
|
|
}
|
|
if reason:
|
|
result["reason"] = reason
|
|
return result
|
|
|
|
def _webdav_capacity(self, source: dict) -> dict | None:
|
|
connector = self.sources.connector(source["id"])
|
|
try:
|
|
response = connector.client.request(
|
|
"PROPFIND",
|
|
connector.base_url,
|
|
headers={"Depth": "0", "Content-Type": "application/xml; charset=utf-8"},
|
|
content=(
|
|
"<?xml version='1.0' encoding='utf-8'?>"
|
|
"<d:propfind xmlns:d='DAV:'><d:prop><d:quota-used-bytes/>"
|
|
"<d:quota-available-bytes/></d:prop></d:propfind>"
|
|
),
|
|
)
|
|
if response.status_code != 207:
|
|
raise RuntimeError(f"WebDAV 配额查询失败：HTTP {response.status_code}")
|
|
root = ET.fromstring(response.content)
|
|
used_text = root.findtext(".//{DAV:}quota-used-bytes")
|
|
free_text = root.findtext(".//{DAV:}quota-available-bytes")
|
|
if used_text is None or free_text is None:
|
|
return None
|
|
used = max(0, int(used_text))
|
|
free = max(0, int(free_text))
|
|
return {"total_bytes": used + free, "used_bytes": used, "available_bytes": free}
|
|
finally:
|
|
connector.close()
|
|
|
|
def _calculate_capacity(self, source_id: str) -> dict:
|
|
if self.sources is None:
|
|
raise RuntimeError("媒体来源服务尚未配置")
|
|
source = self.sources.get(source_id)
|
|
if source["kind"] == "local":
|
|
try:
|
|
disk = shutil.disk_usage(Path(source["config"]["path"]))
|
|
except OSError as exc:
|
|
return self._capacity_result(
|
|
source_id,
|
|
scope="local",
|
|
status="unreachable",
|
|
provider="local",
|
|
reason=f"本地来源不可访问：{type(exc).__name__}",
|
|
cache_seconds=self.cache_seconds,
|
|
)
|
|
return self._capacity_result(
|
|
source_id,
|
|
scope="local",
|
|
status="available",
|
|
provider="local",
|
|
total=disk.total,
|
|
used=disk.used,
|
|
available=disk.free,
|
|
cache_seconds=self.cache_seconds,
|
|
)
|
|
|
|
config = source.get("config", {})
|
|
override_total = config.get("capacity_override_total_bytes")
|
|
override_available = config.get("capacity_override_available_bytes")
|
|
try:
|
|
override_total = int(override_total) if override_total is not None else None
|
|
override_available = int(override_available) if override_available is not None else None
|
|
except (TypeError, ValueError):
|
|
override_total = override_available = None
|
|
if (
|
|
override_total is not None
|
|
and override_available is not None
|
|
and override_total > 0
|
|
and 0 <= override_available <= override_total
|
|
):
|
|
return self._capacity_result(
|
|
source_id,
|
|
scope="remote",
|
|
status="available",
|
|
provider="manual",
|
|
total=override_total,
|
|
used=override_total - override_available,
|
|
available=override_available,
|
|
reason="管理员手工设置；不会自动随云盘变化",
|
|
cache_seconds=self.cache_seconds,
|
|
)
|
|
native_or_alist = bool(
|
|
config.get("storage_backend") == "openlist_native"
|
|
or config.get("mode") == "encrypted"
|
|
or config.get("driver") == "alist"
|
|
)
|
|
try:
|
|
values = self.rclone.about(source) if native_or_alist and self.rclone is not None else None
|
|
provider = "rclone" if values is not None else "webdav"
|
|
if values is None:
|
|
values = self._webdav_capacity(source)
|
|
except NotImplementedError:
|
|
values = None
|
|
provider = "rclone" if native_or_alist else "webdav"
|
|
except Exception as exc:
|
|
return self._capacity_result(
|
|
source_id,
|
|
scope="remote",
|
|
status="unreachable",
|
|
provider="rclone" if native_or_alist else "webdav",
|
|
reason=f"无法连接远端配额接口：{type(exc).__name__}",
|
|
cache_seconds=self.cache_seconds,
|
|
)
|
|
if values is None:
|
|
return self._capacity_result(
|
|
source_id,
|
|
scope="remote",
|
|
status="unsupported",
|
|
provider=provider,
|
|
reason="此媒体来源未提供容量配额",
|
|
cache_seconds=self.cache_seconds,
|
|
)
|
|
return self._capacity_result(
|
|
source_id,
|
|
scope="remote",
|
|
status="available",
|
|
provider=provider,
|
|
total=values["total_bytes"],
|
|
used=values["used_bytes"],
|
|
available=values["available_bytes"],
|
|
cache_seconds=self.cache_seconds,
|
|
)
|
|
|
|
def configure_capacity_override(
|
|
self,
|
|
source_id: str,
|
|
*,
|
|
total_bytes: int | None,
|
|
available_bytes: int | None,
|
|
) -> dict:
|
|
source = self.sources.get(source_id)
|
|
if source["kind"] == "local":
|
|
raise ValueError("本地媒体来源使用文件系统真实容量，不能手工覆盖")
|
|
if (total_bytes is None) != (available_bytes is None):
|
|
raise ValueError("总容量与可用容量必须同时填写或同时清除")
|
|
if total_bytes is not None and (total_bytes <= 0 or not 0 <= available_bytes <= total_bytes):
|
|
raise ValueError("可用容量必须大于等于 0 且不能超过总容量")
|
|
config = dict(source.get("config") or {})
|
|
if total_bytes is None:
|
|
config.pop("capacity_override_total_bytes", None)
|
|
config.pop("capacity_override_available_bytes", None)
|
|
else:
|
|
config["capacity_override_total_bytes"] = int(total_bytes)
|
|
config["capacity_override_available_bytes"] = int(available_bytes)
|
|
with self.db.transaction() as conn:
|
|
conn.execute(
|
|
"UPDATE sources SET config_json=?,updated_at=? WHERE id=?",
|
|
(json.dumps(config), utcnow(), source_id),
|
|
)
|
|
with self._lock:
|
|
self._capacity_cache.pop(source_id, None)
|
|
return self.capacity(source_id, refresh=True)
|
|
|
|
def capacity(self, source_id: str, *, refresh: bool = False) -> dict:
|
|
now = time.monotonic()
|
|
with self._lock:
|
|
cached = self._capacity_cache.get(source_id)
|
|
if not refresh and cached is not None and now - cached[0] < self.cache_seconds:
|
|
return dict(cached[1])
|
|
result = self._calculate_capacity(source_id)
|
|
with self._lock:
|
|
self._capacity_cache[source_id] = (time.monotonic(), dict(result))
|
|
return result
|
|
|
|
def _postgres_relation_bytes(self) -> tuple[int, int]:
|
|
with self.db.read() as conn:
|
|
database = conn.execute("SELECT pg_database_size(current_database())").fetchone()
|
|
vectors = conn.execute(
|
|
"SELECT coalesce(pg_total_relation_size('frames'),0)+"
|
|
"coalesce(pg_total_relation_size('faces'),0)+coalesce(pg_total_relation_size('people'),0)"
|
|
).fetchone()
|
|
return int(database[0] or 0), int(vectors[0] or 0)
|
|
|
|
def _openlist_staging_paths(self) -> list[dict]:
|
|
rows: list[dict] = []
|
|
try:
|
|
with self.db.read() as conn:
|
|
sources = conn.execute("SELECT id,name,config_json FROM sources WHERE kind='webdav'").fetchall()
|
|
except Exception:
|
|
return rows
|
|
for row in sources:
|
|
try:
|
|
config = json.loads(row["config_json"] or "{}")
|
|
except (TypeError, ValueError, json.JSONDecodeError):
|
|
continue
|
|
value = str(config.get("openlist_local_staging_path") or "").strip()
|
|
if value:
|
|
rows.append({"source_id": str(row["id"]), "name": str(row["name"]), "path": Path(value)})
|
|
return rows
|
|
|
|
@staticmethod
|
|
def _child_total(children: dict[str, dict]) -> tuple[int, int, int]:
|
|
total = active = reclaimable = 0
|
|
for value in children.values():
|
|
total += int(value.get("bytes") or 0)
|
|
active += int(value.get("active_bytes") or 0)
|
|
reclaimable += int(value.get("reclaimable_bytes") or 0)
|
|
return total, active, reclaimable
|
|
|
|
def _calculate(self) -> dict:
|
|
disk = shutil.disk_usage(self.settings.data_dir)
|
|
data_total = _path_size(self.settings.data_dir)
|
|
postgres_total, ai_index_bytes = self._postgres_relation_bytes()
|
|
database_bytes = max(0, postgres_total - ai_index_bytes)
|
|
|
|
models, models_ok, models_reason = _root_size(self.settings.models_dir, missing_available=True)
|
|
thumbnails, thumbnails_ok, thumbnails_reason = _root_size(
|
|
self.settings.thumbnails_dir, missing_available=True
|
|
)
|
|
preview, preview_ok, preview_reason = _root_size(self.settings.preview_dir, missing_available=True)
|
|
query_images, query_ok, query_reason = _root_size(
|
|
self.settings.uploads_dir, missing_available=True
|
|
)
|
|
remote_media, remote_media_ok, remote_media_reason = _root_size(
|
|
self.settings.remote_media_cache_dir, missing_available=True
|
|
)
|
|
upload_receive, upload_ok, upload_reason = _root_size(
|
|
self.settings.upload_staging_dir, missing_available=True
|
|
)
|
|
download_staging, download_ok, download_reason = _root_size(
|
|
self.settings.download_staging_dir, missing_available=True
|
|
)
|
|
encryption_work, encryption_ok, encryption_reason = _root_size(
|
|
self.settings.rclone_dir / "local-encryption", missing_available=True
|
|
)
|
|
|
|
rclone_status = self.rclone.status() if self.rclone is not None else {}
|
|
remote_status = self.remote_cache.status() if self.remote_cache is not None else {}
|
|
remote_children = {
|
|
"remote_media_cache": _entry(
|
|
remote_media,
|
|
available=remote_media_ok,
|
|
reason=remote_media_reason,
|
|
active_bytes=int(remote_status.get("active_bytes") or 0),
|
|
reclaimable_bytes=int(remote_status.get("evictable_bytes") or 0),
|
|
entries=int(remote_status.get("entries") or 0),
|
|
),
|
|
"rclone_vfs_cache": _entry(
|
|
int(rclone_status.get("vfs_cache_bytes") or 0),
|
|
active_bytes=int(rclone_status.get("active_cache_bytes") or 0),
|
|
reclaimable_bytes=int(rclone_status.get("reclaimable_cache_bytes") or 0),
|
|
instances=rclone_status.get("instances") or [],
|
|
),
|
|
"rclone_runtime": _entry(int(rclone_status.get("runtime_bytes") or 0)),
|
|
}
|
|
upload_children = {
|
|
"upload_receive": _entry(upload_receive, available=upload_ok, reason=upload_reason),
|
|
"download_staging": _entry(download_staging, available=download_ok, reason=download_reason),
|
|
"encryption_work": _entry(encryption_work, available=encryption_ok, reason=encryption_reason),
|
|
}
|
|
|
|
staging_rows = self._openlist_staging_paths()
|
|
staging_roots: list[Path] = []
|
|
staging_details: list[dict] = []
|
|
for item in staging_rows:
|
|
size, available, reason = _root_size(item["path"])
|
|
staging_details.append(
|
|
{
|
|
"source_id": item["source_id"],
|
|
"name": item["name"],
|
|
"path": str(item["path"]),
|
|
**_entry(size, available=available, reason=reason),
|
|
}
|
|
)
|
|
if available and not item["path"].is_symlink():
|
|
staging_roots.append(item["path"].resolve())
|
|
unique_staging = _dedupe_roots(staging_roots)
|
|
staging_size = sum(_path_size(path) for path in unique_staging)
|
|
upload_children["openlist_local_staging"] = _entry(
|
|
staging_size,
|
|
available=all(item.get("available", False) for item in staging_details)
|
|
if staging_details
|
|
else True,
|
|
reason="部分 OpenList 中转目录当前不可访问"
|
|
if any(not item.get("available", False) for item in staging_details)
|
|
else None,
|
|
locations=staging_details,
|
|
)
|
|
|
|
preview_total, preview_active, preview_reclaimable = self._child_total(
|
|
{
|
|
"preview_cache": _entry(preview, available=preview_ok, reason=preview_reason),
|
|
"query_images": _entry(query_images, available=query_ok, reason=query_reason),
|
|
}
|
|
)
|
|
remote_total, remote_active, remote_reclaimable = self._child_total(remote_children)
|
|
upload_total, upload_active, upload_reclaimable = self._child_total(upload_children)
|
|
|
|
categories: dict[str, dict] = {
|
|
"database": _entry(database_bytes),
|
|
"ai_index": _entry(ai_index_bytes, mode="pgvector"),
|
|
"models": _entry(models, available=models_ok, reason=models_reason),
|
|
"thumbnails": _entry(thumbnails, available=thumbnails_ok, reason=thumbnails_reason),
|
|
"preview_cache": _entry(
|
|
preview_total,
|
|
available=preview_ok and query_ok,
|
|
active_bytes=preview_active,
|
|
reclaimable_bytes=preview_reclaimable,
|
|
children={
|
|
"preview_cache": _entry(preview, available=preview_ok, reason=preview_reason),
|
|
"query_images": _entry(query_images, available=query_ok, reason=query_reason),
|
|
},
|
|
),
|
|
"remote_cache": _entry(
|
|
remote_total,
|
|
active_bytes=remote_active,
|
|
reclaimable_bytes=remote_reclaimable,
|
|
children=remote_children,
|
|
),
|
|
"upload_staging": _entry(
|
|
upload_total,
|
|
available=all(value.get("available", True) for value in upload_children.values()),
|
|
active_bytes=upload_active,
|
|
reclaimable_bytes=upload_reclaimable,
|
|
children=upload_children,
|
|
),
|
|
}
|
|
|
|
managed_roots = _dedupe_roots(
|
|
[
|
|
self.settings.models_dir,
|
|
self.settings.thumbnails_dir,
|
|
self.settings.preview_dir,
|
|
self.settings.uploads_dir,
|
|
self.settings.remote_media_cache_dir,
|
|
self.settings.rclone_dir,
|
|
self.settings.upload_staging_dir,
|
|
self.settings.download_staging_dir,
|
|
]
|
|
)
|
|
known_in_data = sum(
|
|
_path_size(path)
|
|
for path in managed_roots
|
|
if self.settings.data_dir in path.parents or path == self.settings.data_dir
|
|
)
|
|
categories["other"] = _entry(max(0, data_total - known_in_data))
|
|
external_staging = sum(
|
|
_path_size(path)
|
|
for path in unique_staging
|
|
if self.settings.data_dir not in path.parents and path != self.settings.data_dir
|
|
)
|
|
app_bytes = data_total + postgres_total + external_staging
|
|
return {
|
|
"schema_version": 2,
|
|
"collected_at": datetime.now(UTC).isoformat(),
|
|
"disk": {"total_bytes": disk.total, "used_bytes": disk.used, "available_bytes": disk.free},
|
|
"app_bytes": app_bytes,
|
|
"categories": categories,
|
|
"cached_seconds": self.cache_seconds,
|
|
}
|
|
|
|
def usage(self, *, refresh: bool = False) -> dict:
|
|
now = time.monotonic()
|
|
with self._lock:
|
|
if not refresh and self._cached is not None and now - self._cached_at < self.cache_seconds:
|
|
return self._cached
|
|
self._cached = self._calculate()
|
|
self._cached_at = now
|
|
return self._cached
|