bmspy/bmspy/influxdb.py

import atexit
import datetime
import os
import sys
import time
from influxdb_client_3 import InfluxDBClient3, Point
from bmspy import client
from bmspy.classes import UPS
from bmspy.utilities import debugger

DAEMON_UPDATE_PERIOD = 30


def influx_shutdown(influxclient: InfluxDBClient3 | None) -> None:
    """Close the InfluxDB client connection if it is not None."""
    if influxclient is not None:
        influxclient.close()


def influxdb_export(
    bucket: str,
    url: str | None = None,
    org: str | None = None,
    token: str | None = None,
    socket_path: str | None = None,
    ups: str | None = None,
    daemonize: bool = True,
    debug: int = 0,
) -> None:
    """Export BMS data to InfluxDB, either once or as a daemon loop."""
    if not url:
        url = os.environ["INFLUXDB_V2_URL"]
        org = os.environ.get("INFLUXDB_V2_ORG")
        token = os.environ["INFLUXDB_V2_TOKEN"]

    influxclient = InfluxDBClient3(host=url, org=org, database=bucket, token=token)
    atexit.register(influx_shutdown, influxclient)

    if daemonize:
        while True:
            data = client.read_data(socket_path, "influxdb", ups=ups)
            influxdb_write_snapshot(influxclient, bucket, data, debug)
            time.sleep(DAEMON_UPDATE_PERIOD)
    else:
        data = client.read_data(socket_path, "influxdb", ups=ups)
        influxdb_write_snapshot(influxclient, bucket, data, debug)

    influxclient.close()
    atexit.unregister(influx_shutdown)


def influxdb_write_snapshot(influxclient: InfluxDBClient3, bucket: str, ups_data: dict[str, UPS], debug: int = 0) -> None:
    """Create InfluxDB points from ups_data and write them to the database."""
    if debug > 1:
        debugger("influxdb: creating snapshot")
    points = influxdb_create_snapshot(ups_data, debug)
    if debug > 1:
        debugger("influxdb: writing snapshot")
    try:
        influxclient.write(record=points, database=bucket)
    except Exception as e:
        debugger(e)


def influxdb_create_snapshot(ups_data: dict[str, UPS], debug: int = 0) -> list[Point]:
    """Build InfluxDB points from {ups_name: UPS}, tagging each point with the UPS name."""
    points = []
    now = datetime.datetime.now(datetime.timezone.utc)

    for ups_name, data in ups_data.items():
        for kind, contains in data.items():
            helpmsg = contains.get("help") or ""
            units = contains.get("units")

            if contains.get("raw_value") is not None:
                value = contains.get("raw_value")
                if debug > 2:
                    debugger("value: {} [{}] : {}".format(kind, ups_name, value))
                points.append(
                    Point(kind)
                    .tag("ups", ups_name)
                    .tag("units", units)
                    .tag("help", helpmsg)
                    .field("value", value)
                    .time(now)
                )

            if contains.get("raw_values") is not None and isinstance(
                contains.get("raw_values"), dict
            ):
                label = contains.get("label")
                for idx, label_value in contains.get("raw_values").items():
                    if debug > 2:
                        debugger(
                            "labels: {} [{}][{}] : {}".format(
                                kind, ups_name, idx, label_value
                            )
                        )
                    points.append(
                        Point(kind)
                        .tag("ups", ups_name)
                        .tag(label, idx)
                        .tag("units", units)
                        .tag("help", helpmsg)
                        .field("value", label_value)
                        .time(now)
                    )

            if contains.get("info") is not None:
                value = contains.get("info")
                if debug > 2:
                    debugger("info: {} [{}] : {}".format(kind, ups_name, value))
                points.append(
                    Point(kind)
                    .tag("ups", ups_name)
                    .tag("units", units)
                    .tag("help", helpmsg)
                    .field("value", value)
                    .time(now)
                )

    return points


def main():
    import argparse

    parser = argparse.ArgumentParser(
        description="Query JBD BMS and report status",
        add_help=True,
    )
    parser.add_argument(
        "--bucket",
        "-b",
        dest="influx_bucket",
        type=str,
        action="store",
        default="ups",
        help='Set the bucket name when sending data to influxdb (defaults to "ups")',
    )
    parser.add_argument(
        "--url",
        "-u",
        dest="influx_url",
        type=str,
        action="store",
        default=False,
        help="Set the URL when sending data to influxdb (overrides INFLUXDB environment variables)",
    )
    parser.add_argument(
        "--org",
        "-o",
        dest="influx_org",
        type=str,
        action="store",
        default=False,
        help="Set the influx organization when sending data to influxdb (overrides INFLUXDB environment variables)",
    )
    parser.add_argument(
        "--token",
        "-t",
        dest="influx_token",
        type=str,
        action="store",
        default=False,
        help="Set the influx token when sending data to influxdb (overrides INFLUXDB environment variables)",
    )
    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(
        "--verbose",
        "-v",
        action="count",
        default=0,
        help="Print more verbose information (can be specified multiple times)",
    )
    args = parser.parse_args()

    debug = args.verbose

    try:
        if not os.getenv("INFLUXDB_V2_URL") and not args.influx_url:
            raise argparse.ArgumentTypeError("Missing value for --url")
        if not os.getenv("INFLUXDB_V2_ORG") and not args.influx_org:
            raise argparse.ArgumentTypeError("Missing value for --org")
        if not os.getenv("INFLUXDB_V2_TOKEN") and not args.influx_token:
            raise argparse.ArgumentTypeError("Missing value for --token")
    except Exception as e:
        debugger("bmspy-influxdb: {}".format(e))
        sys.exit(1)

    debugger("Running BMS influxdb daemon on socket {}".format(args.socket))

    client.handle_registration(args.socket, "influxdb", debug)
    atexit.register(client.handle_registration, args.socket, "influxdb", debug)

    influxdb_export(
        bucket=args.influx_bucket,
        url=args.influx_url or None,
        org=args.influx_org or None,
        token=args.influx_token or None,
        socket_path=args.socket,
        ups=args.ups,
        daemonize=True,
        debug=debug,
    )


if __name__ == "__main__":
    main()