""" Cluster data-transport helpers for LandSAR distributed processing. Used by _handle_landsar_cluster_item in job_handlers.py to pull input data from the main server and push results back via HTTP. """ from __future__ import annotations import io import json import os import shutil import tempfile import urllib.error import urllib.parse import urllib.request import zipfile from typing import TYPE_CHECKING if TYPE_CHECKING: from ..models.orm import DinsarProductionRunItemORM, DinsarProductionRunORM def _read_cluster_env(name: str) -> str: return str(os.environ.get(name) or "").strip() def _cluster_transfer_timeout() -> int: try: from ..config import read_int_env return read_int_env( "CLUSTER_TRANSFER_TIMEOUT_SECONDS", 3600, minimum=60, maximum=86400, ) except Exception: return 3600 def _cluster_request_headers() -> dict[str, str]: token = _read_cluster_env("CLUSTER_SHARED_TOKEN") if not token: raise RuntimeError("CLUSTER_SHARED_TOKEN is not configured.") return {"X-Cluster-Token": token} def safe_extract_zip(zf: zipfile.ZipFile, target_dir: str) -> None: """Extract a zip after verifying all members stay inside target_dir.""" target_root = os.path.abspath(target_dir) for member in zf.infolist(): member_name = member.filename.replace("\\", "/") if ( not member_name or member_name.startswith("/") or os.path.splitdrive(member_name)[0] or member_name.startswith("../") or "/../" in f"/{member_name}/" ): raise ValueError(f"Unsafe zip member path: {member.filename}") destination = os.path.abspath(os.path.join(target_root, member_name)) if destination != target_root and not destination.startswith( target_root + os.sep ): raise ValueError(f"Unsafe zip member path: {member.filename}") zf.extractall(target_root) def _resolve_cluster_server_url() -> str: """Return the main-server HTTP base URL for cluster data transport. Prefers CLUSTER_MAIN_SERVER_URL; falls back to the DATABASE_URL host. Returns ``http://127.0.0.1`` when nothing is configured (main-server / local worker). """ from ..config import settings explicit = _read_cluster_env("CLUSTER_MAIN_SERVER_URL") or str( getattr(settings, "CLUSTER_MAIN_SERVER_URL", "") or "" ).strip() if explicit: return explicit.rstrip("/") db_url = str(getattr(settings, "DATABASE_URL", "") or "") if "@" in db_url: host_part = db_url.split("@")[1].split("/")[0].split(":")[0] if host_part and host_part not in {"localhost", "127.0.0.1"}: return f"http://{host_part}" return "http://127.0.0.1" def _is_remote_worker() -> bool: """True when this process is configured to talk to a remote main server.""" url = _resolve_cluster_server_url() parsed = urllib.parse.urlparse(url) host = (parsed.hostname or "").strip().lower() return host not in {"", "127.0.0.1", "localhost"} async def materialize_cluster_input( item: DinsarProductionRunItemORM, source_task_dir: str, task_id: str, ) -> None: """Download and extract the input data for a cluster item. Calls ``GET /api/cluster/input-package/{item_id}`` on the main server, retrieves a zip containing the Task_Pool directory tree, and extracts it so that *source_task_dir* exists locally. """ from .task_service import task_service server_url = _resolve_cluster_server_url() download_url = f"{server_url}/api/cluster/input-package/{item.id}" parent_dir = os.path.dirname(source_task_dir) task_name = os.path.basename(source_task_dir) await task_service.add_log( task_id, "INFO", f"[cluster] Downloading input data from {download_url} ...", ) tmp_zip = os.path.join( tempfile.gettempdir(), f"cluster_input_{item.id}_{task_name}.zip", ) try: req = urllib.request.Request( download_url, headers=_cluster_request_headers(), method="GET", ) with urllib.request.urlopen(req, timeout=_cluster_transfer_timeout()) as resp: with open(tmp_zip, "wb") as fh: shutil.copyfileobj(resp, fh, 8 * 1024 * 1024) os.makedirs(parent_dir, exist_ok=True) with zipfile.ZipFile(tmp_zip, "r") as zf: safe_extract_zip(zf, parent_dir) if not os.path.isdir(source_task_dir): raise RuntimeError( f"Extraction did not create expected directory: {source_task_dir}" ) await task_service.add_log( task_id, "INFO", f"[cluster] Input data ready: {source_task_dir}", ) except urllib.error.HTTPError as exc: body_text = "" try: body_text = exc.read().decode("utf-8", errors="replace") except Exception: pass await task_service.add_log( task_id, "WARNING", f"[cluster] Input download HTTP {exc.code}: {body_text[:300]}", ) raise RuntimeError( f"Input download failed HTTP {exc.code}: {body_text[:200]}" ) from exc finally: try: os.unlink(tmp_zip) except Exception: pass async def upload_cluster_result( item: DinsarProductionRunItemORM, run: DinsarProductionRunORM, managed_run_dir: str, run_key: str, task_id: str, ) -> bool: """Package the managed run directory and upload it to the main server. Returns ``True`` when the main server accepted the upload and registered the result in the D-InSAR catalog. """ from .task_service import task_service server_url = _resolve_cluster_server_url() await task_service.add_log( task_id, "INFO", "[cluster] Packaging results for upload ...", ) run_dir_name = os.path.basename(os.path.normpath(managed_run_dir)) tmp_zip = os.path.join( tempfile.gettempdir(), f"cluster_result_{item.id}.zip", ) try: parent = os.path.dirname(managed_run_dir) shutil.make_archive( tmp_zip.replace(".zip", ""), "zip", root_dir=parent, base_dir=run_dir_name, ) await task_service.add_log( task_id, "INFO", "[cluster] Uploading results to main server ...", ) boundary = "----ClusterUploadBoundary" body = io.BytesIO() def _write_field(name, filename, content_type, data): body.write(f"--{boundary}\r\n".encode("utf-8")) if filename: body.write( f'Content-Disposition: form-data; name="{name}"; ' f'filename="{filename}"\r\n'.encode("utf-8") ) body.write(f"Content-Type: {content_type}\r\n".encode("utf-8")) else: body.write( f'Content-Disposition: form-data; name="{name}"\r\n'.encode( "utf-8" ) ) body.write(b"\r\n") body.write(data) body.write(b"\r\n") _write_field("run_id", "", "text/plain", (run.run_id or "").encode("utf-8")) _write_field("run_key", "", "text/plain", str(run_key or "").encode("utf-8")) with open(tmp_zip, "rb") as fh: _write_field( "result_zip", os.path.basename(tmp_zip), "application/zip", fh.read(), ) body.write(f"--{boundary}--\r\n".encode("utf-8")) upload_url = f"{server_url}/api/cluster/upload-result/{item.id}" req = urllib.request.Request( upload_url, data=body.getvalue(), headers={ "Content-Type": f"multipart/form-data; boundary={boundary}", **_cluster_request_headers(), }, method="POST", ) with urllib.request.urlopen(req, timeout=_cluster_transfer_timeout()) as resp: result = json.loads(resp.read().decode("utf-8")) await task_service.add_log( task_id, "INFO", f"[cluster] Results uploaded: " f"registered={result.get('registered', False)} " f"processed={result.get('processed', 0)}", ) return bool(result.get("registered", False)) except urllib.error.HTTPError as exc: body_text = "" try: body_text = exc.read().decode("utf-8", errors="replace") except Exception: pass await task_service.add_log( task_id, "WARNING", f"[cluster] Upload HTTP {exc.code}: {body_text[:300]}", ) raise RuntimeError( f"Upload failed HTTP {exc.code}: {body_text[:200]}" ) from exc finally: try: if os.path.isfile(tmp_zip): os.unlink(tmp_zip) except Exception: pass