diff --git a/spark_executor/server.py b/spark_executor/server.py index 211447f..4e94e7e 100644 --- a/spark_executor/server.py +++ b/spark_executor/server.py @@ -13,6 +13,7 @@ from spark_executor.tools.connections import ( save_connection, ) from spark_executor.tools.generate import generate_job_file +from spark_executor.tools.job_file import read_job_file, update_job_file from spark_executor.tools.kill import kill_job from spark_executor.tools.logs import get_job_logs from spark_executor.tools.requests import ( @@ -23,7 +24,9 @@ from spark_executor.tools.requests import ( JobIdRequest, PendingIdRequest, PrepareSubmitJobRequest, + ReadJobFileRequest, SaveConnectionRequest, + UpdateJobFileRequest, ) from spark_executor.tools.status import get_job_status from spark_executor.tools.submit import ( @@ -226,3 +229,34 @@ def _delete_connection(req: ConnectionNameRequest): ) def _generate_job_file(req: GenerateJobFileRequest): return generate_job_file(req.code) + + +@app.post( + "/read_job_file", + summary="Read the contents of an existing PySpark script", + description=( + "Returns the text content of an existing script file at the given " + "path. Caps reads at 1 MB. Typical use: after generate_job_file " + "returns a path, call read_job_file on that path to inspect what " + "was actually written, before deciding to prepare_submit_job or " + "update_job_file." + ), +) +def _read_job_file(req: ReadJobFileRequest): + return read_job_file(req.script_path) + + +@app.post( + "/update_job_file", + summary="Overwrite an existing PySpark script with new content", + description=( + "Replaces the entire content of an existing script file. Path must " + "be under SPARK_EXECUTOR_JOBS_DIR (the dir generate_job_file writes " + "to) — protects against overwriting host-mounted configs or other " + "non-script files. Caps writes at 1 MB. Typical use: read_job_file, " + "edit the content (LLM or human), update_job_file, then " + "prepare_submit_job with the same path." + ), +) +def _update_job_file(req: UpdateJobFileRequest): + return update_job_file(req.script_path, req.content) diff --git a/spark_executor/tools/job_file.py b/spark_executor/tools/job_file.py new file mode 100644 index 0000000..fd542a3 --- /dev/null +++ b/spark_executor/tools/job_file.py @@ -0,0 +1,111 @@ +# coding=utf-8 +""" +@Time :2026/6/24 +@Author :tao.chen + +Read + update the contents of an existing PySpark script file. These two +tools close the review-and-edit loop: + + generate_job_file(code=...) -> {script_path} + read_job_file(script_path=...) -> {content, path} <-- inspect + update_job_file(path, content) -> {path, bytes_written} <-- edit + prepare_submit_job(path) -> {pending_id, ...} + +Safety: + - read_job_file: any existing regular file. Path-existence only. + - update_job_file: must be under SPARK_EXECUTOR_JOBS_DIR + (settings.jobs_dir) so the agent cannot overwrite host-mounted + configs or arbitrary files on the container FS. + - 1 MB cap on both read and write payloads to keep MCP responses bounded. +""" +import os +from pathlib import Path + +from common.config import settings +from common.logging import logger + +MAX_FILE_BYTES = 1 * 1024 * 1024 # 1 MB + + +class ScriptFileError(ValueError): + """Raised when read/update fails. -> HTTP 400 via the FastAPI ValueError + handler in server.py. + """ + pass + + +def _check_readable(script_path: str) -> None: + if not script_path or not os.path.isfile(script_path): + raise ScriptFileError( + f"script_path does not exist or is not a file: {script_path!r}" + ) + + +def _check_writable(script_path: str) -> None: + """update_job_file is restricted to files under settings.jobs_dir + (the same dir generate_job_file writes to). This prevents the agent + from overwriting arbitrary host-mounted files or the app's own code. + """ + if not script_path or not os.path.isfile(script_path): + raise ScriptFileError( + f"script_path does not exist or is not a file: {script_path!r}. " + f"update_job_file can only edit existing files. " + f"Use generate_job_file to create a new one." + ) + jobs_root = Path(settings.jobs_dir).resolve() + target = Path(script_path).resolve() + try: + target.relative_to(jobs_root) + except ValueError: + raise ScriptFileError( + f"script_path must be under {jobs_root} (the directory " + f"generate_job_file writes to). Got {script_path!r}. " + f"This restriction protects host-mounted configs and other " + f"non-script files from being overwritten by the agent." + ) + + +def read_job_file(script_path: str) -> dict[str, object]: + """Return the text content of an existing script file. + + Caps the read at 1 MB to keep MCP responses bounded; raises + ScriptFileError (-> 400) if the file is missing or too large. + """ + _check_readable(script_path) + size = os.path.getsize(script_path) + if size > MAX_FILE_BYTES: + raise ScriptFileError( + f"Script is too large to read back ({size} bytes > {MAX_FILE_BYTES} " + f"byte cap). Edit it via a host volume mount instead." + ) + logger.debug(f"read_job_file enter script_path={script_path} size={size}") + with open(script_path, encoding="utf-8") as f: + content = f.read() + logger.info(f"read_job_file ok script_path={script_path} size={size}") + return {"path": script_path, "content": content, "size": size} + + +def update_job_file(script_path: str, content: str) -> dict[str, object]: + """Overwrite an existing script file with new content. + + Restricted to paths under settings.jobs_dir. Caps writes at 1 MB. + Raises ScriptFileError (-> 400) if the path is missing, outside + the allowed dir, or the content is too large. + """ + _check_writable(script_path) + encoded_size = len(content.encode("utf-8")) + if encoded_size > MAX_FILE_BYTES: + raise ScriptFileError( + f"content is too large ({encoded_size} bytes > {MAX_FILE_BYTES} " + f"byte cap). Split the script into multiple files." + ) + logger.debug( + f"update_job_file enter script_path={script_path} " + f"new_bytes={encoded_size}" + ) + with open(script_path, "w", encoding="utf-8") as f: + written = f.write(content) + logger.info( + f"update_job_file ok script_path={script_path} bytes_written={written}" + ) + return {"path": script_path, "bytes_written": written} diff --git a/spark_executor/tools/requests.py b/spark_executor/tools/requests.py index 12ce3cc..f7d5eb2 100644 --- a/spark_executor/tools/requests.py +++ b/spark_executor/tools/requests.py @@ -70,3 +70,32 @@ class GenerateJobFileRequest(BaseModel): "returned path." ), ) + + +class ReadJobFileRequest(BaseModel): + script_path: str = Field( + ..., + description=( + "Absolute path to a PySpark script inside the container's " + "filesystem. Must point at an existing regular file." + ), + ) + + +class UpdateJobFileRequest(BaseModel): + script_path: str = Field( + ..., + description=( + "Absolute path to an existing PySpark script inside the " + "container's filesystem. Must be under SPARK_EXECUTOR_JOBS_DIR " + "(the same dir generate_job_file writes to) — protects against " + "overwriting host-mounted configs or other critical files." + ), + ) + content: str = Field( + ..., + description=( + "New file content (replaces the file in full; no merge/diff). " + "Maximum 1 MB to keep the MCP response bounded." + ), + ) diff --git a/tests/integration/test_mcp_routes.py b/tests/integration/test_mcp_routes.py index 9c0995e..973b4e2 100644 --- a/tests/integration/test_mcp_routes.py +++ b/tests/integration/test_mcp_routes.py @@ -33,7 +33,7 @@ def test_health_still_present(): assert r.json() == {"status": "ok"} -def test_thirteen_tool_routes_registered(): +def test_sixteen_tool_routes_registered(): paths = {r.path for r in app.routes} for path in ( # pending-submission flow (5) @@ -51,8 +51,10 @@ def test_thirteen_tool_routes_registered(): "/list_connections", "/get_connection", "/delete_connection", - # LLM-driven PySpark generation (1) — Stage 2 + # LLM-driven PySpark generation (3) — Stage 2 "/generate_job_file", + "/read_job_file", + "/update_job_file", ): assert path in paths, f"missing MCP tool route: {path}" @@ -144,6 +146,52 @@ def test_generate_rejects_sql_policy_violation_with_400(tmp_path): assert list(tmp_path.glob("*.py")) == [] +# --- Stage 2: read_job_file / update_job_file --- + +def test_read_job_file_returns_content_via_mcp(tmp_path): + c = TestClient(app) + script = tmp_path / "demo.py" + script.write_text("print('hello from read_job_file')\n") + r = c.post("/read_job_file", json={"script_path": str(script)}) + assert r.status_code == 200 + body = r.json() + assert body["content"] == "print('hello from read_job_file')\n" + assert body["path"] == str(script) + + +def test_read_job_file_404_for_missing_file(tmp_path): + c = TestClient(app) + r = c.post("/read_job_file", json={"script_path": str(tmp_path / "nope.py")}) + assert r.status_code == 400 + assert "does not exist" in r.json()["detail"] + + +def test_update_job_file_writes_and_reads_back(tmp_path): + from common import config + config.settings.jobs_dir = str(tmp_path) + c = TestClient(app) + script = tmp_path / "edit.py" + script.write_text("v1\n") + # Update + r1 = c.post("/update_job_file", json={"script_path": str(script), "content": "v2\n"}) + assert r1.status_code == 200 + assert r1.json()["bytes_written"] == 3 + # Read back + r2 = c.post("/read_job_file", json={"script_path": str(script)}) + assert r2.json()["content"] == "v2\n" + + +def test_update_job_file_rejects_paths_outside_jobs_dir(tmp_path): + from common import config + config.settings.jobs_dir = str(tmp_path / "jobs") + c = TestClient(app) + other = tmp_path / "elsewhere.py" + other.write_text("x\n") + r = c.post("/update_job_file", json={"script_path": str(other), "content": "y\n"}) + assert r.status_code == 400 + assert "must be under" in r.json()["detail"] + + def test_list_and_get_pending_job_roundtrip_via_body(tmp_path): c = TestClient(app) c.post("/save_connection", json={"name": "prod", "master": "yarn"}) diff --git a/tests/unit/test_job_file.py b/tests/unit/test_job_file.py new file mode 100644 index 0000000..2a56719 --- /dev/null +++ b/tests/unit/test_job_file.py @@ -0,0 +1,162 @@ +# coding=utf-8 +import os +from pathlib import Path + +import pytest + +from common import config +from spark_executor.tools import job_file +from spark_executor.tools.job_file import ( + ScriptFileError, + read_job_file, + update_job_file, +) + + +@pytest.fixture(autouse=True) +def _restore_settings(): + snapshot = config.Settings( + data_dir=config.settings.data_dir, + jobs_dir=config.settings.jobs_dir, + log_dir=config.settings.log_dir, + yarn_resource_manager_url=config.settings.yarn_resource_manager_url, + log_level=config.settings.log_level, + ) + yield + config.settings.data_dir = snapshot.data_dir + config.settings.jobs_dir = snapshot.jobs_dir + config.settings.log_dir = snapshot.log_dir + config.settings.yarn_resource_manager_url = snapshot.yarn_resource_manager_url + config.settings.log_level = snapshot.log_level + + +# --- read_job_file --- + +def test_read_returns_existing_content(tmp_path: Path): + p = tmp_path / "demo.py" + p.write_text("print('hello')\n") + out = read_job_file(str(p)) + assert out["content"] == "print('hello')\n" + assert out["path"] == str(p) + assert out["size"] == len("print('hello')\n") + + +def test_read_works_for_any_existing_file(tmp_path: Path): + """read_job_file has no jobs_dir restriction — it's read-only and the + agent needs to inspect arbitrary files (e.g. logs/, /etc/hadoop/conf).""" + p = tmp_path / "anything.py" + p.write_text("x = 1\n") + out = read_job_file(str(p)) + assert out["content"] == "x = 1\n" + + +def test_read_raises_for_missing_file(tmp_path: Path): + with pytest.raises(ScriptFileError, match="does not exist"): + read_job_file(str(tmp_path / "nope.py")) + + +def test_read_raises_for_directory(tmp_path: Path): + with pytest.raises(ScriptFileError, match="does not exist"): + read_job_file(str(tmp_path)) + + +def test_read_raises_for_empty_path(): + with pytest.raises(ScriptFileError, match="does not exist"): + read_job_file("") + + +# --- update_job_file --- + +def test_update_overwrites_existing_file_in_jobs_dir(tmp_path: Path): + """The file must be under settings.jobs_dir (default './data/jobs').""" + config.settings.jobs_dir = str(tmp_path) + p = tmp_path / "script.py" + p.write_text("v1\n") + out = update_job_file(str(p), "v2\n") + assert out["bytes_written"] == 3 + assert p.read_text() == "v2\n" + + +def test_update_rejects_paths_outside_jobs_dir(tmp_path: Path): + """update_job_file is sandboxed to settings.jobs_dir.""" + config.settings.jobs_dir = str(tmp_path / "jobs") + other = tmp_path / "elsewhere.py" + other.write_text("x\n") + with pytest.raises(ScriptFileError, match="must be under"): + update_job_file(str(other), "y\n") + + +def test_update_rejects_paths_outside_jobs_dir_via_existing_file(tmp_path: Path): + """Path validation: even when the file exists, an absolute path OUTSIDE + settings.jobs_dir is rejected. (In production, /etc/passwd would exist + in the container; here we synthesize the same scenario with a real + file outside the configured jobs_dir.)""" + config.settings.jobs_dir = str(tmp_path / "jobs") + (tmp_path / "jobs").mkdir() + outside = tmp_path / "hadoop-core-site.xml" + outside.write_text("\n") + with pytest.raises(ScriptFileError, match="must be under"): + update_job_file(str(outside), "tampered\n") + + +def test_update_rejects_relative_escape_attempt(tmp_path: Path): + """Path traversal via '..' must not bypass the jobs_dir check.""" + config.settings.jobs_dir = str(tmp_path / "jobs") + (tmp_path / "jobs").mkdir() + (tmp_path / "etc").mkdir() + secret = tmp_path / "etc" / "passwd" + secret.write_text("root:x:0:0:...\n") + # Try to reach the file via '../etc/passwd' relative to jobs_dir + sneaky = str(tmp_path / "jobs" / ".." / "etc" / "passwd") + with pytest.raises(ScriptFileError, match="must be under"): + update_job_file(sneaky, "tampered\n") + + +def test_update_rejects_missing_file(tmp_path: Path): + config.settings.jobs_dir = str(tmp_path) + with pytest.raises(ScriptFileError, match="does not exist"): + update_job_file(str(tmp_path / "nope.py"), "x\n") + + +def test_update_rejects_oversized_content(tmp_path: Path, monkeypatch): + """Cap at 1 MB so the MCP response stays bounded.""" + config.settings.jobs_dir = str(tmp_path) + p = tmp_path / "big.py" + p.write_text("# small\n") + # Synthesize 2 MB of content (don't actually write 2 MB to disk) + big = "x" * (2 * 1024 * 1024) + with pytest.raises(ScriptFileError, match="too large"): + update_job_file(str(p), big) + + +def test_update_rejects_1mb_plus_1_byte(tmp_path: Path): + """Exactly at the boundary: 1 MB + 1 byte must be rejected.""" + config.settings.jobs_dir = str(tmp_path) + p = tmp_path / "x.py" + p.write_text("# t\n") + just_over = "x" * (1024 * 1024 + 1) + with pytest.raises(ScriptFileError, match="too large"): + update_job_file(str(p), just_over) + + +# --- round-trip: write -> read -> update -> read --- + +def test_full_edit_cycle_via_tools(tmp_path: Path): + """Simulate the agent's review-and-edit loop end-to-end.""" + config.settings.jobs_dir = str(tmp_path) + path = tmp_path / "script.py" + # 1. Initial state: file doesn't exist (in real flow, generate_job_file + # would create it; we just create it directly here for the test) + path.write_text("# v1: initial\n") + # 2. Agent reads back + r1 = read_job_file(str(path)) + assert r1["content"] == "# v1: initial\n" + # 3. Agent edits (LLM or human) + edited = "# v2: edited by agent\n" + r1["content"].split("\n", 1)[1] + # 4. Agent writes back + w = update_job_file(str(path), edited) + assert w["bytes_written"] == len(edited) + # 5. Read back to verify + r2 = read_job_file(str(path)) + assert r2["content"] == edited + assert "v2: edited" in r2["content"]