""" Cluster data-transport endpoints for LandSAR distributed processing. Provides HTTP-based input materialize (download) and result upload so that remote Windows workers can pull Task_Pool pair data and push finished products back to the main server without SMB / UNC. """ from __future__ import annotations import os import shutil import tempfile import zipfile from typing import Optional from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from fastapi import ( APIRouter, Depends, File, Form, Header, HTTPException, UploadFile, ) from fastapi.responses import FileResponse from starlette.background import BackgroundTask from ..config import settings from ..database import get_db from ..models.orm import ( DinsarProductionExecutionORM, DinsarProductionRunItemORM, DinsarProductionRunORM, ) from ..services.cluster_transport import safe_extract_zip router = APIRouter() def _cluster_materialize_temp_dir() -> str: explicit = str(os.environ.get("CLUSTER_MATERIALIZE_TEMP_DIR") or "").strip() if not explicit: explicit = str( getattr(settings, "CLUSTER_MATERIALIZE_TEMP_DIR", "") or "" ).strip() if explicit: return os.path.normpath(explicit) task_pool = str(getattr(settings, "TASK_POOL_ROOT", "") or "").strip() if task_pool: return os.path.normpath(os.path.join(task_pool, "_cluster_temp")) return os.path.normpath( os.path.join( os.path.dirname(__file__), "..", "..", "runtime", "_cluster_temp" ) ) def _cluster_shared_token() -> str: return str(os.environ.get("CLUSTER_SHARED_TOKEN") or "").strip() def _require_cluster_token( x_cluster_token: Optional[str] = Header(default=None), ) -> None: expected = _cluster_shared_token() if not expected: raise HTTPException( status_code=503, detail="CLUSTER_SHARED_TOKEN is not configured.", ) if x_cluster_token != expected: raise HTTPException(status_code=403, detail="Invalid cluster token.") # --------------------------------------------------------------------------- # Download input package # --------------------------------------------------------------------------- @router.get("/cluster/input-package/{item_id}") async def download_cluster_input_package( item_id: int, _cluster_token: None = Depends(_require_cluster_token), db: AsyncSession = Depends(get_db), ): """Package a cluster item's source-task directory as a zip and stream it. The remote worker calls this when *source_task_dir* does not exist locally. The zip mirrors the Task_Pool layout: Task_YYYYMMDD_YYYYMMDD/ master/ ... slave/ ... orbit/ ... pair_metadata.json """ item = await db.get(DinsarProductionRunItemORM, int(item_id)) if item is None: raise HTTPException(status_code=404, detail="Cluster item not found.") source_dir = os.path.normpath(str(item.source_task_dir or "")) if not source_dir or not os.path.isdir(source_dir): raise HTTPException( status_code=404, detail=f"Source task directory not found: {source_dir}", ) task_name = os.path.basename(source_dir) parent_dir = os.path.dirname(source_dir) temp_root = _cluster_materialize_temp_dir() os.makedirs(temp_root, exist_ok=True) package_root = tempfile.mkdtemp( prefix=f"cluster_input_{item_id}_", dir=temp_root, ) tmp_base = os.path.join(package_root, task_name) try: zip_path = shutil.make_archive( tmp_base, "zip", root_dir=parent_dir, base_dir=task_name, ) except Exception as exc: try: shutil.rmtree(package_root) except Exception: pass raise HTTPException( status_code=500, detail=f"Failed to create input package: {exc}", ) def _cleanup(): try: if os.path.isdir(package_root): shutil.rmtree(package_root) except Exception: pass return FileResponse( path=zip_path, media_type="application/zip", filename=f"{task_name}.zip", background=BackgroundTask(_cleanup), ) # --------------------------------------------------------------------------- # Upload result package # --------------------------------------------------------------------------- @router.post("/cluster/upload-result/{item_id}") async def upload_cluster_result( item_id: int, run_id: str = Form(...), run_key: str = Form(...), result_zip: UploadFile = File(...), _cluster_token: None = Depends(_require_cluster_token), db: AsyncSession = Depends(get_db), ): """Receive a result zip from a cluster worker and register it in the D-InSAR catalog. The zip is extracted under the cluster item's standard ``results_root_dir/runs/`` directory and then *result_catalog_service.publish_from_sources* is called so the product is immediately visible on the main server. """ item = await db.get(DinsarProductionRunItemORM, int(item_id)) if item is None: raise HTTPException(status_code=404, detail="Cluster item not found.") normalized_run_id = str(run_id or "").strip() if normalized_run_id != str(item.run_id or "").strip(): raise HTTPException( status_code=400, detail="run_id does not match item.", ) publish_root = os.path.normpath( str(getattr(settings, "DINSAR_PRODUCT_DIR", "") or "") ) if not publish_root or not os.path.isdir(publish_root): raise HTTPException( status_code=500, detail="DINSAR_PRODUCT_DIR is not configured.", ) normalized_run_key = str(run_key or "").strip() if not normalized_run_key: raise HTTPException(status_code=400, detail="run_key is required.") item_results_root = os.path.normpath(str(item.results_root_dir or "")) if not item_results_root: raise HTTPException( status_code=500, detail="Cluster item results_root_dir is not configured.", ) publish_root_abs = os.path.abspath(publish_root) item_results_root_abs = os.path.abspath(item_results_root) if item_results_root_abs != publish_root_abs and not item_results_root_abs.startswith( publish_root_abs + os.sep ): raise HTTPException( status_code=400, detail="Cluster item results_root_dir is outside DINSAR_PRODUCT_DIR.", ) extract_dir = os.path.join(item_results_root_abs, "runs", normalized_run_key) os.makedirs(os.path.dirname(extract_dir), exist_ok=True) extract_parent = os.path.dirname(extract_dir) os.makedirs(extract_parent, exist_ok=True) upload_root = tempfile.mkdtemp( prefix=f"cluster_upload_{item_id}_", dir=extract_parent, ) tmp_extract = os.path.join(upload_root, "extract") tmp_zip = os.path.join(upload_root, "upload.zip") backup_dir: Optional[str] = None try: with open(tmp_zip, "wb") as fh: while chunk := await result_zip.read(8 * 1024 * 1024): # 8 MiB fh.write(chunk) os.makedirs(tmp_extract, exist_ok=True) with zipfile.ZipFile(tmp_zip, "r") as zf: safe_extract_zip(zf, tmp_extract) if os.path.isdir(extract_dir): backup_dir = f"{extract_dir}._replace_backup" if os.path.isdir(backup_dir): shutil.rmtree(backup_dir) os.replace(extract_dir, backup_dir) os.replace(tmp_extract, extract_dir) if backup_dir and os.path.isdir(backup_dir): shutil.rmtree(backup_dir) backup_dir = None except zipfile.BadZipFile: raise HTTPException( status_code=400, detail="Uploaded file is not a valid zip archive.", ) except ValueError as exc: raise HTTPException(status_code=400, detail=str(exc)) from exc except Exception: if backup_dir and os.path.isdir(backup_dir) and not os.path.exists(extract_dir): try: os.replace(backup_dir, extract_dir) backup_dir = None except Exception: pass raise finally: try: if os.path.isdir(upload_root): shutil.rmtree(upload_root) except Exception: pass try: if backup_dir and os.path.isdir(backup_dir): shutil.rmtree(backup_dir) except Exception: pass # ----- catalog registration ------------------------------------------------ from ..services.result_catalog_service import result_catalog_service as rcs from ..services.dinsar_production_service import dinsar_production_service try: publish_result = await rcs.publish_from_sources(db, [extract_dir]) processed = int(publish_result.get("processed", 0) or 0) failed = int(publish_result.get("failed", 0) or 0) if processed > 0: await rcs.rebuild_catalog(db, full_rebuild=True) if processed != 1 or failed != 0: raise RuntimeError( f"expected processed=1 failed=0, got processed={processed} failed={failed}" ) details = publish_result.get("details") if isinstance(publish_result, dict) else [] detail = next( ( item_detail for item_detail in (details or []) if str(item_detail.get("run_key") or "").strip() == normalized_run_key ), (details or [{}])[0] if details else {}, ) execution_manifest_path = str( detail.get("execution_manifest_path") or os.path.join(extract_dir, "execution_manifest.json") ) current_pointer_path = str(detail.get("current_pointer_path") or "") if not os.path.isfile(execution_manifest_path): raise RuntimeError( f"execution manifest was not created: {execution_manifest_path}" ) if not current_pointer_path or not os.path.isfile(current_pointer_path): raise RuntimeError( f"current pointer was not created: {current_pointer_path or ''}" ) run = await db.get(DinsarProductionRunORM, str(item.run_id or "")) if run is None: raise RuntimeError(f"Cluster run not found for item {item_id}: {item.run_id}") execution_result = await db.execute( select(DinsarProductionExecutionORM).where( DinsarProductionExecutionORM.item_id == item.id, DinsarProductionExecutionORM.run_key == normalized_run_key, ) ) execution = execution_result.scalar_one_or_none() if execution is None: raise RuntimeError( f"Cluster execution not found for item={item.id} run_key={normalized_run_key}" ) execution.output_dir = extract_dir item.latest_output_dir = extract_dir await dinsar_production_service.mark_item_completed( run=run, item=item, execution=execution, manifest_path=execution_manifest_path, metrics={ "cluster_upload": True, "catalog_processed": processed, "catalog_failed": failed, "current_pointer_path": current_pointer_path, }, db=db, ) final_status = await dinsar_production_service.finalize_cluster_run_if_complete( run, db=db, ) return { "registered": processed > 0, "completed": True, "final_status": final_status, "processed": processed, "failed": failed, "catalog_path": extract_dir, "execution_manifest_path": execution_manifest_path, "current_pointer_path": current_pointer_path, } except Exception as exc: raise HTTPException( status_code=500, detail=f"Catalog registration failed: {exc}", ) from exc