vmware-api-simulator/app/vsphere/inventory.py

"""Persistent vSphere managed-object inventory."""

from __future__ import annotations

import json
from dataclasses import dataclass
from typing import Any

from app.db.pool import Database


@dataclass(frozen=True, slots=True)
class ManagedObject:
    moid: str
    type: str
    name: str
    parent_moid: str | None
    props: dict[str, Any]


def _pool(database: Database) -> Any:
    return database.pool  # type: ignore[attr-defined]


async def count_objects(database: Database) -> int:
    pool = _pool(database)
    async with pool.acquire() as conn:
        return int(await conn.fetchval("SELECT COUNT(*) FROM vsphere_objects") or 0)


async def list_objects(
    database: Database,
    *,
    type_name: str | None = None,
) -> list[ManagedObject]:
    pool = _pool(database)
    async with pool.acquire() as conn:
        if type_name is None:
            rows = await conn.fetch(
                "SELECT moid, type, name, parent_moid, props FROM vsphere_objects ORDER BY type, name"
            )
        else:
            rows = await conn.fetch(
                """
                SELECT moid, type, name, parent_moid, props FROM vsphere_objects
                WHERE type = $1 ORDER BY name
                """,
                type_name,
            )
    return [_row(row) for row in rows]


async def get_object(database: Database, moid: str) -> ManagedObject | None:
    pool = _pool(database)
    async with pool.acquire() as conn:
        row = await conn.fetchrow(
            "SELECT moid, type, name, parent_moid, props FROM vsphere_objects WHERE moid = $1",
            moid,
        )
    return None if row is None else _row(row)


_TYPE_RESOURCE: dict[str, str] = {
    "HostSystem": "host",
    "VirtualMachine": "vm",
    "Datastore": "datastore",
    "Datacenter": "datacenter",
    "Folder": "folder",
    "ClusterComputeResource": "cluster",
    "ResourcePool": "resource_pool",
    "Network": "network",
}


async def require_object(
    database: Database,
    moid: str,
    *,
    type_name: str | None = None,
    resource: str | None = None,
) -> ManagedObject:
    """Return the managed object or raise a native not-found with the bad id."""

    from app.vsphere.errors import no_such

    obj = await get_object(database, moid)
    label = resource or (type_name and _TYPE_RESOURCE.get(type_name)) or "object"
    if obj is None or (type_name is not None and obj.type != type_name):
        raise no_such(label, moid)
    return obj


async def upsert_object(
    database: Database,
    *,
    moid: str,
    type_name: str,
    name: str,
    parent_moid: str | None,
    props: dict[str, Any],
) -> None:
    await upsert_objects_batch(
        database,
        [
            {
                "moid": moid,
                "type": type_name,
                "name": name,
                "parent_moid": parent_moid,
                "props": props,
            }
        ],
    )


async def upsert_objects_batch(database: Database, rows: list[dict[str, Any]]) -> None:
    if not rows:
        return
    pool = _pool(database)
    payload = [
        (
            str(row["moid"]),
            str(row["type"]),
            str(row["name"]),
            None if row.get("parent_moid") is None else str(row["parent_moid"]),
            json.dumps(row.get("props") or {}),
        )
        for row in rows
    ]
    async with pool.acquire() as conn:
        await conn.executemany(
            """
            INSERT INTO vsphere_objects (moid, type, name, parent_moid, props)
            VALUES ($1, $2, $3, $4, $5::jsonb)
            ON CONFLICT (moid) DO UPDATE SET
              type = EXCLUDED.type,
              name = EXCLUDED.name,
              parent_moid = EXCLUDED.parent_moid,
              props = EXCLUDED.props,
              updated_at = now()
            """,
            payload,
        )


async def count_by_type(database: Database) -> dict[str, int]:
    pool = _pool(database)
    async with pool.acquire() as conn:
        rows = await conn.fetch(
            "SELECT type, COUNT(*)::int AS count FROM vsphere_objects GROUP BY type ORDER BY type"
        )
    return {str(row["type"]): int(row["count"]) for row in rows}


async def update_props(database: Database, moid: str, props: dict[str, Any]) -> ManagedObject:
    pool = _pool(database)
    async with pool.acquire() as conn:
        row = await conn.fetchrow(
            """
            UPDATE vsphere_objects
            SET props = $2::jsonb, updated_at = now()
            WHERE moid = $1
            RETURNING moid, type, name, parent_moid, props
            """,
            moid,
            json.dumps(props),
        )
    if row is None:
        raise KeyError(moid)
    return _row(row)


async def delete_object(database: Database, moid: str) -> bool:
    pool = _pool(database)
    async with pool.acquire() as conn:
        result = await conn.execute("DELETE FROM vsphere_objects WHERE moid = $1", moid)
    return result.endswith("1")


async def next_moid(database: Database, prefix: str) -> str:
    """Allocate the next MoID under an advisory lock (safe under concurrent create).

    Inserts a reservation row before releasing the lock so two creators cannot
    compute the same next id between allocate and upsert.
    """
    pool = _pool(database)
    async with pool.acquire() as conn:
        async with conn.transaction():
            await conn.execute("SELECT pg_advisory_xact_lock(hashtext($1))", f"moid:{prefix}")
            rows = await conn.fetch(
                "SELECT moid FROM vsphere_objects WHERE moid LIKE $1",
                f"{prefix}-%",
            )
            numbers: list[int] = []
            for row in rows:
                suffix = str(row["moid"]).removeprefix(f"{prefix}-")
                if suffix.isdigit():
                    numbers.append(int(suffix))
            moid = f"{prefix}-{max(numbers, default=100) + 1}"
            await conn.execute(
                """
                INSERT INTO vsphere_objects (moid, type, name, parent_moid, props)
                VALUES ($1, 'MoIdReservation', $1, NULL, '{}'::jsonb)
                """,
                moid,
            )
            return moid


def _row(row: Any) -> ManagedObject:
    props = row["props"]
    if isinstance(props, str):
        props = json.loads(props)
    return ManagedObject(
        moid=str(row["moid"]),
        type=str(row["type"]),
        name=str(row["name"]),
        parent_moid=None if row["parent_moid"] is None else str(row["parent_moid"]),
        props=dict(props or {}),
    )