feat: add JSON-backed ConnectionStore

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: <DEFAULT_DATA_DIR>/<DEFAULT_FILE_NAME>.
+    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()

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