feat: add JSON-backed PendingStore

This commit is contained in:
Claude
2026-06-24 14:39:13 +08:00
parent 82dafd357f
commit 0aa3c1eaf5
2 changed files with 158 additions and 0 deletions
+80
View File
@@ -0,0 +1,80 @@
# coding=utf-8
"""
@Time :2026/6/24
@Author :tao.chen
"""
import json
import os
import tempfile
from pathlib import Path
from threading import Lock
from spark_executor.models import PendingSubmission
DEFAULT_DATA_DIR = os.environ.get("SPARK_EXECUTOR_DATA_DIR", "./data")
DEFAULT_FILE_NAME = "pending_jobs.json"
class PendingStore:
"""JSON-backed CRUD for PendingSubmission records.
File path: <DEFAULT_DATA_DIR>/<DEFAULT_FILE_NAME>.
Atomic writes (tempfile + os.replace), same as ConnectionStore.
"""
def __init__(self, data_dir: str | None = None, file_name: str = DEFAULT_FILE_NAME) -> None:
self._data_dir = data_dir or DEFAULT_DATA_DIR
self._file_name = file_name
self._lock = Lock()
@property
def path(self) -> Path:
return Path(self._data_dir) / self._file_name
def _load(self) -> dict[str, PendingSubmission]:
if not self.path.exists():
return {}
raw = json.loads(self.path.read_text(encoding="utf-8"))
return {pid: PendingSubmission.model_validate(p) for pid, p in raw.items()}
def _dump(self, records: dict[str, PendingSubmission]) -> None:
os.makedirs(self._data_dir, exist_ok=True)
payload = {pid: p.model_dump() for pid, p in records.items()}
fd, tmp_path = tempfile.mkstemp(
prefix=self._file_name + ".", dir=self._data_dir
)
try:
with os.fdopen(fd, "w", encoding="utf-8") as f:
json.dump(payload, f, indent=2, ensure_ascii=False, default=str)
os.replace(tmp_path, self.path)
except Exception:
if os.path.exists(tmp_path):
os.unlink(tmp_path)
raise
def list_all(self) -> list[PendingSubmission]:
with self._lock:
return list(self._load().values())
def get(self, pending_id: str) -> PendingSubmission | None:
with self._lock:
return self._load().get(pending_id)
def save(self, pending: PendingSubmission) -> None:
with self._lock:
records = self._load()
records[pending.pending_id] = pending
self._dump(records)
def delete(self, pending_id: str) -> bool:
with self._lock:
records = self._load()
if pending_id not in records:
return False
del records[pending_id]
self._dump(records)
return True
# Module-level singleton; replaced in tests.
store = PendingStore()