feat(ops): add managed NE management and batch CLI collection

Introduce workbench operations for multi-vendor NE CRUD/connect-test and Netmiko batch collection with job lifecycle controls, log downloads, and paginated run filters.

Co-authored-by: Cursor <cursoragent@cursor.com>

netx_api/ne_collect_runner.py

@@ -0,0 +1,174 @@
+from __future__ import annotations
+
+import logging
+import re
+from concurrent.futures import ThreadPoolExecutor, TimeoutError as FuturesTimeout
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+from netmiko import ConnectHandler
+
+from .collection_job_state import finalize_collection_job
+from .config import settings
+from .db import SessionLocal
+from .models import ManagedNE, NeCollectionJob, NeCollectionRun
+from .ne_collection_paths import clear_run_output_files, run_output_dir
+from .ne_crypto import CredentialCryptoError
+from .ne_netmiko import normalize_netmiko_device_type
+from .ne_service import get_device_credentials
+
+_log = logging.getLogger("netx.ne.collect")
+_executor: ThreadPoolExecutor | None = None
+
+
+def _executor_pool() -> ThreadPoolExecutor:
+    global _executor
+    if _executor is None:
+        workers = max(1, int(settings.ne_collect_max_workers or 5))
+        _executor = ThreadPoolExecutor(max_workers=workers, thread_name_prefix="ne-collect")
+    return _executor
+
+
+def _safe_filename_part(text: str) -> str:
+    s = re.sub(r'[<>:"/\\|?*]', "_", str(text or "").strip())
+    return s[:80] or "device"
+
+
+def _collect_on_device(creds: dict[str, Any], commands: list[str]) -> str:
+    device_type = normalize_netmiko_device_type(creds["device_type"], creds["protocol"])
+    per_cmd = int(settings.ne_collect_read_timeout_sec or 120)
+    dev: dict[str, Any] = {
+        "device_type": device_type,
+        "host": creds["ip_address"],
+        "username": creds["username"],
+        "password": creds["password"],
+        "port": int(creds["port"] or 22),
+        "conn_timeout": int(settings.ne_connect_timeout_sec or 30),
+        "auth_timeout": int(settings.ne_connect_timeout_sec or 30),
+        "banner_timeout": int(settings.ne_connect_timeout_sec or 30),
+        "session_timeout": per_cmd * max(1, len(commands)) + 60,
+    }
+    secret = str(creds.get("enable_secret") or "").strip()
+    if secret:
+        dev["secret"] = secret
+    chunks: list[str] = []
+    with ConnectHandler(**dev) as conn:
+        prompt = str(conn.find_prompt() or "")
+        for command in commands:
+            ts = datetime.now().isoformat(timespec="seconds")
+            chunks.append(f'>>> [{ts}] {{"String":"{command}", "Match":"{prompt}", "Timeout":0}}\n')
+            out = conn.send_command(command_string=command, read_timeout=per_cmd)
+            chunks.append(str(out or ""))
+            chunks.append("\n")
+    return "".join(chunks)
+
+
+def _collect_with_timeout(creds: dict[str, Any], commands: list[str]) -> str:
+    per_cmd = int(settings.ne_collect_read_timeout_sec or 120)
+    cap = int(settings.ne_collect_run_timeout_cap_sec or 600)
+    budget = min(cap, per_cmd * max(1, len(commands)) + 90)
+    with ThreadPoolExecutor(max_workers=1) as pool:
+        fut = pool.submit(_collect_on_device, creds, commands)
+        try:
+            return fut.result(timeout=budget)
+        except FuturesTimeout as exc:
+            raise TimeoutError(f"collection_timeout ({budget}s)") from exc
+
+
+def _update_run(run_id: str, **fields: Any) -> None:
+    db = SessionLocal()
+    try:
+        row = db.get(NeCollectionRun, run_id)
+        if not row:
+            return
+        for key, val in fields.items():
+            setattr(row, key, val)
+        db.commit()
+    finally:
+        db.close()
+
+
+def _collection_aborted(job_id: str, run_id: str) -> bool:
+    db = SessionLocal()
+    try:
+        job = db.get(NeCollectionJob, job_id)
+        run = db.get(NeCollectionRun, run_id)
+        if not job or not run:
+            return True
+        if str(job.status or "") == "paused":
+            if str(run.status or "") == "pending":
+                run.status = "cancelled"
+                run.message = "paused"
+                run.ended_at = datetime.now()
+                db.commit()
+            return True
+        if str(run.status or "") in ("cancelled", "success", "fail"):
+            return True
+        return False
+    finally:
+        db.close()
+
+
+def _run_single(job_id: str, run_id: str, commands: list[str]) -> None:
+    if _collection_aborted(job_id, run_id):
+        db = SessionLocal()
+        try:
+            finalize_collection_job(db, job_id)
+        finally:
+            db.close()
+        return
+    db = SessionLocal()
+    try:
+        run = db.get(NeCollectionRun, run_id)
+        if not run:
+            return
+        ne = db.get(ManagedNE, str(run.ne_id))
+        if not ne:
+            _update_run(run_id, status="fail", message="managed_ne_not_found", ended_at=datetime.now())
+            return
+        started_at = datetime.now()
+        _update_run(run_id, status="running", message="collecting", started_at=started_at)
+        try:
+            creds = get_device_credentials(ne)
+            output = _collect_with_timeout(creds, commands)
+            finished_at = datetime.now()
+            name_part = _safe_filename_part(str(run.ne_name or creds.get("name") or "ne"))
+            ip_part = _safe_filename_part(str(run.ne_ip or creds.get("ip_address") or "ip"))
+            ts = finished_at.strftime("%Y%m%d%H%M%S")
+            rel_dir = Path(job_id) / run_id
+            out_dir = run_output_dir(job_id, run_id)
+            clear_run_output_files(job_id, run_id)
+            out_dir.mkdir(parents=True, exist_ok=True)
+            filename = f"{name_part}-{ip_part}-{ts}.txt"
+            full_path = out_dir / filename
+            full_path.write_text(output, encoding="utf-8", errors="replace")
+            rel_path = str(rel_dir / filename).replace("\\", "/")
+            _update_run(
+                run_id,
+                status="success",
+                message="collected",
+                output_rel_path=rel_path,
+                ended_at=finished_at,
+            )
+        except CredentialCryptoError as exc:
+            _update_run(run_id, status="fail", message=str(exc)[:1000], ended_at=datetime.now())
+        except Exception as exc:
+            _log.exception("collection failed run=%s", run_id)
+            _update_run(run_id, status="fail", message=str(exc).split("\n")[0][:1000], ended_at=datetime.now())
+    finally:
+        db.close()
+        db2 = SessionLocal()
+        try:
+            finalize_collection_job(db2, job_id)
+        finally:
+            db2.close()
+
+
+def schedule_collection_runs(job_id: str, run_ids: list[str], commands: list[str]) -> int:
+    pool = _executor_pool()
+    submitted = 0
+    for run_id in run_ids:
+        pool.submit(_run_single, job_id, run_id, list(commands))
+        submitted += 1
+    return submitted