From a74934c439ce1be0d743518cdc03de1c532cd04b Mon Sep 17 00:00:00 2001 From: Timothy Allen Date: Sat, 2 May 2026 23:12:40 +0200 Subject: [PATCH] Fix prometheus functionality --- bmspy/prometheus.py | 196 ++++++++++++++++++++++++++------------------ 1 file changed, 117 insertions(+), 79 deletions(-) diff --git a/bmspy/prometheus.py b/bmspy/prometheus.py index 62fe907..3c0f3af 100644 --- a/bmspy/prometheus.py +++ b/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): - global debug - if not can_export_prometheus: - raise ModuleNotFoundError( - "Unable to export to Prometheus. Is prometheus-client installed?" - ) +def prometheus_create_metric( + registry: CollectorRegistry, ups_data: dict[str, UPS] +) -> dict[str, Any]: + """Create one Gauge per scalar/multi field and one Info per info field; skip duplicates.""" + 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 - while bool(data) is False: - data = collect_data() - time.sleep(1) - registry = prometheus_client.CollectorRegistry(auto_describe=True) - # Set up the metric data structure for Prometheus - metric = prometheus_create_metric(registry, data) - # Populate the metric data structure this period - prometheus_populate_metric(metric, data) +def prometheus_populate_metric( + metric: dict[str, Any], ups_data: dict[str, UPS] +) -> 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 = CollectorRegistry(auto_describe=True) + metric = prometheus_create_metric(registry, ups_data) + prometheus_populate_metric(metric, ups_data) if daemonize: prometheus_client.start_http_server(9999, registry=registry) - - while True: - # Delay, collect new data, and start again + while True: # pragma: no cover time.sleep(DAEMON_UPDATE_PERIOD) - # Reset data, so it is re-populated correctly - data = dict() - while bool(data) is False: - data = collect_data() - time.sleep(1) - prometheus_populate_metric(metric, data) + ups_data = {} + while not ups_data: + ups_data = client.read_data(socket_path, "prometheus", ups=ups, debug=debug) + prometheus_populate_metric(metric, ups_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): - metric = dict() - for name, contains in data.items(): - helpmsg = "" - if contains.get("help") is not None: - 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: - metric[name] = prometheus_client.Info(name, helpmsg, registry=registry) - else: - pass - return metric +def main() -> None: + """Entry point for bmspy-prometheus command.""" + parser = argparse.ArgumentParser( + description="Query JBD BMS and report status to Prometheus", + add_help=True, + ) + parser.add_argument( + "--socket", + "-s", + dest="socket", + 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()