openrvdas.logger.writers.mqtt_writer

 1#!/usr/bin/env python3
 2
 3import logging
 4
 5# Don't barf if they don't have paho mqtt installed. Only complain if
 6# they actually try to use it, below. If it *is* installed, check which
 7# version, so we know how to call it.
 8try:
 9    import paho.mqtt.client as mqtt  # import the client | $ pip installing paho-mqtt is necessary
10    PAHO_ENABLED = True
11
12    # Check which paho version is being used
13    from pkg_resources import get_distribution, packaging  # noqa: E402
14    PAHO_VERSION = get_distribution("paho-mqtt").version
15    if packaging.version.parse(PAHO_VERSION) >= packaging.version.parse('2.0.0'):
16        USE_VERSION_FLAG = True
17    else:
18        USE_VERSION_FLAG = False
19except ModuleNotFoundError:
20    PAHO_ENABLED = False
21
22from logger.writers.writer import Writer  # noqa: E402
23
24
25class MQTTWriter(Writer):
26    """Write to paho-mqtt broker channel."""
27
28    def __init__(self, broker, channel, client_name=None, qos=0, **kwargs):
29        """
30        Write text records to a paho-mqtt broker channel.
31        ```
32        broker       MQTT broker to connect, broker format[###.###.#.###]
33        channel      MQTT channel to read from, channel format[@broker/path_of_subscripton]
34        client_name  Deprecated
35        qos          Quality of service: 0 = at most once, 1 = at least once, 2 = exactly once
36
37        ```
38        See /readers/mqtt_reader.py for info on how to start a broker
39        """
40        super().__init__(**kwargs)  # processes 'quiet' and type hints
41
42        if not PAHO_ENABLED:
43            raise ModuleNotFoundError('MQTTReader(): paho-mqtt is not installed. Please '
44                                      'try "pip install paho-mqtt" prior to use.')
45        if qos not in [0, 1, 2]:
46            raise ValueError('MQTTWriter parameter qos must be integer value 0, 1 or 2. '
47                             f'Found type "{type(qos).__name__}", value "{qos}".')
48        self.broker = broker
49        self.channel = channel
50        self.client_name = client_name
51        self.qos = qos
52
53        try:
54            if USE_VERSION_FLAG:
55                self.client = mqtt.Client(mqtt.CallbackAPIVersion.VERSION2)
56            else:
57                self.client = mqtt.Client()
58
59            self.client.connect(broker)
60            self.client.subscribe(channel)
61
62            self.client.loop_start()
63
64        except mqtt.WebsocketConnectionError as e:
65            logging.error('Unable to connect to broker at %s:%s',
66                          self.broker, self.channel)
67            raise e
68
69    ############################
70    def __del__(self):
71        """Clean up the connection once finished"""
72        self.client.loop_stop()
73        self.client.disconnect()
74
75    ############################
76    def write(self, record: str):
77
78        # See if it's something we can process, and if not, try digesting
79        if not self.can_process_record(record):  # inherited from BaseModule()
80            self.digest_record(record)  # inherited from BaseModule()
81            return
82
83        # If record is not a string, try converting to JSON. If we don't know
84        # how, throw a hail Mary and force it into str format
85        # if not type(record) is str:
86        #  if type(record) in [int, float, bool, list, dict]:
87        #    record = json.dumps(record)
88        #  else:
89        #    record = str(record)
90
91        try:
92            self.client.publish(self.channel, record, self.qos)
93        except mqtt.WebsocketConnectionError as e:
94            logging.error('Unable to connect to broker at %s:%d',
95                          self.broker, self.channel)
96            raise e