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
2026-06-26 14:36:24 +08:00
parent c44e9bcc55
commit da9ba657f3
7 changed files with 163 additions and 35 deletions
+86 -31
View File
@@ -6,6 +6,7 @@
import json import json
import os import os
import tempfile import tempfile
from datetime import datetime
from pathlib import Path from pathlib import Path
from threading import Lock from threading import Lock
@@ -14,62 +15,104 @@ from common.logging import logger
from spark_executor.models import PendingSubmission from spark_executor.models import PendingSubmission
DEFAULT_DATA_DIR = settings.data_dir DEFAULT_DATA_DIR = settings.data_dir
DEFAULT_FILE_NAME = "pending_jobs.json" DEFAULT_DIR_NAME = "pending_jobs"
class PendingStore: 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>. Records live in <data_dir>/pending_jobs/YYYY-MM-DD.json, one file per
Atomic writes (tempfile + os.replace), same as ConnectionStore. 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._data_dir = data_dir or DEFAULT_DATA_DIR
self._file_name = file_name self._dir_name = dir_name
self._lock = Lock() self._lock = Lock()
@property @property
def path(self) -> Path: def dir_path(self) -> Path:
return Path(self._data_dir) / self._file_name return Path(self._data_dir) / self._dir_name
def _load(self) -> dict[str, PendingSubmission]: def _date_path(self, created_at: datetime) -> Path:
if not self.path.exists(): # created_at is datetime.utcnow(), so the shard date is a UTC date.
logger.debug(f"PendingStore._load: file {self.path} absent, returning empty") 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 {} return {}
raw = json.loads(self.path.read_text(encoding="utf-8")) raw = json.loads(path.read_text(encoding="utf-8"))
logger.debug(f"PendingStore._load: loaded {len(raw)} records from {self.path}")
return {pid: PendingSubmission.model_validate(p) for pid, p in raw.items()} return {pid: PendingSubmission.model_validate(p) for pid, p in raw.items()}
def _dump(self, records: dict[str, PendingSubmission]) -> None: def _dump_file(self, path: Path, records: dict[str, PendingSubmission]) -> None:
os.makedirs(self._data_dir, exist_ok=True) os.makedirs(path.parent, exist_ok=True)
payload = {pid: p.model_dump() for pid, p in records.items()} payload = {pid: p.model_dump() for pid, p in records.items()}
fd, tmp_path = tempfile.mkstemp( fd, tmp_path = tempfile.mkstemp(prefix=path.name + ".", dir=path.parent)
prefix=self._file_name + ".", dir=self._data_dir
)
try: try:
with os.fdopen(fd, "w", encoding="utf-8") as f: with os.fdopen(fd, "w", encoding="utf-8") as f:
json.dump(payload, f, indent=2, ensure_ascii=False, default=str) json.dump(payload, f, indent=2, ensure_ascii=False, default=str)
os.replace(tmp_path, self.path) os.replace(tmp_path, path)
logger.debug(f"PendingStore._dump: wrote {len(records)} records to {self.path}")
except Exception: except Exception:
if os.path.exists(tmp_path): if os.path.exists(tmp_path):
os.unlink(tmp_path) os.unlink(tmp_path)
raise 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]: def list_all(self) -> list[PendingSubmission]:
with self._lock: with self._lock:
return list(self._load().values()) return list(self._load_all().values())
def get(self, pending_id: str) -> PendingSubmission | None: def get(self, pending_id: str) -> PendingSubmission | None:
with self._lock: 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: def save(self, pending: PendingSubmission) -> None:
with self._lock: 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 records[pending.pending_id] = pending
self._dump(records) self._dump_file(path, records)
logger.info( logger.info(
f"pending saved pending_id={pending.pending_id} " f"pending saved pending_id={pending.pending_id} "
f"status={pending.status} connection={pending.connection}" f"status={pending.status} connection={pending.connection}"
@@ -77,13 +120,25 @@ class PendingStore:
def delete(self, pending_id: str) -> bool: def delete(self, pending_id: str) -> bool:
with self._lock: with self._lock:
records = self._load() for path in self._list_date_files():
if pending_id not in records: records = self._load_file(path)
return False if pending_id in records:
del records[pending_id] del records[pending_id]
self._dump(records) if records:
logger.info(f"pending deleted pending_id={pending_id}") self._dump_file(path, records)
return True 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. # Module-level singleton; replaced in tests.
+1
View File
@@ -86,6 +86,7 @@ class Connection(BaseModel):
class PendingSubmission(BaseModel): class PendingSubmission(BaseModel):
pending_id: str pending_id: str
app_name: str | None = None
connection: str connection: str
master: str master: str
deploy_mode: str deploy_mode: str
+4
View File
@@ -33,6 +33,10 @@ class SaveConnectionRequest(BaseModel):
class PrepareSubmitJobRequest(BaseModel): class PrepareSubmitJobRequest(BaseModel):
connection: str connection: str
app_name: str | None = Field(
default=None,
description="Optional human-readable application name for tracking the pending submission.",
)
script_path: str = Field( script_path: str = Field(
..., ...,
description=( description=(
+2
View File
@@ -68,6 +68,7 @@ def prepare_submit_job(
executor_memory: str = "4G", executor_memory: str = "4G",
executor_cores: int = 2, executor_cores: int = 2,
num_executors: int = 2, num_executors: int = 2,
app_name: str | None = None,
) -> dict[str, object]: ) -> dict[str, object]:
"""Snapshot connection params and persist a PendingSubmission. Does NOT submit.""" """Snapshot connection params and persist a PendingSubmission. Does NOT submit."""
logger.debug( logger.debug(
@@ -109,6 +110,7 @@ def prepare_submit_job(
pending_id = _new_pending_id() pending_id = _new_pending_id()
pending = PendingSubmission( pending = PendingSubmission(
pending_id=pending_id, pending_id=pending_id,
app_name=app_name,
connection=connection, connection=connection,
master=conn.master, master=conn.master,
deploy_mode=conn.deploy_mode, deploy_mode=conn.deploy_mode,
+1
View File
@@ -105,6 +105,7 @@ def test_pending_submission_defaults_to_pending_status():
assert p.job_id is None assert p.job_id is None
assert p.application_id is None assert p.application_id is None
assert p.yarn_rm_url is None # default for connections without one set assert p.yarn_rm_url is None # default for connections without one set
assert p.app_name is None
def test_pending_submission_carries_yarn_rm_url_snapshot(): def test_pending_submission_carries_yarn_rm_url_snapshot():
+51 -4
View File
@@ -8,7 +8,7 @@ from spark_executor.core import pending_store
from spark_executor.models import PendingSubmission from spark_executor.models import PendingSubmission
def _pending(pid: str = "p_1") -> PendingSubmission: def _pending(pid: str = "p_1", created_at: datetime | None = None) -> PendingSubmission:
return PendingSubmission( return PendingSubmission(
pending_id=pid, pending_id=pid,
connection="prod", connection="prod",
@@ -20,7 +20,7 @@ def _pending(pid: str = "p_1") -> PendingSubmission:
executor_cores=2, executor_cores=2,
num_executors=2, num_executors=2,
spark_conf={}, spark_conf={},
created_at=datetime(2026, 6, 24), created_at=created_at or datetime(2026, 6, 24),
) )
@@ -38,9 +38,9 @@ def test_save_then_get_roundtrip():
assert out.master == "yarn" assert out.master == "yarn"
def test_save_persists_to_json(tmp_path: Path): def test_save_persists_to_date_file(tmp_path: Path):
pending_store.store.save(_pending()) pending_store.store.save(_pending())
path = tmp_path / "pending_jobs.json" path = tmp_path / "pending_jobs" / "2026-06-24.json"
assert path.is_file() assert path.is_file()
@@ -58,6 +58,53 @@ def test_list_all_returns_saved():
assert {p.pending_id for p in pending_store.store.list_all()} == {"a", "b"} assert {p.pending_id for p in pending_store.store.list_all()} == {"a", "b"}
def test_list_all_spans_multiple_date_files():
pending_store.store.save(_pending("a", created_at=datetime(2026, 6, 24)))
pending_store.store.save(_pending("b", created_at=datetime(2026, 6, 25)))
assert {p.pending_id for p in pending_store.store.list_all()} == {"a", "b"}
def test_get_searches_date_files():
pending_store.store.save(_pending("a", created_at=datetime(2026, 6, 24)))
assert pending_store.store.get("a") is not None
def test_delete_removes_from_date_file_and_deletes_empty_file():
pending_store.store.save(_pending("a", created_at=datetime(2026, 6, 24)))
date_file = pending_store.store.dir_path / "2026-06-24.json"
assert date_file.is_file()
assert pending_store.store.delete("a") is True
assert pending_store.store.get("a") is None
assert not date_file.exists()
def _legacy_record_text() -> str:
return (
'{"p_legacy": {"pending_id": "p_legacy", "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": "2026-06-23T00:00:00"}}'
)
def test_legacy_file_is_readable(tmp_path: Path):
legacy = tmp_path / "pending_jobs.json"
legacy.write_text(_legacy_record_text(), encoding="utf-8")
assert pending_store.store.get("p_legacy") is not None
assert any(p.pending_id == "p_legacy" for p in pending_store.store.list_all())
def test_save_migrates_legacy_file(tmp_path: Path):
legacy = tmp_path / "pending_jobs.json"
legacy.write_text(_legacy_record_text(), encoding="utf-8")
pending_store.store.save(_pending("p_new", created_at=datetime(2026, 6, 24)))
assert not legacy.exists()
assert pending_store.store.get("p_legacy") is not None
migrated = pending_store.store.dir_path / "2026-06-23.json"
assert migrated.is_file()
def test_save_upserts_existing(): def test_save_upserts_existing():
pending_store.store.save(_pending("a")) pending_store.store.save(_pending("a"))
updated = _pending("a") updated = _pending("a")
+18
View File
@@ -74,6 +74,24 @@ def test_prepare_persists_pending_with_snapshot(monkeypatch, real_script):
assert p.script_path == str(real_script) assert p.script_path == str(real_script)
def test_prepare_persists_app_name(monkeypatch, real_script):
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
submit.prepare_submit_job(
connection="prod", script_path=str(real_script), app_name="my-etl-job"
)
p = submit.pending_store.get(_last_pending_id())
assert p is not None
assert p.app_name == "my-etl-job"
def test_prepare_app_name_defaults_to_none(monkeypatch, real_script):
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
submit.prepare_submit_job(connection="prod", script_path=str(real_script))
p = submit.pending_store.get(_last_pending_id())
assert p is not None
assert p.app_name is None
def test_prepare_raises_for_unknown_connection(real_script): def test_prepare_raises_for_unknown_connection(real_script):
with pytest.raises(KeyError, match="missing"): with pytest.raises(KeyError, match="missing"):
submit.prepare_submit_job(connection="missing", script_path=str(real_script)) submit.prepare_submit_job(connection="missing", script_path=str(real_script))