netx/netx_api/port_traffic_migrate.py

"""Startup migration / backfill for port traffic logical series."""

from __future__ import annotations

import logging
from uuid import uuid4

from sqlalchemy import text
from sqlalchemy.orm import Session

from .models import PortTrafficSample, PortTrafficSeries, PortTrafficTarget

_log = logging.getLogger("netx.port_traffic.migrate")


def ensure_port_traffic_series_schema(conn) -> None:
    """DDL for series + series_id columns (Postgres IF NOT EXISTS)."""
    conn.exec_driver_sql(
        """
        CREATE TABLE IF NOT EXISTS port_traffic_series (
            id VARCHAR(64) PRIMARY KEY,
            task_id VARCHAR(64),
            title VARCHAR(256) DEFAULT '',
            status VARCHAR(32) DEFAULT 'active',
            created_at TIMESTAMP
        )
        """
    )
    conn.exec_driver_sql(
        "CREATE INDEX IF NOT EXISTS ix_port_traffic_series_task_id ON port_traffic_series (task_id)"
    )
    conn.exec_driver_sql(
        "CREATE INDEX IF NOT EXISTS ix_port_traffic_series_status ON port_traffic_series (status)"
    )
    conn.exec_driver_sql(
        "CREATE UNIQUE INDEX IF NOT EXISTS uq_port_traffic_series_title ON port_traffic_series (task_id, title)"
    )
    conn.exec_driver_sql(
        "ALTER TABLE port_traffic_target ADD COLUMN IF NOT EXISTS series_id VARCHAR(64) DEFAULT ''"
    )
    conn.exec_driver_sql(
        "CREATE INDEX IF NOT EXISTS ix_port_traffic_target_series_id ON port_traffic_target (series_id)"
    )
    conn.exec_driver_sql(
        "ALTER TABLE port_traffic_sample ADD COLUMN IF NOT EXISTS series_id VARCHAR(64) DEFAULT ''"
    )
    conn.exec_driver_sql(
        "CREATE INDEX IF NOT EXISTS ix_port_traffic_sample_series_id ON port_traffic_sample (series_id)"
    )


def backfill_port_traffic_series(db: Session) -> int:
    """Create one series per target missing series_id; stamp samples."""
    targets = (
        db.query(PortTrafficTarget)
        .filter((PortTrafficTarget.series_id == None) | (PortTrafficTarget.series_id == ""))  # noqa: E711
        .all()
    )
    created = 0
    for t in targets:
        title = _default_series_title(t.ne_name or "", t.ifname or "")
        title = _unique_series_title(db, str(t.task_id), title)
        sid = uuid4().hex
        db.add(
            PortTrafficSeries(
                id=sid,
                task_id=str(t.task_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, task_id: str, base: str) -> str:
    return _unique_series_title(db, task_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, task_id: str, base: str) -> str:
    title = base[:256]
    exists = (
        db.query(PortTrafficSeries.id)
        .filter(PortTrafficSeries.task_id == task_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.task_id == task_id, PortTrafficSeries.title == candidate)
            .first()
        )
        if not exists:
            return candidate
    return f"{base[:200]}#{uuid4().hex[:8]}"