Prometheus_client custom collector class

kea_exporter/cli.py

@@ -1,23 +1,16 @@
+import logging
 import sys
 import time
 
 import click
-from prometheus_client import REGISTRY, make_wsgi_app, start_http_server
+from prometheus_client import start_http_server
+from prometheus_client.core import REGISTRY
 
 from kea_exporter import __project__, __version__
-from kea_exporter.exporter import Exporter
+from kea_exporter.collector import KeaCollector
+from kea_exporter.exporter import KeaExporter
 
-
-class Timer:
-    def __init__(self):
-        self.reset()
-
-    def reset(self):
-        self.start_time = time.time()
-
-    def time_elapsed(self):
-        now_time = time.time()
-        return now_time - self.start_time
+logger = logging.getLogger(__name__)
 
 
 @click.command()
@@ -38,12 +31,13 @@ def time_elapsed(self):
     help="Port that the exporter binds to.",
 )
 @click.option(
-    "-i",
-    "--interval",
-    envvar="INTERVAL",
-    type=int,
-    default=0,
-    help="Minimal interval between two queries to Kea in seconds.",
+    "-d",
+    "--debug",
+    envvar="DEBUG",
+    type=bool,
+    default=False,
+    is_flag=True,
+    help="Run kea_exporter in debug mode.",
 )
 @click.option(
     "--client-cert",
@@ -61,31 +55,24 @@ def time_elapsed(self):
 )
 @click.argument("targets", envvar="TARGETS", nargs=-1, required=True)
 @click.version_option(prog_name=__project__, version=__version__)
-def cli(port, address, interval, **kwargs):
-    exporter = Exporter(**kwargs)
-
-    if not exporter.targets:
-        sys.exit(1)
-
-    httpd, _ = start_http_server(port, address)
+def cli(port, address, debug, **kwargs):
+    if debug:
+        logging.basicConfig(level=logging.DEBUG)
+    else:
+        logging.basicConfig(level=logging.INFO)
 
-    t = Timer()
+    exporter = KeaExporter(**kwargs)
 
-    def local_wsgi_app(registry):
-        func = make_wsgi_app(registry, False)
+    start_http_server(port, address)
 
-        def app(environ, start_response):
-            if t.time_elapsed() >= interval:
-                exporter.update()
-                t.reset()
-            output_array = func(environ, start_response)
-            return output_array
+    if not exporter.targets:
+        sys.exit(1)
 
-        return app
+    collector = KeaCollector(exporter)
 
-    httpd.set_app(local_wsgi_app(REGISTRY))
+    REGISTRY.register(collector)
 
-    click.echo(f"Listening on http://{address}:{port}")
+    logger.info(f"Listening on http://{address}:{port}")
 
     while True:
         time.sleep(1)

kea_exporter/collector.py

@@ -0,0 +1,21 @@
+from collections.abc import Iterator
+
+from prometheus_client.core import CounterMetricFamily, GaugeMetricFamily
+from prometheus_client.registry import Collector
+
+from kea_exporter.exporter import KeaExporter
+from kea_exporter.metrics import KeaMetrics
+
+
+class KeaCollector(Collector):
+    def __init__(self, exporter: KeaExporter) -> None:
+        super().__init__()
+        self.exporter = exporter
+
+    def collect(self) -> Iterator[GaugeMetricFamily]:
+        # Get all stats from the exporter.
+        all_stats = self.exporter.get_all_stats()
+
+        kea_metrics = KeaMetrics(list(all_stats))
+
+        yield from kea_metrics.collect_metrics()