The agent can now write PySpark that runs DROP/DELETE/UPDATE/etc. on
production tables. Add a static guard that rejects anything other than
SELECT and INSERT at two enforcement points:
1. generate_job_file: validates BEFORE writing to disk. Agent gets
immediate feedback ('rewrite to use only SELECT/INSERT') rather
than learning at submit time.
2. prepare_submit_job: re-validates the script content (reads the
file) as a defense-in-depth check. Catches host-mounted files,
manually-edited files, anything that bypassed generate_job_file.
How it works:
- common/sql_guard.py extracts Python string literals whose first
keyword is a SQL verb (catches spark.sql('...'), f-strings, and any
raw SQL literal)
- sqlparse splits each literal into statements; we check the first
keyword against the policy (SELECT/INSERT/WITH allowed; DROP,
DELETE, UPDATE, TRUNCATE, ALTER, CREATE, REPLACE, MERGE, GRANT,
REVOKE, SET, SHOW, KILL, EXEC, etc. forbidden)
- WITH recurses into the CTE body to catch WITH x AS (DROP ...) ...
- The MCP layer maps ValueError -> HTTP 400 (existing handler)
Test coverage:
- 28 unit tests in test_sql_guard.py cover: extraction (single/double/
f-string, English false positives, multi-literal), statement
classification (SELECT, INSERT, DROP, DELETE, UPDATE, TRUNCATE,
ALTER, CREATE, multi-statement, CTE bodies, comments)
- 2 integration tests verify MCP layer returns 400 with the policy
explanation at both generate_job_file and prepare_submit_job
Limitations (documented in sql_guard.py docstring):
- f-strings where the SQL is built at runtime (e.g. f'SELECT * FROM
{user_input}') look like SELECTs at static-analysis time. The
guard catches the static literal; the runtime substitution is the
caller's responsibility.
- pyspark.sql.functions.expr('...') accepts SQL inline; not currently
caught. (Future work.)
146/146 still pass. Live verified: DROP TABLE -> MCP 400 with policy
explanation; SELECT -> MCP 200 + file written to ./data/jobs/.
The lifespan was hardcoded to mount just one MCP service
(spark_executor_app at /spark-executor-mcp). To add a second MCP
service, you'd have to edit the lifespan function body — error-prone
and not obvious where to add a new entry.
Refactor:
- New McpService dataclass (name, app, mount_path)
- MCP_SERVICES list at module top — adding a new service is one append
- Lifespan iterates MCP_SERVICES; the per-service mount + tool-list
logic is now a helper (_log_mcp_tools)
- Tool-list log now prefixes the service name:
'MCP tools registered (spark_executor, 14):'
instead of just '(14):', so multi-service startups are unambiguous
116/116 still pass. Live verified: gunicorn 2-worker startup logs
each service with its tool list and the source app version.
Future services just need a McpService entry — no lifespan edit. The
MCP_SERVICES list has a docstring example showing the one-liner.
Three problems with the prior Dockerfile:
1. JAVA_HOME and SPARK_HOME were hardcoded at build time (ENV ...),
so docker-compose / .env overrides had no effect.
2. PATH was constructed at build time using those hardcoded values,
so even if you could override the env vars, PATH would still
reference the old literal paths (e.g. /usr/lib/jvm/java-17-openjdk-amd64/bin
hardcoded into PATH at image build).
3. There was no .env.example entry for either, so operators had no
template to follow.
Fix:
- docker-entrypoint.sh: re-derives PATH from the (possibly overridden)
JAVA_HOME and SPARK_HOME at every container start. Strips any stale
JDK/Spark bin dirs from PATH first so a restart with a new override
actually changes which / resolve.
- Dockerfile: COPY + ENTRYPOINT [entrypoint.sh], CMD [gunicorn main:app].
The existing ENV JAVA_HOME and ENV SPARK_HOME stay as the defaults
so the image works out of the box; users override via .env.
- .env.example: new 'Java + Spark install paths' section, default
values match the Dockerfile, comments explain the entrypoint
re-derivation.
- docker-compose.yml: forwards JAVA_HOME and SPARK_HOME with
container-side defaults matching the Dockerfile.
Verified:
- 116/116 tests still pass
- Direct entrypoint run with JAVA_HOME=/fake/jdk shows PATH rebuilt as
/app/.venv/bin:/fake/jdk/bin:/fake/spark/bin:... (override took effect)
Note: uses awk instead of 'paste -sd:' for portability across macOS
(BSD paste doesn't support -s) and Linux (GNU does).
common/logging.py used hardcoded 'data/logs/{debug,info}/' paths,
which forced logs into the same dir as connections.json. Operators
couldn't put logs on a dedicated volume or a different disk for
retention policy reasons.
Now the base log dir comes from settings.log_dir (env var
SPARK_EXECUTOR_LOG_DIR, default <data_dir>/logs). Subdirs debug/ and
info/ are auto-created. Same behavior in dev (./data/logs/), but in
prod you can do:
SPARK_EXECUTOR_LOG_DIR=/var/log/spark-executor # dedicated volume
SPARK_EXECUTOR_LOG_DIR=/mnt/slow-storage/logs # cold storage for old logs
Other changes:
- .env.example: new SPARK_EXECUTOR_LOG_DIR entry, split into its own
'Loguru file sinks' section
- docker-compose.yml: forwards SPARK_EXECUTOR_LOG_DIR with the
container-side default /app/data/logs (the existing ./data:/app/data
volume mount already covers this)
- common/config.py: Settings gets a new log_dir field, defaults
derived from data_dir; reload() also resets it
116/116 still pass. Live smoke verified: with
SPARK_EXECUTOR_LOG_DIR=/tmp/spark-logs, both
/tmp/spark-logs/info/2026-06-25.log
/tmp/spark-logs/debug/2026-06-25.log
are created on the first request.
.env.example now lists every env var read by the app, with a comment
explaining the meaning and the default. The previously-missing ones
(from the common/config.py refactor) are added:
- SPARK_EXECUTOR_DATA_DIR
- SPARK_EXECUTOR_JOBS_DIR
- SPARK_EXECUTOR_LOG_LEVEL
docker-compose.yml forwards all of them with sensible container-side
defaults (SPARK_EXECUTOR_DATA_DIR=/app/data, which is the volume
mount point from the host).
Split the env block in compose into two commented sections
('common/config.py knobs' vs 'gunicorn.conf.py knobs') so it's clear
which file each one is consumed by.
The GUNICORN_LOGLEVEL entry is also added (was missing). All other
GUNICORN_* knobs were already documented.
116/116 still pass.
Before: SPARK_EXECUTOR_DATA_DIR, SPARK_EXECUTOR_JOBS_DIR, and
YARN_RESOURCE_MANAGER_URL were each read directly via os.environ.get()
inside the module that used them. Log level was hardcoded DEBUG in
common/logging.py. There was no single file showing what the full set
of env vars the app reads is.
After: common/config.py defines a single Settings dataclass that
reads all env vars at import time and exposes them as fields on a
module-level singleton. App code uses "from common.config import
settings; settings.data_dir" etc. New SPARK_EXECUTOR_LOG_LEVEL env
var controls stderr + info file verbosity (debug file always gets full
DEBUG).
Improvements:
- One file lists every env var the app reads (was: grep the codebase)
- Tests can monkeypatch fields on the settings singleton directly
instead of monkeypatching the env + reloading
- Adding a new env var means adding one field in config.py, not
editing 3+ call sites
- settings.reload() method for tests that prefer env-var style
Out of scope (kept where they are):
- GUNICORN_* env vars live in gunicorn.conf.py (gunicorn concept)
- PYTHONUNBUFFERED in Dockerfile (Python runtime flag)
- SPARK_SUBMIT_OPTS not in config (JVM flag, not Python)
Test changes:
- test_job_writer.py: settings.jobs_dir instead of monkeypatching
SPARK_EXECUTOR_JOBS_DIR
- test_yarn_client.py: settings.yarn_resource_manager_url instead of
monkeypatching YARN_RESOURCE_MANAGER_URL
- test_generate_tool.py: same as job_writer
- Each test file gets an autouse fixture that snapshots+restores
settings so one test mutation does not leak into the next
116/116 still pass. Live verified: SPARK_EXECUTOR_LOG_LEVEL=INFO
suppresses DEBUG loguru output as expected.
Two problems with the prior prepare_submit_job flow:
1. The agent could pass a script_path that only existed in its own
context (LLM-generated code not yet on disk) or a path on the host
filesystem that's invisible inside the container. The new check
surfaces this as a 400 with a clear remediation hint instead of
letting spark-submit fail later with an opaque FileNotFoundError -> 500.
2. The container-isolation issue: any path the agent gives is interpreted
inside the container. The two ways a file can legitimately exist there
are (a) generate_job_file(code=...) just wrote it to
SPARK_EXECUTOR_JOBS_DIR (the default ./data/jobs/ is the only
gitignored dir that survives restarts), or (b) a host dir was
mounted via -v. The error message spells both out so an agent can
self-correct.
Implementation:
- submit.py: new _check_script_path() that raises ValueError (-> 400)
when the path is missing, empty, or a directory. Called in both
prepare_submit_job and confirm_submit_job (defense in depth).
- prepare_submit_job checks the connection FIRST (KeyError -> 404)
before the script (ValueError -> 400), so an agent with both problems
sees the more fundamental 'unknown connection' error first.
- server.py / requests.py: route description and Pydantic field
description spell out the generate_job_file pattern so an LLM
reading the tool schema learns the right next step.
8 new tests; 8 existing tests adjusted to create real files (they used
synthetic /tmp/*.py paths that don't exist).
Exposes a single new MCP tool: generate_job_file(code) -> {script_path}.
Wiring follows the Stage 1 conventions (Pydantic body model for the
route, loguru DEBUG/INFO logging, summary+description for the startup
log, MCP arg-pattern via the body model so tools/call roundtrips long
code strings without 422-ing on FastAPI query length limits).
Flow:
1. LLM calls generate_job_file(code=...) -> {script_path}
2. LLM calls prepare_submit_job(connection=...,
script_path=...) -> {pending_id}
3. User reviews the file + pending record
4. LLM calls confirm_submit_job(pending_id=...) -> spark-submit runs
The output directory is SPARK_EXECUTOR_JOBS_DIR (default ./data/jobs/,
gitignored, persists across container restarts via the existing
./data volume mount in docker-compose.yml).
3 new tests:
- Unit: env-var override, default fallback in tmp cwd
- Integration: end-to-end body call with a > FastAPI-query-limit code
string (the canary test that would have caught the Stage 1
query-params-422 bug)
Live MCP smoke verified: tools/list shows 14 tools (13 from Stage 1 +
the new one), tools/call generate_job_file returns the absolute path
under ./data/jobs/.
spark_executor/core/job_writer.py provides write_job_file(code, jobs_dir=None)
that resolves the output directory in this order:
1. explicit jobs_dir argument (for tests + programmatic override)
2. SPARK_EXECUTOR_JOBS_DIR environment variable (operator override)
3. DEFAULT_JOBS_DIR = './data/jobs' (gitignored, persists via volume mount)
Filenames are job_<UTC-stamp>_<random-hex>.py — same-second writes still
get distinct names via the random suffix, so concurrent agents won't
clobber each other.
8 tests cover: arg>env>default precedence, dir auto-creation, env-var
override, default fallback in a tmp cwd, and uniqueness under concurrent
calls.
YARN is the dominant use case, and the value of master for YARN is the
literal string 'yarn' (not a URL) — there is no useful variation. Making
it a default saves every YARN user from typing the same value every time:
save_connection(name='prod') # master defaults to 'yarn'
save_connection(name='dev', master='spark://10.0.0.5:7077') # override
Added a field_validator on master that catches common typos
('yarn-cluster', 'http://...', empty string) at save time with a clear
error message instead of letting them reach spark-submit.
Backward compatible: existing Connection records without an explicit
master field default to 'yarn' on load (Pydantic applies defaults on
deserialization), so already-saved data still works.
3 new tests cover the default, the validator, and that other cluster
managers (spark://, k8s://, local[N]) still pass.
docker-compose.yml:
- Builds the Dockerfile, exposes :8000
- Mounts ./data (persistent) and ./hadoop-conf (ro) for cluster configs
- Forwards YARN_RESOURCE_MANAGER_URL via env var (fallback for REST client)
- Healthcheck via /openapi.json (cheap liveness probe; MCP initialize
handshake would be more accurate but requires session-id plumbing)
- restart: unless-stopped for production
- Resource limit section commented out (server itself is lightweight)
.env.example:
- Documents the env vars compose expects
- .env is gitignored
main.py lifespan now enumerates every MCP tool with its description at
startup so operators can see at a glance what the server exposes.
The auto-generated '### Responses:' suffix that fastapi-mcp appends to
each tool description is truncated before logging to keep the startup
log scannable.
Docstrings added to all 12 route handlers in server.py so the MCP tool
descriptions are meaningful (without them, fastapi-mcp would emit empty
descriptions). Each route now carries a summary + description used by
the auto-generated OpenAPI schema.
Spark 4.1.2 requires Java 17 at runtime (it boots a JVM via the
bin/spark-submit shell script). The previous Dockerfile only installed
curl, ca-certificates, and tar — leaving spark-submit unable to find
the 'java' binary.
A JRE is sufficient (not the full JDK) because spark-submit does not
compile Java at runtime — it loads pre-built JARs from $SPARK_HOME/jars/.
JRE-only keeps the image ~160 MB smaller than JDK-headless.
Sets JAVA_HOME=/usr/lib/jvm/java-17-openjdk-amd64 and prepends it to
PATH so spark-submit (which checks JAVA_HOME) finds the right JVM.
Combined with the prior yarn-REST rewrite, the runtime image now needs
exactly three things in addition to python:3.12-slim:
- uv (from ghcr.io/astral-sh/uv)
- openjdk-17-jre-headless
- Spark 4.1.2 from Tsinghua mirror
yarn_client.py no longer invokes the 'yarn' binary via subprocess; it uses
httpx against /ws/v1/cluster/apps/* endpoints. This means the runtime image
no longer needs the Hadoop client installation — the only YARN-side
dependency left in the container is the config dir consumed by
spark-submit itself.
New model field:
- Job.yarn_rm_url: str | None
- PendingSubmission.yarn_rm_url: str | None (snapshotted at prepare)
prepare_submit_job snapshots Connection.yarn_rm_url into the pending
record (consistent with the existing master/deploy_mode/spark_conf
snapshot pattern); confirm_submit_job copies it onto the Job so
status/logs/kill can use it without re-looking-up the connection.
Resolution order for the RM URL at runtime:
1. Job.yarn_rm_url (preferred — survives connection edits/deletes)
2. Connection.yarn_rm_url fallback (if a future tool is added that
doesn't go through a Job)
3. YARN_RESOURCE_MANAGER_URL env var
Errors:
- YarnConfigError (HTTP 4xx semantics) when URL is missing/malformed
- YarnError for HTTP 4xx/5xx from the RM, network failures, missing
state field, or unparseable log responses
10 new tests in test_yarn_client.py cover the REST surface:
success, 404, 5xx, missing state field, env-var fallback, malformed
URL, log 404 with log-aggregation hint, kill PUT body shape, and
httpx connection-error wrapping.
- common/logging.py: improve format to timestamp|LEVEL|module:func:line - message
- core/ layer: DEBUG log every subprocess invocation (cmd, rc, byte counts),
JSON load/dump events, parsed application_id. ERROR log on failures.
- tools/ layer: DEBUG log every public tool entry with key parameters,
INFO log on business outcomes (saved/submitted/killed/...).
- New tests/unit/test_logging.py: capture loguru output via in-memory sink
and assert DEBUG + INFO messages are emitted for representative flows.
Tool functions raise KeyError for unknown ids and ValueError for invalid
state transitions. Without handlers, FastAPI returned a bare 500
"Internal Server Error" to MCP clients — useless for an agent that needs
to know whether the id was wrong vs the cluster was down.
KeyError -> 404 with the message ("Unknown job_id: foo").
ValueError -> 400 with the message (e.g. "pending_id p_xyz is in status
'CANCELLED', not PENDING").
The plan's route signatures used query parameters, which worked for direct
HTTP callers and unit tests, but fastapi-mcp's HTTP transport passes
tools/call arguments as a JSON body. dict-typed parameters like
spark_conf arrived as a string and the route returned 422.
Refactor each route to take a single Pydantic body model (saved in
spark_executor/tools/requests.py). Underlying tool functions unchanged.
Integration tests in tests/integration/test_mcp_routes.py now exercise
the full body-based roundtrip (save_connection with spark_conf, prepare
→ list → get, list_connections with empty body).
Claude