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.
This commit is contained in:
Claude
2026-06-26 16:14:08 +08:00
parent 939f6842d3
commit b8826f8c90
5 changed files with 249 additions and 0 deletions
+74
View File
@@ -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"]
+94
View File
@@ -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")