121 lines
4.0 KiB
Python
121 lines
4.0 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import json
|
||
|
|
import threading
|
||
|
|
import time
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
from yuxi.utils.logging_config import logger
|
||
|
|
|
||
|
|
|
||
|
|
class CredentialQueue:
|
||
|
|
def __init__(self, auth_dir: Path, debounce_ms: int = 500):
|
||
|
|
self._auth_dir = auth_dir
|
||
|
|
self._lock = threading.Lock()
|
||
|
|
self._write_count = 0
|
||
|
|
self._debounce_ms = debounce_ms
|
||
|
|
self._pending_creds: dict[str, Any] | None = None
|
||
|
|
self._last_write_time = 0.0
|
||
|
|
self._debounce_timer: threading.Timer | None = None
|
||
|
|
|
||
|
|
def write_auth(self, creds: dict[str, Any]) -> int:
|
||
|
|
now = time.monotonic()
|
||
|
|
with self._lock:
|
||
|
|
self._write_count += 1
|
||
|
|
count = self._write_count
|
||
|
|
|
||
|
|
if now - self._last_write_time < self._debounce_ms / 1000.0:
|
||
|
|
self._pending_creds = creds
|
||
|
|
self._schedule_debounced_write()
|
||
|
|
return count
|
||
|
|
|
||
|
|
self._last_write_time = now
|
||
|
|
self._pending_creds = None
|
||
|
|
|
||
|
|
return self._do_write(creds, count)
|
||
|
|
|
||
|
|
def _schedule_debounced_write(self) -> None:
|
||
|
|
if self._debounce_timer is not None:
|
||
|
|
self._debounce_timer.cancel()
|
||
|
|
|
||
|
|
self._debounce_timer = threading.Timer(
|
||
|
|
self._debounce_ms / 1000.0,
|
||
|
|
self._flush_debounced,
|
||
|
|
)
|
||
|
|
self._debounce_timer.daemon = True
|
||
|
|
self._debounce_timer.start()
|
||
|
|
|
||
|
|
def _flush_debounced(self) -> None:
|
||
|
|
with self._lock:
|
||
|
|
creds = self._pending_creds
|
||
|
|
self._pending_creds = None
|
||
|
|
self._debounce_timer = None
|
||
|
|
if creds is None:
|
||
|
|
return
|
||
|
|
count = self._write_count
|
||
|
|
|
||
|
|
self._do_write(creds, count)
|
||
|
|
|
||
|
|
def _do_write(self, creds: dict[str, Any], count: int) -> int:
|
||
|
|
path = self._auth_dir / "creds.json"
|
||
|
|
backup = self._auth_dir / "creds.json.bak"
|
||
|
|
tmp_path = self._auth_dir / "creds.json.tmp"
|
||
|
|
try:
|
||
|
|
data = json.dumps(creds, indent=2, ensure_ascii=False)
|
||
|
|
tmp_path.write_text(data, encoding="utf-8")
|
||
|
|
if path.exists():
|
||
|
|
try:
|
||
|
|
path.replace(backup)
|
||
|
|
except OSError:
|
||
|
|
pass
|
||
|
|
tmp_path.replace(path)
|
||
|
|
logger.debug(f"CredentialQueue: wrote creds (#{count}) to {path}")
|
||
|
|
except (OSError, TypeError) as e:
|
||
|
|
logger.error(f"CredentialQueue: write failed (#{count}): {e}")
|
||
|
|
raise
|
||
|
|
return count
|
||
|
|
|
||
|
|
def read_auth(self) -> dict[str, Any] | None:
|
||
|
|
path = self._auth_dir / "creds.json"
|
||
|
|
backup = self._auth_dir / "creds.json.bak"
|
||
|
|
try:
|
||
|
|
if not path.exists():
|
||
|
|
return None
|
||
|
|
return json.loads(path.read_text(encoding="utf-8"))
|
||
|
|
except (OSError, json.JSONDecodeError) as e:
|
||
|
|
logger.error(f"CredentialQueue: read failed: {e}")
|
||
|
|
if backup.exists():
|
||
|
|
try:
|
||
|
|
data = json.loads(backup.read_text(encoding="utf-8"))
|
||
|
|
backup.replace(path)
|
||
|
|
logger.info("CredentialQueue: restored creds from backup")
|
||
|
|
return data
|
||
|
|
except (OSError, json.JSONDecodeError) as e2:
|
||
|
|
logger.error(f"CredentialQueue: backup recovery failed: {e2}")
|
||
|
|
return None
|
||
|
|
|
||
|
|
def clear_auth(self) -> bool:
|
||
|
|
with self._lock:
|
||
|
|
self._pending_creds = None
|
||
|
|
if self._debounce_timer is not None:
|
||
|
|
self._debounce_timer.cancel()
|
||
|
|
self._debounce_timer = None
|
||
|
|
|
||
|
|
path = self._auth_dir / "creds.json"
|
||
|
|
backup = self._auth_dir / "creds.json.bak"
|
||
|
|
try:
|
||
|
|
if path.exists():
|
||
|
|
if backup.exists():
|
||
|
|
backup.unlink()
|
||
|
|
path.rename(backup)
|
||
|
|
logger.info(f"CredentialQueue: cleared creds, backup at {backup}")
|
||
|
|
return True
|
||
|
|
except OSError as e:
|
||
|
|
logger.error(f"CredentialQueue: clear failed: {e}")
|
||
|
|
return False
|
||
|
|
|
||
|
|
@property
|
||
|
|
def write_count(self) -> int:
|
||
|
|
return self._write_count
|