netx/netx_api/ume_alarm_subscription_store.py

from __future__ import annotations

from datetime import datetime, timezone

from sqlalchemy.orm import Session

from .models import UmeAlarmSubscription

DEFAULT_SUBSCRIPTION_KEY = "default"


def _utc_now_naive() -> datetime:
    return datetime.now(timezone.utc).replace(tzinfo=None)


def load_subscription(db: Session, *, cache_key: str = DEFAULT_SUBSCRIPTION_KEY) -> tuple[str, str, str] | None:
    row = db.get(UmeAlarmSubscription, str(cache_key or DEFAULT_SUBSCRIPTION_KEY))
    if row is None:
        return None
    sub_id = str(row.subscription_id or "").strip()
    uri = str(row.wss_uri or "").strip()
    if not sub_id or not uri:
        return None
    return sub_id, uri, str(row.topic or "ALARM")


def save_subscription(
    db: Session,
    *,
    subscription_id: str,
    wss_uri: str,
    topic: str = "ALARM",
    cache_key: str = DEFAULT_SUBSCRIPTION_KEY,
) -> UmeAlarmSubscription:
    key = str(cache_key or DEFAULT_SUBSCRIPTION_KEY)
    now = _utc_now_naive()
    row = db.get(UmeAlarmSubscription, key)
    if row is None:
        row = UmeAlarmSubscription(
            cache_key=key,
            subscription_id=str(subscription_id or "").strip(),
            wss_uri=str(wss_uri or "").strip(),
            topic=str(topic or "ALARM").strip() or "ALARM",
            established_at=now,
            updated_at=now,
        )
        db.add(row)
    else:
        row.subscription_id = str(subscription_id or "").strip()
        row.wss_uri = str(wss_uri or "").strip()
        row.topic = str(topic or "ALARM").strip() or "ALARM"
        row.updated_at = now
    db.flush()
    return row


def clear_subscription(db: Session, *, cache_key: str = DEFAULT_SUBSCRIPTION_KEY) -> None:
    key = str(cache_key or DEFAULT_SUBSCRIPTION_KEY)
    row = db.get(UmeAlarmSubscription, key)
    if row is not None:
        db.delete(row)
        db.flush()