"""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) 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 {}), )