"""Startup migration / backfill for port traffic device-centric model.""" from __future__ import annotations import logging from uuid import uuid4 from sqlalchemy import text from sqlalchemy.orm import Session from .models import PortTrafficDevice, PortTrafficSample, PortTrafficSeries, PortTrafficTarget from .schema_patches import _run_sql _log = logging.getLogger("netx.port_traffic.migrate") def ensure_port_traffic_series_schema(conn) -> None: """DDL for device + series + device_id columns (Postgres IF NOT EXISTS).""" conn.exec_driver_sql( """ CREATE TABLE IF NOT EXISTS port_traffic_device ( id VARCHAR(64) PRIMARY KEY, source VARCHAR(32) DEFAULT 'managed', ne_id VARCHAR(128) DEFAULT '', ne_name VARCHAR(256) DEFAULT '', ne_ip VARCHAR(128) DEFAULT '', vendor VARCHAR(64) DEFAULT '', note VARCHAR(256) DEFAULT '', status VARCHAR(32) DEFAULT 'draft', interval_sec INTEGER DEFAULT 60, retention_days INTEGER DEFAULT 7, concurrency INTEGER DEFAULT 1, collect_running BOOLEAN DEFAULT FALSE, last_collect_started_at TIMESTAMP, last_collect_ended_at TIMESTAMP, last_error VARCHAR(1024) DEFAULT '', created_at TIMESTAMP, updated_at TIMESTAMP ) """ ) conn.exec_driver_sql( "CREATE UNIQUE INDEX IF NOT EXISTS uq_port_traffic_device_ne ON port_traffic_device (source, ne_id)" ) conn.exec_driver_sql( "CREATE INDEX IF NOT EXISTS ix_port_traffic_device_status ON port_traffic_device (status)" ) conn.exec_driver_sql( """ CREATE TABLE IF NOT EXISTS port_traffic_series ( id VARCHAR(64) PRIMARY KEY, device_id VARCHAR(64) DEFAULT '', title VARCHAR(256) DEFAULT '', status VARCHAR(32) DEFAULT 'active', created_at TIMESTAMP ) """ ) # Legacy column may still exist as task_id conn.exec_driver_sql( "ALTER TABLE port_traffic_series ADD COLUMN IF NOT EXISTS device_id VARCHAR(64) DEFAULT ''" ) conn.exec_driver_sql( "ALTER TABLE port_traffic_series ADD COLUMN IF NOT EXISTS task_id VARCHAR(64) DEFAULT ''" ) # Older installs created task_id as NOT NULL before the device-centric rename. # Keep both columns populated and relax NOT NULL so ORM inserts that only set # device_id (or only task_id) do not 500. Each statement uses SAVEPOINT. for sql in ( """ UPDATE port_traffic_series SET device_id = task_id WHERE COALESCE(device_id, '') = '' AND COALESCE(task_id, '') <> '' """, """ UPDATE port_traffic_series SET task_id = device_id WHERE COALESCE(task_id, '') = '' AND COALESCE(device_id, '') <> '' """, "ALTER TABLE port_traffic_series ALTER COLUMN task_id DROP NOT NULL", "ALTER TABLE port_traffic_series ALTER COLUMN task_id SET DEFAULT ''", "CREATE INDEX IF NOT EXISTS ix_port_traffic_series_device_id ON port_traffic_series (device_id)", "CREATE INDEX IF NOT EXISTS ix_port_traffic_series_status ON port_traffic_series (status)", "ALTER TABLE port_traffic_target ADD COLUMN IF NOT EXISTS series_id VARCHAR(64) DEFAULT ''", "ALTER TABLE port_traffic_target ADD COLUMN IF NOT EXISTS device_id VARCHAR(64) DEFAULT ''", "ALTER TABLE port_traffic_target ADD COLUMN IF NOT EXISTS task_id VARCHAR(64) DEFAULT ''", """ UPDATE port_traffic_target SET device_id = task_id WHERE COALESCE(device_id, '') = '' AND COALESCE(task_id, '') <> '' """, """ UPDATE port_traffic_target SET task_id = device_id WHERE COALESCE(task_id, '') = '' AND COALESCE(device_id, '') <> '' """, "ALTER TABLE port_traffic_target ALTER COLUMN task_id DROP NOT NULL", "ALTER TABLE port_traffic_target ALTER COLUMN task_id SET DEFAULT ''", "CREATE INDEX IF NOT EXISTS ix_port_traffic_target_series_id ON port_traffic_target (series_id)", "CREATE INDEX IF NOT EXISTS ix_port_traffic_target_device_id ON port_traffic_target (device_id)", "ALTER TABLE port_traffic_sample ADD COLUMN IF NOT EXISTS series_id VARCHAR(64) DEFAULT ''", "CREATE INDEX IF NOT EXISTS ix_port_traffic_sample_series_id ON port_traffic_sample (series_id)", # One active interface globally per physical port (may fail if duplicates exist) """ CREATE UNIQUE INDEX IF NOT EXISTS uq_port_traffic_target_active_if ON port_traffic_target (source, target_id, ifname) WHERE status = 'active' """, ): _run_sql(conn, sql) conn.exec_driver_sql( """ CREATE TABLE IF NOT EXISTS port_traffic_event ( id VARCHAR(64) PRIMARY KEY, device_id VARCHAR(64) DEFAULT '', target_row_id VARCHAR(64) DEFAULT '', ifname VARCHAR(128) DEFAULT '', level VARCHAR(16) DEFAULT 'error', message TEXT DEFAULT '', created_at TIMESTAMP ) """ ) conn.exec_driver_sql( "CREATE INDEX IF NOT EXISTS ix_port_traffic_event_device_id ON port_traffic_event (device_id)" ) conn.exec_driver_sql( "CREATE INDEX IF NOT EXISTS ix_port_traffic_event_created_at ON port_traffic_event (created_at)" ) conn.exec_driver_sql( "CREATE INDEX IF NOT EXISTS ix_port_traffic_event_level ON port_traffic_event (level)" ) conn.exec_driver_sql( """ CREATE TABLE IF NOT EXISTS port_traffic_board ( id VARCHAR(64) PRIMARY KEY, name VARCHAR(256) DEFAULT '', remark VARCHAR(1024) DEFAULT '', cols INTEGER DEFAULT 2, created_by VARCHAR(64) DEFAULT '', updated_by VARCHAR(64) DEFAULT '', created_at TIMESTAMP, updated_at TIMESTAMP ) """ ) conn.exec_driver_sql( "CREATE INDEX IF NOT EXISTS ix_port_traffic_board_name ON port_traffic_board (name)" ) conn.exec_driver_sql( "CREATE INDEX IF NOT EXISTS ix_port_traffic_board_updated_at ON port_traffic_board (updated_at)" ) conn.exec_driver_sql( """ CREATE TABLE IF NOT EXISTS port_traffic_panel ( id VARCHAR(64) PRIMARY KEY, board_id VARCHAR(64) DEFAULT '', title VARCHAR(256) DEFAULT '', target_id VARCHAR(64) DEFAULT '', range_hours INTEGER DEFAULT 24, baseline VARCHAR(16) DEFAULT 'off', offset_hours INTEGER DEFAULT 0, ahead_hours INTEGER DEFAULT 1, baseline_target_id VARCHAR(64) DEFAULT '', y_mode VARCHAR(16) DEFAULT 'auto', ord INTEGER DEFAULT 0, col_span INTEGER DEFAULT 1, row_span INTEGER DEFAULT 1, created_at TIMESTAMP, updated_at TIMESTAMP ) """ ) conn.exec_driver_sql( "ALTER TABLE port_traffic_panel ADD COLUMN IF NOT EXISTS ahead_hours INTEGER DEFAULT 1" ) conn.exec_driver_sql( "CREATE INDEX IF NOT EXISTS ix_port_traffic_panel_board_id ON port_traffic_panel (board_id)" ) conn.exec_driver_sql( "CREATE INDEX IF NOT EXISTS ix_port_traffic_panel_target_id ON port_traffic_panel (target_id)" ) def migrate_tasks_to_devices(db: Session) -> int: """Collapse legacy port_traffic_task rows into per-NE devices; reassign targets/series.""" # Copy task_id → device_id where missing try: db.execute( text( """ UPDATE port_traffic_target SET device_id = task_id WHERE COALESCE(device_id, '') = '' AND COALESCE(task_id, '') <> '' """ ) ) db.execute( text( """ UPDATE port_traffic_series SET device_id = task_id WHERE COALESCE(device_id, '') = '' AND COALESCE(task_id, '') <> '' """ ) ) db.commit() except Exception: db.rollback() # If legacy task table missing, nothing else to do try: db.execute(text("SELECT 1 FROM port_traffic_task LIMIT 1")) except Exception: db.rollback() return 0 targets = db.query(PortTrafficTarget).all() if not targets: # Still create devices from empty? skip return 0 # Group active+any targets by (source, ne_id) groups: dict[tuple[str, str], list[PortTrafficTarget]] = {} for t in targets: src = str(t.source or "managed").strip().lower() or "managed" ne = str(t.target_id or "").strip() if not ne: continue groups.setdefault((src, ne), []).append(t) created = 0 # Load legacy tasks for interval/status legacy: dict[str, dict] = {} try: rows = db.execute( text( """ SELECT id, title, status, interval_sec, retention_days, concurrency, collect_running, last_collect_started_at, last_collect_ended_at, last_error, created_at, updated_at FROM port_traffic_task """ ) ).mappings().all() for r in rows: legacy[str(r["id"])] = dict(r) except Exception: db.rollback() legacy = {} for (src, ne), members in groups.items(): existing = ( db.query(PortTrafficDevice) .filter(PortTrafficDevice.source == src, PortTrafficDevice.ne_id == ne) .first() ) if existing: device = existing else: # Pick policy from the richest legacy task among members intervals: list[int] = [] retentions: list[int] = [] statuses: list[str] = [] note = "" created_at = None for m in members: tid = str(getattr(m, "device_id", None) or getattr(m, "task_id", None) or "") meta = legacy.get(tid) or {} if meta.get("interval_sec"): intervals.append(int(meta["interval_sec"])) if meta.get("retention_days"): retentions.append(int(meta["retention_days"])) if meta.get("status"): statuses.append(str(meta["status"])) if not note and meta.get("title"): note = str(meta["title"])[:256] if created_at is None and meta.get("created_at"): created_at = meta["created_at"] status = "running" if "running" in statuses else (statuses[0] if statuses else "stopped") sample = members[0] device = PortTrafficDevice( id=uuid4().hex, source=src, ne_id=ne, ne_name=str(sample.ne_name or ""), ne_ip=str(sample.ne_ip or ""), vendor=str(sample.vendor or ""), note=note, status=status if status in ("running", "paused", "stopped", "draft") else "stopped", interval_sec=min(intervals) if intervals else 60, retention_days=max(retentions) if retentions else 7, concurrency=1, collect_running=False, created_at=created_at, ) db.add(device) db.flush() created += 1 # Deduplicate active ifnames: keep earliest active, retire the rest seen_active: set[str] = set() for m in sorted(members, key=lambda x: str(x.created_at or "")): m.device_id = str(device.id) if str(m.status) == "active": key = str(m.ifname or "").strip() if key in seen_active: m.status = "retired" else: seen_active.add(key) # Keep ne metadata in sync if not device.ne_name and m.ne_name: device.ne_name = str(m.ne_name) if not device.ne_ip and m.ne_ip: device.ne_ip = str(m.ne_ip) if not device.vendor and m.vendor: device.vendor = str(m.vendor) # Re-point series series_ids = {str(m.series_id) for m in members if m.series_id} if series_ids: db.query(PortTrafficSeries).filter(PortTrafficSeries.id.in_(list(series_ids))).update( {PortTrafficSeries.device_id: str(device.id)}, synchronize_session=False, ) db.commit() if created: _log.info("port_traffic migrated devices created=%s groups=%s", created, len(groups)) return created def backfill_port_traffic_series(db: Session) -> int: """Create one series per target missing series_id; stamp samples.""" migrate_tasks_to_devices(db) targets = ( db.query(PortTrafficTarget) .filter((PortTrafficTarget.series_id == None) | (PortTrafficTarget.series_id == "")) # noqa: E711 .all() ) created = 0 for t in targets: device_id = str(t.device_id or "") title = _default_series_title(t.ne_name or "", t.ifname or "") title = _unique_series_title(db, device_id, title) sid = uuid4().hex db.add( PortTrafficSeries( id=sid, device_id=device_id, title=title, status="active", ) ) t.series_id = sid db.query(PortTrafficSample).filter( PortTrafficSample.target_row_id == str(t.id), (PortTrafficSample.series_id == None) | (PortTrafficSample.series_id == ""), # noqa: E711 ).update({"series_id": sid}, synchronize_session=False) created += 1 if created: db.commit() _log.info("port_traffic series backfill targets=%s", created) try: orphans = db.execute( text( """ UPDATE port_traffic_sample AS s SET series_id = t.series_id FROM port_traffic_target AS t WHERE s.target_row_id = t.id AND COALESCE(s.series_id, '') = '' AND COALESCE(t.series_id, '') <> '' """ ) ) if orphans.rowcount: db.commit() _log.info("port_traffic series sample stamp rows=%s", orphans.rowcount) except Exception: db.rollback() return created def default_series_title(ne_name: str, ifname: str) -> str: return _default_series_title(ne_name, ifname) def unique_series_title(db: Session, device_id: str, base: str) -> str: return _unique_series_title(db, device_id, base) def _default_series_title(ne_name: str, ifname: str) -> str: ne = (ne_name or "").strip() or "NE" iface = (ifname or "").strip() or "if" return f"{ne}:{iface}"[:256] def _unique_series_title(db: Session, device_id: str, base: str) -> str: title = base[:256] exists = ( db.query(PortTrafficSeries.id) .filter(PortTrafficSeries.device_id == device_id, PortTrafficSeries.title == title) .first() ) if not exists: return title for i in range(2, 1000): candidate = f"{base[:240]}#{i}" exists = ( db.query(PortTrafficSeries.id) .filter(PortTrafficSeries.device_id == device_id, PortTrafficSeries.title == candidate) .first() ) if not exists: return candidate return f"{base[:200]}#{uuid4().hex[:8]}"