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:
@@ -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.
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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=(
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
@@ -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():
|
||||||
|
|||||||
@@ -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")
|
||||||
|
|||||||
@@ -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))
|
||||||
|
|||||||
Reference in New Issue
Block a user