step-parse/backend/app/main.py

"""STEP Parser API — Phase 1.

Endpoints:
  GET  /api/health
  POST /api/upload                      multipart STEP upload -> creates model + job
  GET  /api/jobs/{job_id}               job status / stage / artifacts
  GET  /api/models                      list models
  GET  /api/models/{id}                 model metadata + BOM + artifacts
  GET  /api/models/{id}/artifacts/{name}  download a generated file
  POST /api/models/{id}/query           run a geometry query
"""
import json
import logging
from contextlib import asynccontextmanager
from pathlib import Path

from fastapi import Depends, FastAPI, File, Form, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from pydantic import BaseModel
from sqlmodel import Session, select

from . import processing
from .config import ALLOWED_SUFFIXES, MODELS_DIR
from .db import get_session, init_db
from .models import BomRow, Job, Model, QueryLog
from .worker import submit_job

logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s: %(message)s")
log = logging.getLogger("step_parser.api")


@asynccontextmanager
async def lifespan(app: FastAPI):
    init_db()
    log.info("DB initialized")
    yield


app = FastAPI(title="STEP Parser API", version="0.1.0", lifespan=lifespan)

# Dev-permissive CORS; tighten to the Unraid frontend origin before it leaves the LAN.
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


class QueryRequest(BaseModel):
    query: str


@app.get("/api/health")
def health():
    return {"status": "ok"}


@app.post("/api/upload")
async def upload(
    file: UploadFile = File(...),
    thumbnails: bool = Form(True),
    bom: bool = Form(True),
    diagram: bool = Form(False),
    translate: bool = Form(False),
    diagram_mode: str = Form("enclosure_only"),
    session: Session = Depends(get_session),
):
    safe_name = Path(file.filename or "").name
    suffix = Path(safe_name).suffix.lower()
    if suffix not in ALLOWED_SUFFIXES:
        raise HTTPException(400, f"Unsupported file type '{suffix}'. Expected .step or .stp")

    stem = Path(safe_name).stem
    model = Model(name=stem, original_filename=safe_name, stem=stem)
    session.add(model)
    session.commit()
    session.refresh(model)

    out_dir = MODELS_DIR / str(model.id)
    out_dir.mkdir(parents=True, exist_ok=True)
    (out_dir / safe_name).write_bytes(await file.read())

    options = {
        "thumbnails": thumbnails, "bom": bom, "diagram": diagram,
        "translate": translate, "diagram_mode": diagram_mode,
    }
    job = Job(model_id=model.id, status="pending", options=json.dumps(options))
    session.add(job)
    session.commit()
    session.refresh(job)

    submit_job(job.id)
    return {"model_id": model.id, "job_id": job.id, "status": job.status}


@app.get("/api/jobs/{job_id}")
def get_job(job_id: int, session: Session = Depends(get_session)):
    job = session.get(Job, job_id)
    if job is None:
        raise HTTPException(404, "job not found")
    data = job.model_dump()
    data["artifacts"] = json.loads(job.artifacts) if job.artifacts else []
    data["options"] = json.loads(job.options) if job.options else {}
    return data


@app.get("/api/models")
def list_models(session: Session = Depends(get_session)):
    return session.exec(select(Model).order_by(Model.id.desc())).all()


@app.get("/api/models/{model_id}")
def get_model(model_id: int, session: Session = Depends(get_session)):
    model = session.get(Model, model_id)
    if model is None:
        raise HTTPException(404, "model not found")
    bom = session.exec(
        select(BomRow).where(BomRow.model_id == model_id).order_by(BomRow.level, BomRow.id)
    ).all()
    latest = session.exec(
        select(Job).where(Job.model_id == model_id).order_by(Job.id.desc())
    ).first()
    artifacts = json.loads(latest.artifacts) if (latest and latest.artifacts) else []
    return {
        "model": model,
        "bom": bom,
        "artifacts": artifacts,
        "latest_job": {"id": latest.id, "status": latest.status, "stage": latest.stage} if latest else None,
    }


@app.get("/api/models/{model_id}/artifacts/{name}")
def get_artifact(model_id: int, name: str):
    safe = Path(name).name  # block path traversal
    path = MODELS_DIR / str(model_id) / safe
    if not path.is_file():
        raise HTTPException(404, "artifact not found")
    return FileResponse(path, filename=safe)


@app.post("/api/models/{model_id}/query")
def query_model(model_id: int, req: QueryRequest, session: Session = Depends(get_session)):
    model = session.get(Model, model_id)
    if model is None:
        raise HTTPException(404, "model not found")
    step_path = MODELS_DIR / str(model_id) / model.original_filename
    if not step_path.is_file():
        raise HTTPException(404, "model source file missing")
    try:
        result = processing.run_query(step_path, req.query)
    except Exception as e:  # noqa: BLE001
        raise HTTPException(500, f"query failed: {type(e).__name__}: {e}")
    session.add(QueryLog(model_id=model_id, query=req.query, result=result))
    session.commit()
    return {"query": req.query, "result": result}