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:
@@ -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
|
||||
|
||||
|
||||
@@ -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()}
|
||||
|
||||
Reference in New Issue
Block a user