common/cache/redis.py

0001 #!/usr/bin/env python
0002 #
0003 # Licensed under the Apache License, Version 2.0 (the "License");
0004 # You may not use this file except in compliance with the License.
0005 # You may obtain a copy of the License at
0006 # http://www.apache.org/licenses/LICENSE-2.0OA
0007 #
0008 # Authors:
0009 # - Wen Guan, <wen.guan@cern.ch>, 2022
0010
0011 import logging
0012 import uuid
0013 import redis
0014
0015 from idds.common.constants import Sections
0016 from idds.common.config import config_has_section, config_list_options
0017 from idds.common.utils import json_dumps, json_loads
0018
0019
0020 class Singleton(object):
0021     _instance = None
0022
0023     def __new__(class_, *args, **kwargs):
0024         if not isinstance(class_._instance, class_):
0025             class_._instance = object.__new__(class_, *args, **kwargs)
0026             class_._instance._initialized = False
0027         return class_._instance
0028
0029
0030 class RedisCache(Singleton):
0031     """
0032     Redis cache
0033     """
0034
0035     def __init__(self, logger=None):
0036         if not self._initialized:
0037             self._initialized = True
0038
0039             super(RedisCache, self).__init__()
0040             self._id = str(uuid.uuid4())[:8]
0041             self.logger = logger
0042             self.setup_logger(self.logger)
0043             self.config_section = Sections.Cache
0044             attrs = self.load_attributes()
0045             if 'host' in attrs and attrs['host']:
0046                 self.host = attrs['host']
0047             else:
0048                 self.host = 'localhost'
0049             if 'port' in attrs and attrs['port']:
0050                 self.port = int(attrs['port'])
0051             else:
0052                 self.port = 6379
0053             self.cache = redis.Redis(host=self.host, port=self.port, db=0)
0054
0055     def setup_logger(self, logger=None):
0056         """
0057         Setup logger
0058         """
0059         if logger:
0060             self.logger = logger
0061         else:
0062             self.logger = logging.getLogger(self.get_class_name())
0063
0064     def get_class_name(self):
0065         return self.__class__.__name__
0066
0067     def load_attributes(self):
0068         self.logger.info("Loading config for section: %s" % self.config_section)
0069         attrs = {}
0070         if config_has_section(self.config_section):
0071             options = config_list_options(self.config_section)
0072             for option, value in options:
0073                 if isinstance(value, str) and value.lower() == 'true':
0074                     value = True
0075                 if isinstance(value, str) and value.lower() == 'false':
0076                     value = False
0077                 attrs[option] = value
0078         return attrs
0079
0080     def set(self, key, value, expire_seconds=21600):
0081         value = json_dumps(value)
0082         self.cache.set(key, value, ex=expire_seconds)
0083
0084     def get(self, key, default=None):
0085         value = self.cache.get(key)
0086         if value:
0087             value = json_loads(value)
0088         if not value:
0089             return default
0090         return value
0091
0092     def hset(self, key, value, expire_seconds=21600):
0093         value = json_dumps(value)
0094         self.cache.hset(key, value)
0095         self.cache.expire(key, expire_seconds)
0096
0097     def hget(self, key, default=None):
0098         value = self.cache.hget(key)
0099         if value:
0100             value = json_loads(value)
0101         if not value:
0102             return default
0103         return value
0104
0105
0106 def get_redis_cache():
0107     cache = RedisCache()
0108     return cache