feat(submit): require explicit confirmation for all prepare_submit_job params
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:
@@ -86,7 +86,7 @@ class Connection(BaseModel):
|
|||||||
|
|
||||||
class PendingSubmission(BaseModel):
|
class PendingSubmission(BaseModel):
|
||||||
pending_id: str
|
pending_id: str
|
||||||
app_name: str | None = None
|
app_name: str
|
||||||
connection: str
|
connection: str
|
||||||
master: str
|
master: str
|
||||||
deploy_mode: str
|
deploy_mode: str
|
||||||
@@ -97,6 +97,7 @@ class PendingSubmission(BaseModel):
|
|||||||
executor_cores: int
|
executor_cores: int
|
||||||
num_executors: int
|
num_executors: int
|
||||||
spark_conf: dict[str, str] = Field(default_factory=dict)
|
spark_conf: dict[str, str] = Field(default_factory=dict)
|
||||||
|
extra_args: dict[str, str] = Field(default_factory=dict)
|
||||||
created_at: datetime
|
created_at: datetime
|
||||||
status: str = "PENDING" # PENDING | SUBMITTED | CANCELLED | FAILED
|
status: str = "PENDING" # PENDING | SUBMITTED | CANCELLED | FAILED
|
||||||
error: str | None = None
|
error: str | None = None
|
||||||
|
|||||||
@@ -33,9 +33,9 @@ class SaveConnectionRequest(BaseModel):
|
|||||||
|
|
||||||
class PrepareSubmitJobRequest(BaseModel):
|
class PrepareSubmitJobRequest(BaseModel):
|
||||||
connection: str
|
connection: str
|
||||||
app_name: str | None = Field(
|
app_name: str = Field(
|
||||||
default=None,
|
...,
|
||||||
description="Optional human-readable application name for tracking the pending submission.",
|
description="Human-readable application name for tracking the pending submission.",
|
||||||
)
|
)
|
||||||
script_path: str = Field(
|
script_path: str = Field(
|
||||||
...,
|
...,
|
||||||
@@ -49,10 +49,26 @@ class PrepareSubmitJobRequest(BaseModel):
|
|||||||
"if the path is missing or not a file."
|
"if the path is missing or not a file."
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
queue: str = "default"
|
queue: str = Field(
|
||||||
executor_memory: str = "4G"
|
...,
|
||||||
executor_cores: int = 2
|
description="YARN queue to submit to. Must be explicitly confirmed by the caller.",
|
||||||
num_executors: int = 2
|
)
|
||||||
|
executor_memory: str = Field(
|
||||||
|
...,
|
||||||
|
description="Executor memory, e.g. '4G'. Must be explicitly confirmed by the caller.",
|
||||||
|
)
|
||||||
|
executor_cores: int = Field(
|
||||||
|
...,
|
||||||
|
description="Number of cores per executor. Must be explicitly confirmed by the caller.",
|
||||||
|
)
|
||||||
|
num_executors: int = Field(
|
||||||
|
...,
|
||||||
|
description="Total number of executors. Must be explicitly confirmed by the caller.",
|
||||||
|
)
|
||||||
|
extra_args: dict[str, str] | None = Field(
|
||||||
|
default=None,
|
||||||
|
description="Additional spark-submit flags (e.g. jars, py-files) confirmed at prepare time.",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
class PendingIdRequest(BaseModel):
|
class PendingIdRequest(BaseModel):
|
||||||
|
|||||||
@@ -64,17 +64,18 @@ def prepare_submit_job(
|
|||||||
*,
|
*,
|
||||||
connection: str,
|
connection: str,
|
||||||
script_path: str,
|
script_path: str,
|
||||||
queue: str = "default",
|
queue: str,
|
||||||
executor_memory: str = "4G",
|
executor_memory: str,
|
||||||
executor_cores: int = 2,
|
executor_cores: int,
|
||||||
num_executors: int = 2,
|
num_executors: int,
|
||||||
app_name: str | None = None,
|
app_name: str,
|
||||||
|
extra_args: dict[str, str] | None = None,
|
||||||
) -> dict[str, object]:
|
) -> dict[str, object]:
|
||||||
"""Snapshot connection params and persist a PendingSubmission. Does NOT submit."""
|
"""Snapshot connection params and persist a PendingSubmission. Does NOT submit."""
|
||||||
logger.debug(
|
logger.debug(
|
||||||
f"prepare_submit_job enter connection={connection} script_path={script_path} "
|
f"prepare_submit_job enter connection={connection} script_path={script_path} "
|
||||||
f"queue={queue} executor_memory={executor_memory} executor_cores={executor_cores} "
|
f"queue={queue} executor_memory={executor_memory} executor_cores={executor_cores} "
|
||||||
f"num_executors={num_executors}"
|
f"num_executors={num_executors} app_name={app_name}"
|
||||||
)
|
)
|
||||||
# Order of checks matters for the error the agent sees:
|
# Order of checks matters for the error the agent sees:
|
||||||
# 1. Unknown connection -> 404 (KeyError -> 404 in server.py)
|
# 1. Unknown connection -> 404 (KeyError -> 404 in server.py)
|
||||||
@@ -121,6 +122,7 @@ def prepare_submit_job(
|
|||||||
executor_cores=executor_cores,
|
executor_cores=executor_cores,
|
||||||
num_executors=num_executors,
|
num_executors=num_executors,
|
||||||
spark_conf=dict(conn.spark_conf),
|
spark_conf=dict(conn.spark_conf),
|
||||||
|
extra_args=dict(extra_args or {}),
|
||||||
created_at=datetime.utcnow(),
|
created_at=datetime.utcnow(),
|
||||||
status="PENDING",
|
status="PENDING",
|
||||||
)
|
)
|
||||||
@@ -164,6 +166,7 @@ def confirm_submit_job(*, pending_id: str) -> SubmitResult:
|
|||||||
executor_cores=pending.executor_cores,
|
executor_cores=pending.executor_cores,
|
||||||
num_executors=pending.num_executors,
|
num_executors=pending.num_executors,
|
||||||
spark_conf=pending.spark_conf,
|
spark_conf=pending.spark_conf,
|
||||||
|
extra_args=pending.extra_args,
|
||||||
)
|
)
|
||||||
logger.info(
|
logger.info(
|
||||||
f"confirm_submit_job start pending_id={pending_id} "
|
f"confirm_submit_job start pending_id={pending_id} "
|
||||||
|
|||||||
@@ -86,7 +86,15 @@ def test_prepare_submit_job_works_via_body(tmp_path):
|
|||||||
script.write_text("print('hi')\n")
|
script.write_text("print('hi')\n")
|
||||||
r = c.post(
|
r = c.post(
|
||||||
"/prepare_submit_job",
|
"/prepare_submit_job",
|
||||||
json={"connection": "prod", "script_path": str(script), "queue": "research"},
|
json={
|
||||||
|
"connection": "prod",
|
||||||
|
"script_path": str(script),
|
||||||
|
"queue": "research",
|
||||||
|
"executor_memory": "4G",
|
||||||
|
"executor_cores": 2,
|
||||||
|
"num_executors": 2,
|
||||||
|
"app_name": "test-app",
|
||||||
|
},
|
||||||
)
|
)
|
||||||
assert r.status_code == 200, r.text
|
assert r.status_code == 200, r.text
|
||||||
body = r.json()
|
body = r.json()
|
||||||
@@ -102,7 +110,15 @@ def test_prepare_rejects_nonexistent_script_path_with_400():
|
|||||||
c.post("/save_connection", json={"name": "prod", "master": "yarn"})
|
c.post("/save_connection", json={"name": "prod", "master": "yarn"})
|
||||||
r = c.post(
|
r = c.post(
|
||||||
"/prepare_submit_job",
|
"/prepare_submit_job",
|
||||||
json={"connection": "prod", "script_path": "/nope/does_not_exist.py"},
|
json={
|
||||||
|
"connection": "prod",
|
||||||
|
"script_path": "/nope/does_not_exist.py",
|
||||||
|
"queue": "default",
|
||||||
|
"executor_memory": "4G",
|
||||||
|
"executor_cores": 2,
|
||||||
|
"num_executors": 2,
|
||||||
|
"app_name": "test-app",
|
||||||
|
},
|
||||||
)
|
)
|
||||||
assert r.status_code == 400
|
assert r.status_code == 400
|
||||||
detail = r.json()["detail"]
|
detail = r.json()["detail"]
|
||||||
@@ -121,7 +137,15 @@ def test_prepare_rejects_sql_policy_violation_with_400(tmp_path):
|
|||||||
bad_script.write_text('spark.sql("DROP TABLE users")\n')
|
bad_script.write_text('spark.sql("DROP TABLE users")\n')
|
||||||
r = c.post(
|
r = c.post(
|
||||||
"/prepare_submit_job",
|
"/prepare_submit_job",
|
||||||
json={"connection": "prod", "script_path": str(bad_script)},
|
json={
|
||||||
|
"connection": "prod",
|
||||||
|
"script_path": str(bad_script),
|
||||||
|
"queue": "default",
|
||||||
|
"executor_memory": "4G",
|
||||||
|
"executor_cores": 2,
|
||||||
|
"num_executors": 2,
|
||||||
|
"app_name": "test-app",
|
||||||
|
},
|
||||||
)
|
)
|
||||||
assert r.status_code == 400
|
assert r.status_code == 400
|
||||||
detail = r.json()["detail"]
|
detail = r.json()["detail"]
|
||||||
@@ -199,7 +223,15 @@ def test_list_and_get_pending_job_roundtrip_via_body(tmp_path):
|
|||||||
script.write_text("print('hi')\n")
|
script.write_text("print('hi')\n")
|
||||||
prep = c.post(
|
prep = c.post(
|
||||||
"/prepare_submit_job",
|
"/prepare_submit_job",
|
||||||
json={"connection": "prod", "script_path": str(script)},
|
json={
|
||||||
|
"connection": "prod",
|
||||||
|
"script_path": str(script),
|
||||||
|
"queue": "default",
|
||||||
|
"executor_memory": "4G",
|
||||||
|
"executor_cores": 2,
|
||||||
|
"num_executors": 2,
|
||||||
|
"app_name": "test-app",
|
||||||
|
},
|
||||||
).json()
|
).json()
|
||||||
pid = prep["pending_id"]
|
pid = prep["pending_id"]
|
||||||
|
|
||||||
@@ -261,7 +293,15 @@ def test_unknown_connection_in_prepare_returns_404():
|
|||||||
c = TestClient(app)
|
c = TestClient(app)
|
||||||
r = c.post(
|
r = c.post(
|
||||||
"/prepare_submit_job",
|
"/prepare_submit_job",
|
||||||
json={"connection": "nope", "script_path": "/tmp/x.py"},
|
json={
|
||||||
|
"connection": "nope",
|
||||||
|
"script_path": "/tmp/x.py",
|
||||||
|
"queue": "default",
|
||||||
|
"executor_memory": "4G",
|
||||||
|
"executor_cores": 2,
|
||||||
|
"num_executors": 2,
|
||||||
|
"app_name": "test-app",
|
||||||
|
},
|
||||||
)
|
)
|
||||||
assert r.status_code == 404
|
assert r.status_code == 404
|
||||||
assert "nope" in r.json()["detail"]
|
assert "nope" in r.json()["detail"]
|
||||||
@@ -275,7 +315,15 @@ def test_confirm_non_pending_returns_400(tmp_path):
|
|||||||
script.write_text("print('hi')\n")
|
script.write_text("print('hi')\n")
|
||||||
prep = c.post(
|
prep = c.post(
|
||||||
"/prepare_submit_job",
|
"/prepare_submit_job",
|
||||||
json={"connection": "prod", "script_path": str(script)},
|
json={
|
||||||
|
"connection": "prod",
|
||||||
|
"script_path": str(script),
|
||||||
|
"queue": "default",
|
||||||
|
"executor_memory": "4G",
|
||||||
|
"executor_cores": 2,
|
||||||
|
"num_executors": 2,
|
||||||
|
"app_name": "test-app",
|
||||||
|
},
|
||||||
).json()
|
).json()
|
||||||
pid = prep["pending_id"]
|
pid = prep["pending_id"]
|
||||||
c.post("/cancel_pending_job", json={"pending_id": pid})
|
c.post("/cancel_pending_job", json={"pending_id": pid})
|
||||||
@@ -292,7 +340,15 @@ def test_cancel_already_submitted_returns_400(tmp_path):
|
|||||||
script.write_text("print('hi')\n")
|
script.write_text("print('hi')\n")
|
||||||
prep = c.post(
|
prep = c.post(
|
||||||
"/prepare_submit_job",
|
"/prepare_submit_job",
|
||||||
json={"connection": "prod", "script_path": str(script)},
|
json={
|
||||||
|
"connection": "prod",
|
||||||
|
"script_path": str(script),
|
||||||
|
"queue": "default",
|
||||||
|
"executor_memory": "4G",
|
||||||
|
"executor_cores": 2,
|
||||||
|
"num_executors": 2,
|
||||||
|
"app_name": "test-app",
|
||||||
|
},
|
||||||
).json()
|
).json()
|
||||||
pid = prep["pending_id"]
|
pid = prep["pending_id"]
|
||||||
# Simulate a SUBMITTED state by mutating the pending directly
|
# Simulate a SUBMITTED state by mutating the pending directly
|
||||||
|
|||||||
@@ -47,7 +47,15 @@ def test_prepare_submit_job_emits_debug_and_info(log_capture, tmp_path):
|
|||||||
script = tmp_path / "demo.py"
|
script = tmp_path / "demo.py"
|
||||||
script.write_text("print('hi')\n")
|
script.write_text("print('hi')\n")
|
||||||
log_capture.truncate(0); log_capture.seek(0)
|
log_capture.truncate(0); log_capture.seek(0)
|
||||||
submit.prepare_submit_job(connection="prod", script_path=str(script), queue="research")
|
submit.prepare_submit_job(
|
||||||
|
connection="prod",
|
||||||
|
script_path=str(script),
|
||||||
|
queue="research",
|
||||||
|
executor_memory="4G",
|
||||||
|
executor_cores=2,
|
||||||
|
num_executors=2,
|
||||||
|
app_name="test-app",
|
||||||
|
)
|
||||||
text = log_capture.getvalue()
|
text = log_capture.getvalue()
|
||||||
assert "DEBUG|prepare_submit_job enter" in text
|
assert "DEBUG|prepare_submit_job enter" in text
|
||||||
assert "INFO|prepare_submit_job ok" in text
|
assert "INFO|prepare_submit_job ok" in text
|
||||||
|
|||||||
@@ -89,6 +89,7 @@ def test_connection_rejects_bad_master():
|
|||||||
def test_pending_submission_defaults_to_pending_status():
|
def test_pending_submission_defaults_to_pending_status():
|
||||||
p = PendingSubmission(
|
p = PendingSubmission(
|
||||||
pending_id="p_1",
|
pending_id="p_1",
|
||||||
|
app_name="test-app",
|
||||||
connection="prod",
|
connection="prod",
|
||||||
master="yarn",
|
master="yarn",
|
||||||
deploy_mode="cluster",
|
deploy_mode="cluster",
|
||||||
@@ -105,13 +106,14 @@ def test_pending_submission_defaults_to_pending_status():
|
|||||||
assert p.job_id is None
|
assert p.job_id is None
|
||||||
assert p.application_id is None
|
assert p.application_id is None
|
||||||
assert p.yarn_rm_url is None # default for connections without one set
|
assert p.yarn_rm_url is None # default for connections without one set
|
||||||
assert p.app_name is None
|
assert p.app_name == "test-app"
|
||||||
|
|
||||||
|
|
||||||
def test_pending_submission_carries_yarn_rm_url_snapshot():
|
def test_pending_submission_carries_yarn_rm_url_snapshot():
|
||||||
"""snapshotting yarn_rm_url lets prepare/confirm survive connection edits."""
|
"""snapshotting yarn_rm_url lets prepare/confirm survive connection edits."""
|
||||||
p = PendingSubmission(
|
p = PendingSubmission(
|
||||||
pending_id="p_x",
|
pending_id="p_x",
|
||||||
|
app_name="test-app",
|
||||||
connection="prod",
|
connection="prod",
|
||||||
master="yarn",
|
master="yarn",
|
||||||
deploy_mode="cluster",
|
deploy_mode="cluster",
|
||||||
@@ -122,6 +124,7 @@ def test_pending_submission_carries_yarn_rm_url_snapshot():
|
|||||||
executor_cores=2,
|
executor_cores=2,
|
||||||
num_executors=2,
|
num_executors=2,
|
||||||
spark_conf={},
|
spark_conf={},
|
||||||
|
extra_args={},
|
||||||
created_at=datetime(2026, 6, 24),
|
created_at=datetime(2026, 6, 24),
|
||||||
)
|
)
|
||||||
assert p.yarn_rm_url == "http://rm-prod:8088"
|
assert p.yarn_rm_url == "http://rm-prod:8088"
|
||||||
@@ -130,6 +133,7 @@ def test_pending_submission_carries_yarn_rm_url_snapshot():
|
|||||||
def test_pending_submission_can_record_outcome():
|
def test_pending_submission_can_record_outcome():
|
||||||
p = PendingSubmission(
|
p = PendingSubmission(
|
||||||
pending_id="p_2",
|
pending_id="p_2",
|
||||||
|
app_name="test-app",
|
||||||
connection="prod",
|
connection="prod",
|
||||||
master="yarn",
|
master="yarn",
|
||||||
deploy_mode="cluster",
|
deploy_mode="cluster",
|
||||||
@@ -139,6 +143,7 @@ def test_pending_submission_can_record_outcome():
|
|||||||
executor_cores=2,
|
executor_cores=2,
|
||||||
num_executors=2,
|
num_executors=2,
|
||||||
spark_conf={},
|
spark_conf={},
|
||||||
|
extra_args={},
|
||||||
created_at=datetime(2026, 6, 24),
|
created_at=datetime(2026, 6, 24),
|
||||||
status="SUBMITTED",
|
status="SUBMITTED",
|
||||||
job_id="j_abc",
|
job_id="j_abc",
|
||||||
|
|||||||
@@ -11,6 +11,7 @@ from spark_executor.models import PendingSubmission
|
|||||||
def _pending(pid: str = "p_1", created_at: datetime | None = None) -> PendingSubmission:
|
def _pending(pid: str = "p_1", created_at: datetime | None = None) -> PendingSubmission:
|
||||||
return PendingSubmission(
|
return PendingSubmission(
|
||||||
pending_id=pid,
|
pending_id=pid,
|
||||||
|
app_name="test",
|
||||||
connection="prod",
|
connection="prod",
|
||||||
master="yarn",
|
master="yarn",
|
||||||
deploy_mode="cluster",
|
deploy_mode="cluster",
|
||||||
@@ -20,6 +21,7 @@ def _pending(pid: str = "p_1", created_at: datetime | None = None) -> PendingSub
|
|||||||
executor_cores=2,
|
executor_cores=2,
|
||||||
num_executors=2,
|
num_executors=2,
|
||||||
spark_conf={},
|
spark_conf={},
|
||||||
|
extra_args={},
|
||||||
created_at=created_at or datetime(2026, 6, 24),
|
created_at=created_at or datetime(2026, 6, 24),
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -80,10 +82,10 @@ def test_delete_removes_from_date_file_and_deletes_empty_file():
|
|||||||
|
|
||||||
def _legacy_record_text() -> str:
|
def _legacy_record_text() -> str:
|
||||||
return (
|
return (
|
||||||
'{"p_legacy": {"pending_id": "p_legacy", "connection": "prod", '
|
'{"p_legacy": {"pending_id": "p_legacy", "app_name": "legacy", '
|
||||||
'"master": "yarn", "deploy_mode": "cluster", "script_path": "/tmp/j.py", '
|
'"connection": "prod", "master": "yarn", "deploy_mode": "cluster", '
|
||||||
'"queue": "default", "executor_memory": "4G", "executor_cores": 2, '
|
'"script_path": "/tmp/j.py", "queue": "default", "executor_memory": "4G", '
|
||||||
'"num_executors": 2, "spark_conf": {}, '
|
'"executor_cores": 2, "num_executors": 2, "spark_conf": {}, '
|
||||||
'"created_at": "2026-06-23T00:00:00"}}'
|
'"created_at": "2026-06-23T00:00:00"}}'
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,23 @@
|
|||||||
|
# coding=utf-8
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from spark_executor.tools.requests import PrepareSubmitJobRequest
|
||||||
|
|
||||||
|
|
||||||
|
def test_prepare_submit_job_request_requires_all_confirmed_parameters():
|
||||||
|
with pytest.raises(ValueError):
|
||||||
|
PrepareSubmitJobRequest(connection="prod", script_path="/tmp/x.py")
|
||||||
|
|
||||||
|
|
||||||
|
def test_prepare_submit_job_request_accepts_extra_args():
|
||||||
|
req = PrepareSubmitJobRequest(
|
||||||
|
connection="prod",
|
||||||
|
script_path="/tmp/x.py",
|
||||||
|
queue="default",
|
||||||
|
executor_memory="4G",
|
||||||
|
executor_cores=2,
|
||||||
|
num_executors=2,
|
||||||
|
app_name="test-app",
|
||||||
|
extra_args={"jars": "hdfs:///lib/foo.jar"},
|
||||||
|
)
|
||||||
|
assert req.extra_args == {"jars": "hdfs:///lib/foo.jar"}
|
||||||
+188
-21
@@ -45,7 +45,15 @@ def _last_pending_id() -> str:
|
|||||||
def test_prepare_does_not_invoke_spark_submit(monkeypatch, real_script):
|
def test_prepare_does_not_invoke_spark_submit(monkeypatch, real_script):
|
||||||
called = []
|
called = []
|
||||||
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: called.append(cmd))
|
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 called == [] # never called
|
||||||
assert out["status"] == "PENDING"
|
assert out["status"] == "PENDING"
|
||||||
assert out["pending_id"].startswith("p_")
|
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"},
|
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())
|
p = submit.pending_store.get(_last_pending_id())
|
||||||
assert p is not None
|
assert p is not None
|
||||||
assert p.connection == "prod"
|
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):
|
def test_prepare_persists_app_name(monkeypatch, real_script):
|
||||||
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
|
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
|
||||||
submit.prepare_submit_job(
|
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())
|
p = submit.pending_store.get(_last_pending_id())
|
||||||
assert p is not None
|
assert p is not None
|
||||||
assert p.app_name == "my-etl-job"
|
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)
|
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
|
||||||
|
with pytest.raises(TypeError, match="app_name"):
|
||||||
submit.prepare_submit_job(connection="prod", script_path=str(real_script))
|
submit.prepare_submit_job(connection="prod", script_path=str(real_script))
|
||||||
p = submit.pending_store.get(_last_pending_id())
|
|
||||||
|
|
||||||
|
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 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):
|
def test_prepare_raises_for_unknown_connection(real_script):
|
||||||
with pytest.raises(KeyError, match="missing"):
|
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 ---
|
# --- 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):
|
def test_prepare_rejects_nonexistent_script_path(tmp_path):
|
||||||
missing = str(tmp_path / "does_not_exist.py")
|
missing = str(tmp_path / "does_not_exist.py")
|
||||||
with pytest.raises(ValueError, match="does not exist or is not a file"):
|
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):
|
def test_prepare_rejects_directory_as_script_path(tmp_path):
|
||||||
"""A directory is not a file, even if it exists."""
|
"""A directory is not a file, even if it exists."""
|
||||||
with pytest.raises(ValueError, match="does not exist or is not a file"):
|
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():
|
def test_prepare_rejects_empty_script_path():
|
||||||
with pytest.raises(ValueError, match="does not exist or is not a file"):
|
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):
|
def test_prepare_error_message_mentions_generate_job_file(tmp_path):
|
||||||
"""The agent must be told to call generate_job_file first."""
|
"""The agent must be told to call generate_job_file first."""
|
||||||
with pytest.raises(ValueError, match="generate_job_file"):
|
with pytest.raises(ValueError, match="generate_job_file"):
|
||||||
submit.prepare_submit_job(
|
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."""
|
"""Defense in depth: file present at prepare, gone by confirm -> 400."""
|
||||||
script = tmp_path / "demo.py"
|
script = tmp_path / "demo.py"
|
||||||
script.write_text("print('hi')\n")
|
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()
|
pid = _last_pending_id()
|
||||||
# Simulate the file being deleted (e.g. cleanup job) between prepare and confirm
|
# Simulate the file being deleted (e.g. cleanup job) between prepare and confirm
|
||||||
script.unlink()
|
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):
|
def test_prepare_snapshots_connection_at_prepare_time(monkeypatch, real_script):
|
||||||
"""Editing the connection between prepare and confirm must NOT silently retarget."""
|
"""Editing the connection between prepare and confirm must NOT silently retarget."""
|
||||||
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
|
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())
|
p = submit.pending_store.get(_last_pending_id())
|
||||||
# Now mutate the connection
|
# Now mutate the connection
|
||||||
connection_store.store.save(
|
connection_store.store.save(
|
||||||
@@ -154,7 +257,15 @@ def test_prepare_snapshots_connection_at_prepare_time(monkeypatch, real_script):
|
|||||||
# --- confirm_submit_job ---
|
# --- confirm_submit_job ---
|
||||||
|
|
||||||
def test_confirm_invokes_spark_submit_and_marks_submitted(monkeypatch, real_script):
|
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()
|
pid = _last_pending_id()
|
||||||
fake_proc = type("P", (), {
|
fake_proc = type("P", (), {
|
||||||
"returncode": 0,
|
"returncode": 0,
|
||||||
@@ -182,7 +293,15 @@ def test_confirm_raises_for_unknown_pending_id():
|
|||||||
|
|
||||||
|
|
||||||
def test_confirm_refuses_non_pending_status(monkeypatch, real_script):
|
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()
|
pid = _last_pending_id()
|
||||||
# Mark it CANCELLED first
|
# Mark it CANCELLED first
|
||||||
p = submit.pending_store.get(pid)
|
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):
|
def test_confirm_marks_failed_on_spark_submit_error(monkeypatch, real_script):
|
||||||
from spark_executor.core.spark_submit import SparkSubmitError
|
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()
|
pid = _last_pending_id()
|
||||||
def _raise(_cmd):
|
def _raise(_cmd):
|
||||||
raise SparkSubmitError("boom")
|
raise SparkSubmitError("boom")
|
||||||
@@ -217,8 +344,24 @@ def test_list_pending_jobs_returns_all(monkeypatch, tmp_path):
|
|||||||
b = tmp_path / "b.py"
|
b = tmp_path / "b.py"
|
||||||
a.write_text("a\n"); b.write_text("b\n")
|
a.write_text("a\n"); b.write_text("b\n")
|
||||||
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
|
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
|
||||||
submit.prepare_submit_job(connection="prod", script_path=str(a))
|
submit.prepare_submit_job(
|
||||||
submit.prepare_submit_job(connection="prod", script_path=str(b))
|
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()
|
out = submit.list_pending_jobs()
|
||||||
assert {p["script_path"] for p in out} == {str(a), str(b)}
|
assert {p["script_path"] for p in out} == {str(a), str(b)}
|
||||||
assert all(p["status"] == "PENDING" for p in out)
|
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):
|
def test_get_pending_job_returns_dump(monkeypatch, real_script):
|
||||||
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
|
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()
|
pid = _last_pending_id()
|
||||||
out = submit.get_pending_job(pid)
|
out = submit.get_pending_job(pid)
|
||||||
assert out["pending_id"] == 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):
|
def test_cancel_pending_job_marks_cancelled(monkeypatch, real_script):
|
||||||
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
|
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()
|
pid = _last_pending_id()
|
||||||
out = submit.cancel_pending_job(pid)
|
out = submit.cancel_pending_job(pid)
|
||||||
assert out == {"pending_id": pid, "status": "CANCELLED"}
|
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):
|
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()
|
pid = _last_pending_id()
|
||||||
p = submit.pending_store.get(pid)
|
p = submit.pending_store.get(pid)
|
||||||
p.status = "SUBMITTED"
|
p.status = "SUBMITTED"
|
||||||
|
|||||||
Reference in New Issue
Block a user