utp_service/backend/app/services/object_db.py

"""
Connects to a site object's external PostgreSQL database and syncs devices.
Uses short-lived psycopg3 connections (connect → query → disconnect).

Expected table schema (racks):
  id          TEXT  — external identifier, e.g. "lane-1-1"
  name        TEXT  — short label, e.g. "01"
  type        TEXT  — "rack.lane" | "rack.payment" | ...
  data        JSONB — contains "host" (IP address) and other device-specific fields
"""

from dataclasses import dataclass, field

import asyncssh
import psycopg
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.models.device import Device
from app.models.object import Object
from app.services.crypto import decrypt


@dataclass
class SyncResult:
    created: int = 0
    updated: int = 0
    skipped: int = 0
    errors: list[str] = field(default_factory=list)


# rack.* type → our internal category
_CATEGORY_MAP: dict[str, str] = {
    "rack.lane": "embedded",
    "rack.payment": "other",
    "rack.server": "server",
    "rack.camera": "camera",
    "raspberry": "raspberry",
    "rpi": "raspberry",
    "mikrotik": "mikrotik",
    "camera": "camera",
    "server": "server",
    "embedded": "embedded",
}


def _map_category(raw: str | None) -> str:
    if not raw:
        return "other"
    return _CATEGORY_MAP.get(raw.strip().lower(), "other")


async def _fetch_rows(dsn: str, table: str) -> list[tuple]:
    """Open a short-lived psycopg connection and fetch device rows."""
    async with await psycopg.AsyncConnection.connect(dsn) as conn:
        cursor = await conn.execute(
            f"SELECT id, name, type, data->>'host' AS host FROM \"{table}\""
        )
        return await cursor.fetchall()


async def sync_from_object_db(db: AsyncSession, obj: Object) -> SyncResult:
    result = SyncResult()

    if not all([obj.db_host, obj.db_name, obj.db_user, obj.db_pass_enc, obj.db_table]):
        result.errors.append("Object DB connection is not fully configured")
        return result

    try:
        db_password = decrypt(obj.db_pass_enc)
    except Exception:
        result.errors.append("Failed to decrypt object DB password")
        return result

    if obj.db_via_tunnel:
        # Connect via SSH tunnel through obj.server_ip
        if not obj.server_ip or not obj.ssh_user or not obj.ssh_pass_enc:
            result.errors.append("SSH credentials not configured for tunnel (need server_ip, ssh_user, ssh_password)")
            return result
        try:
            ssh_password = decrypt(obj.ssh_pass_enc)
        except Exception:
            result.errors.append("Failed to decrypt SSH password for tunnel")
            return result

        try:
            async with asyncssh.connect(
                obj.server_ip,
                port=obj.ssh_port or 22,
                username=obj.ssh_user,
                password=ssh_password,
                known_hosts=None,
            ) as tunnel:
                listener = await tunnel.forward_local_port(
                    "127.0.0.1", 0, obj.db_host, obj.db_port or 5432
                )
                local_port = listener.get_port()
                dsn = (
                    f"host=127.0.0.1 port={local_port} dbname={obj.db_name} "
                    f"user={obj.db_user} password={db_password} connect_timeout=10"
                )
                try:
                    remote_rows = await _fetch_rows(dsn, obj.db_table)
                finally:
                    listener.close()
        except Exception as exc:
            result.errors.append(f"SSH tunnel failed: {exc}")
            return result
    else:
        dsn = (
            f"host={obj.db_host} port={obj.db_port or 5432} dbname={obj.db_name} "
            f"user={obj.db_user} password={db_password} connect_timeout=10"
        )
        try:
            remote_rows = await _fetch_rows(dsn, obj.db_table)
        except Exception as exc:
            result.errors.append(f"Cannot connect to object DB: {exc}")
            return result

    for row in remote_rows:
        external_id, name, rack_type, host = row[0], row[1], row[2], row[3]

        ip = str(host).strip() if host else None
        if not ip:
            result.skipped += 1
            continue

        category = _map_category(rack_type)
        hostname = str(name).strip() if name else None

        existing = await db.execute(
            select(Device).where(
                Device.object_id == obj.id,
                Device.external_id == str(external_id),
            )
        )
        device = existing.scalar_one_or_none()

        if device is not None:
            changed = False
            if device.ip != ip:
                device.ip = ip
                changed = True
            if device.hostname != hostname:
                device.hostname = hostname
                changed = True
            if device.category != category:
                device.category = category
                changed = True
            if changed:
                result.updated += 1
            else:
                result.skipped += 1
        else:
            db.add(Device(
                object_id=obj.id,
                ip=ip,
                hostname=hostname,
                category=category,
                source="auto_sync",
                external_id=str(external_id),
                role=rack_type or "unknown",
            ))
            result.created += 1

    return result