Coverage for qbittorrent_exporter/exporter.py: 68%
132 statements
« prev ^ index » next coverage.py v7.3.2, created at 2023-11-21 09:03 +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()