96 lines
3.3 KiB
Python
96 lines
3.3 KiB
Python
import json
|
|
from datetime import datetime, timezone
|
|
from uuid import uuid4
|
|
|
|
from fastapi import APIRouter, Depends, HTTPException, Query
|
|
from pydantic import BaseModel
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.database import get_db
|
|
from app.models import DiscoveryRunModel, JobModel
|
|
from app.schemas import DiscoveryRun, EntityResponse, ListResponse
|
|
|
|
router = APIRouter(prefix='/discoveries', tags=['discoveries'])
|
|
|
|
|
|
class DiscoveryRunCreate(BaseModel):
|
|
profile_id: str
|
|
subnet_ids: list[str] = []
|
|
|
|
|
|
def _to_run(m: DiscoveryRunModel) -> DiscoveryRun:
|
|
return DiscoveryRun(
|
|
id=m.id, profile_id=m.profile_id,
|
|
subnet_ids=json.loads(m.subnet_ids_json or '[]'),
|
|
job_id=m.job_id, status=m.status,
|
|
summary=json.loads(m.summary_json or 'null'),
|
|
created_at=m.created_at,
|
|
started_at=m.started_at,
|
|
finished_at=m.finished_at,
|
|
)
|
|
|
|
|
|
@router.get('', response_model=ListResponse[DiscoveryRun])
|
|
def list_discoveries(
|
|
status: str | None = Query(default=None),
|
|
db: Session = Depends(get_db),
|
|
):
|
|
q = db.query(DiscoveryRunModel)
|
|
if status:
|
|
q = q.filter(DiscoveryRunModel.status == status)
|
|
rows = q.order_by(DiscoveryRunModel.created_at.desc()).all()
|
|
return ListResponse(count=len(rows), items=[_to_run(r) for r in rows])
|
|
|
|
|
|
@router.get('/{discovery_id}', response_model=EntityResponse[DiscoveryRun])
|
|
def get_discovery(discovery_id: str, db: Session = Depends(get_db)):
|
|
row = db.query(DiscoveryRunModel).filter(DiscoveryRunModel.id == discovery_id).first()
|
|
if not row:
|
|
raise HTTPException(status_code=404, detail='Discovery not found')
|
|
return EntityResponse(item=_to_run(row))
|
|
|
|
|
|
@router.post('', response_model=EntityResponse[DiscoveryRun], status_code=201)
|
|
def create_discovery(body: DiscoveryRunCreate, db: Session = Depends(get_db)):
|
|
now = datetime.now(timezone.utc)
|
|
run_id = f'run_{uuid4().hex[:8]}'
|
|
job_id = f'job_{uuid4().hex[:8]}'
|
|
job = JobModel(
|
|
id=job_id, type='discovery', status='pending',
|
|
progress=0, payload_json=json.dumps({'run_id': run_id}),
|
|
created_at=now, updated_at=now,
|
|
)
|
|
run = DiscoveryRunModel(
|
|
id=run_id, profile_id=body.profile_id,
|
|
subnet_ids_json=json.dumps(body.subnet_ids),
|
|
job_id=job_id, status='pending',
|
|
created_at=now, updated_at=now,
|
|
)
|
|
db.add(job)
|
|
db.add(run)
|
|
db.commit()
|
|
db.refresh(run)
|
|
return EntityResponse(item=_to_run(run))
|
|
|
|
|
|
@router.post('/{discovery_id}/cancel', response_model=EntityResponse[DiscoveryRun])
|
|
def cancel_discovery(discovery_id: str, db: Session = Depends(get_db)):
|
|
run = db.query(DiscoveryRunModel).filter(DiscoveryRunModel.id == discovery_id).first()
|
|
if not run:
|
|
raise HTTPException(status_code=404, detail='Discovery not found')
|
|
if run.status not in ('pending', 'running'):
|
|
raise HTTPException(
|
|
status_code=409,
|
|
detail=f'Cannot cancel discovery with status {run.status!r}',
|
|
)
|
|
run.status = 'cancelled'
|
|
run.finished_at = datetime.now(timezone.utc)
|
|
if run.job_id:
|
|
job = db.query(JobModel).filter(JobModel.id == run.job_id).first()
|
|
if job and job.status in ('pending', 'running'):
|
|
job.status = 'cancelled'
|
|
job.finished_at = run.finished_at
|
|
db.commit()
|
|
db.refresh(run)
|
|
return EntityResponse(item=_to_run(run))
|