feat(pending): app_name field and date-sharded pending storage
- Add optional app_name to PendingSubmission and prepare_submit_job, passed through PrepareSubmitJobRequest for user-provided tracking. - Rewrite PendingStore to shard records by UTC date under data/pending_jobs/YYYY-MM-DD.json instead of a single monolithic pending_jobs.json. - Legacy single-file pending_jobs.json remains readable; first save() migrates it and renames the old file to avoid double reads. - Tests cover date sharding, multi-date list, legacy read, legacy migration, and app_name round-trip.
This commit is contained in:
Claude
@@ -6,6 +6,7 @@
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from threading import Lock
|
||||
|
||||
@@ -14,62 +15,104 @@ from common.logging import logger
|
||||
from spark_executor.models import PendingSubmission
|
||||
|
||||
DEFAULT_DATA_DIR = settings.data_dir
|
||||
DEFAULT_FILE_NAME = "pending_jobs.json"
|
||||
DEFAULT_DIR_NAME = "pending_jobs"
|
||||
|
||||
|
||||
class PendingStore:
|
||||
"""JSON-backed CRUD for PendingSubmission records.
|
||||
"""JSON-backed CRUD for PendingSubmission records, sharded by UTC date.
|
||||
|
||||
File path: <DEFAULT_DATA_DIR>/<DEFAULT_FILE_NAME>.
|
||||
Atomic writes (tempfile + os.replace), same as ConnectionStore.
|
||||
Records live in <data_dir>/pending_jobs/YYYY-MM-DD.json, one file per
|
||||
calendar date taken from PendingSubmission.created_at. The legacy single-file
|
||||
pending_jobs.json is still readable and is migrated non-destructively on
|
||||
the first save().
|
||||
"""
|
||||
|
||||
def __init__(self, data_dir: str | None = None, file_name: str = DEFAULT_FILE_NAME) -> None:
|
||||
def __init__(self, data_dir: str | None = None, dir_name: str = DEFAULT_DIR_NAME) -> None:
|
||||
self._data_dir = data_dir or DEFAULT_DATA_DIR
|
||||
self._file_name = file_name
|
||||
self._dir_name = dir_name
|
||||
self._lock = Lock()
|
||||
|
||||
@property
|
||||
def path(self) -> Path:
|
||||
return Path(self._data_dir) / self._file_name
|
||||
def dir_path(self) -> Path:
|
||||
return Path(self._data_dir) / self._dir_name
|
||||
|
||||
def _load(self) -> dict[str, PendingSubmission]:
|
||||
if not self.path.exists():
|
||||
logger.debug(f"PendingStore._load: file {self.path} absent, returning empty")
|
||||
def _date_path(self, created_at: datetime) -> Path:
|
||||
# created_at is datetime.utcnow(), so the shard date is a UTC date.
|
||||
return self.dir_path / f"{created_at.date().isoformat()}.json"
|
||||
|
||||
def _legacy_path(self) -> Path:
|
||||
return Path(self._data_dir) / "pending_jobs.json"
|
||||
|
||||
def _list_date_files(self) -> list[Path]:
|
||||
if not self.dir_path.exists():
|
||||
return []
|
||||
return sorted(self.dir_path.glob("*.json"))
|
||||
|
||||
def _load_file(self, path: Path) -> dict[str, PendingSubmission]:
|
||||
if not path.exists():
|
||||
return {}
|
||||
raw = json.loads(self.path.read_text(encoding="utf-8"))
|
||||
logger.debug(f"PendingStore._load: loaded {len(raw)} records from {self.path}")
|
||||
raw = json.loads(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)
|
||||
def _dump_file(self, path: Path, records: dict[str, PendingSubmission]) -> None:
|
||||
os.makedirs(path.parent, 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
|
||||
)
|
||||
fd, tmp_path = tempfile.mkstemp(prefix=path.name + ".", dir=path.parent)
|
||||
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)
|
||||
logger.debug(f"PendingStore._dump: wrote {len(records)} records to {self.path}")
|
||||
os.replace(tmp_path, path)
|
||||
except Exception:
|
||||
if os.path.exists(tmp_path):
|
||||
os.unlink(tmp_path)
|
||||
raise
|
||||
|
||||
def _load_all(self) -> dict[str, PendingSubmission]:
|
||||
records: dict[str, PendingSubmission] = {}
|
||||
for path in self._list_date_files():
|
||||
records.update(self._load_file(path))
|
||||
legacy = self._legacy_path()
|
||||
if legacy.exists():
|
||||
records.update(self._load_file(legacy))
|
||||
return records
|
||||
|
||||
def list_all(self) -> list[PendingSubmission]:
|
||||
with self._lock:
|
||||
return list(self._load().values())
|
||||
return list(self._load_all().values())
|
||||
|
||||
def get(self, pending_id: str) -> PendingSubmission | None:
|
||||
with self._lock:
|
||||
return self._load().get(pending_id)
|
||||
# Newest first is the cheapest common case.
|
||||
for path in reversed(self._list_date_files()):
|
||||
record = self._load_file(path).get(pending_id)
|
||||
if record is not None:
|
||||
return record
|
||||
legacy = self._legacy_path()
|
||||
if legacy.exists():
|
||||
return self._load_file(legacy).get(pending_id)
|
||||
return None
|
||||
|
||||
def _migrate_legacy(self) -> None:
|
||||
"""Move records out of the monolithic legacy file into date shards."""
|
||||
legacy = self._legacy_path()
|
||||
if not legacy.exists():
|
||||
return
|
||||
for p in self._load_file(legacy).values():
|
||||
path = self._date_path(p.created_at)
|
||||
records = self._load_file(path)
|
||||
records[p.pending_id] = p
|
||||
self._dump_file(path, records)
|
||||
# Preserve the old file but stop reading it as a data file.
|
||||
os.makedirs(self.dir_path, exist_ok=True)
|
||||
legacy.rename(self.dir_path / "pending_jobs.legacy")
|
||||
|
||||
def save(self, pending: PendingSubmission) -> None:
|
||||
with self._lock:
|
||||
records = self._load()
|
||||
self._migrate_legacy()
|
||||
path = self._date_path(pending.created_at)
|
||||
records = self._load_file(path)
|
||||
records[pending.pending_id] = pending
|
||||
self._dump(records)
|
||||
self._dump_file(path, records)
|
||||
logger.info(
|
||||
f"pending saved pending_id={pending.pending_id} "
|
||||
f"status={pending.status} connection={pending.connection}"
|
||||
@@ -77,13 +120,25 @@ class PendingStore:
|
||||
|
||||
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)
|
||||
logger.info(f"pending deleted pending_id={pending_id}")
|
||||
return True
|
||||
for path in self._list_date_files():
|
||||
records = self._load_file(path)
|
||||
if pending_id in records:
|
||||
del records[pending_id]
|
||||
if records:
|
||||
self._dump_file(path, records)
|
||||
else:
|
||||
path.unlink()
|
||||
logger.info(f"pending deleted pending_id={pending_id}")
|
||||
return True
|
||||
legacy = self._legacy_path()
|
||||
if legacy.exists():
|
||||
records = self._load_file(legacy)
|
||||
if pending_id in records:
|
||||
del records[pending_id]
|
||||
self._dump_file(legacy, records)
|
||||
logger.info(f"pending deleted pending_id={pending_id}")
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
# Module-level singleton; replaced in tests.
|
||||
|
||||
Reference in New Issue
Block a user