Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| import asyncio | |
| import logging | |
| import uuid | |
| from datetime import datetime, timezone, timedelta | |
| from typing import Optional | |
| from fastapi import APIRouter, HTTPException, Depends, Request | |
| from pydantic import BaseModel | |
| from app.deps import limiter | |
| from app.services.supabase import get_supabase | |
| from app.services.auth import require_user_id | |
| from app.services.ssrf import validate_url | |
| logger = logging.getLogger(__name__) | |
| router = APIRouter(prefix="/api/ngs", tags=["ngs"]) | |
| _TABLE = "ngs_jobs" | |
| _MAX_JOBS = 200 | |
| _JOB_TTL = 7200 | |
| def _prune_jobs() -> None: | |
| sb = get_supabase() | |
| cutoff = (datetime.now(timezone.utc) - timedelta(seconds=_JOB_TTL)).strftime('%Y-%m-%dT%H:%M:%S') | |
| sb.table(_TABLE).delete().lt("done_at", cutoff).execute() | |
| count = sb.table(_TABLE).select("id", count="exact").execute().count or 0 | |
| if count > _MAX_JOBS: | |
| to_delete = ( | |
| sb.table(_TABLE) | |
| .select("id") | |
| .in_("status", ("complete", "failed")) | |
| .order("created_at", desc=True) | |
| .range(_MAX_JOBS, _MAX_JOBS + 500) | |
| .execute() | |
| .data | |
| ) | |
| ids = [r["id"] for r in to_delete] | |
| if ids: | |
| sb.table(_TABLE).delete().in_("id", ids).execute() | |
| class NGSRequest(BaseModel): | |
| fastq_url: str | |
| reference: str = "sars-cov-2" | |
| class NGSJob(BaseModel): | |
| job_id: str | |
| fastq_url: str | |
| reference: str | |
| status: str = "queued" | |
| result: Optional[dict] = None | |
| error: Optional[str] = None | |
| created_at: str = "" | |
| done_at: Optional[str] = None | |
| def _init(job_id: str, req: NGSRequest, user_id: str) -> None: | |
| try: | |
| _prune_jobs() | |
| except Exception: | |
| pass | |
| _ensure_ngs_table() | |
| get_supabase().table(_TABLE).insert({ | |
| "id": job_id, | |
| "fastq_url": req.fastq_url, | |
| "reference": req.reference, | |
| "status": "queued", | |
| "user_id": user_id, | |
| "result": None, | |
| "error": None, | |
| "done_at": None, | |
| "attempts": 0, | |
| "max_attempts": 3, | |
| }).execute() | |
| _table_initialized = False | |
| def _ensure_ngs_table() -> None: | |
| """Ensure ngs_jobs table has created_at column and the claim RPC exists.""" | |
| global _table_initialized | |
| if _table_initialized: | |
| return | |
| _table_initialized = True | |
| try: | |
| import httpx as _httpx | |
| from app.config import settings | |
| base = settings.SUPABASE_URL.rstrip("/") | |
| key = settings.SUPABASE_SERVICE_ROLE_KEY | |
| headers = {"apikey": key, "Authorization": f"Bearer {key}"} | |
| # Add created_at column if missing | |
| sql = "ALTER TABLE ngs_jobs ADD COLUMN IF NOT EXISTS created_at timestamptz DEFAULT now();" | |
| resp = _httpx.post(f"{base}/rest/v1/rpc/exec_sql", headers=headers, json={"query": sql}, timeout=15) | |
| if resp.status_code == 200: | |
| logger.info("ngs_jobs: added created_at column") | |
| # Create the claim RPC | |
| rpc_sql = """CREATE OR REPLACE FUNCTION claim_next_ngs_job(worker_id text) | |
| RETURNS ngs_jobs LANGUAGE plpgsql SECURITY DEFINER AS $do$ | |
| DECLARE job ngs_jobs; | |
| BEGIN | |
| SELECT * INTO job FROM ngs_jobs WHERE status = 'queued' AND attempts < max_attempts | |
| ORDER BY created_at ASC LIMIT 1 FOR UPDATE SKIP LOCKED; | |
| IF job.id IS NOT NULL THEN | |
| UPDATE ngs_jobs SET status='running', claimed_at=now(), claimed_by=worker_id, | |
| attempts=attempts+1, updated_at=now() WHERE id = job.id RETURNING * INTO job; | |
| END IF; | |
| RETURN job; | |
| END; $do$;""" | |
| resp = _httpx.post(f"{base}/rest/v1/rpc/exec_sql", headers=headers, json={"query": rpc_sql}, timeout=15) | |
| if resp.status_code == 200: | |
| logger.info("ngs_jobs: created claim_next_ngs_job RPC") | |
| except Exception as e: | |
| logger.debug("Auto-migration skipped (exec_sql not available): %s", e) | |
| def _patch(job_id: str, **kw) -> None: | |
| get_supabase().table(_TABLE).update(kw).eq("id", job_id).execute() | |
| def _read(job_id: str, user_id: str | None = None) -> dict | None: | |
| query = get_supabase().table(_TABLE).select("*").eq("id", job_id) | |
| if user_id: | |
| query = query.eq("user_id", user_id) | |
| rows = query.execute().data | |
| if not rows: | |
| return None | |
| job = dict(rows[0]) | |
| if job.get("storage_url") and not job.get("result"): | |
| from app.services.artifact_storage import download_json | |
| result = download_json(job["storage_url"]) | |
| if result: | |
| job["result"] = result | |
| return job | |
| async def _worker(job_id: str) -> None: | |
| job = _read(job_id) | |
| if not job: | |
| return | |
| _patch(job_id, status="running") | |
| from app.tools.ngs import NGSPipeline, PIPELINE_TIMEOUT | |
| tool = NGSPipeline() | |
| try: | |
| result = await asyncio.wait_for( | |
| tool.run({ | |
| "fastq_url": job["fastq_url"], | |
| "reference": job["reference"], | |
| "job_id": job_id, | |
| }), | |
| timeout=PIPELINE_TIMEOUT, | |
| ) | |
| except asyncio.TimeoutError: | |
| _patch(job_id, status="failed", error="Pipeline timed out", done_at=datetime.now(timezone.utc).strftime('%Y-%m-%dT%H:%M:%S')) | |
| return | |
| if "error" in result and not result.get("steps_completed"): | |
| _patch(job_id, status="failed", error=result["error"], done_at=datetime.now(timezone.utc).strftime('%Y-%m-%dT%H:%M:%S')) | |
| else: | |
| from app.services.artifact_storage import upload_json | |
| storage_url = upload_json(job_id, "result", result) | |
| _patch(job_id, status="complete", storage_url=storage_url, result=None, done_at=datetime.now(timezone.utc).strftime('%Y-%m-%dT%H:%M:%S')) | |
| VALID_DEMO = {"synthetic", "demo", "test"} | |
| async def run_ngs(request: Request, req: NGSRequest, user_id: str = Depends(require_user_id)): | |
| if not req.fastq_url.strip(): | |
| raise HTTPException(400, detail="fastq_url is required") | |
| if req.fastq_url.lower() not in VALID_DEMO: | |
| if not req.fastq_url.startswith(("http://", "https://")): | |
| raise HTTPException(400, detail="fastq_url must be a valid URL or 'synthetic' for demo data") | |
| validate_url(req.fastq_url) | |
| job_id = str(uuid.uuid4()) | |
| _init(job_id, req, user_id) | |
| return {"job_id": job_id, "status": "queued"} | |
| async def get_status(job_id: str, user_id: str = Depends(require_user_id)): | |
| job = _read(job_id, user_id) | |
| if not job: | |
| raise HTTPException(404, detail=f"Job {job_id} not found") | |
| return job | |
| async def list_references(): | |
| from app.tools.ngs import REFERENCE_URLS | |
| return { | |
| "references": [ | |
| {"id": k, "name": k.replace("-", " ").title()} | |
| for k in REFERENCE_URLS | |
| ] | |
| } | |