tao.chen/mcp-server
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(submit): require explicit confirmation for all prepare_submit_job params

Browse Source 
Remove defaults from queue, executor_memory, executor_cores,
num_executors, and app_name in prepare_submit_job. All must now be
explicitly confirmed by the caller.

Also add extra_args to the PendingSubmission snapshot so users can
confirm non-conf spark-submit flags (e.g. --jars, --py-files) at
prepare time; confirm_submit_job passes them through to
build_spark_submit_command.

This is an intentional breaking change to the MCP tool contract:
callers can no longer rely on implicit defaults.
This commit is contained in:
Claude
2026-06-26 15:12:14 +08:00
parent da9ba657f3
commit 4b07617af0
9 changed files with 330 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/unit/test_submit_tool.py
+189 -22
View File
@@ -45,7 +45,15 @@ def _last_pending_id() -> str:
def test_prepare_does_not_invoke_spark_submit(monkeypatch, real_script):
called = []
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: called.append(cmd))
out = submit.prepare_submit_job(connection="prod", script_path=str(real_script))
out = 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",
)
assert called == [] # never called
assert out["status"] == "PENDING"
assert out["pending_id"].startswith("p_")
@@ -62,7 +70,15 @@ def test_prepare_persists_pending_with_snapshot(monkeypatch, real_script):
spark_conf={"spark.sql.shuffle.partitions": "200"},
)
)
submit.prepare_submit_job(connection="prod", script_path=str(real_script), queue="research")
submit.prepare_submit_job(
connection="prod",
script_path=str(real_script),
queue="research",
executor_memory="4G",
executor_cores=2,
num_executors=2,
app_name="test-app",
)
p = submit.pending_store.get(_last_pending_id())
assert p is not None
assert p.connection == "prod"
@@ -77,24 +93,65 @@ def test_prepare_persists_pending_with_snapshot(monkeypatch, real_script):
def test_prepare_persists_app_name(monkeypatch, real_script):
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
submit.prepare_submit_job(
connection="prod", script_path=str(real_script), app_name="my-etl-job"
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_cores=2,
num_executors=2,
app_name="my-etl-job",
)
p = submit.pending_store.get(_last_pending_id())
assert p is not None
assert p.app_name == "my-etl-job"
def test_prepare_app_name_defaults_to_none(monkeypatch, real_script):
def test_prepare_requires_explicit_app_name(monkeypatch, real_script):
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
submit.prepare_submit_job(connection="prod", script_path=str(real_script))
p = submit.pending_store.get(_last_pending_id())
with pytest.raises(TypeError, match="app_name"):
submit.prepare_submit_job(connection="prod", script_path=str(real_script))
def test_prepare_snapshots_extra_args(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",
extra_args={"jars": "hdfs:///lib/foo.jar"},
)
pid = _last_pending_id()
p = submit.pending_store.get(pid)
assert p is not None
assert p.app_name is None
assert p.extra_args == {"jars": "hdfs:///lib/foo.jar"}
fake_proc = type("P", (), {
"returncode": 0,
"stderr": "tracking URL: http://rm:8088/proxy/application_17400000002/\n",
})()
with patch("spark_executor.tools.submit.run_spark_submit", return_value=fake_proc) as m:
submit.confirm_submit_job(pending_id=pid)
cmd = m.call_args.args[0]
assert ["--jars", "hdfs:///lib/foo.jar"] in [
cmd[i : i + 2] for i in range(len(cmd) - 1)
]
def test_prepare_raises_for_unknown_connection(real_script):
with pytest.raises(KeyError, match="missing"):
submit.prepare_submit_job(connection="missing", script_path=str(real_script))
submit.prepare_submit_job(
connection="missing",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_cores=2,
num_executors=2,
app_name="test-app",
)
# --- script_path validation ---
@@ -102,25 +159,55 @@ def test_prepare_raises_for_unknown_connection(real_script):
def test_prepare_rejects_nonexistent_script_path(tmp_path):
missing = str(tmp_path / "does_not_exist.py")
with pytest.raises(ValueError, match="does not exist or is not a file"):
submit.prepare_submit_job(connection="prod", script_path=missing)
submit.prepare_submit_job(
connection="prod",
script_path=missing,
queue="default",
executor_memory="4G",
executor_cores=2,
num_executors=2,
app_name="test-app",
)
def test_prepare_rejects_directory_as_script_path(tmp_path):
"""A directory is not a file, even if it exists."""
with pytest.raises(ValueError, match="does not exist or is not a file"):
submit.prepare_submit_job(connection="prod", script_path=str(tmp_path))
submit.prepare_submit_job(
connection="prod",
script_path=str(tmp_path),
queue="default",
executor_memory="4G",
executor_cores=2,
num_executors=2,
app_name="test-app",
)
def test_prepare_rejects_empty_script_path():
with pytest.raises(ValueError, match="does not exist or is not a file"):
submit.prepare_submit_job(connection="prod", script_path="")
submit.prepare_submit_job(
connection="prod",
script_path="",
queue="default",
executor_memory="4G",
executor_cores=2,
num_executors=2,
app_name="test-app",
)
def test_prepare_error_message_mentions_generate_job_file(tmp_path):
"""The agent must be told to call generate_job_file first."""
with pytest.raises(ValueError, match="generate_job_file"):
submit.prepare_submit_job(
connection="prod", script_path=str(tmp_path / "x.py")
connection="prod",
script_path=str(tmp_path / "x.py"),
queue="default",
executor_memory="4G",
executor_cores=2,
num_executors=2,
app_name="test-app",
)
@@ -128,7 +215,15 @@ def test_confirm_rejects_if_script_was_deleted_after_prepare(tmp_path, monkeypat
"""Defense in depth: file present at prepare, gone by confirm -> 400."""
script = tmp_path / "demo.py"
script.write_text("print('hi')\n")
submit.prepare_submit_job(connection="prod", script_path=str(script))
submit.prepare_submit_job(
connection="prod",
script_path=str(script),
queue="default",
executor_memory="4G",
executor_cores=2,
num_executors=2,
app_name="test-app",
)
pid = _last_pending_id()
# Simulate the file being deleted (e.g. cleanup job) between prepare and confirm
script.unlink()
@@ -139,7 +234,15 @@ def test_confirm_rejects_if_script_was_deleted_after_prepare(tmp_path, monkeypat
def test_prepare_snapshots_connection_at_prepare_time(monkeypatch, real_script):
"""Editing the connection between prepare and confirm must NOT silently retarget."""
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
submit.prepare_submit_job(connection="prod", script_path=str(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",
)
p = submit.pending_store.get(_last_pending_id())
# Now mutate the connection
connection_store.store.save(
@@ -154,7 +257,15 @@ def test_prepare_snapshots_connection_at_prepare_time(monkeypatch, real_script):
# --- confirm_submit_job ---
def test_confirm_invokes_spark_submit_and_marks_submitted(monkeypatch, real_script):
submit.prepare_submit_job(connection="prod", script_path=str(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()
fake_proc = type("P", (), {
"returncode": 0,
@@ -182,7 +293,15 @@ def test_confirm_raises_for_unknown_pending_id():
def test_confirm_refuses_non_pending_status(monkeypatch, real_script):
submit.prepare_submit_job(connection="prod", script_path=str(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()
# Mark it CANCELLED first
p = submit.pending_store.get(pid)
@@ -194,7 +313,15 @@ def test_confirm_refuses_non_pending_status(monkeypatch, real_script):
def test_confirm_marks_failed_on_spark_submit_error(monkeypatch, real_script):
from spark_executor.core.spark_submit import SparkSubmitError
submit.prepare_submit_job(connection="prod", script_path=str(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()
def _raise(_cmd):
raise SparkSubmitError("boom")
@@ -217,8 +344,24 @@ def test_list_pending_jobs_returns_all(monkeypatch, tmp_path):
b = tmp_path / "b.py"
a.write_text("a\n"); b.write_text("b\n")
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
submit.prepare_submit_job(connection="prod", script_path=str(a))
submit.prepare_submit_job(connection="prod", script_path=str(b))
submit.prepare_submit_job(
connection="prod",
script_path=str(a),
queue="default",
executor_memory="4G",
executor_cores=2,
num_executors=2,
app_name="test-app",
)
submit.prepare_submit_job(
connection="prod",
script_path=str(b),
queue="default",
executor_memory="4G",
executor_cores=2,
num_executors=2,
app_name="test-app",
)
out = submit.list_pending_jobs()
assert {p["script_path"] for p in out} == {str(a), str(b)}
assert all(p["status"] == "PENDING" for p in out)
@@ -228,7 +371,15 @@ def test_list_pending_jobs_returns_all(monkeypatch, tmp_path):
def test_get_pending_job_returns_dump(monkeypatch, real_script):
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
submit.prepare_submit_job(connection="prod", script_path=str(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()
out = submit.get_pending_job(pid)
assert out["pending_id"] == pid
@@ -245,7 +396,15 @@ def test_get_pending_job_unknown_raises():
def test_cancel_pending_job_marks_cancelled(monkeypatch, real_script):
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
submit.prepare_submit_job(connection="prod", script_path=str(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()
out = submit.cancel_pending_job(pid)
assert out == {"pending_id": pid, "status": "CANCELLED"}
@@ -258,7 +417,15 @@ def test_cancel_pending_job_unknown_raises():
def test_cancel_pending_job_refuses_submitted(monkeypatch, real_script):
submit.prepare_submit_job(connection="prod", script_path=str(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()
p = submit.pending_store.get(pid)
p.status = "SUBMITTED"