adds stat collector
This commit is contained in:
@@ -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
|
||||
Executable
+49
@@ -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()
|
||||
Reference in New Issue
Block a user