openrvdas.logger.readers.redis_reader

 1#!/usr/bin/env python3
 2
 3import logging
 4
 5# Don't barf if they don't have redis installed. Only complain if
 6# they actually try to use it, below
 7try:
 8    import redis
 9    REDIS_ENABLED = True
10except ModuleNotFoundError:
11    REDIS_ENABLED = False
12
13from logger.readers.reader import Reader  # noqa: E402
14
15DEFAULT_HOST = 'localhost'
16DEFAULT_PORT = '6379'
17
18
19################################################################################
20class RedisReader(Reader):
21    """Read messages from a redis pubsub channel."""
22
23    def __init__(self, channel, password=None, **kwargs):
24        """
25        Read text records from a Redis pubsub server channel.
26        ```
27        channel      Redis channel to read from, format channel[@hostname[:port]]
28        ```
29        """
30        super().__init__(**kwargs)
31
32        if not REDIS_ENABLED:
33            raise ModuleNotFoundError('RedisReader(): Redis is not installed. Please '
34                                      'try "pip3 install redis" prior to use.')
35        self.channel = channel
36        self.hostname = DEFAULT_HOST
37        self.port = DEFAULT_PORT
38
39        if channel.find('@') > 0:
40            (self.channel, self.hostname) = channel.split(sep='@', maxsplit=1)
41        if self.hostname.find(':') > 0:
42            (self.hostname, self.port) = self.hostname.split(sep=':', maxsplit=1)
43        self.port = int(self.port)
44
45        # Connect to the specified server and subscribe to channel
46        try:
47            self.redis = redis.StrictRedis(host=self.hostname, port=self.port,
48                                           password=password, decode_responses=True)
49            self.pubsub = self.redis.pubsub()
50            self.pubsub.subscribe(self.channel)
51        except redis.exceptions.ConnectionError as e:
52            logging.error('Unable to connect to server at %s:%d',
53                          self.hostname, self.port)
54            raise e
55
56    ############################
57    def read(self):
58        """Read/wait for message from pubsub channel."""
59
60        while True:
61            message = next(iter(self.pubsub.listen()))
62            logging.debug('Got message "%s"', message)
63            if message.get('type') == 'message':
64                data = message.get('data')
65                if data:
66                    return data
67
68        # Alternatively, we could use
69        # while True:
70        #  message = self.pubsub.get_message(timeout=10)
71        #  if message:
72        #    record = message.get('data')
73        #    if record:
74        #      return record