"""Collect runner: expand task items → CLI → match → parse → batch rows.""" from __future__ import annotations import logging import threading from concurrent.futures import ThreadPoolExecutor from datetime import datetime from typing import Any from uuid import uuid4 from fastapi import HTTPException from sqlalchemy.exc import DBAPIError, OperationalError from ..cli_creds import cli_creds_skip_reason from ..cli_resolve import resolve_cli_target from ..cli_timeout import run_cli_with_timeout from ..config import settings from ..db import SessionLocal from ..lldp_shared import resolve_vendor_key from ..models import ( BizStateBatch, BizStateBatchCommand, BizStateEvent, BizStateLldpNeighbor, BizStateMetricRow, BizStateTask, BizStateTaskItem, BizStateTaskItemBinding, ) from ..ne_netmiko import disable_target_paging, send_show_command from ..ne_session_factory import close_netmiko_connection, open_netmiko_connection from .collect_session import ( CollectSession, build_parse_bundle, resolve_aux_command, run_primary_with_bundle, ) from .command_match import ( EXPAND_ALL_COMMAND, expand_bindings_from_discover_records, expand_from_bindings, match_command, normalize_command, ) from .parsers import get_parser from .profiles import get_profile _log = logging.getLogger("netx.biz_state.runner") # concrete, params, profile_id, item_id, mode WorkItem = tuple[str, dict[str, str], str, str, str] _heavy_pool_lock = threading.Lock() _heavy_pool: ThreadPoolExecutor | None = None def _heavy_cli_pool() -> ThreadPoolExecutor: global _heavy_pool with _heavy_pool_lock: if _heavy_pool is None: n = max(1, int(getattr(settings, "biz_state_heavy_workers", 4) or 4)) _heavy_pool = ThreadPoolExecutor(max_workers=n, thread_name_prefix="biz-heavy") return _heavy_pool def work_item_lane(profile_id: str, mode: str) -> str: """Return collect lane for a work item (custom_raw always light).""" if str(mode or "").strip() == "custom": return "light" profile = get_profile(profile_id) or _resolve_collect_profile(profile_id) lane = str(getattr(profile, "collect_lane", "light") or "light").strip().lower() if profile else "light" return "heavy" if lane == "heavy" else "light" def partition_work(work: list[WorkItem]) -> tuple[list[WorkItem], list[WorkItem]]: """Split work into (light, heavy) lanes.""" light: list[WorkItem] = [] heavy: list[WorkItem] = [] for item in work: if work_item_lane(item[2], item[4]) == "heavy": heavy.append(item) else: light.append(item) return light, heavy def _utcnow() -> datetime: return datetime.utcnow() def _format_error(exc: BaseException) -> str: return f"{type(exc).__name__}: {exc}"[:1020] def _append_event(db, *, task_id: str, message: str, level: str = "error") -> None: msg = str(message or "").strip() if not msg or not task_id: return db.add( BizStateEvent( id=uuid4().hex, task_id=task_id, level=str(level or "error")[:16], message=msg[:4000], created_at=_utcnow(), ) ) def _bindings_for_item(db, item_id: str) -> list[dict[str, str]]: rows = ( db.query(BizStateTaskItemBinding) .filter(BizStateTaskItemBinding.item_id == item_id) .all() ) # Keep one row per binding value (same placeholder may appear many times). return [ {"placeholder": str(r.placeholder or "").strip(), "value": str(r.value or "").strip()} for r in rows if str(r.placeholder or "").strip() and str(r.value or "").strip() ] def _resolve_collect_profile(profile_id: str): """Resolve task-item profile; remap disabled if_intf → config_interface.""" pid = str(profile_id or "").strip() if not pid: return None profile = get_profile(pid) if profile is None: return None if profile.enabled: return profile # Legacy IF VRF check merged into Config Interface Intent — VRF is a subset. if profile.metric_id == "if_intf" or pid.endswith(".if_intf"): vk = str(profile.vendor_key or "zte").strip() or "zte" remapped = get_profile(f"{vk}.config_interface") or get_profile("zte.config_interface") if remapped and remapped.enabled: return remapped return None def _persist_lldp_rows( db, *, batch: BizStateBatch, cmd_row: BizStateBatchCommand, records: list[dict[str, Any]], ) -> int: n = 0 seen: set[tuple[str, str, str]] = set() for rec in records: local_if = str(rec.get("local_if") or "").strip()[:128] remote_sys = str(rec.get("remote_sys") or "").strip()[:256] remote_if = str(rec.get("remote_if") or "").strip()[:128] if not local_if and not remote_sys and not remote_if: continue key = (local_if, remote_sys, remote_if) if key in seen: continue seen.add(key) db.add( BizStateLldpNeighbor( id=uuid4().hex, batch_id=batch.id, batch_command_id=cmd_row.id, task_id=batch.task_id, ne_id=batch.ne_id, local_if=local_if, remote_sys=remote_sys, remote_if=remote_if, remote_ip=str(rec.get("remote_ip") or "")[:128], protocol=str(rec.get("protocol") or "lldp")[:32], collected_at=_utcnow(), ) ) n += 1 return n _GENERIC_METRICS = { "isis_adjacency", "interface_brief", "interface_detail", "arp", "if_intf", "nd6_cache", "bgp_peer", "ospf_neighbor", "vrrp", "optical_brief", "bgp_route", "ip_route", "ipv6_route", "l2vpn_pw", "l2vpn_pw_detail", "l2vpn_mac", "evpn_mac", "config_vrf", "config_interface", "config_bgp_peer", "config_l2vpn_pw", "config_static_route", "config_ospf", "config_isis", } _METRIC_CHUNK = 2000 def _emit_task_event(*, task_id: str, message: str, level: str = "error") -> None: """Short-lived session for lane events (no long-held DB during CLI).""" if not task_id or not str(message or "").strip(): return db = SessionLocal() try: _append_event(db, task_id=task_id, message=message, level=level) db.commit() except Exception: _log.exception("biz_state emit event failed task=%s", task_id) try: db.rollback() except Exception: pass finally: db.close() def _flush_spooled_commands( batch_id: str, pending: list[Any], ) -> tuple[int, int]: """Insert SpooledCommand rows (+ metric/lldp) in one transaction. Returns (cmds, rows).""" from .spool import SpooledCommand, raw_max_bytes, read_raw_text, read_records if not pending: return 0, 0 items: list[SpooledCommand] = list(pending) pending.clear() def _write(db) -> tuple[int, int]: batch = db.get(BizStateBatch, batch_id) if not batch: return 0, 0 rows_n = 0 max_raw = raw_max_bytes() for item in items: raw = "" if item.raw_rel_path: raw = read_raw_text(item.raw_rel_path, max_bytes=max_raw) cmd_row = BizStateBatchCommand( id=item.id, batch_id=batch_id, task_item_id=item.task_item_id, profile_id=item.profile_id, parser_id=item.parser_id, metric_id=item.metric_id, raw_command=str(item.raw_command or "")[:512], params_json=dict(item.params_json or {}), parse_status=item.parse_status, message=str(item.message or "")[:1020], raw_text=raw, row_count=int(item.row_count or 0), created_at=_utcnow(), ) db.add(cmd_row) if item.persist_kind == "metric" and item.records_rel_path: records = read_records(item.records_rel_path) mid = str(item.metric_id or "").strip() if mid and records: n = _persist_metric_rows( db, batch=batch, cmd_row=cmd_row, metric_id=mid, records=records, ) cmd_row.row_count = n rows_n += n elif item.persist_kind == "lldp" and item.records_rel_path: records = read_records(item.records_rel_path) if records: n = _persist_lldp_rows( db, batch=batch, cmd_row=cmd_row, records=records ) cmd_row.row_count = n rows_n += n db.commit() return len(items), rows_n try: cmds, rows = _run_db_with_reconnect(_write, label="biz_state_flush_spool") if cmds or rows: _bump_batch_progress(batch_id, add_cmds=cmds, add_rows=rows) return int(cmds or 0), int(rows or 0) except Exception: _log.exception("biz_state flush spool failed batch=%s n=%s", batch_id, len(items)) # Put back so a later flush / finalize can retry. pending.extend(items) raise def _persist_metric_rows( db, *, batch: BizStateBatch, cmd_row: BizStateBatchCommand, metric_id: str, records: list[dict[str, Any]], ) -> int: """Bulk-insert generic metric rows (JSON payload per row).""" mid = str(metric_id or "").strip() if not mid or not records: return 0 buf: list[dict[str, Any]] = [] n = 0 for i, rec in enumerate(records): if not isinstance(rec, dict) or not rec: continue buf.append( { "id": uuid4().hex, "batch_id": batch.id, "batch_command_id": cmd_row.id, "task_id": batch.task_id, "ne_id": batch.ne_id, "metric_id": mid, "seq": i, "data_json": dict(rec), "collected_at": _utcnow(), } ) n += 1 if len(buf) >= _METRIC_CHUNK: db.bulk_insert_mappings(BizStateMetricRow, buf) buf.clear() if buf: db.bulk_insert_mappings(BizStateMetricRow, buf) return n def _finish_task(task_id: str, *, error: str = "") -> None: db = SessionLocal() try: task = db.get(BizStateTask, task_id) if not task: return task.collect_running = False task.last_collect_ended_at = _utcnow() task.last_error = str(error or "")[:1020] task.updated_at = _utcnow() if error: _append_event(db, task_id=task_id, message=error, level="error") db.commit() finally: db.close() def dispatch_collect(task_id: str, *, manual: bool = False) -> None: """Claim and run one collect round. Scheduler calls with ``manual=False`` (only when task status is ``running``). Collect-now calls with ``manual=True`` (any status, as long as not already collecting). """ db = SessionLocal() batch_id = "" try: task = db.get(BizStateTask, task_id) if not task: return if task.collect_running: return st = str(task.status or "").strip() if manual: # Idle manual trigger: allow scheduled / paused / draft / stopped if st in ("", "deleted"): return else: if st != "running": return items = ( db.query(BizStateTaskItem) .filter( BizStateTaskItem.task_id == task_id, BizStateTaskItem.enabled.is_(True), ) .order_by(BizStateTaskItem.sort_order.asc()) .all() ) if not items: task.last_error = "no enabled task items" task.updated_at = _utcnow() db.commit() return task.collect_running = True task.last_collect_started_at = _utcnow() task.last_error = "" task.updated_at = _utcnow() batch = BizStateBatch( id=uuid4().hex, task_id=task.id, source=task.source, ne_id=task.ne_id, ne_name=task.ne_name, vendor=task.vendor, status="running", started_at=_utcnow(), ) db.add(batch) db.commit() batch_id = batch.id vendor = str(task.vendor or "") device_type = str(task.device_type or "") source = str(task.source or "managed").strip().lower() ne_id = str(task.ne_id or "").strip() finally: db.close() if not batch_id: return error = "" try: _run_collect_session( task_id=task_id, batch_id=batch_id, source=source, ne_id=ne_id, vendor=vendor, device_type=device_type, ) except Exception as exc: _log.exception("biz_state collect failed task=%s", task_id) error = _format_error(exc) try: _fail_batch_status(batch_id, error) except Exception: _log.exception("biz_state fail-batch after collect error failed batch=%s", batch_id) finally: _finish_task(task_id, error=error) def _run_collect_lane( *, work: list[WorkItem], batch_id: str, creds: dict[str, Any], vendor_eff: str, device_type_eff: str, vendor_key: str, per_cmd: int, cap: int, label: str, shared_cache: dict[str, Any] | None = None, cache_lock: Any | None = None, cmd_locks: dict[str, Any] | None = None, aux_persisted: set[tuple[str, str]] | None = None, ) -> tuple[int, int, bool, bool]: """Run one SSH lane: collect+parse to spool, flush to DB in batches.""" if not work: return 0, 0, False, False from .spool import ( SpooledCommand, persist_every_cmds, write_meta, write_raw_text, write_records, ) budget = min(int(cap), int(per_cmd) * max(1, len(work)) + 90) holder: dict[str, Any] = {} flush_every = persist_every_cmds() def _session() -> tuple[int, int, bool, bool]: from ..ne_netmiko import drain_read_channel conn = open_netmiko_connection(creds, session_timeout=budget) holder["conn"] = conn total_rows = 0 cmd_count = 0 any_fail = False any_ok = False pending: list[SpooledCommand] = [] task_id = "" def _queue(item: SpooledCommand, *, records: list[dict[str, Any]] | None = None) -> None: nonlocal cmd_count, total_rows if records is not None and item.persist_kind: item.records_rel_path = write_records(batch_id, item.id, records) item.row_count = len(records) try: write_meta(batch_id, item.id, item.to_meta()) except Exception: _log.exception("biz_state write meta failed cmd=%s", item.id) pending.append(item) cmd_count += 1 if len(pending) >= flush_every: try: _c, _r = _flush_spooled_commands(batch_id, pending) total_rows += int(_r or 0) except Exception: # Keep collecting to spool; retry flush at lane end. _log.exception( "biz_state mid-lane flush failed batch=%s pending=%s", batch_id, len(pending), ) try: try: disable_target_paging( conn, vendor=str(creds.get("vendor") or vendor_eff or ""), device_type=str(creds.get("device_type") or device_type_eff or ""), ) except Exception: pass try: drain_read_channel(conn) except Exception: pass # Lightweight lookup for task_id / batch existence (no long hold). sdb = SessionLocal() try: batch_row = sdb.get(BizStateBatch, batch_id) if not batch_row: return 0, 0, True, False task_id = str(batch_row.task_id or "") finally: sdb.close() session = CollectSession( conn, vendor=vendor_eff, device_type=device_type_eff, vendor_key=vendor_key, read_timeout=per_cmd, shared_cache=shared_cache, cache_lock=cache_lock, cmd_locks=cmd_locks, ) # Resolve expand_all → concrete per-VRF commands via discover profile. flat_work: list[WorkItem] = [] persisted = aux_persisted if aux_persisted is not None else set() for concrete, params, profile_id, item_id, mode in work: if mode != "expand_all": flat_work.append((concrete, params, profile_id, item_id, mode)) continue profile = get_profile(profile_id) if profile is None or not profile.placeholders: any_fail = True _emit_task_event( task_id=task_id, message=f"expand_all missing profile {profile_id}", level="error", ) continue ph = profile.placeholders[0] disc = get_profile(str(ph.discover_profile_id or "").strip()) if disc is None: any_fail = True _emit_task_event( task_id=task_id, message=f"expand_all discover profile missing for {profile_id}", level="error", ) continue disc_cmd = normalize_command(disc.command_template) entry, _ = session.fetch_and_parse( disc_cmd, parser_id=disc.parser_id, textfsm_command=disc.textfsm_command or disc_cmd, params={}, ) if not entry.ok: any_fail = True _emit_task_event( task_id=task_id, message=f"expand_all discover failed: {entry.error}", level="error", ) continue try: pairs = expand_bindings_from_discover_records( profile=profile, records=entry.records, ) except ValueError as exc: any_fail = True _emit_task_event(task_id=task_id, message=str(exc), level="error") continue for cmd, p in pairs: flat_work.append((cmd, p, profile_id, item_id, "normal")) for concrete, params, profile_id, item_id, mode in flat_work: if holder.get("timed_out"): raise TimeoutError(f"{label}_aborted") cmd_id = uuid4().hex raw_text = "" cache_hit_primary = False try: cached = session.get_cached(concrete) if cached is not None and str(cached.raw or "").strip(): raw_text = str(cached.raw or "") cache_hit_primary = True else: raw_text = str( send_show_command(conn, concrete, read_timeout=per_cmd) or "" ) except Exception as exc: any_fail = True sp = SpooledCommand( id=cmd_id, batch_id=batch_id, task_item_id=item_id, profile_id=profile_id, raw_command=concrete[:512], params_json=dict(params or {}), parse_status="failed", message=_format_error(exc), ) _queue(sp) continue raw_rel = "" try: raw_rel = write_raw_text(batch_id, cmd_id, raw_text) except Exception: _log.exception("biz_state spool raw failed cmd=%s", cmd_id) if mode == "custom": any_ok = True _queue( SpooledCommand( id=cmd_id, batch_id=batch_id, task_item_id=item_id, profile_id=profile_id, raw_command=concrete[:512], params_json=dict(params or {}), parse_status="skipped_custom", message="custom_raw", raw_rel_path=raw_rel, ) ) continue hit = match_command(vendor_key=vendor_key, command=concrete) if not hit: any_fail = True _queue( SpooledCommand( id=cmd_id, batch_id=batch_id, task_item_id=item_id, profile_id=profile_id, raw_command=concrete[:512], params_json=dict(params or {}), parse_status="unmatched", message="no profile matched concrete command", raw_rel_path=raw_rel, ) ) continue merged = {**params, **hit.params} if not get_parser(hit.profile.parser_id): any_fail = True _queue( SpooledCommand( id=cmd_id, batch_id=batch_id, task_item_id=item_id, profile_id=hit.profile.profile_id, parser_id=hit.profile.parser_id, metric_id=hit.profile.metric_id, raw_command=concrete[:512], params_json=merged, parse_status="failed", message=f"unknown parser {hit.profile.parser_id}", raw_rel_path=raw_rel, ) ) continue if not cache_hit_primary: session.remember( concrete, raw=raw_text, ok=True, cmd_row_id=cmd_id, ) resolved_aux = [] aux_results: dict[str, Any] = {} for aux in list(hit.profile.aux_commands or []): try: ra = resolve_aux_command(aux, params=merged) except ValueError as exc: _queue( SpooledCommand( id=uuid4().hex, batch_id=batch_id, task_item_id=item_id, profile_id=str(aux.profile_id or "")[:128], raw_command=str(aux.key or "")[:512], parse_status="aux_failed", message=f"aux_for={cmd_id};resolve:{exc}"[:1020], ) ) continue resolved_aux.append(ra) aux_id = uuid4().hex entry, cache_hit = session.fetch_and_parse( ra.command, parser_id=ra.parser_id, textfsm_command=ra.textfsm_command, params=merged, cmd_row_id=aux_id, ) aux_results[ra.key] = entry aux_mid = str(getattr(ra.profile, "metric_id", "") or "").strip() aux_sp = SpooledCommand( id=aux_id, batch_id=batch_id, task_item_id=item_id, profile_id=ra.profile_id, parser_id=ra.parser_id, metric_id=aux_mid, raw_command=ra.command[:512], params_json={}, ) if cache_hit: aux_sp.parse_status = "aux_cached" aux_sp.message = ( f"aux_for={cmd_id};cache_hit;src={entry.cmd_row_id}" )[:1020] aux_sp.row_count = len(entry.records or []) elif not entry.ok: aux_sp.parse_status = "aux_failed" aux_sp.message = f"aux_for={cmd_id};{entry.error}"[:1020] try: aux_sp.raw_rel_path = write_raw_text( batch_id, aux_id, entry.raw or "" ) except Exception: pass else: aux_sp.parse_status = "aux" aux_sp.message = f"aux_for={cmd_id}"[:1020] try: aux_sp.raw_rel_path = write_raw_text( batch_id, aux_id, entry.raw or "" ) except Exception: pass aux_sp.row_count = len(entry.records or []) entry.cmd_row_id = aux_id persist_recs: list[dict[str, Any]] | None = None if entry.ok and entry.records and aux_mid in _GENERIC_METRICS: persist_key = (normalize_command(ra.command), aux_mid) do_persist = False if cache_lock is not None: with cache_lock: if persist_key not in persisted: persisted.add(persist_key) do_persist = True elif persist_key not in persisted: persisted.add(persist_key) do_persist = True if do_persist: aux_sp.persist_kind = "metric" persist_recs = list(entry.records) _queue(aux_sp, records=persist_recs) bundle = build_parse_bundle( primary_raw=raw_text, primary_parser_id=hit.profile.parser_id, aux_results=aux_results, resolved_aux=resolved_aux, ) primary = SpooledCommand( id=cmd_id, batch_id=batch_id, task_item_id=item_id, profile_id=hit.profile.profile_id, parser_id=hit.profile.parser_id, metric_id=hit.profile.metric_id, raw_command=concrete[:512], params_json=merged, raw_rel_path=raw_rel, ) try: records, fsm_tables, rule_keys = run_primary_with_bundle( hit.profile.parser_id, bundle=bundle, vendor=vendor_eff, device_type=device_type_eff, command=hit.profile.textfsm_command or concrete, textfsm_command=hit.profile.textfsm_command or "", params=merged, enrich_joins=list(hit.profile.enrich_joins or []), ) session.remember( concrete, raw=raw_text, fsm_tables=fsm_tables, records=records, ok=True, cmd_row_id=cmd_id, ) hints = [] if rule_keys: nonempty = [k for k in rule_keys if fsm_tables.get(k)] hints.append( f"fsm_keys={','.join(rule_keys)};hit={','.join(nonempty)}" ) if bundle.aux_records: hints.append( "aux=" + ",".join( f"{k}:{len(v)}" for k, v in bundle.aux_records.items() ) ) if hit.profile.enrich_joins: hints.append( "enrich=" + ",".join(j.from_aux for j in hit.profile.enrich_joins) ) if hints: primary.message = ";".join(hints)[:1020] primary.parse_status = "ok" any_ok = True persist_recs = None if hit.profile.metric_id == "lldp_neighbor": primary.persist_kind = "lldp" persist_recs = list(records or []) elif hit.profile.metric_id in _GENERIC_METRICS: primary.persist_kind = "metric" persist_recs = list(records or []) _queue(primary, records=persist_recs) except Exception as exc: any_fail = True primary.parse_status = "failed" primary.message = f"parse: {_format_error(exc)}" _queue(primary) # Final flush for this lane. if pending: _c, _r = _flush_spooled_commands(batch_id, pending) total_rows += int(_r or 0) return total_rows, cmd_count, any_fail, any_ok finally: # Best-effort: persist whatever was collected before timeout/abort. if pending: try: _c, _r = _flush_spooled_commands(batch_id, pending) total_rows += int(_r or 0) except Exception: _log.exception( "biz_state flush on lane exit failed batch=%s", batch_id ) holder.pop("conn", None) close_netmiko_connection(conn) try: return run_cli_with_timeout( _session, timeout_sec=budget, conn_holder=holder, label=label, acquire_budget=True, ) except TimeoutError as exc: raise RuntimeError(str(exc)[:1020]) from exc def _absorb_lane_result( result: tuple[int, int, bool, bool] | BaseException, *, total_rows: int, cmd_count: int, any_fail: bool, any_ok: bool, lane_errors: list[str], ) -> tuple[int, int, bool, bool]: if isinstance(result, BaseException): lane_errors.append(_format_error(result)) return total_rows, cmd_count, True, any_ok rows, cmds, fail, ok = result return ( total_rows + int(rows or 0), cmd_count + int(cmds or 0), any_fail or bool(fail), any_ok or bool(ok), ) def _bump_batch_progress(batch_id: str, *, add_cmds: int = 0, add_rows: int = 0) -> None: """Atomically bump batch counters so UI can show progress while lanes still run.""" cmds = int(add_cmds or 0) rows = int(add_rows or 0) if not batch_id or (cmds <= 0 and rows <= 0): return db = SessionLocal() try: batch = ( db.query(BizStateBatch) .filter(BizStateBatch.id == batch_id) .with_for_update() .one_or_none() ) if not batch: return if cmds > 0: batch.command_count = int(batch.command_count or 0) + cmds if rows > 0: batch.row_count = int(batch.row_count or 0) + rows db.commit() except Exception: _log.exception("biz_state bump batch progress failed batch=%s", batch_id) try: db.rollback() except Exception: pass finally: db.close() def _is_stale_db_connection(exc: BaseException) -> bool: """True when PG/middleware closed an idle connection mid-collect.""" if isinstance(exc, OperationalError): return True if isinstance(exc, DBAPIError) and bool(getattr(exc, "connection_invalidated", False)): return True msg = str(exc or "").lower() return ( "server closed the connection" in msg or "connection not open" in msg or "connection already closed" in msg or "ssl connection has been closed" in msg ) def _invalidate_session(db) -> None: try: conn = db.connection() conn.invalidate() except Exception: pass try: db.close() except Exception: pass def _run_db_with_reconnect(fn, *, label: str = "biz_state_db"): """Run ``fn(db)`` on a fresh Session; retry once after disconnect/OperationalError.""" last: BaseException | None = None for attempt in range(2): db = SessionLocal() try: result = fn(db) try: db.close() except Exception: pass return result except Exception as exc: last = exc try: db.rollback() except Exception: pass _invalidate_session(db) if attempt == 0 and _is_stale_db_connection(exc): _log.warning("%s reconnect after stale connection: %s", label, exc) continue raise assert last is not None raise last def _batch_has_progress(batch_id: str) -> bool: """Whether progressive bumps already persisted cmds/rows for this batch.""" def _read(db) -> bool: live = db.get(BizStateBatch, batch_id) if not live: return False return int(live.command_count or 0) > 0 or int(live.row_count or 0) > 0 return bool(_run_db_with_reconnect(_read, label="biz_state_batch_progress")) def _finalize_batch_status( *, batch_id: str, task_id: str, cmd_count: int, total_rows: int, any_fail: bool, any_ok: bool, lane_errors: list[str], ) -> str: """Write terminal batch status on a fresh Session (retry once on disconnect).""" def _write(db) -> str: batch = db.get(BizStateBatch, batch_id) if not batch: return "" # Prefer progressive counters (survive lane timeout) over in-memory lane totals. batch.command_count = max(int(batch.command_count or 0), int(cmd_count or 0)) batch.row_count = max(int(batch.row_count or 0), int(total_rows or 0)) batch.ended_at = _utcnow() if any_fail and any_ok: batch.status = "partial" if lane_errors: batch.message = "; ".join(lane_errors)[:1020] elif any_fail and not any_ok: batch.status = "failed" batch.message = ( "; ".join(lane_errors)[:1020] if lane_errors else "all commands failed" ) else: batch.status = "success" batch.message = "" status = str(batch.status or "") db.commit() if status in ("success", "partial"): try: from .compare_service import try_auto_compare_for_task try_auto_compare_for_task(db, task_id, batch_id) except Exception: _log.exception("biz_state auto compare hook failed task=%s", task_id) return status return str( _run_db_with_reconnect(_write, label="biz_state_finalize") or "" ) def _fail_batch_status(batch_id: str, error: str) -> None: """Mark batch failed on a fresh Session (retry once on disconnect).""" msg = str(error or "")[:1020] def _write(db) -> None: batch = db.get(BizStateBatch, batch_id) if not batch: return if str(batch.status or "") != "running": return batch.status = "failed" batch.message = msg batch.ended_at = _utcnow() db.commit() _run_db_with_reconnect(_write, label="biz_state_fail_batch") def _run_collect_session( *, task_id: str, batch_id: str, source: str, ne_id: str, vendor: str, device_type: str, ) -> None: light_per = int(settings.ne_collect_read_timeout_sec or 120) light_cap = int(settings.ne_collect_run_timeout_cap_sec or 600) heavy_per = int(getattr(settings, "biz_state_heavy_read_timeout_sec", 300) or 300) heavy_cap = int(getattr(settings, "biz_state_heavy_run_timeout_cap_sec", 900) or 900) # Phase 1: resolve target + build work list, then release the DB connection. # Holding one Session across heavy CLI (up to ~2400s) lets PG/middleware close # the idle connection; finalize would then hit OperationalError. creds: dict[str, Any] vendor_eff: str device_type_eff: str vendor_key: str light_work: list[WorkItem] heavy_work: list[WorkItem] db = SessionLocal() try: task = db.get(BizStateTask, task_id) batch = db.get(BizStateBatch, batch_id) if not task or not batch: return try: if source == "managed": creds, info = resolve_cli_target(db, managed_ne_id=ne_id) elif source == "ume": creds, info = resolve_cli_target(db, ume_ne_id=ne_id) else: raise RuntimeError("invalid_source") except HTTPException as exc: raise RuntimeError(str(exc.detail or "resolve_failed")) from exc skip = cli_creds_skip_reason(creds, interactive=False) if skip: raise RuntimeError(skip) vendor_eff = str(info.get("vendor") or vendor or "") device_type_eff = str(info.get("device_type") or device_type or "") if vendor_eff and vendor_eff != task.vendor: task.vendor = vendor_eff if device_type_eff and device_type_eff != task.device_type: task.device_type = device_type_eff db.commit() vendor_key = resolve_vendor_key(vendor_eff, device_type_eff) items = ( db.query(BizStateTaskItem) .filter( BizStateTaskItem.task_id == task_id, BizStateTaskItem.enabled.is_(True), ) .order_by(BizStateTaskItem.sort_order.asc()) .all() ) # Build work list before opening session work: list[WorkItem] = [] # Dedupe same CLI → same metric (e.g. legacy if_intf + config_interface). seen_work: set[tuple[str, str]] = set() for item in items: if item.kind == "custom_raw": cmd = normalize_command(item.command_override) if cmd: key = (cmd, "__custom__") if key in seen_work: continue seen_work.add(key) work.append((cmd, {}, "", item.id, "custom")) continue profile = _resolve_collect_profile(item.source_profile_id) if profile is None: _append_event( db, task_id=task_id, message=f"skip profile {item.source_profile_id} (missing or disabled)", level="info", ) continue binds = _bindings_for_item(db, item.id) try: pairs = expand_from_bindings( profile=profile, bindings=binds, command_override=item.command_override, ) except ValueError as exc: _append_event(db, task_id=task_id, message=str(exc), level="error") continue for concrete, params in pairs: if concrete == EXPAND_ALL_COMMAND: # Defer VRF list expansion until CollectSession is open. work.append(("", dict(params or {}), profile.profile_id, item.id, "expand_all")) continue cmd = normalize_command(concrete) # Prefer match_command metric so remapped if_intf shares key with config_interface hit = match_command(vendor_key=vendor_key, command=cmd) mid = str((hit.profile.metric_id if hit else profile.metric_id) or "").strip() pid = str((hit.profile.profile_id if hit else profile.profile_id) or "").strip() key = (cmd, mid or pid) if key in seen_work: continue seen_work.add(key) work.append((cmd, params, pid or profile.profile_id, item.id, "normal")) if not work: batch.status = "failed" batch.message = "no commands to run" batch.ended_at = _utcnow() db.commit() raise RuntimeError("no commands to run") light_work, heavy_work = partition_work(work) db.commit() finally: db.close() # Fresh spool dir for this batch (collect → disk, then flush to DB). try: from .spool import clear_batch_spool clear_batch_spool(batch_id) except Exception: _log.exception("biz_state clear spool failed batch=%s", batch_id) # Phase 2: CLI lanes — no outer Session held across long timeouts. shared_cache: dict[str, Any] = {} cache_lock = threading.RLock() cmd_locks: dict[str, Any] = {} aux_persisted: set[tuple[str, str]] = set() lane_kwargs = dict( batch_id=batch_id, creds=creds, vendor_eff=vendor_eff, device_type_eff=device_type_eff, vendor_key=vendor_key, shared_cache=shared_cache, cache_lock=cache_lock, cmd_locks=cmd_locks, aux_persisted=aux_persisted, ) def _run_light() -> tuple[int, int, bool, bool]: return _run_collect_lane( work=light_work, per_cmd=light_per, cap=light_cap, label="biz_state_light", **lane_kwargs, ) def _run_heavy() -> tuple[int, int, bool, bool]: return _run_collect_lane( work=heavy_work, per_cmd=heavy_per, cap=heavy_cap, label="biz_state_heavy", **lane_kwargs, ) total_rows = 0 cmd_count = 0 any_fail = False any_ok = False lane_errors: list[str] = [] if light_work and heavy_work: heavy_fut = _heavy_cli_pool().submit(_run_heavy) light_res: tuple[int, int, bool, bool] | BaseException try: light_res = _run_light() except Exception as exc: light_res = exc heavy_res: tuple[int, int, bool, bool] | BaseException try: heavy_res = heavy_fut.result() except Exception as exc: heavy_res = exc total_rows, cmd_count, any_fail, any_ok = _absorb_lane_result( light_res, total_rows=total_rows, cmd_count=cmd_count, any_fail=any_fail, any_ok=any_ok, lane_errors=lane_errors, ) total_rows, cmd_count, any_fail, any_ok = _absorb_lane_result( heavy_res, total_rows=total_rows, cmd_count=cmd_count, any_fail=any_fail, any_ok=any_ok, lane_errors=lane_errors, ) elif heavy_work: try: total_rows, cmd_count, any_fail, any_ok = _run_heavy() except Exception as exc: total_rows, cmd_count, any_fail, any_ok = _absorb_lane_result( exc, total_rows=0, cmd_count=0, any_fail=False, any_ok=False, lane_errors=lane_errors, ) else: try: total_rows, cmd_count, any_fail, any_ok = _run_light() except Exception as exc: total_rows, cmd_count, any_fail, any_ok = _absorb_lane_result( exc, total_rows=0, cmd_count=0, any_fail=False, any_ok=False, lane_errors=lane_errors, ) if lane_errors and not any_ok and cmd_count == 0: # Progressive bumps may already have cmds; only hard-fail if nothing landed. if not _batch_has_progress(batch_id): raise RuntimeError("; ".join(lane_errors)[:1020]) any_fail = True _finalize_batch_status( batch_id=batch_id, task_id=task_id, cmd_count=cmd_count, total_rows=total_rows, any_fail=any_fail, any_ok=any_ok, lane_errors=lane_errors, ) def _purge(db) -> None: _purge_task_retention(db, task_id=task_id) try: _run_db_with_reconnect(_purge, label="biz_state_purge") except Exception: _log.exception("biz_state retention purge wrapper failed task=%s", task_id) def _purge_task_retention(db, *, task_id: str) -> None: from .retention import purge_task_batches task = db.get(BizStateTask, task_id) if not task: return try: info = purge_task_batches(db, task) if info.get("dropped"): _log.info( "biz_state retention purged task=%s dropped=%s days=%s daily=%s", task_id, info.get("dropped"), info.get("retention_days"), info.get("daily_keep_enabled"), ) except Exception: _log.exception("biz_state retention purge failed task=%s", task_id) def trigger_collect_now(task_id: str) -> dict[str, Any]: db = SessionLocal() try: task = db.get(BizStateTask, task_id) if not task: raise HTTPException(status_code=404, detail="task not found") if task.collect_running: raise HTTPException(status_code=409, detail="collect already running") finally: db.close() # Manual: allow even when schedule is on (status=running) or paused/stopped. dispatch_collect(task_id, manual=True) return {"ok": True, "task_id": task_id}