+ Coverage for qbittorrent_exporter/exporter.py: + 40% +
+ ++ 132 statements + + + +
++ « prev + ^ index + » next + + coverage.py v7.3.2, + created at 2023-11-20 12:15 +0000 +
+ +1import faulthandler
+2import logging
+3import os
+4import signal
+5import sys
+6import time
+7from dataclasses import dataclass, field
+8from enum import StrEnum, auto
+9from typing import Any, Iterable
+ +11from prometheus_client import start_http_server
+12from prometheus_client.core import REGISTRY, CounterMetricFamily, GaugeMetricFamily
+13from pythonjsonlogger import jsonlogger
+14from qbittorrentapi import Client, TorrentStates
+ +16# Enable dumps on stderr in case of segfault
+17faulthandler.enable()
+18logger = logging.getLogger()
+ + +21class MetricType(StrEnum):
+22 """
+23 Represents possible metric types (used in this project).
+24 """
+ +26 GAUGE = auto()
+27 COUNTER = auto()
+ + +30@dataclass
+31class Metric:
+32 """
+33 Contains data and metadata about a single counter or gauge.
+34 """
+ +36 name: str
+37 value: Any
+38 labels: dict[str, str] = field(default_factory=lambda: {}) # Default to empty dict
+39 help_text: str = ""
+40 metric_type: MetricType = MetricType.GAUGE
+ + +43class QbittorrentMetricsCollector:
+44 def __init__(self, config: dict) -> None:
+45 self.config = config
+46 self.client = Client(
+47 host=config["host"],
+48 port=config["port"],
+49 username=config["username"],
+50 password=config["password"],
+51 VERIFY_WEBUI_CERTIFICATE=config["verify_webui_certificate"],
+52 )
+ +54 def collect(self) -> Iterable[GaugeMetricFamily | CounterMetricFamily]:
+55 """
+56 Yields Prometheus gauges and counters from metrics collected from qbittorrent.
+57 """
+58 metrics: list[Metric] = self.get_qbittorrent_metrics()
+ +60 for metric in metrics:
+61 if metric.metric_type == MetricType.COUNTER:
+62 prom_metric = CounterMetricFamily(
+63 metric.name, metric.help_text, labels=list(metric.labels.keys())
+64 )
+65 else:
+66 prom_metric = GaugeMetricFamily(
+67 metric.name, metric.help_text, labels=list(metric.labels.keys())
+68 )
+69 prom_metric.add_metric(
+70 value=metric.value, labels=list(metric.labels.values())
+71 )
+72 yield prom_metric
+ +74 def get_qbittorrent_metrics(self) -> list[Metric]:
+75 """
+76 Calls and combines qbittorrent state metrics with torrent metrics.
+77 """
+78 metrics: list[Metric] = []
+79 metrics.extend(self._get_qbittorrent_status_metrics())
+80 metrics.extend(self._get_qbittorrent_torrent_tags_metrics())
+ +82 return metrics
+ +84 def _get_qbittorrent_status_metrics(self) -> list[Metric]:
+85 """
+86 Returns metrics about the state of the qbittorrent server.
+87 """
+88 response: dict[str, Any] = {}
+89 version: str = ""
+ +91 # Fetch data from API
+92 try:
+93 response = self.client.transfer.info
+94 version = self.client.app.version
+95 except Exception as e:
+96 logger.error(f"Couldn't get server info: {e}")
+ +98 return [
+99 Metric(
+100 name=f"{self.config['metrics_prefix']}_up",
+101 value=bool(response),
+102 labels={"version": version},
+103 help_text=(
+104 "Whether the qBittorrent server is answering requests from this"
+105 " exporter. A `version` label with the server version is added."
+106 ),
+107 ),
+108 Metric(
+109 name=f"{self.config['metrics_prefix']}_connected",
+110 value=response.get("connection_status", "") == "connected",
+111 labels={}, # no labels in the example
+112 help_text=(
+113 "Whether the qBittorrent server is connected to the Bittorrent"
+114 " network."
+115 ),
+116 ),
+117 Metric(
+118 name=f"{self.config['metrics_prefix']}_firewalled",
+119 value=response.get("connection_status", "") == "firewalled",
+120 labels={}, # no labels in the example
+121 help_text=(
+122 "Whether the qBittorrent server is connected to the Bittorrent"
+123 " network but is behind a firewall."
+124 ),
+125 ),
+126 Metric(
+127 name=f"{self.config['metrics_prefix']}_dht_nodes",
+128 value=response.get("dht_nodes", 0),
+129 labels={}, # no labels in the example
+130 help_text="Number of DHT nodes connected to.",
+131 ),
+132 Metric(
+133 name=f"{self.config['metrics_prefix']}_dl_info_data",
+134 value=response.get("dl_info_data", 0),
+135 labels={}, # no labels in the example
+136 help_text="Data downloaded since the server started, in bytes.",
+137 metric_type=MetricType.COUNTER,
+138 ),
+139 Metric(
+140 name=f"{self.config['metrics_prefix']}_up_info_data",
+141 value=response.get("up_info_data", 0),
+142 labels={}, # no labels in the example
+143 help_text="Data uploaded since the server started, in bytes.",
+144 metric_type=MetricType.COUNTER,
+145 ),
+146 ]
+ +148 def _fetch_categories(self) -> dict:
+149 """Fetches all categories in use from qbittorrent."""
+150 try:
+151 categories = dict(self.client.torrent_categories.categories)
+152 for key, value in categories.items():
+153 categories[key] = dict(value) # type: ignore
+154 return categories
+155 except Exception as e:
+156 logger.error(f"Couldn't fetch categories: {e}")
+157 return {}
+ +159 def _fetch_torrents(self) -> list[dict]:
+160 """Fetches torrents from qbittorrent"""
+161 try:
+162 return [dict(_attr_dict) for _attr_dict in self.client.torrents.info()]
+163 except Exception as e:
+164 logger.error(f"Couldn't fetch torrents: {e}")
+165 return []
+ +167 def _filter_torrents_by_category(
+168 self, category: str, torrents: list[dict]
+169 ) -> list[dict]:
+170 """Filters torrents by the given category."""
+171 return [
+172 torrent
+173 for torrent in torrents
+174 if torrent["category"] == category
+175 or (category == "Uncategorized" and torrent["category"] == "")
+176 ]
+ +178 def _filter_torrents_by_state(
+179 self, state: TorrentStates, torrents: list[dict]
+180 ) -> list[dict]:
+181 """Filters torrents by the given state."""
+182 return [torrent for torrent in torrents if torrent["state"] == state.value]
+ +184 def _construct_metric(self, state: str, category: str, count: int) -> Metric:
+185 """Constructs and returns a metric object with a torrent count and appropriate
+186 labels."""
+187 return Metric(
+188 name=f"{self.config['metrics_prefix']}_torrents_count",
+189 value=count,
+190 labels={
+191 "status": state,
+192 "category": category,
+193 },
+194 help_text=f"Number of torrents in status {state} under category {category}",
+195 )
+ +197 def _get_qbittorrent_torrent_tags_metrics(self) -> list[Metric]:
+198 categories = self._fetch_categories()
+199 torrents = self._fetch_torrents()
+ +201 metrics: list[Metric] = []
+202 categories["Uncategorized"] = {"name": "Uncategorized", "savePath": ""}
+ +204 for category in categories:
+205 category_torrents = self._filter_torrents_by_category(category, torrents)
+206 for state in TorrentStates:
+207 state_torrents = self._filter_torrents_by_state(
+208 state, category_torrents
+209 )
+210 metric = self._construct_metric(
+211 state.value, category, len(state_torrents)
+212 )
+213 metrics.append(metric)
+ +215 return metrics
+ + +218class ShutdownSignalHandler:
+219 def __init__(self):
+220 self.shutdown_count: int = 0
+ +222 # Register signal handler
+223 signal.signal(signal.SIGINT, self._on_signal_received)
+224 signal.signal(signal.SIGTERM, self._on_signal_received)
+ +226 def is_shutting_down(self):
+227 return self.shutdown_count > 0
+ +229 def _on_signal_received(self, signal, frame):
+230 if self.shutdown_count > 1:
+231 logger.warn("Forcibly killing exporter")
+232 sys.exit(1)
+233 logger.info("Exporter is shutting down")
+234 self.shutdown_count += 1
+ + +237def _get_config_value(key: str, default: str = "") -> str:
+238 input_path = os.environ.get("FILE__" + key, None)
+239 if input_path is not None:
+240 try:
+241 with open(input_path, "r") as input_file:
+242 return input_file.read().strip()
+243 except IOError as e:
+244 logger.error(f"Unable to read value for {key} from {input_path}: {str(e)}")
+ +246 return os.environ.get(key, default)
+ + +249def get_config() -> dict:
+250 """Loads all config values."""
+251 return {
+252 "host": _get_config_value("QBITTORRENT_HOST", ""),
+253 "port": _get_config_value("QBITTORRENT_PORT", ""),
+254 "username": _get_config_value("QBITTORRENT_USER", ""),
+255 "password": _get_config_value("QBITTORRENT_PASS", ""),
+256 "exporter_port": int(_get_config_value("EXPORTER_PORT", "8000")),
+257 "log_level": _get_config_value("EXPORTER_LOG_LEVEL", "INFO"),
+258 "metrics_prefix": _get_config_value("METRICS_PREFIX", "qbittorrent"),
+259 "verify_webui_certificate": (
+260 _get_config_value("VERIFY_WEBUI_CERTIFICATE", "True") == "True"
+261 ),
+262 }
+ + +265def main():
+266 # Init logger so it can be used
+267 logHandler = logging.StreamHandler()
+268 formatter = jsonlogger.JsonFormatter(
+269 "%(asctime) %(levelname) %(message)", datefmt="%Y-%m-%d %H:%M:%S"
+270 )
+271 logHandler.setFormatter(formatter)
+272 logger.addHandler(logHandler)
+273 logger.setLevel("INFO") # default until config is loaded
+ +275 config = get_config()
+ +277 # set level once config has been loaded
+278 logger.setLevel(config["log_level"])
+ +280 # Register signal handler
+281 signal_handler = ShutdownSignalHandler()
+ +283 if not config["host"]:
+284 logger.error(
+285 "No host specified, please set QBITTORRENT_HOST environment variable"
+286 )
+287 sys.exit(1)
+288 if not config["port"]:
+289 logger.error(
+290 "No port specified, please set QBITTORRENT_PORT environment variable"
+291 )
+292 sys.exit(1)
+ +294 # Register our custom collector
+295 logger.info("Exporter is starting up")
+296 REGISTRY.register(QbittorrentMetricsCollector(config)) # type: ignore
+ +298 # Start server
+299 start_http_server(config["exporter_port"])
+300 logger.info(f"Exporter listening on port {config['exporter_port']}")
+ +302 while not signal_handler.is_shutting_down():
+303 time.sleep(1)
+ +305 logger.info("Exporter has shutdown")
+ + +308if __name__ == "__main__":
+309 main()
+