# coding=utf-8 """ @Time :2026/6/24 @Author :tao.chen Read + update the contents of an existing PySpark script file. These two tools close the review-and-edit loop: 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, ...} Safety: - read_job_file: any existing regular file. Path-existence only. - update_job_file: must be under SPARK_EXECUTOR_JOBS_DIR (settings.jobs_dir) so the agent cannot overwrite host-mounted configs or arbitrary files on the container FS. - 1 MB cap on both read and write payloads to keep MCP responses bounded. """ import os from pathlib import Path from common.config import settings from common.logging import logger MAX_FILE_BYTES = 1 * 1024 * 1024 # 1 MB class ScriptFileError(ValueError): """Raised when read/update fails. -> HTTP 400 via the FastAPI ValueError handler in server.py. """ pass def _check_readable(script_path: str) -> None: 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}" ) def _check_writable(script_path: str) -> None: """update_job_file is restricted to files under settings.jobs_dir (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 write_job_file to create a new one." ) jobs_root = Path(settings.jobs_dir).resolve() target = Path(script_path).resolve() try: target.relative_to(jobs_root) except ValueError: raise ScriptFileError( f"script_path must be under {jobs_root} (the directory " 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." ) def read_job_file(script_path: str) -> dict[str, object]: """Return the text content of an existing script file. Caps the read at 1 MB to keep MCP responses bounded; raises ScriptFileError (-> 400) if the file is missing or too large. """ _check_readable(script_path) size = os.path.getsize(script_path) if size > MAX_FILE_BYTES: raise ScriptFileError( f"Script is too large to read back ({size} bytes > {MAX_FILE_BYTES} " f"byte cap). Edit it via a host volume mount instead." ) logger.debug(f"read_job_file enter script_path={script_path} size={size}") with open(script_path, encoding="utf-8") as f: content = f.read() logger.info(f"read_job_file ok script_path={script_path} size={size}") return {"path": script_path, "content": content, "size": size} def update_job_file(script_path: str, content: str) -> dict[str, object]: """Overwrite an existing script file with new content. Restricted to paths under settings.jobs_dir. Caps writes at 1 MB. Raises ScriptFileError (-> 400) if the path is missing, outside the allowed dir, or the content is too large. """ _check_writable(script_path) encoded_size = len(content.encode("utf-8")) if encoded_size > MAX_FILE_BYTES: raise ScriptFileError( f"content is too large ({encoded_size} bytes > {MAX_FILE_BYTES} " f"byte cap). Split the script into multiple files." ) logger.debug( f"update_job_file enter script_path={script_path} " f"new_bytes={encoded_size}" ) with open(script_path, "w", encoding="utf-8") as f: written = f.write(content) logger.info( f"update_job_file ok script_path={script_path} bytes_written={written}" ) return {"path": script_path, "bytes_written": written}