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

refactor(mcp): rename generate_job_file to write_job_file to match what it does

Browse Source 
The MCP tool was named `generate_job_file` from Stage 2 but it does
NOT generate PySpark code — the calling LLM writes the code in its own
context, and this tool only persists it to a file under
SPARK_EXECUTOR_JOBS_DIR so `spark-submit` can see it. The misleading
`generate_` prefix sent agents (and humans) looking for a code
generator that doesn't exist.

This commit folds three related polish changes into one (split later
with rebase -i if you want them as separate history):

  1. The rename itself:
     - `tools/generate.py`  →  `tools/write_job.py`
     - `generate_job_file`  →  `write_job_file`
     - `GenerateJobFileRequest`  →  `WriteJobFileRequest`
     - `/generate_job_file` route  →  `/write_job_file`
     - `operation_id="generate_job_file"`  →  `operation_id="write_job_file"`
     The internal helper `core.job_writer.write_job_file` (which just
     writes bytes to disk with no SQL guard) is imported with an
     `_write_to_disk` alias to avoid the name collision with the
     MCP-exposed function in the same module.
     The description for the tool now explicitly states 'this tool
     does NOT generate PySpark code. The calling LLM is expected to
     have already written the code; this tool only persists it.'

  2. Skill for LLM agents operating the service
     (`docs/superpowers/skills/spark-executor-mcp-operate/SKILL.md`,
     449 lines). Covers the 16 tools, the two-step prepare/confirm
     flow, the dual-ID contract (job_id vs application_id), the
     PendingSubmission state machine, the Connection profile, the
     job-file workflow, the error reference, common pitfalls, and a
     full end-to-end word-count example.

  3. Default `executor_memory` lowered 4G → 2G
     (`_DEFAULTS_TO_CONFIRM` in `server.py`). Mirrors the matching
     change in `test_mcp_routes.py` and the 5 unit tests that
     reference the default. Aligns with the lighter workloads the
     service is sized for in its current container profile.

Also tracked in git for the first time:
  - `docs/superpowers/plans/2026-06-24-spark-executor-mcp.md`
    (the original Stage 1/2/3 design plan, updated to use the new
    tool name throughout).

Test rename:
  - `tests/unit/test_generate_tool.py`  →  `test_write_job_tool.py`
  - the new test file picks up an extra assertion that the SQL guard
    rejects a `DROP TABLE` statement at write time.

