From 50ca1e4162daef1215591e572ebdefffd1e4d41f Mon Sep 17 00:00:00 2001 From: Joel Heaps Date: Thu, 28 Sep 2023 23:10:21 -0500 Subject: [PATCH] Refactor - add enums, dataclasses, more type hints --- qbittorrent_exporter/exporter.py | 142 +++++++++++++++++-------------- 1 file changed, 80 insertions(+), 62 deletions(-) diff --git a/qbittorrent_exporter/exporter.py b/qbittorrent_exporter/exporter.py index 9c36d41..72eb233 100644 --- a/qbittorrent_exporter/exporter.py +++ b/qbittorrent_exporter/exporter.py @@ -10,7 +10,8 @@ from prometheus_client.core import GaugeMetricFamily, CounterMetricFamily, REGIS import logging from pythonjsonlogger import jsonlogger from enum import StrEnum, auto - +from typing import Iterable, Any +from dataclasses import dataclass, field # Enable dumps on stderr in case of segfault faulthandler.enable() @@ -25,8 +26,22 @@ class TorrentStatus(StrEnum): UPLOADING = auto() +class MetricType(StrEnum): + GAUGE = auto() + COUNTER = auto() + + +@dataclass +class Metric: + name: str + value: int | float | bool + labels: dict[str, str] = field(default_factory={}) + help_text: str = "" + metric_type: MetricType = MetricType.GAUGE + + class QbittorrentMetricsCollector: - def __init__(self, config): + def __init__(self, config: dict[str, str | int | bool]) -> None: self.config = config self.client = Client( host=config["host"], @@ -36,33 +51,31 @@ class QbittorrentMetricsCollector: VERIFY_WEBUI_CERTIFICATE=config["verify_webui_certificate"], ) - def collect(self): - metrics = self.get_qbittorrent_metrics() + def collect(self) -> Iterable[GaugeMetricFamily | CounterMetricFamily]: + metrics: list[Metric] = self.get_qbittorrent_metrics() for metric in metrics: - name = metric["name"] - value = metric["value"] - help_text = metric.get("help", "") - labels = metric.get("labels", {}) - metric_type = metric.get("type", "gauge") - - if metric_type == "counter": - prom_metric = CounterMetricFamily(name, help_text, labels=labels.keys()) + if metric.metric_type == MetricType.COUNTER: + prom_metric = CounterMetricFamily( + metric.name, metric.help_text, labels=metric.labels.keys() + ) else: - prom_metric = GaugeMetricFamily(name, help_text, labels=labels.keys()) - prom_metric.add_metric(value=value, labels=labels.values()) + prom_metric = GaugeMetricFamily( + metric.name, metric.help_text, labels=metric.labels.keys() + ) + prom_metric.add_metric(value=metric.value, labels=metric.labels.values()) yield prom_metric - def get_qbittorrent_metrics(self): - metrics = [] - metrics.extend(self.get_qbittorrent_status_metrics()) - metrics.extend(self.get_qbittorrent_torrent_tags_metrics()) + def get_qbittorrent_metrics(self) -> list[Metric]: + metrics: list[Metric] = [] + metrics.extend(self._get_qbittorrent_status_metrics()) + metrics.extend(self._get_qbittorrent_torrent_tags_metrics()) return metrics - def get_qbittorrent_status_metrics(self): - response = {} - version = "" + def _get_qbittorrent_status_metrics(self) -> list[dict]: + response: dict[str, Any] = {} + version: str = "" # Fetch data from API try: @@ -72,42 +85,47 @@ class QbittorrentMetricsCollector: logger.error(f"Couldn't get server info: {e}") return [ - { - "name": f"{self.config['metrics_prefix']}_up", - "value": bool(response), - "labels": {"version": version}, - "help": "Whether server is reachable", - }, - { - "name": f"{self.config['metrics_prefix']}_connected", - "value": response.get("connection_status", "") == "connected", - "help": "Whether server is currently connected", - }, - { - "name": f"{self.config['metrics_prefix']}_firewalled", - "value": response.get("connection_status", "") == "firewalled", - "help": "Whether server is behind a firewall", - }, - { - "name": f"{self.config['metrics_prefix']}_dht_nodes", - "value": response.get("dht_nodes", 0), - "help": "Number of connected DHT nodes", - }, - { - "name": f"{self.config['metrics_prefix']}_dl_info_data", - "value": response.get("dl_info_data", 0), - "help": "Data downloaded this session (bytes)", - "type": "counter", - }, - { - "name": f"{self.config['metrics_prefix']}_up_info_data", - "value": response.get("up_info_data", 0), - "help": "Data uploaded this session (bytes)", - "type": "counter", - }, + Metric( + name=f"{self.config['metrics_prefix']}_up", + value=bool(response), + labels={"version": version}, + help_text="Whether server is reachable", + ), + Metric( + name=f"{self.config['metrics_prefix']}_connected", + value=response.get("connection_status", "") == "connected", + labels={}, # no labels in the example + help_text="Whether server is currently connected", + ), + Metric( + name=f"{self.config['metrics_prefix']}_firewalled", + value=response.get("connection_status", "") == "firewalled", + labels={}, # no labels in the example + help_text="Whether server is behind a firewall", + ), + Metric( + name=f"{self.config['metrics_prefix']}_dht_nodes", + value=response.get("dht_nodes", 0), + labels={}, # no labels in the example + help_text="Number of connected DHT nodes", + ), + Metric( + name=f"{self.config['metrics_prefix']}_dl_info_data", + value=response.get("dl_info_data", 0), + labels={}, # no labels in the example + help_text="Data downloaded this session (bytes)", + metric_type=MetricType.COUNTER, + ), + Metric( + name=f"{self.config['metrics_prefix']}_up_info_data", + value=response.get("up_info_data", 0), + labels={}, # no labels in the example + help_text="Data uploaded this session (bytes)", + metric_type=MetricType.COUNTER, + ), ] - def get_qbittorrent_torrent_tags_metrics(self): + def _get_qbittorrent_torrent_tags_metrics(self) -> list[Metric]: try: categories = self.client.torrent_categories.categories torrents = self.client.torrents.info() @@ -115,7 +133,7 @@ class QbittorrentMetricsCollector: logger.error(f"Couldn't fetch torrent info: {e}") return [] - metrics = [] + metrics: list[Metric] = [] categories.Uncategorized = attridict({"name": "Uncategorized", "savePath": ""}) for category in categories: category_torrents = [ @@ -135,18 +153,18 @@ class QbittorrentMetricsCollector: ) ] metrics.append( - { - "name": f"{self.config['metrics_prefix']}_torrents_count", - "value": len(status_torrents), - "labels": { + Metric( + name=f"{self.config['metrics_prefix']}_torrents_count", + value=len(status_torrents), + labels={ "status": status.value, "category": category, }, - "help": ( + help_text=( f"Number of torrents in status {status.value} under" f" category {category}" ), - } + ) ) return metrics