mo
244 lines
8.9 KiB
Python
244 lines
8.9 KiB
Python
"""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)
|