From 7f08f0590a2bbf42f96d4af8eb63ca0dba559de7 Mon Sep 17 00:00:00 2001
From: Claude <claude@anthropic.com>
Date: Wed, 24 Jun 2026 14:39:31 +0800
Subject: [PATCH] feat: add prepare_submit_job (no spark-submit invocation)

---
 spark_executor/tools/submit.py | 57 ++++++++++++++++++++++++
 tests/unit/test_submit_tool.py | 79 ++++++++++++++++++++++++++++++++++
 2 files changed, 136 insertions(+)
 create mode 100644 spark_executor/tools/submit.py
 create mode 100644 tests/unit/test_submit_tool.py

diff --git a/spark_executor/tools/submit.py b/spark_executor/tools/submit.py
new file mode 100644
index 0000000..5e717d9
--- /dev/null
+++ b/spark_executor/tools/submit.py
@@ -0,0 +1,57 @@
+# coding=utf-8
+"""
+@Time   :2026/6/24
+@Author :tao.chen
+"""
+import secrets
+from datetime import datetime
+
+from common.logging import logger
+from spark_executor.core.connection_store import store as conn_store
+from spark_executor.core.pending_store import store as pending_store
+from spark_executor.core.spark_submit import run_spark_submit  # re-exported for monkeypatch in tests
+from spark_executor.models import PendingSubmission
+
+
+def _new_pending_id() -> str:
+    return "p_" + secrets.token_hex(6)
+
+
+def prepare_submit_job(
+    *,
+    connection: str,
+    script_path: str,
+    queue: str = "default",
+    executor_memory: str = "4G",
+    executor_cores: int = 2,
+    num_executors: int = 2,
+) -> dict[str, object]:
+    """Snapshot connection params and persist a PendingSubmission. Does NOT submit."""
+    conn = conn_store.get(connection)
+    if conn is None:
+        raise KeyError(f"Unknown connection: {connection}")
+
+    pending_id = _new_pending_id()
+    pending = PendingSubmission(
+        pending_id=pending_id,
+        connection=connection,
+        master=conn.master,
+        deploy_mode=conn.deploy_mode,
+        script_path=script_path,
+        queue=queue,
+        executor_memory=executor_memory,
+        executor_cores=executor_cores,
+        num_executors=num_executors,
+        spark_conf=dict(conn.spark_conf),
+        created_at=datetime.utcnow(),
+        status="PENDING",
+    )
+    pending_store.save(pending)
+    logger.info(
+        f"prepare_submit_job pending_id={pending_id} connection={connection} master={conn.master}"
+    )
+    return {
+        "pending_id": pending_id,
+        "status": "PENDING",
+        "parameters": pending.model_dump(),
+    }
diff --git a/tests/unit/test_submit_tool.py b/tests/unit/test_submit_tool.py
new file mode 100644
index 0000000..8d50f42
--- /dev/null
+++ b/tests/unit/test_submit_tool.py
@@ -0,0 +1,79 @@
+# coding=utf-8
+from pathlib import Path
+from unittest.mock import patch
+
+import pytest
+
+from spark_executor.core import connection_store, pending_store
+from spark_executor.core.pending_store import PendingStore
+from spark_executor.models import Connection
+from spark_executor.tools import submit
+
+
+@pytest.fixture(autouse=True)
+def _fresh(tmp_path: Path, monkeypatch):
+    monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path))
+    monkeypatch.setattr(connection_store, "store", connection_store.ConnectionStore())
+    monkeypatch.setattr(pending_store, "DEFAULT_DATA_DIR", str(tmp_path))
+    monkeypatch.setattr(pending_store, "store", PendingStore())
+    submit.conn_store = connection_store.store
+    submit.pending_store = pending_store.store
+    connection_store.store.save(Connection(name="prod", master="yarn", deploy_mode="cluster"))
+
+
+def _last_pending_id() -> str:
+    pending = submit.pending_store.list_all()
+    assert pending, "no pending submission was created"
+    return pending[-1].pending_id
+
+
+# --- prepare_submit_job ---
+
+def test_prepare_does_not_invoke_spark_submit(monkeypatch):
+    called = []
+    monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: called.append(cmd))
+    out = submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py")
+    assert called == []  # never called
+    assert out["status"] == "PENDING"
+    assert out["pending_id"].startswith("p_")
+
+
+def test_prepare_persists_pending_with_snapshot(monkeypatch):
+    monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
+    connection_store.store.save(
+        Connection(
+            name="prod",
+            master="yarn",
+            deploy_mode="cluster",
+            spark_conf={"spark.sql.shuffle.partitions": "200"},
+        )
+    )
+    submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py", queue="research")
+    p = submit.pending_store.get(_last_pending_id())
+    assert p is not None
+    assert p.connection == "prod"
+    assert p.master == "yarn"
+    assert p.deploy_mode == "cluster"
+    assert p.spark_conf == {"spark.sql.shuffle.partitions": "200"}
+    assert p.queue == "research"
+    assert p.script_path == "/tmp/j.py"
+
+
+def test_prepare_raises_for_unknown_connection():
+    with pytest.raises(KeyError, match="missing"):
+        submit.prepare_submit_job(connection="missing", script_path="/tmp/j.py")
+
+
+def test_prepare_snapshots_connection_at_prepare_time(monkeypatch):
+    """Editing the connection between prepare and confirm must NOT silently retarget."""
+    monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
+    submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py")
+    p = submit.pending_store.get(_last_pending_id())
+    # Now mutate the connection
+    connection_store.store.save(
+        Connection(name="prod", master="spark://attacker:7077", deploy_mode="client")
+    )
+    # Snapshot is unchanged
+    p2 = submit.pending_store.get(p.pending_id)
+    assert p2.master == "yarn"
+    assert p2.deploy_mode == "cluster"