44 lines
1.5 KiB
Python
44 lines
1.5 KiB
Python
# 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}
|