diff --git a/common/config.py b/common/config.py new file mode 100644 index 0000000..55f7bcb --- /dev/null +++ b/common/config.py @@ -0,0 +1,90 @@ +# coding=utf-8 +""" +@Time :2026/6/24 +@Author :tao.chen + +Single source of truth for environment-driven application configuration. + +Why this module exists: every prior `os.environ.get("SPARK_EXECUTOR_*")` +was scattered across `connection_store.py`, `pending_store.py`, +`yarn_client.py`, `job_writer.py`. The same env var name appeared in +multiple files, and there was no single place to see "what's +configurable?". This file is that place. + +All application code should: + from common.config import settings + ... settings.data_dir / settings.yarn_resource_manager_url / ... + +...rather than `os.environ.get(...)` directly. This way: + - The env var name appears in exactly one place. + - Ops can audit the full config surface from one file. + - Tests can monkeypatch fields on the `settings` singleton directly + (more reliable than monkeypatching env + reloading). + +NOT in this module: + - Gunicorn-specific env vars (workers, threads, bind, timeout) live in + gunicorn.conf.py because they configure gunicorn, not the app. + - `PYTHONUNBUFFERED` — Python runtime flag, set in Dockerfile. + - `SPARK_SUBMIT_OPTS` — JVM flags, not Python config; users pass them + to spark-submit directly. +""" +import os +from dataclasses import dataclass + + +@dataclass +class Settings: + """Application settings loaded from environment variables at import time. + + Mutable so tests can reassign fields directly. For tests that use + `monkeypatch.setenv(...)`, call `settings.reload()` to re-read from + the (now-patched) environment. + """ + + # --- Data persistence (./data/ in dev, /var/lib/... in prod) --- + data_dir: str = "./data" + + # --- Job files (LLM-generated PySpark) --- + # Defaults to /jobs; can be pointed at a larger disk via + # SPARK_EXECUTOR_JOBS_DIR (e.g. /var/spark-jobs on a big-disk host). + jobs_dir: str = "" + + # --- YARN REST client (fallback for Job.yarn_rm_url snapshot) --- + # Set in the env OR per-Connection via save_connection. + yarn_resource_manager_url: str | None = None + + # --- Loguru --- + # Controls stderr verbosity and the info-level file sink. The debug + # file sink always captures DEBUG (full audit trail). + log_level: str = "DEBUG" + + @classmethod + def from_env(cls) -> "Settings": + data_dir = os.environ.get("SPARK_EXECUTOR_DATA_DIR", "./data") + jobs_dir = os.environ.get( + "SPARK_EXECUTOR_JOBS_DIR", os.path.join(data_dir, "jobs") + ) + return cls( + data_dir=data_dir, + jobs_dir=jobs_dir, + # `or None` collapses empty string to None for the URL fallback + yarn_resource_manager_url=os.environ.get("YARN_RESOURCE_MANAGER_URL") or None, + log_level=os.environ.get("SPARK_EXECUTOR_LOG_LEVEL", "DEBUG"), + ) + + def reload(self) -> "Settings": + """Re-read from environment, mutate in place, return self (chainable). + + Useful in tests that `monkeypatch.setenv(...)` and want the + singleton to pick up the new value without a re-import. + """ + fresh = self.from_env() + self.data_dir = fresh.data_dir + self.jobs_dir = fresh.jobs_dir + self.yarn_resource_manager_url = fresh.yarn_resource_manager_url + self.log_level = fresh.log_level + return self + + +# Module-level singleton. Loaded once at import time. +settings = Settings.from_env() diff --git a/common/logging.py b/common/logging.py index 7122a46..1d828cc 100644 --- a/common/logging.py +++ b/common/logging.py @@ -6,6 +6,10 @@ Process-wide loguru configuration. Import `logger` from here in every module instead of instantiating new loggers. +Log level is controlled via `SPARK_EXECUTOR_LOG_LEVEL` (see common/config.py): + DEBUG - default; full verbosity + INFO - quieter; recommended for production + Levels used in this project: DEBUG - entry/exit of public tools, subprocess commands, file I/O paths INFO - business events (job submitted, status changed, connection saved) @@ -15,13 +19,16 @@ Levels used in this project: import sys from pathlib import Path from loguru import logger + +from common.config import settings + Path("data/logs/debug").mkdir(parents=True, exist_ok=True) Path("data/logs/info").mkdir(parents=True, exist_ok=True) logger.remove() logger.add( sys.stderr, - level="DEBUG", + level=settings.log_level, format=( "{time:HH:mm:ss.SSS} | " "{level: <7} | " @@ -34,7 +41,7 @@ logger.add( logger.add( "data/logs/debug/{time:YYYY-MM-DD}.log", - level="DEBUG", + level="DEBUG", # always full at the file level for audit enqueue=True, retention="30 days", compression="gz", @@ -43,7 +50,7 @@ logger.add( logger.add( "data/logs/info/{time:YYYY-MM-DD}.log", - level="INFO", + level=settings.log_level, enqueue=True, retention="30 days", compression="gz", diff --git a/spark_executor/core/connection_store.py b/spark_executor/core/connection_store.py index 15729a5..6f54bb5 100644 --- a/spark_executor/core/connection_store.py +++ b/spark_executor/core/connection_store.py @@ -9,10 +9,11 @@ import tempfile from pathlib import Path from threading import Lock +from common.config import settings from common.logging import logger from spark_executor.models import Connection -DEFAULT_DATA_DIR = os.environ.get("SPARK_EXECUTOR_DATA_DIR", "./data") +DEFAULT_DATA_DIR = settings.data_dir DEFAULT_FILE_NAME = "connections.json" diff --git a/spark_executor/core/job_writer.py b/spark_executor/core/job_writer.py index f174209..3ecbf56 100644 --- a/spark_executor/core/job_writer.py +++ b/spark_executor/core/job_writer.py @@ -16,20 +16,21 @@ import os import secrets from datetime import datetime +from common.config import settings from common.logging import logger -DEFAULT_JOBS_DIR = "./data/jobs" ENV_JOBS_DIR = "SPARK_EXECUTOR_JOBS_DIR" def resolve_jobs_dir(jobs_dir: str | None = None) -> str: - """Pick the effective jobs directory in priority order: arg > env > default.""" + """Pick the effective jobs directory in priority order: arg > settings.jobs_dir. + + `settings.jobs_dir` itself is computed in `common/config.py` as: + SPARK_EXECUTOR_JOBS_DIR env var (if set) > /jobs + """ if jobs_dir is not None: return jobs_dir - env = os.environ.get(ENV_JOBS_DIR) - if env: - return env - return DEFAULT_JOBS_DIR + return settings.jobs_dir def write_job_file(code: str, jobs_dir: str | None = None) -> str: diff --git a/spark_executor/core/pending_store.py b/spark_executor/core/pending_store.py index 8487881..7a01983 100644 --- a/spark_executor/core/pending_store.py +++ b/spark_executor/core/pending_store.py @@ -9,10 +9,11 @@ import tempfile from pathlib import Path from threading import Lock +from common.config import settings from common.logging import logger from spark_executor.models import PendingSubmission -DEFAULT_DATA_DIR = os.environ.get("SPARK_EXECUTOR_DATA_DIR", "./data") +DEFAULT_DATA_DIR = settings.data_dir DEFAULT_FILE_NAME = "pending_jobs.json" diff --git a/spark_executor/core/yarn_client.py b/spark_executor/core/yarn_client.py index afbae3a..7ce94f4 100644 --- a/spark_executor/core/yarn_client.py +++ b/spark_executor/core/yarn_client.py @@ -18,10 +18,10 @@ var if unset. This matches the pattern the original Connection.yarn_rm_url field was designed for, but no longer requires the `yarn` CLI to interpret it. """ import json -import os import httpx +from common.config import settings from common.logging import logger @@ -35,7 +35,7 @@ class YarnConfigError(YarnError): def _base_url(yarn_rm_url: str | None) -> str: """Resolve and validate the RM URL. Raises YarnConfigError if unusable.""" - url = yarn_rm_url or os.environ.get("YARN_RESOURCE_MANAGER_URL") + url = yarn_rm_url or settings.yarn_resource_manager_url if not url: raise YarnConfigError( "YARN ResourceManager URL is not configured. " diff --git a/tests/unit/test_generate_tool.py b/tests/unit/test_generate_tool.py index e95fc80..833efd9 100644 --- a/tests/unit/test_generate_tool.py +++ b/tests/unit/test_generate_tool.py @@ -2,11 +2,29 @@ import os from pathlib import Path +import pytest + +from common import config from spark_executor.tools import generate +@pytest.fixture(autouse=True) +def _restore_settings(): + snapshot = config.Settings( + data_dir=config.settings.data_dir, + jobs_dir=config.settings.jobs_dir, + yarn_resource_manager_url=config.settings.yarn_resource_manager_url, + log_level=config.settings.log_level, + ) + yield + config.settings.data_dir = snapshot.data_dir + config.settings.jobs_dir = snapshot.jobs_dir + config.settings.yarn_resource_manager_url = snapshot.yarn_resource_manager_url + config.settings.log_level = snapshot.log_level + + def test_generate_writes_code_and_returns_path(monkeypatch, tmp_path: Path): - monkeypatch.chdir(tmp_path) + config.settings.jobs_dir = str(tmp_path / "data" / "jobs") out = generate.generate_job_file( "from pyspark.sql import SparkSession\n" "spark = SparkSession.builder.getOrCreate()\n" @@ -20,8 +38,8 @@ def test_generate_writes_code_and_returns_path(monkeypatch, tmp_path: Path): assert "SparkSession.builder.getOrCreate()" in f.read() -def test_generate_uses_env_var_when_set(monkeypatch, tmp_path: Path): - monkeypatch.setenv("SPARK_EXECUTOR_JOBS_DIR", str(tmp_path / "custom")) +def test_generate_uses_settings_jobs_dir(tmp_path: Path): + config.settings.jobs_dir = str(tmp_path / "custom") out = generate.generate_job_file("x = 1\n") assert out["script_path"].startswith(str(tmp_path / "custom")) assert os.path.isfile(out["script_path"]) diff --git a/tests/unit/test_job_writer.py b/tests/unit/test_job_writer.py index 526b086..24cb3fe 100644 --- a/tests/unit/test_job_writer.py +++ b/tests/unit/test_job_writer.py @@ -2,29 +2,48 @@ import os from pathlib import Path +import pytest + +from common import config from spark_executor.core.job_writer import ( - DEFAULT_JOBS_DIR, ENV_JOBS_DIR, resolve_jobs_dir, write_job_file, ) -# --- resolve_jobs_dir (priority: arg > env > default) --- +@pytest.fixture(autouse=True) +def _restore_settings(): + """job_writer uses settings.jobs_dir which is captured at module import. + Tests that set env vars then call settings.reload() — this fixture saves + and restores the original settings so one test doesn't leak into another.""" + snapshot = config.Settings( + data_dir=config.settings.data_dir, + jobs_dir=config.settings.jobs_dir, + yarn_resource_manager_url=config.settings.yarn_resource_manager_url, + log_level=config.settings.log_level, + ) + yield + config.settings.data_dir = snapshot.data_dir + config.settings.jobs_dir = snapshot.jobs_dir + config.settings.yarn_resource_manager_url = snapshot.yarn_resource_manager_url + config.settings.log_level = snapshot.log_level + + +# --- resolve_jobs_dir (priority: arg > settings.jobs_dir) --- def test_resolve_explicit_arg_wins(tmp_path: Path, monkeypatch): - monkeypatch.setenv(ENV_JOBS_DIR, "/from/env") + config.settings.jobs_dir = "/from/settings" assert resolve_jobs_dir(str(tmp_path)) == str(tmp_path) -def test_resolve_env_var_used_when_no_arg(monkeypatch, tmp_path: Path): - monkeypatch.setenv(ENV_JOBS_DIR, str(tmp_path)) +def test_resolve_uses_settings_jobs_dir_when_no_arg(monkeypatch, tmp_path: Path): + config.settings.jobs_dir = str(tmp_path) assert resolve_jobs_dir() == str(tmp_path) -def test_resolve_default_when_neither_set(monkeypatch): - monkeypatch.delenv(ENV_JOBS_DIR, raising=False) - assert resolve_jobs_dir() == DEFAULT_JOBS_DIR +def test_resolve_default_when_settings_unset(monkeypatch): + config.settings.jobs_dir = "./data/jobs" assert resolve_jobs_dir() == "./data/jobs" @@ -48,27 +67,28 @@ def test_write_creates_jobs_dir_if_missing(tmp_path: Path): assert os.path.isfile(out) -def test_write_uses_env_var_when_no_arg(monkeypatch, tmp_path: Path): - monkeypatch.setenv(ENV_JOBS_DIR, str(tmp_path)) - out = write_job_file("env-driven\n") +def test_write_uses_settings_jobs_dir_when_no_arg(tmp_path: Path): + config.settings.jobs_dir = str(tmp_path) + out = write_job_file("settings-driven\n") assert out.startswith(str(tmp_path)) with open(out) as f: - assert f.read() == "env-driven\n" + assert f.read() == "settings-driven\n" -def test_write_uses_default_when_neither_set(monkeypatch, tmp_path: Path): - """Default is ./data/jobs relative to cwd. Run in a tmp dir so the test - doesn't pollute the real project and stays self-contained.""" - monkeypatch.delenv(ENV_JOBS_DIR, raising=False) - monkeypatch.chdir(tmp_path) # pytest built-in: chdir for this test only - out = write_job_file("default\n") - assert os.path.isabs(out) - # Should resolve to /data/jobs/job_*.py - expected_dir = tmp_path / "data" / "jobs" - assert expected_dir.is_dir() - assert str(out).startswith(str(expected_dir)) - with open(out) as f: - assert f.read() == "default\n" +def test_write_uses_default_jobs_dir(tmp_path: Path): + """When settings.jobs_dir is the default './data/jobs', writes resolve + relative to cwd. chdir to tmp so the test doesn't pollute the project.""" + config.settings.jobs_dir = "./data/jobs" + old_cwd = os.getcwd() + try: + os.chdir(tmp_path) + out = write_job_file("default\n") + assert os.path.isabs(out) + expected_dir = tmp_path / "data" / "jobs" + assert expected_dir.is_dir() + assert str(out).startswith(str(expected_dir)) + finally: + os.chdir(old_cwd) def test_write_returns_unique_paths_for_concurrent_calls(tmp_path: Path): diff --git a/tests/unit/test_yarn_client.py b/tests/unit/test_yarn_client.py index 566bcec..6a170c3 100644 --- a/tests/unit/test_yarn_client.py +++ b/tests/unit/test_yarn_client.py @@ -4,6 +4,7 @@ from unittest.mock import patch import httpx import pytest +from common import config from spark_executor.core.yarn_client import ( YarnConfigError, YarnError, @@ -16,6 +17,23 @@ from spark_executor.core.yarn_client import ( RM = "http://rm:8088" +@pytest.fixture(autouse=True) +def _restore_settings(): + """Each test gets a clean copy of `settings` so env-var-style overrides + in one test don't leak into the next.""" + snapshot = config.Settings( + data_dir=config.settings.data_dir, + jobs_dir=config.settings.jobs_dir, + yarn_resource_manager_url=config.settings.yarn_resource_manager_url, + log_level=config.settings.log_level, + ) + yield + config.settings.data_dir = snapshot.data_dir + config.settings.jobs_dir = snapshot.jobs_dir + config.settings.yarn_resource_manager_url = snapshot.yarn_resource_manager_url + config.settings.log_level = snapshot.log_level + + def _resp(status: int, *, json_data=None, text: str | None = None) -> httpx.Response: if json_data is not None: return httpx.Response(status, json=json_data) @@ -57,16 +75,16 @@ def test_status_raises_when_state_field_missing(): def test_status_requires_rm_url(): - with patch.dict("os.environ", {}, clear=True): - with pytest.raises(YarnConfigError): - get_application_status("application_1", None) + config.settings.yarn_resource_manager_url = None + with pytest.raises(YarnConfigError): + get_application_status("application_1", None) -def test_status_falls_back_to_env_var(): +def test_status_falls_back_to_settings_yarn_rm_url(): + config.settings.yarn_resource_manager_url = "http://env-rm:8088" fake = _resp(200, json_data={"app": {"state": "FINISHED"}}) - with patch.dict("os.environ", {"YARN_RESOURCE_MANAGER_URL": "http://env-rm:8088"}, clear=False): - with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake) as m: - state, _ = get_application_status("application_1", None) + with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake) as m: + state, _ = get_application_status("application_1", None) assert state == "FINISHED" assert "env-rm:8088" in m.call_args.args[1]