feat: add JSON-backed PendingStore
This commit is contained in:
@@ -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()
|
||||
@@ -0,0 +1,78 @@
|
||||
# coding=utf-8
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from spark_executor.core import pending_store
|
||||
from spark_executor.models import PendingSubmission
|
||||
|
||||
|
||||
def _pending(pid: str = "p_1") -> PendingSubmission:
|
||||
return PendingSubmission(
|
||||
pending_id=pid,
|
||||
connection="prod",
|
||||
master="yarn",
|
||||
deploy_mode="cluster",
|
||||
script_path="/tmp/j.py",
|
||||
queue="default",
|
||||
executor_memory="4G",
|
||||
executor_cores=2,
|
||||
num_executors=2,
|
||||
spark_conf={},
|
||||
created_at=datetime(2026, 6, 24),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _fresh(tmp_path: Path, monkeypatch):
|
||||
monkeypatch.setattr(pending_store, "DEFAULT_DATA_DIR", str(tmp_path))
|
||||
monkeypatch.setattr(pending_store, "store", pending_store.PendingStore())
|
||||
|
||||
|
||||
def test_save_then_get_roundtrip():
|
||||
pending_store.store.save(_pending())
|
||||
out = pending_store.store.get("p_1")
|
||||
assert out is not None
|
||||
assert out.connection == "prod"
|
||||
assert out.master == "yarn"
|
||||
|
||||
|
||||
def test_save_persists_to_json(tmp_path: Path):
|
||||
pending_store.store.save(_pending())
|
||||
path = tmp_path / "pending_jobs.json"
|
||||
assert path.is_file()
|
||||
|
||||
|
||||
def test_get_missing_returns_none():
|
||||
assert pending_store.store.get("missing") is None
|
||||
|
||||
|
||||
def test_list_all_empty_when_file_absent():
|
||||
assert pending_store.store.list_all() == []
|
||||
|
||||
|
||||
def test_list_all_returns_saved():
|
||||
pending_store.store.save(_pending("a"))
|
||||
pending_store.store.save(_pending("b"))
|
||||
assert {p.pending_id for p in pending_store.store.list_all()} == {"a", "b"}
|
||||
|
||||
|
||||
def test_save_upserts_existing():
|
||||
pending_store.store.save(_pending("a"))
|
||||
updated = _pending("a")
|
||||
updated.status = "SUBMITTED"
|
||||
updated.job_id = "j_1"
|
||||
pending_store.store.save(updated)
|
||||
assert pending_store.store.get("a").status == "SUBMITTED"
|
||||
assert pending_store.store.get("a").job_id == "j_1"
|
||||
|
||||
|
||||
def test_delete_returns_true_when_present():
|
||||
pending_store.store.save(_pending())
|
||||
assert pending_store.store.delete("p_1") is True
|
||||
assert pending_store.store.get("p_1") is None
|
||||
|
||||
|
||||
def test_delete_returns_false_when_absent():
|
||||
assert pending_store.store.delete("nope") is False
|
||||
Reference in New Issue
Block a user