mirror of
https://github.com/hansjone/netx.git
synced 2026-10-08 23:33:21 +08:00
Topology MCP adds pull/compact/bundle/suggest-hubs with a no-template skill path; API gains fabric level and NE collection policy; web list pages get paging and denser collect/network workflows. Co-authored-by: Cursor <cursoragent@cursor.com>
742 lines
25 KiB
Python
742 lines
25 KiB
Python
from __future__ import annotations
|
|
|
|
import io
|
|
import logging
|
|
import re
|
|
import shutil
|
|
import zipfile
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Any, TypedDict
|
|
|
|
from fastapi import HTTPException
|
|
from sqlalchemy import func, or_
|
|
from sqlalchemy.orm import Session
|
|
|
|
from .models import ManagedNE, NeCollectionJob, NeCollectionRun, UmeCliOverride, UmeInventoryNE
|
|
from .collection_job_state import (
|
|
finalize_collection_job,
|
|
reconcile_stale_collection_job,
|
|
sync_job_progress,
|
|
_sync_job_counts,
|
|
)
|
|
from .collection_policy import (
|
|
ensure_policy,
|
|
expand_policy_targets,
|
|
history_keep_value,
|
|
next_due_at,
|
|
policy_to_out,
|
|
prune_collection_jobs,
|
|
)
|
|
from .collection_schemas import (
|
|
CollectionDashboardOut,
|
|
CollectionJobCreate,
|
|
CollectionJobOut,
|
|
CollectionJobSummary,
|
|
CollectionRunOut,
|
|
)
|
|
from .ne_collection_paths import clear_run_output_files, collection_data_root
|
|
|
|
_log = logging.getLogger("netx.collection")
|
|
|
|
|
|
class CollectionSchedulePayload(TypedDict):
|
|
job_id: str
|
|
run_ids: list[str]
|
|
commands: list[str]
|
|
|
|
|
|
def _now() -> datetime:
|
|
return datetime.now()
|
|
|
|
|
|
def _parse_commands(text: str) -> list[str]:
|
|
lines = []
|
|
for raw in str(text or "").splitlines():
|
|
line = raw.strip()
|
|
if not line or line.startswith("#"):
|
|
continue
|
|
lines.append(line)
|
|
return lines
|
|
|
|
|
|
def _output_counts_for_jobs(db: Session, job_ids: list[str]) -> dict[str, int]:
|
|
if not job_ids:
|
|
return {}
|
|
rows = (
|
|
db.query(NeCollectionRun.job_id, func.count())
|
|
.filter(
|
|
NeCollectionRun.job_id.in_(job_ids),
|
|
NeCollectionRun.output_rel_path != "",
|
|
NeCollectionRun.output_rel_path.isnot(None),
|
|
)
|
|
.group_by(NeCollectionRun.job_id)
|
|
.all()
|
|
)
|
|
return {str(job_id): int(count) for job_id, count in rows}
|
|
|
|
|
|
def _output_count_for_job(db: Session, job_id: str) -> int:
|
|
return int(
|
|
db.query(func.count())
|
|
.select_from(NeCollectionRun)
|
|
.filter(
|
|
NeCollectionRun.job_id == job_id,
|
|
NeCollectionRun.output_rel_path != "",
|
|
NeCollectionRun.output_rel_path.isnot(None),
|
|
)
|
|
.scalar()
|
|
or 0
|
|
)
|
|
|
|
|
|
def job_to_out(row: NeCollectionJob, *, output_count: int | None = None) -> CollectionJobOut:
|
|
return CollectionJobOut(
|
|
id=str(row.id),
|
|
title=str(row.title or ""),
|
|
commands=str(row.commands or ""),
|
|
trigger_mode=str(getattr(row, "trigger_mode", None) or "manual"),
|
|
status=str(row.status or "pending"),
|
|
ne_count=int(row.ne_count or 0),
|
|
success_count=int(row.success_count or 0),
|
|
fail_count=int(row.fail_count or 0),
|
|
output_count=int(output_count if output_count is not None else 0),
|
|
error_message=str(row.error_message or "")[:1000],
|
|
created_at=row.created_at,
|
|
started_at=row.started_at,
|
|
ended_at=row.ended_at,
|
|
last_run_at=row.last_run_at,
|
|
)
|
|
|
|
|
|
def job_to_summary(row: NeCollectionJob | None) -> CollectionJobSummary | None:
|
|
if row is None:
|
|
return None
|
|
return CollectionJobSummary(
|
|
id=str(row.id),
|
|
title=str(row.title or "").strip() or str(row.id)[:8],
|
|
status=str(row.status or "pending"),
|
|
trigger_mode=str(getattr(row, "trigger_mode", None) or "manual"),
|
|
ne_count=int(row.ne_count or 0),
|
|
success_count=int(row.success_count or 0),
|
|
fail_count=int(row.fail_count or 0),
|
|
created_at=row.created_at,
|
|
started_at=row.started_at,
|
|
ended_at=row.ended_at,
|
|
last_run_at=row.last_run_at,
|
|
)
|
|
|
|
|
|
def has_active_collection_job(db: Session) -> NeCollectionJob | None:
|
|
return (
|
|
db.query(NeCollectionJob)
|
|
.filter(NeCollectionJob.status.in_(("pending", "running", "paused")))
|
|
.order_by(NeCollectionJob.created_at.desc())
|
|
.first()
|
|
)
|
|
|
|
|
|
def last_finished_collection_job(db: Session) -> NeCollectionJob | None:
|
|
return (
|
|
db.query(NeCollectionJob)
|
|
.filter(NeCollectionJob.status.in_(("done", "failed")))
|
|
.order_by(NeCollectionJob.created_at.desc())
|
|
.first()
|
|
)
|
|
|
|
|
|
def get_collection_dashboard(db: Session) -> CollectionDashboardOut:
|
|
job_count = int(db.query(NeCollectionJob).count() or 0)
|
|
active_count = int(
|
|
db.query(NeCollectionJob)
|
|
.filter(NeCollectionJob.status.in_(("pending", "running", "paused")))
|
|
.count()
|
|
or 0
|
|
)
|
|
running = has_active_collection_job(db)
|
|
if running is not None:
|
|
jid = str(running.id)
|
|
if str(running.status or "") not in ("done", "failed", "paused"):
|
|
reconcile_stale_collection_job(db, jid)
|
|
db.refresh(running)
|
|
if str(running.status or "") == "running":
|
|
sync_job_progress(db, jid)
|
|
finalize_collection_job(db, jid)
|
|
db.refresh(running)
|
|
if str(running.status or "") not in ("pending", "running", "paused"):
|
|
running = has_active_collection_job(db)
|
|
active_count = int(
|
|
db.query(NeCollectionJob)
|
|
.filter(NeCollectionJob.status.in_(("pending", "running", "paused")))
|
|
.count()
|
|
or 0
|
|
)
|
|
last = last_finished_collection_job(db)
|
|
policy = ensure_policy(db)
|
|
return CollectionDashboardOut(
|
|
job_count=job_count,
|
|
active_count=active_count,
|
|
running_job=job_to_summary(running),
|
|
last_job=job_to_summary(last),
|
|
next_due_at=next_due_at(db, policy),
|
|
policy=policy_to_out(policy),
|
|
)
|
|
|
|
|
|
def run_to_out(row: NeCollectionRun) -> CollectionRunOut:
|
|
rel = str(row.output_rel_path or "").strip()
|
|
return CollectionRunOut(
|
|
id=str(row.id),
|
|
job_id=str(row.job_id),
|
|
ne_id=str(row.ne_id),
|
|
ne_source=str(getattr(row, "ne_source", None) or "managed"),
|
|
ne_name=str(row.ne_name or ""),
|
|
ne_ip=str(row.ne_ip or ""),
|
|
status=str(row.status or "pending"),
|
|
message=str(row.message or ""),
|
|
output_rel_path=rel,
|
|
has_output=bool(rel),
|
|
started_at=row.started_at,
|
|
ended_at=row.ended_at,
|
|
)
|
|
|
|
|
|
def list_eligible_ne(
|
|
db: Session,
|
|
*,
|
|
page: int = 1,
|
|
page_size: int = 200,
|
|
keyword: str = "",
|
|
) -> dict[str, Any]:
|
|
from .device_types import WEBCRT_NE_SOURCE
|
|
|
|
page = max(1, int(page or 1))
|
|
page_size = max(1, min(500, int(page_size or 200)))
|
|
kw = str(keyword or "").strip()
|
|
like = f"%{kw}%" if kw else ""
|
|
|
|
# Managed inventory (any connect_status); exclude WebCRT session hosts.
|
|
m_stmt = db.query(ManagedNE).filter(ManagedNE.source != WEBCRT_NE_SOURCE)
|
|
if like:
|
|
m_stmt = m_stmt.filter(
|
|
or_(
|
|
ManagedNE.name.ilike(like),
|
|
ManagedNE.ip_address.ilike(like),
|
|
ManagedNE.vendor.ilike(like),
|
|
ManagedNE.device_type.ilike(like),
|
|
)
|
|
)
|
|
managed_items = [
|
|
{
|
|
"id": str(x.id),
|
|
"source": "managed",
|
|
"name": str(x.name or ""),
|
|
"vendor": str(x.vendor or ""),
|
|
"device_type": str(x.device_type or ""),
|
|
"ip_address": str(x.ip_address or ""),
|
|
"connect_status": str(x.connect_status or ""),
|
|
"connect_tested_at": x.connect_tested_at.isoformat() if x.connect_tested_at else None,
|
|
}
|
|
for x in m_stmt.order_by(ManagedNE.name.asc()).all()
|
|
]
|
|
|
|
u_stmt = db.query(UmeInventoryNE)
|
|
if like:
|
|
u_stmt = u_stmt.filter(
|
|
or_(
|
|
UmeInventoryNE.ne_id.ilike(like),
|
|
UmeInventoryNE.ne_name.ilike(like),
|
|
UmeInventoryNE.user_label.ilike(like),
|
|
UmeInventoryNE.host_name.ilike(like),
|
|
UmeInventoryNE.ip_address.ilike(like),
|
|
UmeInventoryNE.vendor.ilike(like),
|
|
UmeInventoryNE.ne_type.ilike(like),
|
|
)
|
|
)
|
|
ume_rows = u_stmt.order_by(UmeInventoryNE.host_name.asc(), UmeInventoryNE.ne_name.asc()).all()
|
|
overrides = {
|
|
str(o.ume_ne_id): o
|
|
for o in db.query(UmeCliOverride).filter(
|
|
UmeCliOverride.ume_ne_id.in_([str(x.ne_id) for x in ume_rows] or ["__none__"])
|
|
).all()
|
|
}
|
|
ume_items = []
|
|
for x in ume_rows:
|
|
uid = str(x.ne_id)
|
|
ov = overrides.get(uid)
|
|
ume_items.append(
|
|
{
|
|
"id": uid,
|
|
"source": "ume",
|
|
"name": str(x.host_name or x.user_label or x.ne_name or x.ip_address or uid),
|
|
"vendor": str(x.vendor or "ZTE"),
|
|
"device_type": str(x.ne_type or ""),
|
|
"ip_address": str(x.ip_address or ""),
|
|
"connect_status": str(ov.connect_status or "unknown") if ov else "unknown",
|
|
"connect_tested_at": ov.connect_tested_at.isoformat()
|
|
if ov and ov.connect_tested_at
|
|
else None,
|
|
}
|
|
)
|
|
|
|
items = managed_items + ume_items
|
|
items.sort(key=lambda r: (str(r.get("name") or "").lower(), str(r.get("id") or "")))
|
|
total = len(items)
|
|
start = (page - 1) * page_size
|
|
page_items = items[start : start + page_size]
|
|
return {"total": total, "page": page, "page_size": page_size, "items": page_items}
|
|
|
|
|
|
def create_collection(db: Session, body: CollectionJobCreate) -> CollectionJobOut:
|
|
commands = _parse_commands(body.commands)
|
|
if not commands:
|
|
raise HTTPException(status_code=400, detail="commands_empty")
|
|
ne_ids = [str(x).strip() for x in (body.ne_ids or []) if str(x).strip()]
|
|
ume_ids = [str(x).strip() for x in (body.ume_ne_ids or []) if str(x).strip()]
|
|
if not ne_ids and not ume_ids:
|
|
raise HTTPException(status_code=400, detail="ne_ids_required")
|
|
|
|
targets: list[tuple[str, str, str, str]] = [] # source, id, name, ip
|
|
missing_m: list[str] = []
|
|
for ne_id in ne_ids:
|
|
row = db.get(ManagedNE, ne_id)
|
|
if not row:
|
|
missing_m.append(ne_id)
|
|
continue
|
|
targets.append(
|
|
(
|
|
"managed",
|
|
str(row.id),
|
|
str(row.name or row.ip_address or ""),
|
|
str(row.ip_address or ""),
|
|
)
|
|
)
|
|
if missing_m:
|
|
raise HTTPException(status_code=404, detail=f"managed_ne_not_found: {','.join(missing_m[:5])}")
|
|
|
|
missing_u: list[str] = []
|
|
for uid in ume_ids:
|
|
inv = db.get(UmeInventoryNE, uid)
|
|
if not inv:
|
|
missing_u.append(uid)
|
|
continue
|
|
name = str(inv.host_name or inv.user_label or inv.ne_name or inv.ip_address or uid)
|
|
targets.append(("ume", uid, name, str(inv.ip_address or "")))
|
|
if missing_u:
|
|
raise HTTPException(status_code=404, detail=f"ume_ne_not_found: {','.join(missing_u[:5])}")
|
|
if not targets:
|
|
raise HTTPException(status_code=400, detail="no_eligible_ne")
|
|
|
|
now = _now()
|
|
job = NeCollectionJob(
|
|
title=str(body.title or "").strip() or f"collect-{now.strftime('%Y%m%d-%H%M%S')}",
|
|
commands="\n".join(commands),
|
|
trigger_mode="manual",
|
|
status="pending",
|
|
ne_count=len(targets),
|
|
created_at=now,
|
|
started_at=None,
|
|
last_run_at=None,
|
|
)
|
|
db.add(job)
|
|
db.flush()
|
|
|
|
for source, tid, name, ip in targets:
|
|
run = NeCollectionRun(
|
|
job_id=str(job.id),
|
|
ne_id=tid,
|
|
ne_source=source,
|
|
ne_name=name,
|
|
ne_ip=ip,
|
|
status="pending",
|
|
)
|
|
db.add(run)
|
|
db.commit()
|
|
db.refresh(job)
|
|
return job_to_out(job, output_count=0)
|
|
|
|
|
|
def create_and_start_from_policy(
|
|
db: Session,
|
|
*,
|
|
trigger_mode: str = "manual",
|
|
) -> tuple[CollectionJobOut, CollectionSchedulePayload]:
|
|
"""Create a job from the singleton policy and start it immediately."""
|
|
if has_active_collection_job(db) is not None:
|
|
raise HTTPException(status_code=409, detail="collection_job_running")
|
|
policy = ensure_policy(db)
|
|
commands = _parse_commands(str(policy.commands or ""))
|
|
if not commands:
|
|
raise HTTPException(status_code=400, detail="commands_empty")
|
|
targets = expand_policy_targets(db, policy)
|
|
if not targets:
|
|
raise HTTPException(status_code=400, detail="no_eligible_ne")
|
|
mode = str(trigger_mode or "manual").strip().lower() or "manual"
|
|
if mode not in {"manual", "schedule"}:
|
|
mode = "manual"
|
|
now = _now()
|
|
title = str(policy.title or "").strip() or f"collect-{now.strftime('%Y%m%d-%H%M%S')}"
|
|
job = NeCollectionJob(
|
|
title=title,
|
|
commands="\n".join(commands),
|
|
trigger_mode=mode,
|
|
status="pending",
|
|
ne_count=len(targets),
|
|
created_at=now,
|
|
started_at=None,
|
|
last_run_at=None,
|
|
)
|
|
db.add(job)
|
|
db.flush()
|
|
for source, tid, name, ip in targets:
|
|
db.add(
|
|
NeCollectionRun(
|
|
job_id=str(job.id),
|
|
ne_id=tid,
|
|
ne_source=source,
|
|
ne_name=name,
|
|
ne_ip=ip,
|
|
status="pending",
|
|
)
|
|
)
|
|
db.commit()
|
|
db.refresh(job)
|
|
out, payload = start_collection_job(db, str(job.id))
|
|
try:
|
|
prune_collection_jobs(db, keep=history_keep_value(policy))
|
|
except Exception: # noqa: BLE001
|
|
_log.exception("prune_collection_jobs after start failed")
|
|
return out, payload
|
|
|
|
|
|
def list_collection_jobs(
|
|
db: Session,
|
|
*,
|
|
page: int = 1,
|
|
page_size: int = 20,
|
|
status: str = "",
|
|
keyword: str = "",
|
|
) -> dict[str, Any]:
|
|
stmt = db.query(NeCollectionJob)
|
|
st = str(status or "").strip()
|
|
if st:
|
|
stmt = stmt.filter(NeCollectionJob.status == st)
|
|
kw = str(keyword or "").strip()
|
|
if kw:
|
|
like = f"%{kw}%"
|
|
stmt = stmt.filter(
|
|
or_(
|
|
NeCollectionJob.title.ilike(like),
|
|
NeCollectionJob.id.ilike(like),
|
|
NeCollectionJob.error_message.ilike(like),
|
|
)
|
|
)
|
|
total = int(stmt.count())
|
|
rows = stmt.order_by(NeCollectionJob.created_at.desc()).offset((page - 1) * page_size).limit(page_size).all()
|
|
for row in rows:
|
|
job_id = str(row.id)
|
|
if str(row.status or "") not in ("done", "failed", "paused"):
|
|
reconcile_stale_collection_job(db, job_id)
|
|
db.refresh(row)
|
|
if str(row.status or "") == "running":
|
|
sync_job_progress(db, job_id)
|
|
finalize_collection_job(db, job_id)
|
|
db.refresh(row)
|
|
job_ids = [str(x.id) for x in rows]
|
|
output_counts = _output_counts_for_jobs(db, job_ids)
|
|
return {
|
|
"total": total,
|
|
"page": page,
|
|
"page_size": page_size,
|
|
"items": [job_to_out(x, output_count=output_counts.get(str(x.id), 0)).model_dump() for x in rows],
|
|
}
|
|
|
|
|
|
def get_collection_job(db: Session, job_id: str) -> dict[str, Any]:
|
|
job = db.get(NeCollectionJob, job_id)
|
|
if not job:
|
|
raise HTTPException(status_code=404, detail="collection_job_not_found")
|
|
if str(job.status or "") not in ("done", "failed", "paused"):
|
|
reconcile_stale_collection_job(db, job_id)
|
|
db.refresh(job)
|
|
if str(job.status or "") == "running":
|
|
sync_job_progress(db, job_id)
|
|
finalize_collection_job(db, job_id)
|
|
db.refresh(job)
|
|
return {
|
|
"job": job_to_out(job, output_count=_output_count_for_job(db, job_id)).model_dump(),
|
|
}
|
|
|
|
|
|
def list_collection_runs(
|
|
db: Session,
|
|
job_id: str,
|
|
*,
|
|
page: int = 1,
|
|
page_size: int = 20,
|
|
status: str = "",
|
|
keyword: str = "",
|
|
) -> dict[str, Any]:
|
|
job = db.get(NeCollectionJob, job_id)
|
|
if not job:
|
|
raise HTTPException(status_code=404, detail="collection_job_not_found")
|
|
stmt = db.query(NeCollectionRun).filter(NeCollectionRun.job_id == job_id)
|
|
st = str(status or "").strip()
|
|
if st:
|
|
stmt = stmt.filter(NeCollectionRun.status == st)
|
|
kw = str(keyword or "").strip()
|
|
if kw:
|
|
like = f"%{kw}%"
|
|
stmt = stmt.filter(or_(NeCollectionRun.ne_name.ilike(like), NeCollectionRun.ne_ip.ilike(like)))
|
|
total = int(stmt.count())
|
|
rows = (
|
|
stmt.order_by(NeCollectionRun.ne_name.asc(), NeCollectionRun.ne_ip.asc())
|
|
.offset((page - 1) * page_size)
|
|
.limit(page_size)
|
|
.all()
|
|
)
|
|
return {
|
|
"total": total,
|
|
"page": page,
|
|
"page_size": page_size,
|
|
"items": [run_to_out(x).model_dump() for x in rows],
|
|
}
|
|
|
|
|
|
def _active_runs(runs: list[NeCollectionRun]) -> bool:
|
|
return any(str(r.status or "") in ("pending", "running") for r in runs)
|
|
|
|
|
|
def _runs_in_progress(runs: list[NeCollectionRun]) -> bool:
|
|
"""True when a worker is actively collecting (not merely queued as pending)."""
|
|
return any(str(r.status or "") == "running" for r in runs)
|
|
|
|
|
|
def _blocking_running_job_ids(db: Session, ne_ids: list[str], *, exclude_job_id: str) -> list[str]:
|
|
if not ne_ids:
|
|
return []
|
|
other_jobs = (
|
|
db.query(NeCollectionJob)
|
|
.filter(NeCollectionJob.status == "running", NeCollectionJob.id != exclude_job_id)
|
|
.all()
|
|
)
|
|
blocked: list[str] = []
|
|
ne_set = set(ne_ids)
|
|
for other in other_jobs:
|
|
other_id = str(other.id)
|
|
overlap = (
|
|
db.query(NeCollectionRun.ne_id)
|
|
.filter(NeCollectionRun.job_id == other_id, NeCollectionRun.ne_id.in_(list(ne_set)))
|
|
.first()
|
|
)
|
|
if overlap:
|
|
blocked.append(other_id)
|
|
return blocked
|
|
|
|
|
|
def pause_collection_job(db: Session, job_id: str) -> CollectionJobOut:
|
|
job = db.get(NeCollectionJob, job_id)
|
|
if not job:
|
|
raise HTTPException(status_code=404, detail="collection_job_not_found")
|
|
if str(job.status or "") != "running":
|
|
raise HTTPException(status_code=400, detail="collection_job_not_running")
|
|
now = _now()
|
|
runs = db.query(NeCollectionRun).filter(NeCollectionRun.job_id == job_id).all()
|
|
for run in runs:
|
|
if str(run.status or "") != "pending":
|
|
continue
|
|
run.status = "cancelled"
|
|
run.message = "paused"
|
|
run.ended_at = now
|
|
job.status = "paused"
|
|
_sync_job_counts(job, runs)
|
|
if not any(str(r.status or "") not in {"success", "fail", "cancelled"} for r in runs):
|
|
job.ended_at = now
|
|
db.commit()
|
|
db.refresh(job)
|
|
finalize_collection_job(db, job_id)
|
|
db.refresh(job)
|
|
return job_to_out(job, output_count=_output_count_for_job(db, job_id))
|
|
|
|
|
|
def _reset_runs_for_retry(
|
|
db: Session,
|
|
job_id: str,
|
|
runs: list[NeCollectionRun],
|
|
*,
|
|
only_statuses: frozenset[str],
|
|
) -> list[str]:
|
|
retry_ids: list[str] = []
|
|
for run in runs:
|
|
st = str(run.status or "")
|
|
if st in ("pending", "running") or st not in only_statuses:
|
|
continue
|
|
clear_run_output_files(job_id, str(run.id))
|
|
run.status = "pending"
|
|
run.message = ""
|
|
run.output_rel_path = ""
|
|
run.started_at = None
|
|
run.ended_at = None
|
|
retry_ids.append(str(run.id))
|
|
return retry_ids
|
|
|
|
|
|
def _start_job_retry(
|
|
db: Session,
|
|
job: NeCollectionJob,
|
|
job_id: str,
|
|
retry_ids: list[str],
|
|
commands: list[str],
|
|
*,
|
|
reset_all_counts: bool,
|
|
) -> tuple[CollectionJobOut, CollectionSchedulePayload]:
|
|
if not retry_ids:
|
|
raise HTTPException(status_code=400, detail="collection_nothing_to_retry")
|
|
now = _now()
|
|
job.status = "running"
|
|
job.ended_at = None
|
|
job.error_message = ""
|
|
if reset_all_counts:
|
|
job.success_count = 0
|
|
job.fail_count = 0
|
|
else:
|
|
runs = db.query(NeCollectionRun).filter(NeCollectionRun.job_id == job_id).all()
|
|
_sync_job_counts(job, runs)
|
|
job.started_at = now
|
|
job.last_run_at = now
|
|
db.commit()
|
|
db.refresh(job)
|
|
payload: CollectionSchedulePayload = {
|
|
"job_id": job_id,
|
|
"run_ids": retry_ids,
|
|
"commands": commands,
|
|
}
|
|
return job_to_out(job, output_count=_output_count_for_job(db, job_id)), payload
|
|
|
|
|
|
def start_collection_job(db: Session, job_id: str) -> tuple[CollectionJobOut, CollectionSchedulePayload]:
|
|
"""Start a draft job or re-run all NEs after pause/completion."""
|
|
job = db.get(NeCollectionJob, job_id)
|
|
if not job:
|
|
raise HTTPException(status_code=404, detail="collection_job_not_found")
|
|
runs = db.query(NeCollectionRun).filter(NeCollectionRun.job_id == job_id).all()
|
|
if not runs:
|
|
raise HTTPException(status_code=400, detail="collection_no_runs")
|
|
if str(job.status or "") == "running" or _runs_in_progress(runs):
|
|
raise HTTPException(status_code=400, detail="collection_job_running")
|
|
commands = _parse_commands(str(job.commands or ""))
|
|
if not commands:
|
|
raise HTTPException(status_code=400, detail="commands_empty")
|
|
|
|
ne_ids = list({str(r.ne_id) for r in runs if str(r.ne_id or "").strip()})
|
|
blocking = _blocking_running_job_ids(db, ne_ids, exclude_job_id=job_id)
|
|
if blocking:
|
|
raise HTTPException(status_code=409, detail=f"collection_ne_busy: {blocking[0][:12]}")
|
|
|
|
if str(job.status or "") == "pending":
|
|
run_ids = [str(r.id) for r in runs if str(r.status or "") == "pending"]
|
|
if not run_ids:
|
|
raise HTTPException(status_code=400, detail="collection_nothing_to_start")
|
|
return _start_job_retry(db, job, job_id, run_ids, commands, reset_all_counts=True)
|
|
|
|
retry_ids = _reset_runs_for_retry(
|
|
db,
|
|
job_id,
|
|
runs,
|
|
only_statuses=frozenset({"success", "fail", "cancelled"}),
|
|
)
|
|
return _start_job_retry(db, job, job_id, retry_ids, commands, reset_all_counts=True)
|
|
|
|
|
|
def restart_collection_job(db: Session, job_id: str) -> tuple[CollectionJobOut, CollectionSchedulePayload]:
|
|
return start_collection_job(db, job_id)
|
|
|
|
|
|
def retry_failed_collection_job(
|
|
db: Session, job_id: str
|
|
) -> tuple[CollectionJobOut, CollectionSchedulePayload]:
|
|
job = db.get(NeCollectionJob, job_id)
|
|
if not job:
|
|
raise HTTPException(status_code=404, detail="collection_job_not_found")
|
|
runs = db.query(NeCollectionRun).filter(NeCollectionRun.job_id == job_id).all()
|
|
if not runs:
|
|
raise HTTPException(status_code=400, detail="collection_no_runs")
|
|
if str(job.status or "") == "running" or _runs_in_progress(runs):
|
|
raise HTTPException(status_code=400, detail="collection_job_running")
|
|
commands = _parse_commands(str(job.commands or ""))
|
|
if not commands:
|
|
raise HTTPException(status_code=400, detail="commands_empty")
|
|
ne_ids = list({str(r.ne_id) for r in runs if str(r.ne_id or "").strip()})
|
|
blocking = _blocking_running_job_ids(db, ne_ids, exclude_job_id=job_id)
|
|
if blocking:
|
|
raise HTTPException(status_code=409, detail=f"collection_ne_busy: {blocking[0][:12]}")
|
|
retry_ids = _reset_runs_for_retry(db, job_id, runs, only_statuses=frozenset({"fail"}))
|
|
return _start_job_retry(db, job, job_id, retry_ids, commands, reset_all_counts=False)
|
|
|
|
|
|
def delete_collection_job(db: Session, job_id: str) -> dict[str, bool]:
|
|
job = db.get(NeCollectionJob, job_id)
|
|
if not job:
|
|
raise HTTPException(status_code=404, detail="collection_job_not_found")
|
|
runs = db.query(NeCollectionRun).filter(NeCollectionRun.job_id == job_id).all()
|
|
if str(job.status or "") == "running" or _active_runs(runs):
|
|
raise HTTPException(status_code=400, detail="collection_job_running")
|
|
db.query(NeCollectionRun).filter(NeCollectionRun.job_id == job_id).delete()
|
|
db.delete(job)
|
|
db.commit()
|
|
job_dir = (collection_data_root() / job_id).resolve()
|
|
root = collection_data_root().resolve()
|
|
if str(job_dir).startswith(str(root)) and job_dir.is_dir():
|
|
shutil.rmtree(job_dir, ignore_errors=True)
|
|
return {"ok": True}
|
|
|
|
|
|
def _safe_archive_part(text: str, fallback: str = "device") -> str:
|
|
s = re.sub(r'[<>:"/\\|?*]', "_", str(text or "").strip())
|
|
return (s[:80] or fallback).strip("._") or fallback
|
|
|
|
|
|
def build_collection_job_zip(db: Session, job_id: str) -> tuple[str, bytes]:
|
|
job = db.get(NeCollectionJob, job_id)
|
|
if not job:
|
|
raise HTTPException(status_code=404, detail="collection_job_not_found")
|
|
runs = (
|
|
db.query(NeCollectionRun)
|
|
.filter(NeCollectionRun.job_id == job_id)
|
|
.order_by(NeCollectionRun.ne_name.asc())
|
|
.all()
|
|
)
|
|
files: list[tuple[str, Path]] = []
|
|
used_names: set[str] = set()
|
|
for run in runs:
|
|
rel = str(run.output_rel_path or "").strip()
|
|
if not rel:
|
|
continue
|
|
path = resolve_run_output_path(rel)
|
|
arcname = path.name
|
|
if arcname in used_names:
|
|
arcname = f"{str(run.id)[:8]}_{path.name}"
|
|
used_names.add(arcname)
|
|
files.append((arcname, path))
|
|
if not files:
|
|
raise HTTPException(status_code=404, detail="collection_outputs_not_found")
|
|
buf = io.BytesIO()
|
|
with zipfile.ZipFile(buf, mode="w", compression=zipfile.ZIP_DEFLATED) as zf:
|
|
for arcname, path in files:
|
|
zf.write(path, arcname=arcname)
|
|
title = _safe_archive_part(str(job.title or "collect"), "collect")
|
|
zip_name = f"{title}_{job_id[:8]}.zip"
|
|
return zip_name, buf.getvalue()
|
|
|
|
|
|
def resolve_run_output_path(rel_path: str) -> Path:
|
|
rel = str(rel_path or "").strip().replace("\\", "/")
|
|
if not rel or ".." in rel.split("/"):
|
|
raise HTTPException(status_code=400, detail="invalid_output_path")
|
|
full = (collection_data_root() / rel).resolve()
|
|
root = collection_data_root()
|
|
if not str(full).startswith(str(root)):
|
|
raise HTTPException(status_code=400, detail="invalid_output_path")
|
|
if not full.is_file():
|
|
raise HTTPException(status_code=404, detail="output_file_not_found")
|
|
return full
|