243 tests pass (was 242; +1 new SQL-guard assertion). Zero regressions.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Claude
2026-06-29 19:13:29 +08:00
co-authored by Claude Fable 5
parent 0fef77c0b8
commit 523e6a9c76
18 changed files with 3661 additions and 154 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CLAUDE.md
+1 -1
View File
@@ -94,5 +94,5 @@ docs/superpowers/plans/ # implementation plans
## Stage Status
- **Stage 1: complete** — 12 MCP tools, 2-step submit, connection management, structured logging, Pydantic body models, exception handlers. 22 commits on `feat/stage-1`.
- **Stage 2: not started** — `generate_job_file` for LLM-written PySpark code (`job_writer` + one new tool).
- **Stage 2: not started** — `write_job_file` for LLM-written PySpark code (`job_writer` + one new tool).
- **Stage 3: deferred** — async submission, SQLite `JobRegistry`, status poller, offset-based `LogCache`, multi-tenant `owner`. Plan lives at `docs/superpowers/plans/2026-06-24-spark-executor-mcp.md`.
docs/superpowers/plans/2026-06-24-spark-executor-mcp.md
+3027
View File
File diff suppressed because it is too large Load Diff
docs/superpowers/skills/spark-executor-mcp-operate/SKILL.md
+449
View File
@@ -0,0 +1,449 @@
---
name: spark-executor-mcp-operate
description: Operation guide for the spark-executor-mcp service — 16 MCP tools that submit, monitor, fetch logs from, and kill PySpark jobs on YARN, plus connection-profile and script-file management. USE THIS SKILL whenever the user wants to do anything with the spark-executor-mcp MCP service: submit a Spark/PySpark job to YARN, check job status or result, fetch container logs, kill a running job, save/list/get/delete a YARN connection, or write/read/update a PySpark script file. Also use it when the user asks generally about Spark, YARN, pyspark, cluster jobs, or distributed compute via this service. Do NOT guess tool names or invent a submit flow — read this skill first to learn the two-step prepare/confirm flow, the dual-ID contract (job_id vs application_id), and the SQL guard / path restrictions. The MCP service does NOT generate PySpark code — the calling LLM writes the code, then uses `write_job_file` to persist it on disk before submitting.
---
# spark-executor-mcp 操作指南
Spark-executor-mcp is an MCP (Model Context Protocol) service that exposes
**16 tools** for managing PySpark jobs on YARN. It runs as a FastAPI app
mounted at `/spark-executor-mcp`; clients (typically LLM agents) call the
tools via the standard MCP `tools/call` flow.
This skill is the canonical reference for an LLM agent using the service.
It does **not** describe implementation details — for that, see the source
in `spark_executor/`.
---
## 0. Mindset before you start
Three things you MUST internalize before calling any tool:
1. **The submit flow is two-step: `prepare_submit_job` then `confirm_submit_job`.**
Prepare is cheap and reversible (just snapshots params + writes JSON to
disk). Confirm is irreversible — it actually invokes `spark-submit` and
a YARN application starts running. Never skip the prepare step, even
when the user is in a hurry.
2. **Every Job has two IDs that are NOT interchangeable everywhere.**
- `job_id` — local, 12-char hex (`a1b2c3d4e5f6`), generated by the
service. Used for: looking up the Job record, all MCP tool calls.
- `application_id` — YARN's own ID (`application_17400000001_0001`).
Used for: YARN UI, `tracking_url`, the YARN REST API.
The service's 4 job-lifecycle tools accept **either** ID (after the
2026-06-29 fix), but other systems (YARN UI, kubectl-style scripts)
only know `application_id`. **Always store both** from every
`confirm_submit_job` response.
3. **The LLM writes the PySpark code. The MCP service writes the file.**
`write_job_file` is a file writer with SQL-safety checks, not a
code generator. You compose the code in your context, then call
`write_job_file(code=...)` to drop it on the container's
filesystem where `spark-submit` can see it.
---
## 1. Tool index
| Operation ID | HTTP | Group | Purpose |
|---|---|---|---|
| `save_connection` | POST | Connection | Upsert a named YARN connection (master, deploy_mode, spark_conf, …) |
| `list_connections` | POST | Connection | List all saved connections |
| `get_connection` | POST | Connection | Fetch one connection by name |
| `delete_connection` | POST | Connection | Delete a connection by name |
| `write_job_file` | POST | Job file | Write LLM-written PySpark code to disk |
| `read_job_file` | POST | Job file | Read a script file back (capped at 1 MB) |
| `update_job_file` | POST | Job file | Overwrite a script file (capped at 1 MB) |
| `prepare_submit_job` | POST | Submit | Snapshot a job to `PendingSubmission`; **does NOT submit** |
| `confirm_submit_job` | POST | Submit | Actually invoke `spark-submit` for a `pending_id` |
| `update_pending_job` | POST | Submit | Edit a `PENDING` submission (script_path, queue, …) |
| `cancel_pending_job` | POST | Submit | Cancel a `PENDING` submission |
| `list_pending_jobs` | POST | Submit | List every PendingSubmission regardless of status |
| `get_pending_job` | POST | Submit | Fetch one PendingSubmission by id |
| `get_job_status` | POST | Job lifecycle | Query YARN for a job's live state (RUNNING / SUCCEEDED / …) |
| `get_job_result` | POST | Job lifecycle | Terminal-oriented view: final_status, diagnostics, timing |
| `get_job_logs` | POST | Job lifecycle | Aggregated container logs (default last 5000 chars) |
| `kill_job` | POST | Job lifecycle | PUT state=KILLED to YARN REST for the app |
(Total: 16 MCP tools + `/health` GET, which is not exposed as a tool.)
---
## 2. The canonical happy path: submit a PySpark job
```
1. save_connection # one-time, per cluster
2. write_job_file # write the LLM's PySpark to disk
3. read_job_file # (optional) sanity-check what was written
4. prepare_submit_job # snapshot + persist; no YARN call yet
5. confirm_submit_job # NOW spark-submit runs
6. get_job_status / logs # poll while running
7. get_job_result # terminal view when done
```
### 2.1 Save a connection (one-time per cluster)
```python
save_connection(
name="prod",
master="yarn",
deploy_mode="cluster",
yarn_rm_url="http://rm.example.com:8088",
spark_conf={"spark.sql.shuffle.partitions": "200"},
# optional: auth_type="kerberos", ssl_verify=False, ...
)
```
**`name`** is the lookup key — keep it short and stable. Re-calling
`save_connection` with the same `name` overwrites; the original `yarn_rm_url`
and `spark_conf` are **snapshot into the PendingSubmission at prepare time**
and survive later edits to the connection (this is intentional — the user
can re-point the connection at a different cluster without retargeting
in-flight submissions).
### 2.2 Write the script to disk
The LLM writes the PySpark code. Then:
```python
script_path = write_job_file(code=PYSPARK_SOURCE)["script_path"]
# Returns {"script_path": "/absolute/path/in/container.py", "bytes": N}
```
The path lives inside the container's `SPARK_EXECUTOR_JOBS_DIR` (default
`./data/jobs/`). For pre-existing scripts on the host, mount the host
directory into the container and skip this step — pass the in-container
path directly to `prepare_submit_job`.
### 2.3 Prepare (snapshot, no YARN call)
```python
prepare_submit_job(
connection="prod",
script_path=script_path,
app_name="daily-aggregation", # REQUIRED, no default
queue="default", # these 4 must be explicit
executor_memory="2G",
executor_cores=2,
num_executors=2,
# extra_args={"jars": "hdfs://...jar"}, # optional
)
# Returns {"pending_id": "p_a1b2c3", "status": "PENDING", "parameters": {...}}
```
**You MUST pass `queue`, `executor_memory`, `executor_cores`, `num_executors`
explicitly every time** — even though defaults exist server-side, the
service rejects implicit defaults with a 400 listing the assumed values.
This forces the agent (or the human reviewer) to see and approve what
will actually be submitted.
The response is the only chance to capture the `pending_id` you need for
`confirm_submit_job`. Store it.
### 2.4 (Optional) Edit before submit
If something needs to change between prepare and confirm:
```python
update_pending_job(
pending_id="p_a1b2c3",
queue="research",
# script_path can also be changed — but the new file MUST exist
# and pass the SQL guard
)
```
Refuses to edit anything that is no longer `PENDING`.
### 2.5 Confirm (the irreversible step)
```python
result = confirm_submit_job(pending_id="p_a1b2c3")
# Returns {"job_id": "a1b2c3d4e5f6",
# "application_id": "application_17400000001_0001",
# "tracking_url": "http://rm:8088/proxy/application_17400000001_0001/"}
```
**STORE BOTH IDS NOW.** Future calls to `get_job_status` / `get_job_logs`
/ `kill_job` accept either, but `application_id` is the one YARN itself
recognizes.
`confirm_submit_job` is **idempotent** on `SUBMITTED` (re-calling returns
the same record) and **retries** on transient `SparkSubmitError` (3
attempts with 5 s delay by default). It transitions to `FAILED` only
after all retries are exhausted.
### 2.6 Monitor / fetch / kill
```python
# While running
state = get_job_status(job_id="a1b2c3d4e5f6").state
logs = get_job_logs(job_id="a1b2c3d4e5f6", tail_chars=10000)
# When done
result = get_job_result(job_id="a1b2c3d4e5f6")
# result.state -> "SUCCEEDED" / "FAILED" / "KILLED" / ...
# result.final_status -> "SUCCEEDED" / "FAILED" / "KILLED" / "UNDEFINED"
# result.diagnostics -> YARN's "why" string (often empty on success)
# result.tracking_url -> the same proxy URL confirm_submit_job returned
# To stop it
kill_job(job_id="a1b2c3d4e5f6")
```
All four accept **either** `job_id` or `application_id`. Pass whichever
you have on hand; the service tries `job_id` first, falls back to
`application_id`.
---
## 3. The PendingSubmission state machine
```
prepare_submit_job
|
v
+-------+ cancel_pending_job
|PENDING+----------> CANCELLED (terminal)
+---+---+
| confirm_submit_job
v
+----------+
| SUBMITTED| (terminal — kill via kill_job, not cancel)
+----------+
(or)
+--------+
| FAILED | <-- confirm_submit_job exhausted retries
+--------+
| confirm_submit_job (resets to PENDING, retries)
v
PENDING
```
- `cancel_pending_job` **refuses** to cancel `SUBMITTED` or `FAILED`
records — those are terminal; use `kill_job` instead.
- `update_pending_job` **refuses** to edit anything that is not
`PENDING`.
- `confirm_submit_job` on a `SUBMITTED` record is a no-op (returns the
same `job_id` / `application_id`). On a `FAILED` record, it resets
to `PENDING` and retries. On a `CANCELLED` record, it raises 400.
---
## 4. The dual-ID contract
| ID | Where it comes from | What it's good for |
|---|---|---|
| `job_id` | service-generated, 12-char hex (`a1b2c3d4e5f6`) | Looking up the Job in this service's store; all 4 lifecycle tools |
| `application_id` | parsed from `spark-submit` stderr (`application_<ts>_<n>`) | YARN UI, YARN REST, `tracking_url`, talking to ops scripts |
**Both IDs are returned by `confirm_submit_job`.** Both are accepted by
`get_job_status` / `get_job_result` / `get_job_logs` / `kill_job`. **You
don't need to remember which to pass — but DO remember to store both
from the confirm response** because:
- If the Job store has been wiped (rare; only on data dir loss), only
`application_id` is still meaningful to YARN — you can't poll a YARN
app you can't identify.
- If you're going to paste a link to a teammate, `tracking_url` (derived
from `application_id`) is the human-friendly form.
If a lifecycle tool returns `KeyError("No Job found for id='X' …")`,
check that `X` came from a successful `confirm_submit_job` response.
The error message intentionally mentions both ID forms and shows an
example of each so you can self-diagnose.
---
## 5. Connection profiles
`Connection` is the cluster profile — it's the "where to submit" record.
Fields:
| Field | Required | Default | Notes |
|---|---|---|---|
| `name` | yes | — | The lookup key |
| `master` | yes | `"yarn"` | `"yarn"` / `"spark://…"` / `"k8s://…"` / `"local[N]"` |
| `deploy_mode` | yes | `"cluster"` | `"cluster"` or `"client"` |
| `yarn_rm_url` | no | env `YARN_RESOURCE_MANAGER_URL` | Required if you'll call `get_job_status` / `get_job_logs` / `kill_job` (the lifecycle tools need a YARN endpoint) |
| `spark_conf` | no | `{}` | Merged into `spark-submit --conf` at submit time |
| `ssl_verify` | no | env `SPARK_EXECUTOR_SSL_VERIFY_DEFAULT` | bool |
| `ssl_ca_bundle` | no | env `SPARK_EXECUTOR_SSL_CA_BUNDLE_DEFAULT` | Path to CA file |
| `auth_type` | no | `"none"` | `"none"` / `"simple"` / `"basic"` / `"kerberos"` |
| `auth_user` / `auth_password` | for `basic` | — | |
| `auth_principal` / `auth_keytab` | for `kerberos` (display only) | — | Real auth uses the system cache |
**Snapshot semantics:** when `prepare_submit_job` runs, it copies
`master` / `deploy_mode` / `spark_conf` / `yarn_rm_url` from the
**current** Connection into the `PendingSubmission`. Editing the
Connection later does **not** retarget an in-flight submission.
**Lifecycle tools use the Job's snapshotted `yarn_rm_url`** (set at
confirm time) — not the Connection's current value. If the RM moves
between confirm and a follow-up `get_job_logs`, you have a problem;
re-pointing the Connection does not help.
---
## 6. Job file workflow (`write_job_file` / `read_job_file` / `update_job_file`)
These are **filesystem I/O tools, not code generators**. The LLM in
context writes the PySpark code; these tools just persist it where
`spark-submit` can find it, with safety guards.
**Typical lifecycle:**
```
LLM composes code --> write_job_file(code=...) # write
LLM reviews code --> read_job_file(script_path=...) # verify
LLM revises code --> update_job_file(script_path=..., # overwrite
content=...)
LLM submits --> prepare_submit_job(script_path=...)
```
**Guards (these are why the tools exist instead of just `cat > file.sh`):**
- **SQL guard** — `write_job_file` and `update_job_file` both reject
code containing forbidden SQL (DROP, DELETE, UPDATE, TRUNCATE, ALTER,
GRANT, REVOKE, …). The check is conservative; if it false-positives,
rewrite to use only SELECT/INSERT or restructure.
- **Path restriction** — `update_job_file` only writes under
`SPARK_EXECUTOR_JOBS_DIR`. This blocks accidentally overwriting
host-mounted configs or the service's own files.
- **Size cap** — 1 MB on both `read_job_file` and `update_job_file`.
Use `read_job_file` for small scripts; for huge ones, mount the host
directory and skip these tools.
**For pre-existing scripts on the host:** mount the host dir into the
container (e.g. `-v /host/scripts:/app/scripts:ro` in `docker run`) and
pass the in-container path directly to `prepare_submit_job` — no need
to round-trip through the file tools.
---
## 7. Error reference
| Status | When | What to do |
|---|---|---|
| 400 | `prepare_submit_job` without explicit defaults; `cancel_pending_job` on terminal state; `update_pending_job` on non-PENDING; `prepare_submit_job` with missing/non-file `script_path`; SQL policy violation | Read the message — it almost always names the fix. Re-submit with corrected args. |
| 404 | `KeyError` from tool layer: unknown `connection`, `pending_id`, `job_id`, or `application_id` | The ID you sent doesn't match anything in the store. For lifecycle tools, the error message explicitly suggests both ID forms. |
| 422 | Pydantic validation: missing field, wrong type | Usually a typo or missing required arg. The response body lists every field error. |
| 500 | spark-submit subprocess failed and was not a `SparkSubmitError` (e.g. binary missing); YARN REST call failed | The tool logs the full traceback to `data/logs/error/...` — check there. |
The service translates exceptions at the FastAPI layer:
- `KeyError` → 404 (with the KeyError's string as `detail`)
- `ValueError` → 400
- Pydantic validation → 422
So when the server returns 404, the body `detail` field is the same
"Unknown job_id: …" / "Unknown connection: …" string the tool layer
raised.
---
## 8. Common pitfalls
| Pitfall | Why it bites | What to do |
|---|---|---|
| Skipping the prepare step | `confirm_submit_job` requires a `pending_id`; calling it without one returns 400 | Always prepare first |
| Passing implicit defaults | Server returns 400 listing what you didn't confirm | Pass `queue`, `executor_memory`, `executor_cores`, `num_executors` every time |
| `script_path` doesn't exist in the container | `prepare_submit_job` returns 400; the agent generated the path in its own context but the container's filesystem is different | Either call `write_job_file` first, or mount a host dir |
| SQL injection in the script | `write_job_file` / `update_job_file` / `prepare_submit_job` all 400 | Only use SELECT/INSERT; never DROP/DELETE/UPDATE |
| `cancel_pending_job` on a running YARN app | Returns 400 — submissions are terminal once submitted | Use `kill_job(job_id=...)` instead, which talks to YARN REST |
| Re-pointing a Connection and expecting in-flight jobs to follow | They don't — the Job snapshots `yarn_rm_url` at confirm time | Re-issue with a new `prepare_submit_job` if you actually need a different RM |
| Calling lifecycle tools with the wrong ID type | Used to fail; **fixed in the 2026-06-29 update** — both IDs are accepted | Just pass whichever you have; the service tries `job_id` first |
| "Session not found" at the MCP layer | Multi-worker gunicorn + in-process session state | Out of scope for this skill; the service warns on startup if `GUNICORN_WORKERS>1`; set `GUNICORN_WORKERS=1` or front with a sticky LB |
---
## 9. End-to-end example (LLM agent's perspective)
User: "Run a word count on `s3a://my-bucket/books/*.txt` and write the
result back to `s3a://my-bucket/wc/`."
```
# 1. Cluster profile (assumes this is a fresh setup)
save_connection(
name="prod",
master="yarn",
deploy_mode="cluster",
yarn_rm_url="http://rm.prod.example.com:8088",
spark_conf={
"spark.sql.shuffle.partitions": "200",
"spark.hadoop.fs.s3a.access.key": "...",
"spark.hadoop.fs.s3a.secret.key": "...",
},
)
# 2. Compose the PySpark code IN YOUR CONTEXT
pyspark_code = '''
from pyspark.sql import SparkSession
spark = SparkSession.builder.appName("wordcount").getOrCreate()
sc = spark.sparkContext
files = sc.wholeTextFiles("s3a://my-bucket/books/")
words = files.flatMap(lambda f: f[1].lower().split()) \\
.map(lambda w: (w, 1)) \\
.reduceByKey(lambda a, b: a + b)
df = words.toDF(["word", "count"])
df.write.mode("overwrite").parquet("s3a://my-bucket/wc/")
'''
# 3. Persist to disk
script_path = write_job_file(code=pyspark_code)["script_path"]
# 4. Sanity-check (optional but recommended)
contents = read_job_file(script_path=script_path)["content"]
# ... review in context ...
# 5. Prepare (snapshot — no YARN call yet)
pending = prepare_submit_job(
connection="prod",
script_path=script_path,
app_name="wordcount-books",
queue="default",
executor_memory="2G",
executor_cores=2,
num_executors=2,
)
pending_id = pending["pending_id"]
# 6. Confirm (NOW spark-submit runs)
result = confirm_submit_job(pending_id=pending_id)
job_id = result["job_id"] # STORE
application_id = result["application_id"] # STORE
tracking_url = result["tracking_url"] # STORE
# 7. Poll
import time
while True:
st = get_job_status(job_id=job_id).state
if st in ("SUCCEEDED", "FAILED", "KILLED", "FINISHED"):
break
time.sleep(30)
# 8. Final result
res = get_job_result(job_id=job_id)
print(res.final_status, res.diagnostics)
# If something went wrong and you want to see why:
# logs = get_job_logs(job_id=job_id, tail_chars=20000)
```
---
## 10. Reference paths (for the curious)
This skill is the **operation** guide. The architectural / implementation
details live in:
- `CLAUDE.md` — project conventions, persistence layout, two-step submit
rationale, MCP arg pattern, MCP session affinity caveat.
- `docs/superpowers/plans/2026-06-24-spark-executor-mcp.md` — the
original implementation plan with Stage 1 / 2 / 3 breakdown.
- `spark_executor/server.py` — the route + exception-handler layer.
- `spark_executor/tools/requests.py` — the Pydantic body models; useful
to read when the MCP 422 response confuses you.
- `spark_executor/core/job_store.py` — Job persistence; explains the
dual-ID contract and the file-backed JSON layout.
spark_executor/server.py
+19 -14
View File
@@ -12,14 +12,14 @@ from spark_executor.tools.connections import (
list_connections,
save_connection,
)
from spark_executor.tools.generate import generate_job_file
from spark_executor.tools.write_job import write_job_file
from spark_executor.tools.job_file import read_job_file, update_job_file
from spark_executor.tools.kill import kill_job
from spark_executor.tools.logs import get_job_logs
from spark_executor.tools.requests import (
ConnectionNameRequest,
EmptyRequest,
GenerateJobFileRequest,
WriteJobFileRequest,
GetJobLogsRequest,
JobIdRequest,
PendingIdRequest,
@@ -46,7 +46,7 @@ app = FastAPI(title="Spark Executor MCP", version="0.0.1", description="Spark Ex
_DEFAULTS_TO_CONFIRM = {
"queue": "default",
"executor_memory": "4G",
"executor_memory": "2G",
"executor_cores": 2,
"num_executors": 2,
}
@@ -90,7 +90,7 @@ def health_check():
"yarn_rm_url into a PendingSubmission record and persist it. "
"Does NOT invoke spark-submit. Returns pending_id for use with "
"confirm_submit_job (the user-second-confirmation step).\n\n"
"REQUIRED PATTERN for LLM-generated code: call generate_job_file(code=...) "
"REQUIRED PATTERN for LLM-generated code: call write_job_file(code=...) "
"first, then pass the returned script_path here. Direct submission with a "
"synthetic path (one that only exists in the agent's context) will be "
"rejected with HTTP 400 — the script must exist inside the container's "
@@ -297,19 +297,24 @@ def _delete_connection(req: ConnectionNameRequest):
# --- LLM-driven PySpark generation (Stage 2) ---
@app.post(
"/generate_job_file",
operation_id="generate_job_file",
summary="Write LLM-generated PySpark code to disk",
"/write_job_file",
operation_id="write_job_file",
summary="Write LLM-authored PySpark code to disk",
description=(
"Takes a PySpark code string and writes it to a timestamped file under "
"Takes a PySpark code string the LLM has already composed in its "
"context and writes it to a timestamped file under "
"SPARK_EXECUTOR_JOBS_DIR (default ./data/jobs/). Returns the absolute "
"path for use as the script_path argument of prepare_submit_job — the "
"two-step pattern means the LLM can produce code, the user can review "
"the resulting file, and only then is the job submitted."
"two-step pattern means the LLM writes the file, the user can review "
"it (via read_job_file), and only then is the job submitted.\n\n"
"Note: this tool does NOT generate PySpark code. The calling LLM is "
"expected to have already written the code; this tool only persists "
"it. Code is also run through the SQL safety policy (SELECT/INSERT "
"only) before being written — forbidden statements cause a 400."
),
)
def _generate_job_file(req: GenerateJobFileRequest):
return generate_job_file(req.code)
def _write_job_file(req: WriteJobFileRequest):
return write_job_file(req.code)
@app.post(
@@ -318,7 +323,7 @@ def _generate_job_file(req: GenerateJobFileRequest):
summary="Read the contents of an existing PySpark script",
description=(
"Returns the text content of an existing script file at the given "
"path. Caps reads at 1 MB. Typical use: after generate_job_file "
"path. Caps reads at 1 MB. Typical use: after write_job_file "
"returns a path, call read_job_file on that path to inspect what "
"was actually written, before deciding to prepare_submit_job or "
"update_job_file."
@@ -334,7 +339,7 @@ def _read_job_file(req: ReadJobFileRequest):
summary="Overwrite an existing PySpark script with new content",
description=(
"Replaces the entire content of an existing script file. Path must "
"be under SPARK_EXECUTOR_JOBS_DIR (the dir generate_job_file writes "
"be under SPARK_EXECUTOR_JOBS_DIR (the dir write_job_file writes "
"to) — protects against overwriting host-mounted configs or other "
"non-script files. Caps writes at 1 MB. Typical use: read_job_file, "
"edit the content (LLM or human), update_job_file, then "
spark_executor/tools/generate.py
-43
View File
@@ -1,43 +0,0 @@
# coding=utf-8
"""
@Time :2026/6/24
@Author :tao.chen
"""
from common.logging import logger
from common.sql_guard import validate_pyspark_code
from spark_executor.core.job_writer import write_job_file
class SqlGuardViolation(ValueError):
"""Raised when generate_job_file receives code that violates the SQL
safety policy. -> HTTP 400 via the FastAPI ValueError handler.
"""
pass
def generate_job_file(code: str) -> dict[str, str]:
"""Write a PySpark code string to disk; return its absolute path.
Validates the code against the SQL safety policy (SELECT/INSERT only)
BEFORE writing, so the agent gets immediate feedback rather than
learning at prepare_submit_job time. Use the returned path as the
`script_path` argument of prepare_submit_job.
The output directory is controlled by the SPARK_EXECUTOR_JOBS_DIR env var
(default: ./data/jobs/).
"""
logger.debug(f"generate_job_file enter code_bytes={len(code)}")
offenses = validate_pyspark_code(code)
if offenses:
logger.warning(
f"generate_job_file rejected: SQL policy violation(s): {offenses}"
)
raise SqlGuardViolation(
f"PySpark code violates SQL safety policy. "
f"Only SELECT and INSERT statements are allowed. "
f"Found forbidden statement(s): {offenses}. "
f"Rewrite the code to use only SELECT/INSERT."
)
path = write_job_file(code)
logger.info(f"generate_job_file ok script_path={path}")
return {"script_path": path}
spark_executor/tools/job_file.py
+4 -4
View File
@@ -6,7 +6,7 @@
Read + update the contents of an existing PySpark script file. These two
tools close the review-and-edit loop:
generate_job_file(code=...) -> {script_path}
write_job_file(code=...) -> {script_path}
read_job_file(script_path=...) -> {content, path} <-- inspect
update_job_file(path, content) -> {path, bytes_written} <-- edit
prepare_submit_job(path) -> {pending_id, ...}
@@ -43,14 +43,14 @@ def _check_readable(script_path: str) -> None:
def _check_writable(script_path: str) -> None:
"""update_job_file is restricted to files under settings.jobs_dir
(the same dir generate_job_file writes to). This prevents the agent
(the same dir write_job_file writes to). This prevents the agent
from overwriting arbitrary host-mounted files or the app's own code.
"""
if not script_path or not os.path.isfile(script_path):
raise ScriptFileError(
f"script_path does not exist or is not a file: {script_path!r}. "
f"update_job_file can only edit existing files. "
f"Use generate_job_file to create a new one."
f"Use write_job_file to create a new one."
)
jobs_root = Path(settings.jobs_dir).resolve()
target = Path(script_path).resolve()
@@ -59,7 +59,7 @@ def _check_writable(script_path: str) -> None:
except ValueError:
raise ScriptFileError(
f"script_path must be under {jobs_root} (the directory "
f"generate_job_file writes to). Got {script_path!r}. "
f"write_job_file writes to). Got {script_path!r}. "
f"This restriction protects host-mounted configs and other "
f"non-script files from being overwritten by the agent."
)
spark_executor/tools/requests.py
+10 -8
View File
@@ -42,7 +42,7 @@ class PrepareSubmitJobRequest(BaseModel):
description=(
"Absolute path to the PySpark script inside the container's "
"filesystem. Must point at an existing regular file. For "
"LLM-generated code, call generate_job_file(code=...) first "
"LLM-generated code, call write_job_file(code=...) first "
"and pass the returned script_path here. For pre-existing "
"files on the host, mount them via a docker volume and pass "
"the in-container path. Returns 400 with a remediation hint "
@@ -54,8 +54,8 @@ class PrepareSubmitJobRequest(BaseModel):
description="YARN queue to submit to. Must be explicitly confirmed by the caller.",
)
executor_memory: str = Field(
default="4G",
description="Executor memory, e.g. '4G'. Must be explicitly confirmed by the caller.",
default="2G",
description="Executor memory, e.g. '2G'. Must be explicitly confirmed by the caller.",
)
executor_cores: int = Field(
default=2,
@@ -102,13 +102,15 @@ class ConnectionNameRequest(BaseModel):
name: str
class GenerateJobFileRequest(BaseModel):
class WriteJobFileRequest(BaseModel):
code: str = Field(
...,
description=(
"Full PySpark source code to write to disk. Will be passed verbatim "
"to spark-submit after the agent calls prepare_submit_job on the "
"returned path."
"Full PySpark source code to write to disk. The LLM is expected "
"to have already written this code in its own context; this tool "
"persists it to a file under SPARK_EXECUTOR_JOBS_DIR so "
"spark-submit can see it. The returned script_path is what you "
"pass to prepare_submit_job."
),
)
@@ -129,7 +131,7 @@ class UpdateJobFileRequest(BaseModel):
description=(
"Absolute path to an existing PySpark script inside the "
"container's filesystem. Must be under SPARK_EXECUTOR_JOBS_DIR "
"(the same dir generate_job_file writes to) — protects against "
"(the same dir write_job_file writes to) — protects against "
"overwriting host-mounted configs or other critical files."
),
)
spark_executor/tools/submit.py
+5 -5
View File
@@ -30,7 +30,7 @@ def _script_path_error(path: str) -> ValueError:
return ValueError(
f"script_path does not exist or is not a file: {path!r}. "
f"Two ways to fix this:\n"
f" 1. (Recommended) Call generate_job_file(code=...) first to write "
f" 1. (Recommended) Call write_job_file(code=...) first to write "
f"the PySpark source to disk, then pass the returned script_path.\n"
f" 2. If the file already exists on the host, mount it into the "
f"container (e.g. -v /host/path:/app/scripts:ro in docker run) and "
@@ -52,7 +52,7 @@ def _sql_guard_error(offenses: list[str]) -> ValueError:
f"Only SELECT and INSERT statements are allowed. "
f"Found forbidden statement(s): {offenses}. "
f"Edit the script and try again. (If the code was generated via "
f"generate_job_file, the generator should have caught this — check "
f"write_job_file, the generator should have caught this — check "
f"for f-string-based SQL injection where the static analysis can't "
f"see the runtime value.)"
)
@@ -67,7 +67,7 @@ def prepare_submit_job(
connection: str,
script_path: str,
queue: str = "default",
executor_memory: str = "4G",
executor_memory: str = "2G",
executor_cores: int = 2,
num_executors: int = 2,
app_name: str,
@@ -92,12 +92,12 @@ def prepare_submit_job(
raise KeyError(f"Unknown connection: {connection}")
# Fail fast: a non-existent path is the most common agent mistake (it
# generated the code in its own context but forgot to call
# generate_job_file first, or its path refers to the host filesystem
# write_job_file first, or its path refers to the host filesystem
# which is invisible inside the container). Better to surface this with
# a 400 + clear remediation than to let spark-submit fail later with
# an opaque FileNotFoundError -> 500.
_check_script_path(script_path)
# SQL safety: re-validate the script even though generate_job_file
# SQL safety: re-validate the script even though write_job_file
# already guards its own output. Catches files written by other means
# (host volume mounts, manual edits).
with open(script_path, encoding="utf-8") as f:
spark_executor/tools/write_job.py
+57
View File
@@ -0,0 +1,57 @@
# coding=utf-8
"""
@Time :2026/6/24
@Author :tao.chen
The MCP-exposed `write_job_file` tool. Despite the historical name
`generate_job_file`, this tool does NOT generate PySpark code — the
calling LLM writes the code, and this tool persists it to disk so
`spark-submit` can see it. The rename (2026-06-29) tightens the naming
to match what the tool actually does, eliminating the "I thought this
would generate code for me" confusion.
It also runs the SQL safety policy (SELECT/INSERT only) on the code
BEFORE writing, so the agent gets immediate feedback if it slipped a
DROP / DELETE / UPDATE past its own code-generation step.
"""
from common.logging import logger
from common.sql_guard import validate_pyspark_code
# Internal helper already named write_job_file (it just writes bytes to
# disk; no SQL guard). Import with an alias to avoid the name collision
# with the MCP-exposed function below.
from spark_executor.core.job_writer import write_job_file as _write_to_disk
class SqlGuardViolation(ValueError):
"""Raised when write_job_file receives code that violates the SQL
safety policy. -> HTTP 400 via the FastAPI ValueError handler."""
pass
def write_job_file(code: str) -> dict[str, str]:
"""Persist a PySpark code string to disk and return its absolute path.
Validates the code against the SQL safety policy (SELECT/INSERT only)
BEFORE writing, so the agent gets immediate feedback rather than
learning at prepare_submit_job time. Use the returned path as the
`script_path` argument of prepare_submit_job.
The output directory is controlled by the SPARK_EXECUTOR_JOBS_DIR env var
(default: ./data/jobs/). Each call writes to a fresh timestamped file
under that directory, so repeated calls never overwrite one another.
"""
logger.debug(f"write_job_file enter code_bytes={len(code)}")
offenses = validate_pyspark_code(code)
if offenses:
logger.warning(
f"write_job_file rejected: SQL policy violation(s): {offenses}"
)
raise SqlGuardViolation(
f"PySpark code violates SQL safety policy. "
f"Only SELECT and INSERT statements are allowed. "
f"Found forbidden statement(s): {offenses}. "
f"Rewrite the code to use only SELECT/INSERT."
)
path = _write_to_disk(code)
logger.info(f"write_job_file ok script_path={path}")
return {"script_path": path}
tests/integration/test_mcp_routes.py
+22 -22
View File
@@ -52,7 +52,7 @@ def test_sixteen_tool_routes_registered():
"/get_connection",
"/delete_connection",
# LLM-driven PySpark generation (3) — Stage 2
"/generate_job_file",
"/write_job_file",
"/read_job_file",
"/update_job_file",
):
@@ -146,7 +146,7 @@ def test_prepare_submit_job_works_via_body(tmp_path):
"connection": "prod",
"script_path": str(script),
"queue": "research",
"executor_memory": "4G",
"executor_memory": "2G",
"executor_cores": 2,
"num_executors": 2,
"app_name": "test-app",
@@ -177,7 +177,7 @@ def test_prepare_rejects_unconfirmed_defaults(tmp_path):
detail = r.json()["detail"]
assert "Please confirm default values" in detail
assert "queue='default'" in detail
assert "executor_memory='4G'" in detail
assert "executor_memory='2G'" in detail
assert "executor_cores=2" in detail
assert "num_executors=2" in detail
@@ -193,7 +193,7 @@ def test_prepare_accepts_confirmed_defaults(tmp_path):
"connection": "prod",
"script_path": str(script),
"queue": "default",
"executor_memory": "4G",
"executor_memory": "2G",
"executor_cores": 2,
"num_executors": 2,
"app_name": "test-app",
@@ -203,7 +203,7 @@ def test_prepare_accepts_confirmed_defaults(tmp_path):
body = r.json()
assert body["status"] == "PENDING"
assert body["parameters"]["queue"] == "default"
assert body["parameters"]["executor_memory"] == "4G"
assert body["parameters"]["executor_memory"] == "2G"
assert body["parameters"]["executor_cores"] == 2
assert body["parameters"]["num_executors"] == 2
@@ -218,7 +218,7 @@ def test_prepare_rejects_nonexistent_script_path_with_400():
"connection": "prod",
"script_path": "/nope/does_not_exist.py",
"queue": "default",
"executor_memory": "4G",
"executor_memory": "2G",
"executor_cores": 2,
"num_executors": 2,
"app_name": "test-app",
@@ -228,12 +228,12 @@ def test_prepare_rejects_nonexistent_script_path_with_400():
detail = r.json()["detail"]
# The message must guide the agent to the right next step
assert "does not exist" in detail
assert "generate_job_file" in detail
assert "write_job_file" in detail
def test_prepare_rejects_sql_policy_violation_with_400(tmp_path):
"""SQL guard: even if the file exists, prepare_submit_job must reject
code containing DROP/DELETE/etc. (defense in depth — generate_job_file
code containing DROP/DELETE/etc. (defense in depth — write_job_file
also enforces this, but a host-mounted file might bypass that)."""
c = TestClient(app)
c.post("/save_connection", json={"name": "prod", "master": "yarn"})
@@ -245,7 +245,7 @@ def test_prepare_rejects_sql_policy_violation_with_400(tmp_path):
"connection": "prod",
"script_path": str(bad_script),
"queue": "default",
"executor_memory": "4G",
"executor_memory": "2G",
"executor_cores": 2,
"num_executors": 2,
"app_name": "test-app",
@@ -257,14 +257,14 @@ def test_prepare_rejects_sql_policy_violation_with_400(tmp_path):
assert "SELECT" in detail and "INSERT" in detail # the policy is explained
def test_generate_rejects_sql_policy_violation_with_400(tmp_path):
"""Same policy at generate_job_file — agent gets immediate feedback
def test_write_rejects_sql_policy_violation_with_400(tmp_path):
"""Same policy at write_job_file — agent gets immediate feedback
before the file is even written to disk."""
from common import config
config.settings.jobs_dir = str(tmp_path / "jobs")
c = TestClient(app)
r = c.post(
"/generate_job_file",
"/write_job_file",
json={"code": 'spark.sql("DELETE FROM events")\n'},
)
assert r.status_code == 400
@@ -331,7 +331,7 @@ def test_list_and_get_pending_job_roundtrip_via_body(tmp_path):
"connection": "prod",
"script_path": str(script),
"queue": "default",
"executor_memory": "4G",
"executor_memory": "2G",
"executor_cores": 2,
"num_executors": 2,
"app_name": "test-app",
@@ -354,15 +354,15 @@ def test_list_connections_works_with_empty_body():
assert r.json() == []
# --- Stage 2: generate_job_file ---
# --- Stage 2: write_job_file ---
def test_generate_job_file_accepts_code_string_in_body(tmp_path, monkeypatch):
def test_write_job_file_accepts_code_string_in_body(tmp_path, monkeypatch):
"""End-to-end: a Pydantic body model lets tools/call pass a code string
that FastAPI would reject if it were a query parameter (length limits)."""
monkeypatch.chdir(tmp_path)
c = TestClient(app)
code = "print('from MCP integration test')\n" * 100 # > FastAPI query limit
r = c.post("/generate_job_file", json={"code": code})
r = c.post("/write_job_file", json={"code": code})
assert r.status_code == 200, r.text
p = r.json()["script_path"]
assert os.path.isfile(p)
@@ -401,7 +401,7 @@ def test_unknown_connection_in_prepare_returns_404():
"connection": "nope",
"script_path": "/tmp/x.py",
"queue": "default",
"executor_memory": "4G",
"executor_memory": "2G",
"executor_cores": 2,
"num_executors": 2,
"app_name": "test-app",
@@ -423,7 +423,7 @@ def test_confirm_non_pending_returns_400(tmp_path):
"connection": "prod",
"script_path": str(script),
"queue": "default",
"executor_memory": "4G",
"executor_memory": "2G",
"executor_cores": 2,
"num_executors": 2,
"app_name": "test-app",
@@ -448,7 +448,7 @@ def test_cancel_already_submitted_returns_400(tmp_path):
"connection": "prod",
"script_path": str(script),
"queue": "default",
"executor_memory": "4G",
"executor_memory": "2G",
"executor_cores": 2,
"num_executors": 2,
"app_name": "test-app",
@@ -477,7 +477,7 @@ def test_update_pending_job_route(tmp_path):
"connection": "prod",
"script_path": str(script),
"queue": "default",
"executor_memory": "4G",
"executor_memory": "2G",
"executor_cores": 2,
"num_executors": 2,
"app_name": "test-app",
@@ -496,7 +496,7 @@ def test_update_pending_job_route(tmp_path):
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["executor_memory"] == "2G"
assert got["app_name"] == "test-app"
@@ -511,7 +511,7 @@ def test_update_pending_job_route_rejects_submitted(tmp_path, monkeypatch):
"connection": "prod",
"script_path": str(script),
"queue": "default",
"executor_memory": "4G",
"executor_memory": "2G",
"executor_cores": 2,
"num_executors": 2,
"app_name": "test-app",
tests/unit/test_job_file.py
+1 -1
View File
@@ -145,7 +145,7 @@ def test_full_edit_cycle_via_tools(tmp_path: Path):
"""Simulate the agent's review-and-edit loop end-to-end."""
config.settings.jobs_dir = str(tmp_path)
path = tmp_path / "script.py"
# 1. Initial state: file doesn't exist (in real flow, generate_job_file
# 1. Initial state: file doesn't exist (in real flow, write_job_file
# would create it; we just create it directly here for the test)
path.write_text("# v1: initial\n")
# 2. Agent reads back
tests/unit/test_logging.py
+1 -1
View File
@@ -51,7 +51,7 @@ def test_prepare_submit_job_emits_debug_and_info(log_capture, tmp_path):
connection="prod",
script_path=str(script),
queue="research",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
tests/unit/test_models.py
+3 -3
View File
@@ -95,7 +95,7 @@ def test_pending_submission_defaults_to_pending_status():
deploy_mode="cluster",
script_path="/tmp/j.py",
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
spark_conf={},
@@ -120,7 +120,7 @@ def test_pending_submission_carries_yarn_rm_url_snapshot():
yarn_rm_url="http://rm-prod:8088",
script_path="/tmp/j.py",
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
spark_conf={},
@@ -139,7 +139,7 @@ def test_pending_submission_can_record_outcome():
deploy_mode="cluster",
script_path="/tmp/j.py",
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
spark_conf={},
tests/unit/test_pending_store.py
+3 -3
View File
@@ -18,7 +18,7 @@ def _pending(pid: str = "p_1", created_at: datetime | None = None) -> PendingSub
deploy_mode="cluster",
script_path="/tmp/j.py",
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
spark_conf={},
@@ -85,7 +85,7 @@ def _legacy_record_text() -> str:
return (
'{"p_legacy": {"pending_id": "p_legacy", "app_name": "legacy", '
'"connection": "prod", "master": "yarn", "deploy_mode": "cluster", '
'"script_path": "/tmp/j.py", "queue": "default", "executor_memory": "4G", '
'"script_path": "/tmp/j.py", "queue": "default", "executor_memory": "2G", '
'"executor_cores": 2, "num_executors": 2, "spark_conf": {}, '
'"created_at": "2026-06-23T00:00:00"}}'
)
@@ -139,7 +139,7 @@ def test_loads_records_with_missing_or_null_app_name(tmp_path: Path):
"deploy_mode": "cluster",
"script_path": "/tmp/j.py",
"queue": "default",
"executor_memory": "4G",
"executor_memory": "2G",
"executor_cores": 2,
"num_executors": 2,
"spark_conf": {},
tests/unit/test_requests.py
+2 -2
View File
@@ -16,7 +16,7 @@ def test_prepare_submit_job_request_restores_defaults_for_queue_and_resources():
app_name="test-app",
)
assert req.queue == "default"
assert req.executor_memory == "4G"
assert req.executor_memory == "2G"
assert req.executor_cores == 2
assert req.num_executors == 2
@@ -26,7 +26,7 @@ def test_prepare_submit_job_request_accepts_extra_args():
connection="prod",
script_path="/tmp/x.py",
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
tests/unit/test_spark_submit.py
+9 -9
View File
@@ -36,7 +36,7 @@ def test_build_command_uses_provided_master_and_deploy_mode():
deploy_mode="cluster",
script_path="/tmp/jobs/job_001.py",
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
)
@@ -44,7 +44,7 @@ def test_build_command_uses_provided_master_and_deploy_mode():
assert "yarn" in cmd
assert "cluster" in cmd
assert "--queue" in cmd and "default" in cmd
assert "--executor-memory" in cmd and "4G" in cmd
assert "--executor-memory" in cmd and "2G" in cmd
assert "--executor-cores" in cmd and "2" in cmd
assert "--num-executors" in cmd and "2" in cmd
assert cmd[-1] == "/tmp/jobs/job_001.py"
@@ -56,7 +56,7 @@ def test_extra_args_none_default():
deploy_mode="cluster",
script_path="/tmp/jobs/job_001.py",
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
)
@@ -70,7 +70,7 @@ def test_extra_args_single_pair():
deploy_mode="cluster",
script_path="/tmp/jobs/job_001.py",
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
extra_args={"jars": "hdfs:///libs/foo.jar"},
@@ -88,7 +88,7 @@ def test_extra_args_multiple_pairs_preserve_order():
deploy_mode="cluster",
script_path="/tmp/jobs/job_001.py",
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
extra_args={
@@ -112,7 +112,7 @@ def test_extra_args_does_not_collide_with_spark_conf():
deploy_mode="cluster",
script_path="/tmp/jobs/job_001.py",
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
spark_conf={"spark.executor.memory": "8G"},
@@ -146,7 +146,7 @@ def test_build_command_appends_spark_conf_entries():
deploy_mode="cluster",
script_path="/tmp/j.py",
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
spark_conf={"spark.sql.shuffle.partitions": "200", "spark.executor.memoryOverhead": "1G"},
@@ -169,7 +169,7 @@ def test_build_command_uses_settings_spark_submit_bin():
deploy_mode="cluster",
script_path="/tmp/j.py",
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
)
@@ -183,7 +183,7 @@ def test_build_command_supports_pyspark_binary():
deploy_mode="cluster",
script_path="/tmp/j.py",
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
)
tests/unit/test_submit_tool.py
+33 -33
View File
@@ -56,7 +56,7 @@ def test_prepare_does_not_invoke_spark_submit(monkeypatch, real_script):
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -81,7 +81,7 @@ def test_prepare_persists_pending_with_snapshot(monkeypatch, real_script):
connection="prod",
script_path=str(real_script),
queue="research",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -103,7 +103,7 @@ def test_prepare_persists_app_name(monkeypatch, real_script):
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="my-etl-job",
@@ -125,7 +125,7 @@ def test_prepare_accepts_restored_defaults(monkeypatch, real_script):
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -133,7 +133,7 @@ def test_prepare_accepts_restored_defaults(monkeypatch, real_script):
assert out["status"] == "PENDING"
p = submit.pending_store.get(_last_pending_id())
assert p.queue == "default"
assert p.executor_memory == "4G"
assert p.executor_memory == "2G"
assert p.executor_cores == 2
assert p.num_executors == 2
@@ -144,7 +144,7 @@ def test_prepare_snapshots_extra_args(monkeypatch, real_script):
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -173,7 +173,7 @@ def test_prepare_raises_for_unknown_connection(real_script):
connection="missing",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -189,7 +189,7 @@ def test_prepare_rejects_nonexistent_script_path(tmp_path):
connection="prod",
script_path=missing,
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -203,7 +203,7 @@ def test_prepare_rejects_directory_as_script_path(tmp_path):
connection="prod",
script_path=str(tmp_path),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -216,21 +216,21 @@ def test_prepare_rejects_empty_script_path():
connection="prod",
script_path="",
queue="default",
executor_memory="4G",
executor_memory="2G",
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"):
def test_prepare_error_message_mentions_write_job_file(tmp_path):
"""The agent must be told to call write_job_file first."""
with pytest.raises(ValueError, match="write_job_file"):
submit.prepare_submit_job(
connection="prod",
script_path=str(tmp_path / "x.py"),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -245,7 +245,7 @@ def test_confirm_rejects_if_script_was_deleted_after_prepare(tmp_path, monkeypat
connection="prod",
script_path=str(script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -264,7 +264,7 @@ def test_prepare_snapshots_connection_at_prepare_time(monkeypatch, real_script):
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -287,7 +287,7 @@ def test_confirm_invokes_spark_submit_and_marks_submitted(monkeypatch, real_scri
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -323,7 +323,7 @@ def test_confirm_refuses_non_pending_status(monkeypatch, real_script):
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -342,7 +342,7 @@ def test_confirm_marks_failed_on_spark_submit_error(monkeypatch, real_script):
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -368,7 +368,7 @@ def test_confirm_retries_spark_submit_failure_then_succeeds(monkeypatch, real_sc
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -395,7 +395,7 @@ def test_confirm_exhausts_retries_and_sets_failed(monkeypatch, real_script):
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -419,7 +419,7 @@ def test_confirm_idempotent_when_submitted(monkeypatch, real_script):
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -446,7 +446,7 @@ def test_confirm_resets_failed_and_retries(monkeypatch, real_script):
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -476,7 +476,7 @@ def test_confirm_updates_pending_before_job_store(monkeypatch, real_script):
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -510,7 +510,7 @@ def test_list_pending_jobs_returns_all(monkeypatch, tmp_path):
connection="prod",
script_path=str(a),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -519,7 +519,7 @@ def test_list_pending_jobs_returns_all(monkeypatch, tmp_path):
connection="prod",
script_path=str(b),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -537,7 +537,7 @@ def test_get_pending_job_returns_dump(monkeypatch, real_script):
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -562,7 +562,7 @@ def test_cancel_pending_job_marks_cancelled(monkeypatch, real_script):
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -583,7 +583,7 @@ def test_cancel_pending_job_refuses_submitted(monkeypatch, real_script):
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -604,7 +604,7 @@ def test_update_pending_updates_resource_params(monkeypatch, real_script):
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -639,7 +639,7 @@ def test_update_pending_rejects_non_pending_status(monkeypatch, real_script):
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -657,7 +657,7 @@ def test_update_pending_rejects_bad_script_path(tmp_path, real_script):
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -673,7 +673,7 @@ def test_update_pending_rejects_sql_violation_script(tmp_path, real_script):
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
tests/unit/test_generate_tool.py → tests/unit/test_write_job_tool.py
+15 -5
View File
@@ -5,7 +5,7 @@ from pathlib import Path
import pytest
from common import config
from spark_executor.tools import generate
from spark_executor.tools import write_job
@pytest.fixture(autouse=True)
@@ -23,9 +23,9 @@ def _restore_settings():
config.settings.log_level = snapshot.log_level
def test_generate_writes_code_and_returns_path(monkeypatch, tmp_path: Path):
def test_write_writes_code_and_returns_path(monkeypatch, tmp_path: Path):
config.settings.jobs_dir = str(tmp_path / "data" / "jobs")
out = generate.generate_job_file(
out = write_job.write_job_file(
"from pyspark.sql import SparkSession\n"
"spark = SparkSession.builder.getOrCreate()\n"
)
@@ -38,8 +38,18 @@ def test_generate_writes_code_and_returns_path(monkeypatch, tmp_path: Path):
assert "SparkSession.builder.getOrCreate()" in f.read()
def test_generate_uses_settings_jobs_dir(tmp_path: Path):
def test_write_uses_settings_jobs_dir(tmp_path: Path):
config.settings.jobs_dir = str(tmp_path / "custom")
out = generate.generate_job_file("x = 1\n")
out = write_job.write_job_file("x = 1\n")
assert out["script_path"].startswith(str(tmp_path / "custom"))
assert os.path.isfile(out["script_path"])
def test_write_rejects_sql_violation(tmp_path: Path):
"""The SQL guard is the whole reason this isn't a bare file-write
tool: the LLM might have produced a script with a forbidden
statement and we want to fail at write time, not at prepare time."""
from spark_executor.tools.write_job import SqlGuardViolation
config.settings.jobs_dir = str(tmp_path / "data" / "jobs")
with pytest.raises(SqlGuardViolation):
write_job.write_job_file("spark.sql('DROP TABLE users')\n")