From b8826f8c90c20c45863e3695ce7c09789116fe29 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 26 Jun 2026 16:14:08 +0800 Subject: [PATCH] feat(submit): update_pending_job tool for editing PENDING submissions Add an update_pending_job MCP tool that lets callers modify parameters of a PENDING submission before confirm_submit_job: - queue, executor_memory, executor_cores, num_executors - app_name - extra_args - script_path (re-validates file existence and SQL guard) Only PENDING submissions can be updated; SUBMITTED/CANCELLED/FAILED are rejected with HTTP 400. --- spark_executor/server.py | 15 +++++ spark_executor/tools/requests.py | 14 +++++ spark_executor/tools/submit.py | 52 +++++++++++++++ tests/integration/test_mcp_routes.py | 74 ++++++++++++++++++++++ tests/unit/test_submit_tool.py | 94 ++++++++++++++++++++++++++++ 5 files changed, 249 insertions(+) diff --git a/spark_executor/server.py b/spark_executor/server.py index a60a08e..e1bfe29 100644 --- a/spark_executor/server.py +++ b/spark_executor/server.py @@ -27,6 +27,7 @@ from spark_executor.tools.requests import ( ReadJobFileRequest, SaveConnectionRequest, UpdateJobFileRequest, + UpdatePendingJobRequest, ) from spark_executor.tools.status import get_job_status from spark_executor.tools.submit import ( @@ -35,6 +36,7 @@ from spark_executor.tools.submit import ( get_pending_job, list_pending_jobs, prepare_submit_job, + update_pending_job, ) from spark_executor.tools.result import get_job_result @@ -138,6 +140,19 @@ def _get_pending_job(req: PendingIdRequest): return get_pending_job(req.pending_id) +@app.post( + "/update_pending_job", + summary="Update an unsubmitted pending submission", + description=( + "Modify parameters of a PENDING submission before confirm_submit_job. " + "Only the provided fields are changed. If script_path is changed, the " + "new file must exist and pass the SQL guard." + ), +) +def _update_pending_job(req: UpdatePendingJobRequest): + return update_pending_job(**req.model_dump(exclude_none=True)) + + @app.post( "/cancel_pending_job", summary="Cancel a pending submission", diff --git a/spark_executor/tools/requests.py b/spark_executor/tools/requests.py index 1529b9f..54ca2a7 100644 --- a/spark_executor/tools/requests.py +++ b/spark_executor/tools/requests.py @@ -75,6 +75,20 @@ class PendingIdRequest(BaseModel): pending_id: str +class UpdatePendingJobRequest(BaseModel): + pending_id: str + script_path: str | None = Field( + default=None, + description="Optional new absolute path to the PySpark script. If provided, the file must exist and pass SQL guard.", + ) + queue: str | None = None + executor_memory: str | None = None + executor_cores: int | None = None + num_executors: int | None = None + app_name: str | None = None + extra_args: dict[str, str] | None = None + + class JobIdRequest(BaseModel): job_id: str diff --git a/spark_executor/tools/submit.py b/spark_executor/tools/submit.py index 75aa74f..5adfdde 100644 --- a/spark_executor/tools/submit.py +++ b/spark_executor/tools/submit.py @@ -287,3 +287,55 @@ def cancel_pending_job(pending_id: str) -> dict[str, str]: pending_store.save(p) logger.info(f"cancel_pending_job ok pending_id={pending_id}") return {"pending_id": pending_id, "status": "CANCELLED"} + + +def update_pending_job( + *, + pending_id: str, + script_path: str | None = None, + queue: str | None = None, + executor_memory: str | None = None, + executor_cores: int | None = None, + num_executors: int | None = None, + app_name: str | None = None, + extra_args: dict[str, str] | None = None, +) -> dict[str, object]: + logger.debug(f"update_pending_job enter pending_id={pending_id}") + pending = pending_store.get(pending_id) + if pending is None: + raise KeyError(f"Unknown pending_id: {pending_id}") + if pending.status != "PENDING": + raise ValueError( + f"pending_id {pending_id} is in status {pending.status!r}; " + f"only PENDING submissions can be updated" + ) + + if script_path is not None: + _check_script_path(script_path) + with open(script_path, encoding="utf-8") as f: + script_body = f.read() + offenses = validate_pyspark_code(script_body) + if offenses: + logger.warning( + f"update_pending_job rejected: SQL policy violation(s) in " + f"{script_path}: {offenses}" + ) + raise _sql_guard_error(offenses) + pending.script_path = script_path + + if queue is not None: + pending.queue = queue + if executor_memory is not None: + pending.executor_memory = executor_memory + if executor_cores is not None: + pending.executor_cores = executor_cores + if num_executors is not None: + pending.num_executors = num_executors + if app_name is not None: + pending.app_name = app_name + if extra_args is not None: + pending.extra_args = dict(extra_args) + + pending_store.save(pending) + logger.info(f"update_pending_job ok pending_id={pending_id}") + return {"pending_id": pending_id, "status": pending.status, "parameters": pending.model_dump()} diff --git a/tests/integration/test_mcp_routes.py b/tests/integration/test_mcp_routes.py index 1d6a6dd..1d13527 100644 --- a/tests/integration/test_mcp_routes.py +++ b/tests/integration/test_mcp_routes.py @@ -59,6 +59,11 @@ def test_sixteen_tool_routes_registered(): assert path in paths, f"missing MCP tool route: {path}" +def test_seventeen_tool_routes_registered(): + paths = {r.path for r in app.routes} + assert "/update_pending_job" in paths + + # --- End-to-end body-based calls (the gap the route-registration test missed) --- def test_save_connection_accepts_dict_spark_conf_in_body(): @@ -406,3 +411,72 @@ def test_cancel_already_submitted_returns_400(tmp_path): r = c.post("/cancel_pending_job", json={"pending_id": pid}) assert r.status_code == 400 assert "SUBMITTED" in r.json()["detail"] + + +# --- update_pending_job route --- + +def test_update_pending_job_route(tmp_path): + c = TestClient(app) + c.post("/save_connection", json={"name": "prod", "master": "yarn"}) + script = tmp_path / "x.py" + script.write_text("print('hi')\n") + prep = c.post( + "/prepare_submit_job", + json={ + "connection": "prod", + "script_path": str(script), + "queue": "default", + "executor_memory": "4G", + "executor_cores": 2, + "num_executors": 2, + "app_name": "test-app", + }, + ).json() + pid = prep["pending_id"] + + r = c.post( + "/update_pending_job", + json={"pending_id": pid, "queue": "research"}, + ) + assert r.status_code == 200, r.text + assert r.json()["status"] == "PENDING" + assert r.json()["parameters"]["queue"] == "research" + + got = c.post("/get_pending_job", json={"pending_id": pid}).json() + assert got["queue"] == "research" + # untouched fields are preserved + assert got["executor_memory"] == "4G" + assert got["app_name"] == "test-app" + + +def test_update_pending_job_route_rejects_submitted(tmp_path, monkeypatch): + c = TestClient(app) + c.post("/save_connection", json={"name": "prod", "master": "yarn"}) + script = tmp_path / "x.py" + script.write_text("print('hi')\n") + prep = c.post( + "/prepare_submit_job", + json={ + "connection": "prod", + "script_path": str(script), + "queue": "default", + "executor_memory": "4G", + "executor_cores": 2, + "num_executors": 2, + "app_name": "test-app", + }, + ).json() + pid = prep["pending_id"] + fake_proc = type("P", (), { + "returncode": 0, + "stderr": "tracking URL: http://rm:8088/proxy/application_17400000001/\n", + })() + monkeypatch.setattr(submit, "run_spark_submit", lambda _cmd: fake_proc) + c.post("/confirm_submit_job", json={"pending_id": pid}) + + r = c.post( + "/update_pending_job", + json={"pending_id": pid, "queue": "research"}, + ) + assert r.status_code == 400 + assert "only PENDING submissions can be updated" in r.json()["detail"] diff --git a/tests/unit/test_submit_tool.py b/tests/unit/test_submit_tool.py index a761a10..377bda5 100644 --- a/tests/unit/test_submit_tool.py +++ b/tests/unit/test_submit_tool.py @@ -594,3 +594,97 @@ def test_cancel_pending_job_refuses_submitted(monkeypatch, real_script): submit.pending_store.save(p) with pytest.raises(ValueError, match="SUBMITTED"): submit.cancel_pending_job(pid) + + +# --- update_pending_job --- + +def test_update_pending_updates_resource_params(monkeypatch, real_script): + monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None) + submit.prepare_submit_job( + connection="prod", + script_path=str(real_script), + queue="default", + executor_memory="4G", + executor_cores=2, + num_executors=2, + app_name="test-app", + ) + pid = _last_pending_id() + out = submit.update_pending_job( + pending_id=pid, + queue="research", + executor_memory="8G", + executor_cores=4, + num_executors=8, + app_name="renamed-app", + extra_args={"jars": "hdfs:///lib/bar.jar"}, + ) + assert out["pending_id"] == pid + assert out["status"] == "PENDING" + params = out["parameters"] + assert params["queue"] == "research" + assert params["executor_memory"] == "8G" + assert params["executor_cores"] == 4 + assert params["num_executors"] == 8 + assert params["app_name"] == "renamed-app" + assert params["extra_args"] == {"jars": "hdfs:///lib/bar.jar"} + # unchanged fields are preserved + assert params["connection"] == "prod" + assert params["script_path"] == str(real_script) + + +def test_update_pending_rejects_non_pending_status(monkeypatch, real_script): + monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None) + submit.prepare_submit_job( + connection="prod", + script_path=str(real_script), + queue="default", + executor_memory="4G", + executor_cores=2, + num_executors=2, + app_name="test-app", + ) + pid = _last_pending_id() + p = submit.pending_store.get(pid) + p.status = "SUBMITTED" + submit.pending_store.save(p) + with pytest.raises(ValueError, match="only PENDING submissions can be updated"): + submit.update_pending_job(pending_id=pid, queue="research") + + +def test_update_pending_rejects_bad_script_path(tmp_path, real_script): + submit.prepare_submit_job( + connection="prod", + script_path=str(real_script), + queue="default", + executor_memory="4G", + executor_cores=2, + num_executors=2, + app_name="test-app", + ) + pid = _last_pending_id() + missing = str(tmp_path / "does_not_exist.py") + with pytest.raises(ValueError, match="does not exist or is not a file"): + submit.update_pending_job(pending_id=pid, script_path=missing) + + +def test_update_pending_rejects_sql_violation_script(tmp_path, real_script): + submit.prepare_submit_job( + connection="prod", + script_path=str(real_script), + queue="default", + executor_memory="4G", + executor_cores=2, + num_executors=2, + app_name="test-app", + ) + pid = _last_pending_id() + bad_script = tmp_path / "evil.py" + bad_script.write_text('spark.sql("DROP TABLE users")\n') + with pytest.raises(ValueError, match="DROP"): + submit.update_pending_job(pending_id=pid, script_path=str(bad_script)) + + +def test_update_pending_unknown_id_raises(): + with pytest.raises(KeyError, match="missing"): + submit.update_pending_job(pending_id="missing", queue="research")