generated from nathanwoodburn/python-webserver-template
feat: Initial code
This commit is contained in:
305
database.py
Normal file
305
database.py
Normal file
@@ -0,0 +1,305 @@
|
||||
import json
|
||||
import sqlite3
|
||||
from contextlib import closing
|
||||
from datetime import datetime, timezone
|
||||
from typing import Dict, Iterable, List, Optional
|
||||
|
||||
|
||||
def utc_now() -> str:
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
|
||||
class InventoryStore:
|
||||
def __init__(self, database_path: str):
|
||||
self.database_path = database_path
|
||||
|
||||
def _connect(self) -> sqlite3.Connection:
|
||||
conn = sqlite3.connect(self.database_path, check_same_thread=False)
|
||||
conn.row_factory = sqlite3.Row
|
||||
return conn
|
||||
|
||||
def init(self) -> None:
|
||||
with closing(self._connect()) as conn:
|
||||
conn.executescript(
|
||||
"""
|
||||
PRAGMA journal_mode=WAL;
|
||||
|
||||
CREATE TABLE IF NOT EXISTS sources (
|
||||
name TEXT PRIMARY KEY,
|
||||
enabled INTEGER NOT NULL,
|
||||
last_status TEXT,
|
||||
last_error TEXT,
|
||||
last_success TEXT,
|
||||
updated_at TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS collection_runs (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
started_at TEXT NOT NULL,
|
||||
finished_at TEXT,
|
||||
status TEXT NOT NULL,
|
||||
error_summary TEXT
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS assets (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
source TEXT NOT NULL,
|
||||
asset_type TEXT NOT NULL,
|
||||
external_id TEXT NOT NULL,
|
||||
name TEXT NOT NULL,
|
||||
hostname TEXT,
|
||||
status TEXT,
|
||||
ip_addresses TEXT,
|
||||
subnet TEXT,
|
||||
public_ip TEXT,
|
||||
node TEXT,
|
||||
parent_id TEXT,
|
||||
cpu REAL,
|
||||
memory_mb REAL,
|
||||
disk_gb REAL,
|
||||
metadata_json TEXT,
|
||||
last_seen TEXT NOT NULL,
|
||||
updated_at TEXT NOT NULL,
|
||||
UNIQUE(source, asset_type, external_id)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_assets_source ON assets(source);
|
||||
CREATE INDEX IF NOT EXISTS idx_assets_status ON assets(status);
|
||||
CREATE INDEX IF NOT EXISTS idx_assets_subnet ON assets(subnet);
|
||||
"""
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
def seed_sources(self, source_states: Dict[str, bool]) -> None:
|
||||
now = utc_now()
|
||||
with closing(self._connect()) as conn:
|
||||
valid_sources = set(source_states.keys())
|
||||
for source, enabled in source_states.items():
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO sources(name, enabled, updated_at)
|
||||
VALUES(?, ?, ?)
|
||||
ON CONFLICT(name) DO UPDATE SET
|
||||
enabled=excluded.enabled,
|
||||
updated_at=excluded.updated_at
|
||||
""",
|
||||
(source, int(enabled), now),
|
||||
)
|
||||
|
||||
if valid_sources:
|
||||
placeholders = ",".join(["?"] * len(valid_sources))
|
||||
conn.execute(
|
||||
f"DELETE FROM sources WHERE name NOT IN ({placeholders})",
|
||||
tuple(valid_sources),
|
||||
)
|
||||
|
||||
conn.commit()
|
||||
|
||||
def run_start(self) -> int:
|
||||
with closing(self._connect()) as conn:
|
||||
cursor = conn.execute(
|
||||
"INSERT INTO collection_runs(started_at, status) VALUES(?, ?)",
|
||||
(utc_now(), "running"),
|
||||
)
|
||||
conn.commit()
|
||||
row_id = cursor.lastrowid
|
||||
if row_id is None:
|
||||
raise RuntimeError("Failed to create collection run")
|
||||
return int(row_id)
|
||||
|
||||
def run_finish(self, run_id: int, status: str, error_summary: str = "") -> None:
|
||||
with closing(self._connect()) as conn:
|
||||
conn.execute(
|
||||
"""
|
||||
UPDATE collection_runs
|
||||
SET finished_at=?, status=?, error_summary=?
|
||||
WHERE id=?
|
||||
""",
|
||||
(utc_now(), status, error_summary.strip(), run_id),
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
def set_source_status(self, source: str, status: str, error: str = "") -> None:
|
||||
now = utc_now()
|
||||
success_ts = now if status == "ok" else None
|
||||
with closing(self._connect()) as conn:
|
||||
conn.execute(
|
||||
"""
|
||||
UPDATE sources
|
||||
SET last_status=?,
|
||||
last_error=?,
|
||||
last_success=COALESCE(?, last_success),
|
||||
updated_at=?
|
||||
WHERE name=?
|
||||
""",
|
||||
(status, error.strip(), success_ts, now, source),
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
def upsert_assets(self, source: str, assets: Iterable[Dict]) -> None:
|
||||
now = utc_now()
|
||||
with closing(self._connect()) as conn:
|
||||
for asset in assets:
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO assets(
|
||||
source, asset_type, external_id, name, hostname, status,
|
||||
ip_addresses, subnet, public_ip, node, parent_id,
|
||||
cpu, memory_mb, disk_gb, metadata_json, last_seen, updated_at
|
||||
) VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
ON CONFLICT(source, asset_type, external_id)
|
||||
DO UPDATE SET
|
||||
name=excluded.name,
|
||||
hostname=excluded.hostname,
|
||||
status=excluded.status,
|
||||
ip_addresses=excluded.ip_addresses,
|
||||
subnet=excluded.subnet,
|
||||
public_ip=excluded.public_ip,
|
||||
node=excluded.node,
|
||||
parent_id=excluded.parent_id,
|
||||
cpu=excluded.cpu,
|
||||
memory_mb=excluded.memory_mb,
|
||||
disk_gb=excluded.disk_gb,
|
||||
metadata_json=excluded.metadata_json,
|
||||
last_seen=excluded.last_seen,
|
||||
updated_at=excluded.updated_at
|
||||
""",
|
||||
(
|
||||
source,
|
||||
asset.get("asset_type", "unknown"),
|
||||
str(asset.get("external_id", asset.get("name", "unknown"))),
|
||||
asset.get("name", "unknown"),
|
||||
asset.get("hostname"),
|
||||
asset.get("status"),
|
||||
json.dumps(asset.get("ip_addresses", [])),
|
||||
asset.get("subnet"),
|
||||
asset.get("public_ip"),
|
||||
asset.get("node"),
|
||||
asset.get("parent_id"),
|
||||
asset.get("cpu"),
|
||||
asset.get("memory_mb"),
|
||||
asset.get("disk_gb"),
|
||||
json.dumps(asset.get("metadata", {})),
|
||||
now,
|
||||
now,
|
||||
),
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
def list_assets(self) -> List[Dict]:
|
||||
with closing(self._connect()) as conn:
|
||||
rows = conn.execute(
|
||||
"""
|
||||
SELECT source, asset_type, external_id, name, hostname, status,
|
||||
ip_addresses, subnet, public_ip, node, parent_id,
|
||||
cpu, memory_mb, disk_gb, metadata_json, last_seen, updated_at
|
||||
FROM assets
|
||||
ORDER BY source, asset_type, name
|
||||
"""
|
||||
).fetchall()
|
||||
return [self._row_to_asset(row) for row in rows]
|
||||
|
||||
def source_health(self) -> List[Dict]:
|
||||
with closing(self._connect()) as conn:
|
||||
rows = conn.execute(
|
||||
"""
|
||||
SELECT name, enabled, last_status, last_error, last_success, updated_at
|
||||
FROM sources
|
||||
ORDER BY name
|
||||
"""
|
||||
).fetchall()
|
||||
return [dict(row) for row in rows]
|
||||
|
||||
def last_run(self) -> Optional[Dict]:
|
||||
with closing(self._connect()) as conn:
|
||||
row = conn.execute(
|
||||
"""
|
||||
SELECT id, started_at, finished_at, status, error_summary
|
||||
FROM collection_runs
|
||||
ORDER BY id DESC
|
||||
LIMIT 1
|
||||
"""
|
||||
).fetchone()
|
||||
return dict(row) if row else None
|
||||
|
||||
def summary(self) -> Dict:
|
||||
with closing(self._connect()) as conn:
|
||||
totals = conn.execute(
|
||||
"""
|
||||
SELECT
|
||||
COUNT(*) AS total_assets,
|
||||
SUM(CASE WHEN status IN ('running', 'online', 'healthy', 'up') THEN 1 ELSE 0 END) AS online_assets,
|
||||
SUM(CASE WHEN status IN ('stopped', 'offline', 'down', 'error', 'unhealthy') THEN 1 ELSE 0 END) AS offline_assets,
|
||||
COUNT(DISTINCT source) AS source_count,
|
||||
COUNT(DISTINCT subnet) AS subnet_count
|
||||
FROM assets
|
||||
"""
|
||||
).fetchone()
|
||||
|
||||
by_type = conn.execute(
|
||||
"""
|
||||
SELECT asset_type, COUNT(*) AS count
|
||||
FROM assets
|
||||
GROUP BY asset_type
|
||||
ORDER BY count DESC
|
||||
"""
|
||||
).fetchall()
|
||||
|
||||
return {
|
||||
"total_assets": int(totals["total_assets"] or 0),
|
||||
"online_assets": int(totals["online_assets"] or 0),
|
||||
"offline_assets": int(totals["offline_assets"] or 0),
|
||||
"source_count": int(totals["source_count"] or 0),
|
||||
"subnet_count": int(totals["subnet_count"] or 0),
|
||||
"asset_breakdown": [{"asset_type": row["asset_type"], "count": row["count"]} for row in by_type],
|
||||
}
|
||||
|
||||
def topology(self) -> Dict:
|
||||
with closing(self._connect()) as conn:
|
||||
rows = conn.execute(
|
||||
"""
|
||||
SELECT
|
||||
COALESCE(subnet, 'unassigned') AS subnet,
|
||||
COUNT(*) AS asset_count,
|
||||
COUNT(DISTINCT source) AS source_count,
|
||||
GROUP_CONCAT(DISTINCT public_ip) AS public_ips
|
||||
FROM assets
|
||||
GROUP BY COALESCE(subnet, 'unassigned')
|
||||
ORDER BY subnet
|
||||
"""
|
||||
).fetchall()
|
||||
|
||||
networks = []
|
||||
for row in rows:
|
||||
ips = [value for value in (row["public_ips"] or "").split(",") if value]
|
||||
networks.append(
|
||||
{
|
||||
"subnet": row["subnet"],
|
||||
"asset_count": row["asset_count"],
|
||||
"source_count": row["source_count"],
|
||||
"public_ips": ips,
|
||||
}
|
||||
)
|
||||
return {"networks": networks}
|
||||
|
||||
@staticmethod
|
||||
def _row_to_asset(row: sqlite3.Row) -> Dict:
|
||||
return {
|
||||
"source": row["source"],
|
||||
"asset_type": row["asset_type"],
|
||||
"external_id": row["external_id"],
|
||||
"name": row["name"],
|
||||
"hostname": row["hostname"],
|
||||
"status": row["status"],
|
||||
"ip_addresses": json.loads(row["ip_addresses"] or "[]"),
|
||||
"subnet": row["subnet"],
|
||||
"public_ip": row["public_ip"],
|
||||
"node": row["node"],
|
||||
"parent_id": row["parent_id"],
|
||||
"cpu": row["cpu"],
|
||||
"memory_mb": row["memory_mb"],
|
||||
"disk_gb": row["disk_gb"],
|
||||
"metadata": json.loads(row["metadata_json"] or "{}"),
|
||||
"last_seen": row["last_seen"],
|
||||
"updated_at": row["updated_at"],
|
||||
}
|
||||
Reference in New Issue
Block a user