Commit Graph
32 Commits
Author SHA1 Message Date
Claude ad557b5984 feat(yarn_client): SSL/TLS config for YARN REST connections
Add per-Connection ssl_verify / ssl_ca_bundle plus global defaults so
CDH 5 / on-prem clusters with self-signed certs or custom CA bundles can
be queried without patching code.

- Connection gets ssl_verify (bool|None) and ssl_ca_bundle (str|None)
- Settings gets ssl_verify_default and ssl_ca_bundle_default
- New YarnClientConfig dataclass carries the resolved verify= value
- _request passes verify= through to httpx.request
- All public yarn_client functions now take YarnClientConfig instead of
  a bare yarn_rm_url string; tool call sites resolve the Connection
- SaveConnectionRequest exposes the two new fields

Tests cover per-connection CA bundle, per-connection verify=False,
global default fallback, and connection-not-found error.
2026-06-26 13:50:30 +08:00
Claude 5566d55a7c feat(spark_submit): extra_args kwarg for non-conf spark-submit flags
The function had hardcoded --master / --deploy-mode / --queue /
--executor-memory / --executor-cores / --num-executors plus a
`spark_conf` dict for --conf, with no way to pass other flags like
--jars, --py-files, --files, --driver-memory, --name, etc.

Add `extra_args: dict[str, str] | None = None` that emits `--{key}
{value}` pairs, placed after the --conf loop and before the script
path. Default None preserves the existing cmd shape exactly.

Structured (dict) instead of raw list[str] so LLM agents writing
tool-call payloads can't accidentally pack flag+value into a single
string. Key order in the dict is preserved.

Tests: 4 cases — None default, single pair, multiple pairs in order,
no collision with spark_conf. uv run pytest -> 175 passed.
2026-06-26 11:22:00 +08:00
Claude 70400d3ed1 feat(result): get_job_result tool for terminal view of Spark job
`get_job_status` returns YARN state + the raw response blob, so the
terminal fields (finalStatus, diagnostics, trackingUrl, startedTime,
finishedTime) are buried inside `raw` and not surfaced in a structured
form. Add a new tool that parses them.

`get_job_result(job_id)` reuses `yarn_client.get_application_status` and
extracts:
  - finalStatus  (SUCCEEDED / FAILED / KILLED / UNDEFINED)
  - diagnostics  (YARN final message)
  - tracking_url (Spark Web UI)
  - started_time / finished_time (epoch ms)

All five fields are optional: running jobs have no `finishedTime`, and
older YARN versions (CDH 5 / H2) may omit some fields. Missing fields
stay None — never raise.

Coexistence with `get_job_status` is intentional: the latter is for
polling the running YARN state, the former is the terminal view.

Tests: 4 cases — happy path, running job (no finishedTime), bare-minimum
raw (all optionals None), unknown job_id raises KeyError. uv run pytest
-> 171 passed.
2026-06-26 11:21:53 +08:00
Claude a5c587dc5f feat(yarn_client): H2 fallback to NM container logs when /aggregated-logs 404/501
`/ws/v1/cluster/apps/{id}/aggregated-logs` is a Hadoop 3+ endpoint; on
CDH 5 / Hadoop 2.6 it returns 404 and the old code surfaced a hard error
to the user. Add capability detection in `get_application_logs`: on 404 or
501 from the aggregated-logs endpoint, walk `appattempts` -> `containers`
-> NodeManager `/node/containerlogs/{id}/{user}/` and concatenate the
results.

5xx on the aggregated-logs endpoint still raises immediately (no fallback
attempted for genuine server errors). The H3 fast path is unchanged.

Constraints honored:
- No `subprocess` / `yarn` CLI. The whole point of yarn_client.py is to
  avoid that dependency, so the fallback stays on `httpx` + REST.
- No new pyproject.toml / uv.lock dependencies.
- Public signature of `get_application_logs` unchanged; `tools/logs.py`
  needs no changes.
- No version detection — pure capability detection (try, react to status).

Tests: replaced the old "raise on 404" test with 5 new cases (H3 fast
path, NM fallback, NM path empty, multiple containers, 5xx immediate
raise). `uv run pytest` -> 167 passed.
2026-06-26 11:05:52 +08:00
Claude 1c9e4a321d feat: env-configurable spark-submit binary name
Hardcoding 'spark-submit' as cmd[0] in build_spark_submit_command
breaks for hosts where:
  - both Spark 1.x and 2.x/3.x are installed and 'spark-submit' resolves
    to the wrong one (use 'spark2-submit' or 'spark3-submit' explicitly)
  - the user wants to launch via the PySpark entrypoint ('pyspark')
  - a custom wrapper script sits on PATH (e.g. a credentials-injecting
    'spark-submit-wrapper')

New env var SPARK_EXECUTOR_SPARK_SUBMIT_BIN. Default is 'spark-submit'
(preserves the current behavior for everyone). Override in .env /
docker-compose.yml to change.

common/config.py:
  - new Settings.spark_submit_bin field
  - env-var resolution in from_env() with default 'spark-submit'
  - included in reload() so tests work

spark_executor/core/spark_submit.py:
  - cmd[0] reads settings.spark_submit_bin (was hardcoded 'spark-submit')

.env.example: new section with the override and example values.
docker-compose.yml: forwards the var with the standard 'spark-submit'
default.

