From ec6615a007810b8e7f95d98d549ffb979b812348 Mon Sep 17 00:00:00 2001 From: Mark Hoekveen Date: Thu, 5 Dec 2024 22:29:15 +0100 Subject: [PATCH] adds stat collector --- .gitignore | 2 + bin/lib/statscollector.py | 137 ++++++++++++++++++++++++++++++++++++++ bin/stats-collector | 49 ++++++++++++++ 3 files changed, 188 insertions(+) create mode 100644 bin/lib/statscollector.py create mode 100755 bin/stats-collector diff --git a/.gitignore b/.gitignore index b8a6a92..e75c798 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,8 @@ /* */ !/bin/ +!/bin/lib/ +!/bin/res/ !/utils/ /*.* !/bin/*.* diff --git a/bin/lib/statscollector.py b/bin/lib/statscollector.py new file mode 100644 index 0000000..d94e097 --- /dev/null +++ b/bin/lib/statscollector.py @@ -0,0 +1,137 @@ +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 diff --git a/bin/stats-collector b/bin/stats-collector new file mode 100755 index 0000000..268152c --- /dev/null +++ b/bin/stats-collector @@ -0,0 +1,49 @@ +#!/usr/bin/env python3 +import time +import configparser +import os +from lib.statscollector import StatsCollector + +def get_config(): + config = configparser.ConfigParser() + config_home = os.environ.get("XDG_CONFIG_HOME", os.path.expanduser("~/.config")) + config_path = os.path.join(config_home, "stats-collector", "config.ini") + + if not os.path.exists(config_path): + os.makedirs(os.path.dirname(config_path), exist_ok=True) + config['DEFAULT'] = { + 'port': '5747', + 'interval': '60' + } + config['localhost'] = { + 'host': 'localhost', + 'displayname': 'Local System' + } + with open(config_path, 'w') as f: + config.write(f) + + config.read(config_path) + return config + +def main(): + config = get_config() + collector = StatsCollector() + default_port = int(config['DEFAULT']['port']) + interval = int(config['DEFAULT']['interval']) + + def collect_all(): + for section in config.sections(): + host = config[section]['host'] + port = int(config[section].get('port', default_port)) + collector.collect_from_host(host, port) + + if interval > 0: + print(f"Starting continuous collection every {interval} seconds...") + while True: + collect_all() + time.sleep(interval) + else: + collect_all() + +if __name__ == '__main__': + main()