Fix prometheus functionality

bmspy/prometheus.py

@@ -1,40 +1,86 @@
+import argparse
 import time
+from typing import Any
+
 import prometheus_client
+from prometheus_client import CollectorRegistry, Gauge, Info
+
+from bmspy import client
+from bmspy.classes import UPS, BMSScalarField, BMSMultiField, BMSInfoField
 from bmspy.utilities import debugger
-from bmspy.server import collect_data
+
+DAEMON_UPDATE_PERIOD = 30
 
 
-def prometheus_export(daemonize=True, filename=None):
+def prometheus_create_metric(
-    global debug
+    registry: CollectorRegistry, ups_data: dict[str, UPS]
-    if not can_export_prometheus:
+) -> dict[str, Any]:
-        raise ModuleNotFoundError(
+    """Create one Gauge per scalar/multi field and one Info per info field; skip duplicates."""
-            "Unable to export to Prometheus. Is prometheus-client installed?"
+    metric: dict[str, Any] = {}
+    for ups_name, ups_obj in ups_data.items():
+        for name, field in ups_obj.items():
+            if name in metric:
+                continue
+            helpmsg = field.get("help") or ""
+            units = field.get("units")
+            if units:
+                helpmsg = "{} ({})".format(helpmsg, units)
+            if isinstance(field, BMSScalarField):
+                metric[name] = Gauge(name, helpmsg, ["ups"], registry=registry)
+            elif isinstance(field, BMSMultiField):
+                label = field.label
+                metric[name] = Gauge(
+                    name, helpmsg, ["ups", label], registry=registry
                 )
+            elif isinstance(field, BMSInfoField):
+                metric[name] = Info(name, helpmsg, ["ups"], registry=registry)
+    return metric
 
-    data = dict()
+
-    # Initialize data structure, to fill in help values
+def prometheus_populate_metric(
-    while bool(data) is False:
+    metric: dict[str, Any], ups_data: dict[str, UPS]
-        data = collect_data()
+) -> None:
+    """Populate metric values from UPS data using isinstance checks on field types."""
+    for ups_name, ups_obj in ups_data.items():
+        for name, field in ups_obj.items():
+            if name not in metric:
+                continue
+            if isinstance(field, BMSScalarField):
+                metric[name].labels(ups=ups_name).set(field.raw_value)
+            elif isinstance(field, BMSMultiField):
+                label = field.label
+                for idx, value in field.raw_values.items():
+                    metric[name].labels(**{"ups": ups_name, label: str(idx)}).set(value)
+            elif isinstance(field, BMSInfoField):
+                metric[name].labels(ups=ups_name).info({name: field.info})
+
+
+def prometheus_export(
+    daemonize: bool = True,
+    filename: str | None = None,
+    socket_path: str = "/run/bmspy/bms",
+    ups: str | None = None,
+    debug: int = 0,
+) -> bool | None:
+    """Export BMS data to Prometheus; daemonize or write to textfile."""
+    ups_data: dict[str, UPS] = {}
+    while not ups_data:
+        ups_data = client.read_data(socket_path, "prometheus", ups=ups, debug=debug)
+        if not ups_data:
             time.sleep(1)
 
-    registry = prometheus_client.CollectorRegistry(auto_describe=True)
+    registry = CollectorRegistry(auto_describe=True)
-    # Set up the metric data structure for Prometheus
+    metric = prometheus_create_metric(registry, ups_data)
-    metric = prometheus_create_metric(registry, data)
+    prometheus_populate_metric(metric, ups_data)
-    # Populate the metric data structure this period
-    prometheus_populate_metric(metric, data)
 
     if daemonize:
         prometheus_client.start_http_server(9999, registry=registry)
-
+        while True:  # pragma: no cover
-        while True:
-            # Delay, collect new data, and start again
             time.sleep(DAEMON_UPDATE_PERIOD)
-            # Reset data, so it is re-populated correctly
+            ups_data = {}
-            data = dict()
+            while not ups_data:
-            while bool(data) is False:
+                ups_data = client.read_data(socket_path, "prometheus", ups=ups, debug=debug)
-                data = collect_data()
+            prometheus_populate_metric(metric, ups_data)
-            time.sleep(1)
-            prometheus_populate_metric(metric, data)
             prometheus_client.generate_latest(registry)
     else:
         if filename is None:
@@ -44,61 +90,53 @@ def prometheus_export(daemonize=True, filename=None):
         return True
 
 
-def prometheus_create_metric(registry, data):
+def main() -> None:
-    metric = dict()
+    """Entry point for bmspy-prometheus command."""
-    for name, contains in data.items():
+    parser = argparse.ArgumentParser(
-        helpmsg = ""
+        description="Query JBD BMS and report status to Prometheus",
-        if contains.get("help") is not None:
+        add_help=True,
-            helpmsg = contains.get("help")
-            if contains.get("units"):
-                helpmsg += " (" + contains.get("units") + ")"
-        if contains.get("value") is not None:
-            metric[name] = prometheus_client.Gauge(name, helpmsg, registry=registry)
-        # Has multiple values, each a different label
-        elif contains.get("values") is not None:
-            if contains.get("label") is None:
-                debugger("ERROR: no label for {0} specified".format(name))
-            label = contains.get("label")
-            metric[name] = prometheus_client.Gauge(
-                name, helpmsg, [label], registry=registry
     )
-        elif contains.get("info") is not None:
+    parser.add_argument(
-            metric[name] = prometheus_client.Info(name, helpmsg, registry=registry)
+        "--socket",
-        else:
+        "-s",
-            pass
+        dest="socket",
-    return metric
+        action="store",
+        default="/run/bmspy/bms",
+        help="Socket to communicate with daemon",
+    )
+    parser.add_argument(
+        "--ups",
+        dest="ups",
+        action="store",
+        default=None,
+        help="Only export data for this UPS name (default: all)",
+    )
+    parser.add_argument(
+        "--file",
+        "-f",
+        dest="filename",
+        type=str,
+        action="store",
+        default=None,
+        help="Write Prometheus textfile to this path (non-daemonized)",
+    )
+    parser.add_argument(
+        "--verbose",
+        "-v",
+        action="count",
+        default=0,
+        help="Print more verbose information (can be specified multiple times)",
+    )
+    args = parser.parse_args()
 
+    prometheus_export(
+        daemonize=args.filename is None,
+        filename=args.filename,
+        socket_path=args.socket,
+        ups=args.ups,
+        debug=args.verbose,
+    )
 
-def prometheus_populate_metric(metric, data):
-    for name, contains in data.items():
-        if contains.get("value") is not None:
-            value = contains.get("value")
-            metric[name].set(value)
-        # doesn't have a value, but has [1-4]:
-        if contains.get("values") is not None and isinstance(
-            contains.get("values"), dict
-        ):
-            for idx, label_value in contains.get("values").items():
-                metric[name].labels(idx).set(label_value)
-        if contains.get("info"):
-            value = contains.get("info")
-            metric[name].info({name: value})
-        else:
-            pass
-
-
-# TODO fork bms daemon if need be?
-
-
-def main():
-    debugger("TODO. At present, run from bmspy directly.")
-
-
-#    influxdb_export(bucket=args.influx_bucket, \
-#                    url=args.influx_url, \
-#                    org=args.influx_org, \
-#                    token=args.influx_token, \
-#                    daemonize=True)
 
 if __name__ == "__main__":
     main()