Tests: 2 new (settings.spark_submit_bin='spark2-submit', ='pyspark')
plus existing tests updated to use the settings-restore fixture so
mutations don't leak between tests.

165/163 still pass.
2026-06-25 16:46:59 +08:00
Claude 35bff11edf feat: add read_job_file and update_job_file MCP tools
Closes the review-and-edit loop for LLM-generated PySpark code:

  generate_job_file(code=...)      ->  {script_path}
  read_job_file(script_path=...)   ->  {content, path, size}
  update_job_file(path, content)   ->  {path, bytes_written}
  prepare_submit_job(path)         ->  {pending_id, ...}

Or, in a single edit cycle:
  1. generate    (LLM writes initial draft)
  2. read        (LLM or human inspects)
  3. update      (overwrite with edited version)
  4. prepare     (submit for two-step confirmation)

Safety:
  - read_job_file has no path restriction (read-only; useful for
    inspecting any file the agent can see: scripts, logs/, hadoop-conf/)
  - update_job_file is sandboxed to settings.jobs_dir (the same dir
    generate_job_file writes to). Rejects paths outside that tree,
    including ../-traversal attempts. This protects host-mounted
    configs (/etc/passwd, hadoop-conf/*) from being overwritten by
    the agent.
  - 1 MB cap on both reads and writes so MCP responses stay bounded.

Pydantic body models (ReadJobFileRequest, UpdateJobFileRequest) follow
the Stage 1 pattern so tools/call roundtrips long code strings without
the FastAPI query-length 422.

Tests (13 new):
  - 9 unit tests: read success/missing/empty/dir, update success/outside/
    relative-escape/missing/oversize/1mb+1, full edit cycle round-trip
  - 4 integration tests: read via MCP, missing file 400, write+readback
    via MCP, outside-jobs_dir rejection via MCP

163/146 still pass. Live verified end-to-end: generate -> read v1
-> update -> read v2; update /etc/passwd correctly 400'd with
'script_path must be under ... data/jobs/'.
2026-06-25 13:18:39 +08:00
Claude 34e6208f54 feat: SQL safety policy (SELECT/INSERT only) at submit time
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/.
2026-06-25 12:52:16 +08:00
Claude f8e367d43c refactor: unify env-var config in common/config.py
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.
2026-06-25 10:52:53 +08:00
Claude dd197019f9 fix: validate script_path exists + tell agent to call generate_job_file
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).
2026-06-25 10:41:19 +08:00
Claude 994c2d67ab feat: add generate_job_file MCP tool (Stage 2 Task 22)
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/.
2026-06-25 10:19:59 +08:00
Claude 9340061b30 feat(job_writer): write PySpark code to disk with env-var override
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.
2026-06-25 10:17:12 +08:00
Claude 516541a4df feat: Connection.master defaults to 'yarn'
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.
2026-06-25 10:05:27 +08:00
Claude b3564737f7 refactor: replace yarn CLI shell-out with YARN REST API
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.
2026-06-24 17:33:36 +08:00
Claude 5c7dcf57cc feat: structured DEBUG/INFO logging via loguru
- 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.
2026-06-24 15:02:54 +08:00
Claude fad296591c fix: map KeyError->404 and ValueError->400 for MCP clients
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").
2026-06-24 14:59:00 +08:00
Claude 6b9f278294 fix: switch FastAPI routes to Pydantic body models for MCP tools/call
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).
2026-06-24 14:56:31 +08:00
Claude cc7eafaa16 feat: register 12 MCP tools on FastAPI app 2026-06-24 14:50:51 +08:00
Claude 0d60f9b246 feat: add kill_job tool 2026-06-24 14:44:10 +08:00
Claude 1cd3128522 feat: add get_job_logs tool with tail 2026-06-24 14:43:56 +08:00
Claude e443f96756 feat: add get_job_status tool 2026-06-24 14:43:41 +08:00
Claude 63dcae8d75 test: cover cancel_pending_job 2026-06-24 14:43:26 +08:00
Claude c734823e79 feat: add list_pending_jobs and get_pending_job 2026-06-24 14:40:54 +08:00
Claude 2f57c98aa6 feat: add confirm_submit_job (two-step submit flow) 2026-06-24 14:40:33 +08:00
Claude 7f08f0590a feat: add prepare_submit_job (no spark-submit invocation) 2026-06-24 14:40:01 +08:00
Claude 0aa3c1eaf5 feat: add JSON-backed PendingStore 2026-06-24 14:39:13 +08:00
Claude 00592f5cf8 feat: add save_connection tool 2026-06-24 14:37:58 +08:00
Claude 8dd39367ed feat: add JSON-backed ConnectionStore 2026-06-24 14:37:09 +08:00
Claude 31d28baace feat: add in-memory JobStore 2026-06-24 14:36:05 +08:00
Claude a9f64d9942 feat: add yarn client (status, logs, kill) 2026-06-24 14:35:30 +08:00
Claude b4b2ad534a feat: add spark-submit command builder and runner 2026-06-24 14:35:07 +08:00
Claude e6902b00a0 feat: add spark-submit output parser 2026-06-24 14:34:40 +08:00
Claude 864093f1fb chore: add pytest dev-deps and pydantic models (Job, Connection, PendingSubmission) 2026-06-24 14:34:09 +08:00