tao.chen/mcp-server
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: drop fetch_url body cap and modernize datetime usage

Browse Source 
- fetch_url: remove the 1MB body cap and the streaming helper. The
  manual redirect loop with allowlist re-check (the SSRF fix) is kept
  intact; the body is now read in full via resp.text. Drop the now-
  meaningless `truncated` field from FetchUrlResult and the tests that
  asserted on it. Switched from httpx.stream() back to httpx.get()
  for the redirect loop — cleaner without the body cap.

- datetime: replace deprecated datetime.utcnow() with
  datetime.now(timezone.utc) in submit.py (3 sites) and
  core/job_writer.py (1 site). Update the stale comment in
  core/pending_store.py that referenced the old call.

- Clean up an unused `from common.config import settings` import in
  submit.py that ruff flagged.

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
Claude
2026-07-09 14:35:16 +08:00
co-authored by Claude
parent 7d4e512cb0
commit deafb5a26b
7 changed files with 50 additions and 158 deletions
Download Patch File Download Diff File Expand all files Collapse all files
spark_executor/core/job_writer.py
+2 -2
View File
@@ -14,7 +14,7 @@ Resolution order for the output directory:
"""
import os
import secrets
from datetime import datetime
from datetime import datetime, timezone
from common.config import settings
from common.logging import logger
@@ -63,7 +63,7 @@ def write_job_file(code: str, jobs_dir: str | None = None) -> str:
effective_dir = resolve_jobs_dir(jobs_dir)
os.makedirs(effective_dir, exist_ok=True)
stamp = datetime.utcnow().strftime("%Y%m%d%H%M%S")
stamp = datetime.now(timezone.utc).strftime("%Y%m%d%H%M%S")
name = f"job_{stamp}_{secrets.token_hex(3)}.py"
path = os.path.join(effective_dir, name)
abs_path = os.path.abspath(path)
spark_executor/core/pending_store.py
+1 -1
View File
@@ -37,7 +37,7 @@ class PendingStore:
return Path(self._data_dir) / self._dir_name
def _date_path(self, created_at: datetime) -> Path:
# created_at is datetime.utcnow(), so the shard date is a UTC date.
# created_at is timezone-aware UTC, so the shard date is a UTC date.
return self.dir_path / f"{created_at.date().isoformat()}.json"
def _legacy_path(self) -> Path:
spark_executor/models.py
-1
View File
@@ -45,7 +45,6 @@ class FetchUrlResult(BaseModel):
status_code: int
content_type: str
body: str
truncated: bool = False
class ApplicationSummary(BaseModel):
spark_executor/server.py
-1
View File
@@ -549,7 +549,6 @@ def _update_job_file(req: UpdateJobFileRequest):
"guardrails — the allowlist is the only gate — so keep it tight. The "
"Connection's saved auth is reused, so the agent does not need cluster "
"credentials.\n\n"
"**Limits:** response body capped at 1 MB (truncated=true if larger), "
"30s timeout, redirects followed."
),
)
spark_executor/tools/fetch_url.py
+27 -57
View File
@@ -21,7 +21,6 @@ from spark_executor.core.yarn_client import YarnClientConfig
from spark_executor.models import FetchUrlResult
from spark_executor.tools.connections import store as conn_store
_MAX_BODY_BYTES = 1_000_000 # 1 MB cap on response body
_REQUEST_TIMEOUT_SECONDS = 30
_MAX_REDIRECTS = 3
@@ -70,29 +69,6 @@ def _validate_url_host(url: str, allowlist: list[str] | None) -> None:
)
def _read_streamed_body(resp: httpx.Response) -> tuple[bytes, bool]:
"""Stream the response body, keeping at most ``_MAX_BODY_BYTES`` bytes.
Returns ``(body_bytes, truncated)``. Stops reading as soon as the cap is
exceeded so that a multi-gigabyte response from an allowlisted host cannot
OOM the service.
"""
chunks: list[bytes] = []
total = 0
truncated = False
for chunk in resp.iter_bytes():
if total + len(chunk) <= _MAX_BODY_BYTES:
chunks.append(chunk)
total += len(chunk)
else:
if total < _MAX_BODY_BYTES:
chunks.append(chunk[: _MAX_BODY_BYTES - total])
total = _MAX_BODY_BYTES
truncated = True
break
return b"".join(chunks), truncated
def fetch_url(url: str, connection_name: str) -> FetchUrlResult:
"""Proxy an HTTP GET to url using the auth/SSL settings of connection_name."""
logger.debug(f"fetch_url enter url={url} connection_name={connection_name}")
@@ -107,45 +83,39 @@ def fetch_url(url: str, connection_name: str) -> FetchUrlResult:
verify = config.verify_for_httpx()
for hop in range(_MAX_REDIRECTS + 1):
with httpx.stream(
"GET",
resp = httpx.get(
url,
auth=auth,
verify=verify,
timeout=_REQUEST_TIMEOUT_SECONDS,
follow_redirects=False,
) as resp:
if resp.status_code not in _REDIRECT_STATUSES:
body_bytes, truncated = _read_streamed_body(resp)
encoding = resp.encoding or "utf-8"
body = body_bytes.decode(encoding, errors="replace")
logger.info(
f"fetch_url ok url={url} connection_name={connection_name} "
f"status_code={resp.status_code} "
f"content_type={resp.headers.get('content-type', '')} "
f"body_bytes={len(body)} truncated={truncated}"
)
return FetchUrlResult(
url=url,
status_code=resp.status_code,
content_type=resp.headers.get("content-type", ""),
body=body,
truncated=truncated,
)
location = resp.headers.get("Location") or resp.headers.get("location")
if not location:
raise ValueError(
f"HTTP GET returned {resp.status_code} with no Location header at {url}"
)
next_url = urljoin(url, location)
logger.debug(
f"fetch_url redirect url={url} status={resp.status_code} to={next_url}"
)
if resp.status_code not in _REDIRECT_STATUSES:
logger.info(
f"fetch_url ok url={url} connection_name={connection_name} "
f"status_code={resp.status_code} "
f"content_type={resp.headers.get('content-type', '')} "
f"body_chars={len(resp.text)}"
)
_validate_url_host(next_url, conn.url_allowlist)
url = next_url
return FetchUrlResult(
url=url,
status_code=resp.status_code,
content_type=resp.headers.get("content-type", ""),
body=resp.text,
)
location = resp.headers.get("Location") or resp.headers.get("location")
if not location:
raise ValueError(
f"HTTP GET returned {resp.status_code} with no Location header at {url}"
)
next_url = urljoin(url, location)
logger.debug(
f"fetch_url redirect url={url} status={resp.status_code} to={next_url}"
)
_validate_url_host(next_url, conn.url_allowlist)
url = next_url
raise ValueError(
f"fetch_url exceeded {_MAX_REDIRECTS} redirects (last url: {url})"
spark_executor/tools/submit.py
+4 -5
View File
@@ -6,9 +6,8 @@
import os
import secrets
import uuid
from datetime import datetime
from datetime import datetime, timezone
from common.config import settings
from common.logging import logger
from common.sql_guard import validate_pyspark_code
from spark_executor.core.connection_store import store as conn_store
@@ -124,7 +123,7 @@ def prepare_submit_job(
num_executors=num_executors,
spark_conf=dict(conn.spark_conf),
extra_args=dict(extra_args or {}),
created_at=datetime.utcnow(),
created_at=datetime.now(timezone.utc),
status="PENDING",
)
pending_store.save(pending)
@@ -248,7 +247,7 @@ def confirm_submit_job(*, pending_id: str) -> SubmitResult:
application_id=application_id,
script_path=pending.script_path,
queue=pending.queue,
submit_time=datetime.utcnow(),
submit_time=datetime.now(timezone.utc),
connection=pending.connection,
yarn_rm_url=pending.yarn_rm_url,
)
@@ -282,7 +281,7 @@ def confirm_submit_job(*, pending_id: str) -> SubmitResult:
application_id=application_id,
script_path=pending.script_path,
queue=pending.queue,
submit_time=datetime.utcnow(),
submit_time=datetime.now(timezone.utc),
connection=pending.connection,
yarn_rm_url=pending.yarn_rm_url,
)