138 lines
5.3 KiB
Python
138 lines
5.3 KiB
Python
import requests
|
|
import time
|
|
import os
|
|
import sys
|
|
from sqlalchemy import create_engine, ForeignKey, UniqueConstraint, Index
|
|
from sqlalchemy.orm import sessionmaker, relationship, DeclarativeBase, mapped_column, Mapped
|
|
from dataclasses import dataclass
|
|
from typing import List
|
|
|
|
class Base(DeclarativeBase):
|
|
pass
|
|
|
|
@dataclass
|
|
class Host(Base):
|
|
__tablename__ = 'hosts'
|
|
id: Mapped[int] = mapped_column(primary_key=True)
|
|
hostname: Mapped[str] = mapped_column(unique=True)
|
|
current_stats: Mapped[List["CurrentStat"]] = relationship("CurrentStat", back_populates="host")
|
|
historical_stats: Mapped[List["HistoricalStat"]] = relationship("HistoricalStat", back_populates="host")
|
|
|
|
@dataclass
|
|
class CurrentStat(Base):
|
|
__tablename__ = 'current_stats'
|
|
id: Mapped[int] = mapped_column(primary_key=True)
|
|
host_id: Mapped[int] = mapped_column(ForeignKey('hosts.id'))
|
|
category: Mapped[str] = mapped_column()
|
|
name: Mapped[str] = mapped_column()
|
|
value: Mapped[str] = mapped_column()
|
|
unit: Mapped[str | None] = mapped_column(nullable=True)
|
|
timestamp: Mapped[float] = mapped_column()
|
|
host: Mapped["Host"] = relationship("Host", back_populates="current_stats")
|
|
__table_args__ = (
|
|
UniqueConstraint('host_id', 'category', 'name'),
|
|
Index('idx_current_host', 'host_id'),
|
|
)
|
|
|
|
@dataclass
|
|
class HistoricalStat(Base):
|
|
__tablename__ = 'historical_stats'
|
|
id: Mapped[int] = mapped_column(primary_key=True)
|
|
host_id: Mapped[int] = mapped_column(ForeignKey('hosts.id'))
|
|
category: Mapped[str] = mapped_column()
|
|
name: Mapped[str] = mapped_column()
|
|
value: Mapped[str] = mapped_column()
|
|
unit: Mapped[str | None] = mapped_column(nullable=True)
|
|
timestamp: Mapped[float] = mapped_column()
|
|
host: Mapped["Host"] = relationship("Host", back_populates="historical_stats")
|
|
__table_args__ = (
|
|
Index('idx_historical_host', 'host_id'),
|
|
Index('idx_historical_timestamp', 'timestamp'),
|
|
)
|
|
|
|
class StatsCollector:
|
|
def __init__(self, db_path="stats.db"):
|
|
state_home = os.environ.get("XDG_STATE_HOME", os.path.expanduser("~/.local/state"))
|
|
stats_dir = os.path.join(state_home, "stats")
|
|
os.makedirs(stats_dir, exist_ok=True)
|
|
self.db_path = os.path.join(stats_dir, db_path)
|
|
if os.path.exists(self.db_path):
|
|
os.remove(self.db_path)
|
|
self.engine = create_engine(f'sqlite:///{self.db_path}')
|
|
self.setup_database()
|
|
Session = sessionmaker(bind=self.engine)
|
|
self.session = Session()
|
|
|
|
def setup_database(self):
|
|
"""Initialize the SQLite database schema"""
|
|
Base.metadata.create_all(self.engine)
|
|
|
|
def get_or_create_host_id(self, hostname):
|
|
"""Get host ID from database or create if not exists"""
|
|
host = self.session.query(Host).filter_by(hostname=hostname).first()
|
|
if not host:
|
|
host = Host(hostname=hostname)
|
|
self.session.add(host)
|
|
self.session.commit()
|
|
return host.id
|
|
|
|
def store_stats(self, hostname, stats_data):
|
|
"""Store stats in the database"""
|
|
host_id = self.get_or_create_host_id(hostname)
|
|
timestamp = stats_data.get('timestamp', time.time())
|
|
|
|
for category, category_stats in stats_data['stats'].items():
|
|
for name, data in category_stats.items():
|
|
value = data['value']
|
|
unit = data.get('unit', None)
|
|
is_historical = data.get('track', False)
|
|
|
|
if is_historical:
|
|
hist_stat = HistoricalStat(
|
|
host_id=host_id,
|
|
category=category,
|
|
name=name,
|
|
value=value,
|
|
unit=unit,
|
|
timestamp=timestamp
|
|
)
|
|
self.session.add(hist_stat)
|
|
else:
|
|
# Only update/create current stat if not historical
|
|
curr_stat = self.session.query(CurrentStat).filter_by(
|
|
host_id=host_id,
|
|
category=category,
|
|
name=name
|
|
).first()
|
|
|
|
if curr_stat:
|
|
curr_stat.value = value
|
|
curr_stat.unit = unit
|
|
curr_stat.timestamp = timestamp
|
|
else:
|
|
curr_stat = CurrentStat(
|
|
host_id=host_id,
|
|
category=category,
|
|
name=name,
|
|
value=value,
|
|
unit=unit,
|
|
timestamp=timestamp
|
|
)
|
|
self.session.add(curr_stat)
|
|
|
|
self.session.commit()
|
|
|
|
def collect_from_host(self, hostname, port=5747):
|
|
"""Collect stats from a single host"""
|
|
try:
|
|
url = f"http://{hostname}:{port}"
|
|
response = requests.get(url, timeout=5)
|
|
response.raise_for_status()
|
|
stats_data = response.json()
|
|
self.store_stats(hostname, stats_data)
|
|
print(f"Successfully collected stats from {hostname}")
|
|
return True
|
|
except Exception as e:
|
|
print(f"Error collecting stats from {hostname}: {str(e)}", file=sys.stderr)
|
|
return False
|