"""SQLite database for tracking monitors and hosts.""" import os import sqlite3 import logging from contextlib import contextmanager from dataclasses import dataclass, asdict from datetime import datetime from typing import Optional from config import get_config logger = logging.getLogger(__name__) SCHEMA = """ -- Hosts that have been scanned CREATE TABLE IF NOT EXISTS hosts ( id INTEGER PRIMARY KEY AUTOINCREMENT, hostname TEXT UNIQUE NOT NULL, last_scan_at TIMESTAMP, last_sync_at TIMESTAMP, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ); -- Monitors created by Kuma Strapper CREATE TABLE IF NOT EXISTS monitors ( id INTEGER PRIMARY KEY AUTOINCREMENT, kuma_monitor_id INTEGER UNIQUE NOT NULL, host_id INTEGER NOT NULL, name TEXT NOT NULL, type TEXT NOT NULL, target TEXT, port INTEGER, interval_seconds INTEGER DEFAULT 60, push_metric TEXT, status TEXT DEFAULT 'active', created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, synced_at TIMESTAMP, FOREIGN KEY (host_id) REFERENCES hosts(id) ON DELETE CASCADE ); -- Push script deployments CREATE TABLE IF NOT EXISTS deployments ( id INTEGER PRIMARY KEY AUTOINCREMENT, monitor_id INTEGER NOT NULL, host_id INTEGER NOT NULL, script_path TEXT NOT NULL, scheduling_method TEXT, scheduling_info TEXT, deployed_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, status TEXT DEFAULT 'deployed', FOREIGN KEY (monitor_id) REFERENCES monitors(id) ON DELETE CASCADE, FOREIGN KEY (host_id) REFERENCES hosts(id) ON DELETE CASCADE ); -- Indexes CREATE INDEX IF NOT EXISTS idx_monitors_host_id ON monitors(host_id); CREATE INDEX IF NOT EXISTS idx_monitors_kuma_id ON monitors(kuma_monitor_id); CREATE INDEX IF NOT EXISTS idx_hosts_hostname ON hosts(hostname); """ @dataclass class Host: id: Optional[int] hostname: str last_scan_at: Optional[datetime] = None last_sync_at: Optional[datetime] = None created_at: Optional[datetime] = None def to_dict(self) -> dict: d = asdict(self) for k, v in d.items(): if isinstance(v, datetime): d[k] = v.isoformat() return d @dataclass class TrackedMonitor: id: Optional[int] kuma_monitor_id: int host_id: int name: str type: str target: Optional[str] = None port: Optional[int] = None interval_seconds: int = 60 push_metric: Optional[str] = None status: str = "active" created_at: Optional[datetime] = None synced_at: Optional[datetime] = None def to_dict(self) -> dict: d = asdict(self) for k, v in d.items(): if isinstance(v, datetime): d[k] = v.isoformat() return d @dataclass class Deployment: id: Optional[int] monitor_id: int host_id: int script_path: str scheduling_method: Optional[str] = None scheduling_info: Optional[str] = None deployed_at: Optional[datetime] = None status: str = "deployed" def to_dict(self) -> dict: d = asdict(self) for k, v in d.items(): if isinstance(v, datetime): d[k] = v.isoformat() return d class Database: def __init__(self, db_path: Optional[str] = None): self.db_path = db_path or get_config().database_path self._ensure_directory() self._init_schema() def _ensure_directory(self): dir_path = os.path.dirname(self.db_path) if dir_path: os.makedirs(dir_path, exist_ok=True) @contextmanager def get_connection(self): conn = sqlite3.connect(self.db_path) conn.row_factory = sqlite3.Row conn.execute("PRAGMA foreign_keys = ON") try: yield conn conn.commit() except Exception: conn.rollback() raise finally: conn.close() def _init_schema(self): with self.get_connection() as conn: conn.executescript(SCHEMA) logger.info(f"Database initialized at {self.db_path}") # Host operations def get_or_create_host(self, hostname: str) -> Host: with self.get_connection() as conn: row = conn.execute( "SELECT * FROM hosts WHERE hostname = ?", (hostname,) ).fetchone() if row: return self._row_to_host(row) cursor = conn.execute( "INSERT INTO hosts (hostname) VALUES (?)", (hostname,) ) return Host(id=cursor.lastrowid, hostname=hostname) def get_host_by_hostname(self, hostname: str) -> Optional[Host]: with self.get_connection() as conn: row = conn.execute( "SELECT * FROM hosts WHERE hostname = ?", (hostname,) ).fetchone() return self._row_to_host(row) if row else None def get_all_hosts(self) -> list[Host]: with self.get_connection() as conn: rows = conn.execute("SELECT * FROM hosts ORDER BY hostname").fetchall() return [self._row_to_host(row) for row in rows] def update_host_scan_time(self, host_id: int): with self.get_connection() as conn: conn.execute( "UPDATE hosts SET last_scan_at = ? WHERE id = ?", (datetime.utcnow(), host_id), ) def update_host_sync_time(self, host_id: int): with self.get_connection() as conn: conn.execute( "UPDATE hosts SET last_sync_at = ? WHERE id = ?", (datetime.utcnow(), host_id), ) def _row_to_host(self, row: sqlite3.Row) -> Host: return Host( id=row["id"], hostname=row["hostname"], last_scan_at=self._parse_timestamp(row["last_scan_at"]), last_sync_at=self._parse_timestamp(row["last_sync_at"]), created_at=self._parse_timestamp(row["created_at"]), ) # Monitor operations def add_monitor(self, monitor: TrackedMonitor) -> int: with self.get_connection() as conn: cursor = conn.execute( """INSERT INTO monitors (kuma_monitor_id, host_id, name, type, target, port, interval_seconds, push_metric, status, synced_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", ( monitor.kuma_monitor_id, monitor.host_id, monitor.name, monitor.type, monitor.target, monitor.port, monitor.interval_seconds, monitor.push_metric, monitor.status, datetime.utcnow(), ), ) return cursor.lastrowid def get_monitor_by_kuma_id(self, kuma_monitor_id: int) -> Optional[TrackedMonitor]: with self.get_connection() as conn: row = conn.execute( "SELECT * FROM monitors WHERE kuma_monitor_id = ?", (kuma_monitor_id,) ).fetchone() return self._row_to_monitor(row) if row else None def get_monitors_for_host(self, host_id: int) -> list[TrackedMonitor]: with self.get_connection() as conn: rows = conn.execute( "SELECT * FROM monitors WHERE host_id = ? ORDER BY name", (host_id,), ).fetchall() return [self._row_to_monitor(row) for row in rows] def get_monitors_for_hostname(self, hostname: str) -> list[TrackedMonitor]: with self.get_connection() as conn: rows = conn.execute( """SELECT m.* FROM monitors m JOIN hosts h ON m.host_id = h.id WHERE h.hostname = ? ORDER BY m.name""", (hostname,), ).fetchall() return [self._row_to_monitor(row) for row in rows] def get_all_monitors(self) -> list[TrackedMonitor]: with self.get_connection() as conn: rows = conn.execute("SELECT * FROM monitors ORDER BY name").fetchall() return [self._row_to_monitor(row) for row in rows] def update_monitor(self, monitor: TrackedMonitor): with self.get_connection() as conn: conn.execute( """UPDATE monitors SET name = ?, type = ?, target = ?, port = ?, interval_seconds = ?, push_metric = ?, status = ?, synced_at = ? WHERE id = ?""", ( monitor.name, monitor.type, monitor.target, monitor.port, monitor.interval_seconds, monitor.push_metric, monitor.status, datetime.utcnow(), monitor.id, ), ) def update_monitor_status(self, monitor_id: int, status: str): with self.get_connection() as conn: conn.execute( "UPDATE monitors SET status = ?, synced_at = ? WHERE id = ?", (status, datetime.utcnow(), monitor_id), ) def mark_monitor_synced(self, monitor_id: int): with self.get_connection() as conn: conn.execute( "UPDATE monitors SET synced_at = ? WHERE id = ?", (datetime.utcnow(), monitor_id), ) def delete_monitor(self, monitor_id: int): with self.get_connection() as conn: conn.execute("DELETE FROM monitors WHERE id = ?", (monitor_id,)) def _row_to_monitor(self, row: sqlite3.Row) -> TrackedMonitor: return TrackedMonitor( id=row["id"], kuma_monitor_id=row["kuma_monitor_id"], host_id=row["host_id"], name=row["name"], type=row["type"], target=row["target"], port=row["port"], interval_seconds=row["interval_seconds"], push_metric=row["push_metric"], status=row["status"], created_at=self._parse_timestamp(row["created_at"]), synced_at=self._parse_timestamp(row["synced_at"]), ) # Deployment operations def add_deployment(self, deployment: Deployment) -> int: with self.get_connection() as conn: cursor = conn.execute( """INSERT INTO deployments (monitor_id, host_id, script_path, scheduling_method, scheduling_info, status) VALUES (?, ?, ?, ?, ?, ?)""", ( deployment.monitor_id, deployment.host_id, deployment.script_path, deployment.scheduling_method, deployment.scheduling_info, deployment.status, ), ) return cursor.lastrowid def get_deployment_for_monitor(self, monitor_id: int) -> Optional[Deployment]: with self.get_connection() as conn: row = conn.execute( "SELECT * FROM deployments WHERE monitor_id = ?", (monitor_id,) ).fetchone() return self._row_to_deployment(row) if row else None def update_deployment_status(self, deployment_id: int, status: str): with self.get_connection() as conn: conn.execute( "UPDATE deployments SET status = ? WHERE id = ?", (status, deployment_id), ) def _row_to_deployment(self, row: sqlite3.Row) -> Deployment: return Deployment( id=row["id"], monitor_id=row["monitor_id"], host_id=row["host_id"], script_path=row["script_path"], scheduling_method=row["scheduling_method"], scheduling_info=row["scheduling_info"], deployed_at=self._parse_timestamp(row["deployed_at"]), status=row["status"], ) def _parse_timestamp(self, value) -> Optional[datetime]: if value is None: return None if isinstance(value, datetime): return value try: return datetime.fromisoformat(value) except (ValueError, TypeError): return None # Global database instance _database: Optional[Database] = None def get_database() -> Database: """Get the global database instance.""" global _database if _database is None: _database = Database() return _database