diff --git a/spark_executor/core/connection_store.py b/spark_executor/core/connection_store.py new file mode 100644 index 0000000..503f67b --- /dev/null +++ b/spark_executor/core/connection_store.py @@ -0,0 +1,85 @@ +# 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 Connection + +DEFAULT_DATA_DIR = os.environ.get("SPARK_EXECUTOR_DATA_DIR", "./data") +DEFAULT_FILE_NAME = "connections.json" + + +class ConnectionNotFound(KeyError): + """Raised when a named connection does not exist.""" + + +class ConnectionStore: + """JSON-backed CRUD for Connection records. + + File path: /. + Writes are atomic: tempfile + os.replace. + """ + + 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, Connection]: + if not self.path.exists(): + return {} + raw = json.loads(self.path.read_text(encoding="utf-8")) + return {name: Connection.model_validate(c) for name, c in raw.items()} + + def _dump(self, records: dict[str, Connection]) -> None: + os.makedirs(self._data_dir, exist_ok=True) + payload = {name: c.model_dump() for name, c in records.items()} + # atomic write: tempfile in same dir, then replace + 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) + os.replace(tmp_path, self.path) + except Exception: + if os.path.exists(tmp_path): + os.unlink(tmp_path) + raise + + def list_all(self) -> list[Connection]: + with self._lock: + return list(self._load().values()) + + def get(self, name: str) -> Connection | None: + with self._lock: + return self._load().get(name) + + def save(self, conn: Connection) -> None: + with self._lock: + records = self._load() + records[conn.name] = conn + self._dump(records) + + def delete(self, name: str) -> bool: + with self._lock: + records = self._load() + if name not in records: + return False + del records[name] + self._dump(records) + return True + + +# Module-level singleton; replaced in tests. +store = ConnectionStore() diff --git a/tests/unit/test_connection_store.py b/tests/unit/test_connection_store.py new file mode 100644 index 0000000..798b62c --- /dev/null +++ b/tests/unit/test_connection_store.py @@ -0,0 +1,74 @@ +# coding=utf-8 +import json +from pathlib import Path + +from spark_executor.core import connection_store +from spark_executor.models import Connection + + +def _conn(name: str) -> Connection: + return Connection(name=name, master="yarn") + + +def test_save_then_get_roundtrip(tmp_path: Path, monkeypatch): + monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path)) + monkeypatch.setattr(connection_store, "store", connection_store.ConnectionStore()) + connection_store.store.save(_conn("prod")) + assert connection_store.store.get("prod") is not None + assert connection_store.store.get("prod").master == "yarn" + + +def test_save_persists_to_json_file(tmp_path: Path, monkeypatch): + monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path)) + monkeypatch.setattr(connection_store, "store", connection_store.ConnectionStore()) + connection_store.store.save(_conn("prod")) + path = tmp_path / "connections.json" + assert path.is_file() + payload = json.loads(path.read_text()) + assert "prod" in payload + assert payload["prod"]["master"] == "yarn" + + +def test_get_missing_returns_none(tmp_path: Path, monkeypatch): + monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path)) + monkeypatch.setattr(connection_store, "store", connection_store.ConnectionStore()) + assert connection_store.store.get("missing") is None + + +def test_list_all_returns_empty_when_file_absent(tmp_path: Path, monkeypatch): + monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path)) + monkeypatch.setattr(connection_store, "store", connection_store.ConnectionStore()) + assert connection_store.store.list_all() == [] + + +def test_list_all_returns_saved_connections(tmp_path: Path, monkeypatch): + monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path)) + monkeypatch.setattr(connection_store, "store", connection_store.ConnectionStore()) + connection_store.store.save(_conn("a")) + connection_store.store.save(_conn("b")) + names = {c.name for c in connection_store.store.list_all()} + assert names == {"a", "b"} + + +def test_save_upserts_existing(tmp_path: Path, monkeypatch): + monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path)) + monkeypatch.setattr(connection_store, "store", connection_store.ConnectionStore()) + connection_store.store.save(_conn("prod")) + updated = Connection(name="prod", master="spark://new:7077", deploy_mode="client") + connection_store.store.save(updated) + assert connection_store.store.get("prod").master == "spark://new:7077" + assert connection_store.store.get("prod").deploy_mode == "client" + + +def test_delete_returns_true_when_present(tmp_path: Path, monkeypatch): + monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path)) + monkeypatch.setattr(connection_store, "store", connection_store.ConnectionStore()) + connection_store.store.save(_conn("prod")) + assert connection_store.store.delete("prod") is True + assert connection_store.store.get("prod") is None + + +def test_delete_returns_false_when_absent(tmp_path: Path, monkeypatch): + monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path)) + monkeypatch.setattr(connection_store, "store", connection_store.ConnectionStore()) + assert connection_store.store.delete("nope") is False