From da9ba657f3219001f38b9a01749f0d6bdeebb772 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 26 Jun 2026 14:36:24 +0800 Subject: [PATCH] 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. --- spark_executor/core/pending_store.py | 117 ++++++++++++++++++++------- spark_executor/models.py | 1 + spark_executor/tools/requests.py | 4 + spark_executor/tools/submit.py | 2 + tests/unit/test_models.py | 1 + tests/unit/test_pending_store.py | 55 ++++++++++++- tests/unit/test_submit_tool.py | 18 +++++ 7 files changed, 163 insertions(+), 35 deletions(-) diff --git a/spark_executor/core/pending_store.py b/spark_executor/core/pending_store.py index 7a01983..5bfccde 100644 --- a/spark_executor/core/pending_store.py +++ b/spark_executor/core/pending_store.py @@ -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: /. - Atomic writes (tempfile + os.replace), same as ConnectionStore. + Records live in /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. diff --git a/spark_executor/models.py b/spark_executor/models.py index 50b05eb..b0e8427 100644 --- a/spark_executor/models.py +++ b/spark_executor/models.py @@ -86,6 +86,7 @@ class Connection(BaseModel): class PendingSubmission(BaseModel): pending_id: str + app_name: str | None = None connection: str master: str deploy_mode: str diff --git a/spark_executor/tools/requests.py b/spark_executor/tools/requests.py index 104db0a..e316981 100644 --- a/spark_executor/tools/requests.py +++ b/spark_executor/tools/requests.py @@ -33,6 +33,10 @@ class SaveConnectionRequest(BaseModel): class PrepareSubmitJobRequest(BaseModel): connection: str + app_name: str | None = Field( + default=None, + description="Optional human-readable application name for tracking the pending submission.", + ) script_path: str = Field( ..., description=( diff --git a/spark_executor/tools/submit.py b/spark_executor/tools/submit.py index e7ba0f0..f1c8b2a 100644 --- a/spark_executor/tools/submit.py +++ b/spark_executor/tools/submit.py @@ -68,6 +68,7 @@ def prepare_submit_job( executor_memory: str = "4G", executor_cores: int = 2, num_executors: int = 2, + app_name: str | None = None, ) -> dict[str, object]: """Snapshot connection params and persist a PendingSubmission. Does NOT submit.""" logger.debug( @@ -109,6 +110,7 @@ def prepare_submit_job( pending_id = _new_pending_id() pending = PendingSubmission( pending_id=pending_id, + app_name=app_name, connection=connection, master=conn.master, deploy_mode=conn.deploy_mode, diff --git a/tests/unit/test_models.py b/tests/unit/test_models.py index 4b9b053..52ed2b0 100644 --- a/tests/unit/test_models.py +++ b/tests/unit/test_models.py @@ -105,6 +105,7 @@ def test_pending_submission_defaults_to_pending_status(): assert p.job_id is None assert p.application_id is None 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(): diff --git a/tests/unit/test_pending_store.py b/tests/unit/test_pending_store.py index 43a42c4..3aedaf7 100644 --- a/tests/unit/test_pending_store.py +++ b/tests/unit/test_pending_store.py @@ -8,7 +8,7 @@ from spark_executor.core import pending_store 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( pending_id=pid, connection="prod", @@ -20,7 +20,7 @@ def _pending(pid: str = "p_1") -> PendingSubmission: executor_cores=2, num_executors=2, 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" -def test_save_persists_to_json(tmp_path: Path): +def test_save_persists_to_date_file(tmp_path: Path): pending_store.store.save(_pending()) - path = tmp_path / "pending_jobs.json" + path = tmp_path / "pending_jobs" / "2026-06-24.json" 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"} +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(): pending_store.store.save(_pending("a")) updated = _pending("a") diff --git a/tests/unit/test_submit_tool.py b/tests/unit/test_submit_tool.py index d44eb74..f5349dd 100644 --- a/tests/unit/test_submit_tool.py +++ b/tests/unit/test_submit_tool.py @@ -74,6 +74,24 @@ def test_prepare_persists_pending_with_snapshot(monkeypatch, 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): with pytest.raises(KeyError, match="missing"): submit.prepare_submit_job(connection="missing", script_path=str(real_script))