Add Command Center v2: DQ/RAG integration, S3 browser, Jupyter, GPU matrix.
Mirror mo/atc-GPU layout with config/, docs/, scripts/ for Gitea deploy.
This commit is contained in:
mo
@@ -0,0 +1,243 @@
|
||||
"""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)
|
||||
Reference in New Issue
Block a user