"""Live database inventory — sizes, row counts, schemas for LLM context.""" from __future__ import annotations import asyncio import os from concurrent.futures import ThreadPoolExecutor, TimeoutError as FuturesTimeout from typing import Any DB_HOST = os.getenv("DB_VAULT_HOST", "10.0.21.51") PG_USER = os.getenv("PG_USER", "mo") PG_PASS = os.getenv("PG_PASSWORD", "Dell2026!") MYSQL_USER = os.getenv("MYSQL_USER", "mo") MYSQL_PASS = os.getenv("MYSQL_PASSWORD", "Dell2026!") NEO4J_USER = os.getenv("NEO4J_USER", "neo4j") NEO4J_PASS = os.getenv("NEO4J_PASSWORD", "testpwd") ENGINE_TIMEOUT = float(os.getenv("DB_INVENTORY_TIMEOUT", "20")) _executor = ThreadPoolExecutor(max_workers=4) def _fmt_bytes(n: int | float | None) -> str: if n is None: return "?" n = float(n) for unit in ("B", "KB", "MB", "GB", "TB"): if n < 1024 or unit == "TB": return f"{n:.1f} {unit}" if unit != "B" else f"{int(n)} B" n /= 1024 return f"{n:.1f} TB" def _inventory_postgres() -> dict[str, Any]: import psycopg2 out: dict[str, Any] = {"engine": "PostgreSQL", "host": DB_HOST, "database": "postgres", "ok": False} try: conn = psycopg2.connect( host=DB_HOST, user=PG_USER, password=PG_PASS, dbname="postgres", connect_timeout=5, ) cur = conn.cursor() cur.execute("SELECT pg_database_size(current_database())") out["size_bytes"] = cur.fetchone()[0] out["size_human"] = _fmt_bytes(out["size_bytes"]) cur.execute( "SELECT table_name FROM information_schema.tables " "WHERE table_schema='public' AND table_type='BASE TABLE' ORDER BY table_name", ) tables = [] for (tname,) in cur.fetchall(): cur.execute(f'SELECT reltuples::bigint FROM pg_class WHERE relname = %s', (tname,)) est = cur.fetchone() rows = int(est[0]) if est and est[0] else None cur.execute( "SELECT column_name, data_type FROM information_schema.columns " "WHERE table_schema='public' AND table_name=%s ORDER BY ordinal_position", (tname,), ) cols = [f"{c} ({dt})" for c, dt in cur.fetchall()] tbl: dict[str, Any] = {"name": tname, "rows": rows, "rows_estimated": True, "columns": cols} if tname == "sales_orders" and rows: cur.execute( "SELECT region, COUNT(*) FROM sales_orders TABLESAMPLE SYSTEM (0.1) " "GROUP BY region ORDER BY COUNT(*) DESC LIMIT 5", ) sample = cur.fetchall() if sample: tbl["sample_regions"] = {r: c for r, c in sample} tables.append(tbl) out["tables"] = tables out["ok"] = True conn.close() except Exception as exc: out["error"] = str(exc) return out def _inventory_mysql() -> dict[str, Any]: import pymysql out: dict[str, Any] = {"engine": "MySQL", "host": DB_HOST, "database": "hr", "ok": False} try: conn = pymysql.connect( host=DB_HOST, user=MYSQL_USER, password=MYSQL_PASS, database="hr", connect_timeout=5, ) cur = conn.cursor() cur.execute( "SELECT table_name, data_length+index_length, table_rows " "FROM information_schema.tables WHERE table_schema='hr'", ) tables = [] total_bytes = 0 for tname, tbytes, trows in cur.fetchall(): total_bytes += tbytes or 0 cur.execute(f"SHOW COLUMNS FROM `{tname}`") cols = [f"{r[0]} ({r[1]})" for r in cur.fetchall()] tbl: dict[str, Any] = { "name": tname, "rows": int(trows) if trows else None, "rows_estimated": True, "size_bytes": tbytes, "columns": cols, } if tname == "employee_events": tbl["note"] = "HR employee lifecycle events (promotions, transfers, salary changes, etc.)" tables.append(tbl) out["tables"] = tables out["size_bytes"] = total_bytes out["size_human"] = _fmt_bytes(total_bytes) out["ok"] = True conn.close() except Exception as exc: out["error"] = str(exc) return out def _inventory_mongo() -> dict[str, Any]: from pymongo import MongoClient out: dict[str, Any] = {"engine": "MongoDB", "host": DB_HOST, "ok": False} try: client = MongoClient(f"mongodb://{DB_HOST}:27017/", serverSelectionTimeoutMS=5000) db = client["supplychain"] collections = [] for cname in db.list_collection_names(): if cname.startswith("__"): continue col = db[cname] docs = col.estimated_document_count() sample = col.find_one() or {} fields = sorted(k for k in sample if k != "_id") coll: dict[str, Any] = {"name": cname, "documents": docs, "fields": fields} if cname == "events" and docs: try: pipe = [ {"$sample": {"size": 5000}}, {"$group": {"_id": "$type", "count": {"$sum": 1}}}, {"$sort": {"count": -1}}, {"$limit": 5}, ] coll["sample_types"] = {r["_id"]: r["count"] for r in col.aggregate(pipe, maxTimeMS=5000)} except Exception: pass collections.append(coll) out["database"] = "supplychain" out["collections"] = collections out["ok"] = True client.close() except Exception as exc: out["error"] = str(exc) return out def _inventory_cassandra() -> dict[str, Any]: out: dict[str, Any] = {"engine": "Cassandra", "host": DB_HOST, "ok": False} try: from cassandra.cluster import Cluster cluster = Cluster([DB_HOST], connect_timeout=5) session = cluster.connect() keyspaces = [ r.keyspace_name for r in session.execute("SELECT keyspace_name FROM system_schema.keyspaces") if r.keyspace_name not in ( "system", "system_schema", "system_traces", "system_distributed", "system_virtual_schema", "system_auth", "system_views", ) ] tables_out = [] for ks in keyspaces: for row in session.execute( "SELECT table_name FROM system_schema.tables WHERE keyspace_name=%s", (ks,), ): tables_out.append({ "keyspace": ks, "name": row.table_name, "rows": None, "note": "COUNT skipped (large table; use Trino/Iceberg for analytics)", }) out["keyspaces"] = keyspaces out["tables"] = tables_out out["ok"] = True cluster.shutdown() except Exception as exc: out["error"] = str(exc) return out def _inventory_neo4j() -> dict[str, Any]: out: dict[str, Any] = {"engine": "Neo4j", "host": DB_HOST, "ok": False} try: from neo4j import GraphDatabase driver = GraphDatabase.driver(f"bolt://{DB_HOST}:7687", auth=(NEO4J_USER, NEO4J_PASS)) with driver.session() as session: nodes = [ {"label": r["lbl"], "count": r["c"]} for r in session.run( "MATCH (n) RETURN labels(n)[0] AS lbl, count(*) AS c ORDER BY c DESC LIMIT 10", ) ] rels = [ {"type": r["t"], "count": r["c"]} for r in session.run( "MATCH ()-[r]->() RETURN type(r) AS t, count(*) AS c ORDER BY c DESC LIMIT 10", ) ] out["nodes"] = nodes out["relationships"] = rels out["ok"] = True driver.close() except Exception as exc: out["error"] = str(exc) return out def _run_with_timeout(fn, timeout: float) -> dict[str, Any]: future = _executor.submit(fn) try: return future.result(timeout=timeout) except FuturesTimeout: return {"engine": fn.__name__.replace("_inventory_", ""), "ok": False, "error": f"timeout after {timeout}s"} except Exception as exc: return {"ok": False, "error": str(exc)} def collect_database_inventory_sync() -> dict[str, Any]: fns = { "postgresql": _inventory_postgres, "mysql": _inventory_mysql, "mongodb": _inventory_mongo, "cassandra": _inventory_cassandra, "neo4j": _inventory_neo4j, } engines = {k: _run_with_timeout(fn, ENGINE_TIMEOUT) for k, fn in fns.items()} ok_count = sum(1 for e in engines.values() if e.get("ok")) return {"host": DB_HOST, "engines_ok": ok_count, "engines_total": len(engines), "engines": engines} async def collect_database_inventory() -> dict[str, Any]: loop = asyncio.get_event_loop() return await loop.run_in_executor(_executor, collect_database_inventory_sync)