2
0

redisstorage.py 3.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105
  1. #!/usr/bin/python
  2. # -*- coding: utf-8 -*-
  3. import json
  4. import logging
  5. import redis
  6. from .basestorage import BaseStorage
  7. class RedisStorage(BaseStorage):
  8. """Redis-backed storage"""
  9. FIELDSUFFIX_TYPE = '__TYPE'
  10. def __init__(self,
  11. redis_host='127.0.0.1', redis_port=6379, redis_password=None):
  12. self.logger = logging.getLogger('RedisStorage')
  13. self.logger.debug('Connecting to REDIS database at %s on port %d.',
  14. redis_host, redis_port)
  15. self.db = redis.StrictRedis(host=redis_host, port=redis_port,
  16. password=redis_password)
  17. self.logger.info('Connected to REDIS database with %d entries.',
  18. self.db.dbsize())
  19. def save(self):
  20. self.db.save()
  21. def get_node(self, id):
  22. return self.get_node_data(id)
  23. def get_all_nodes_raw(self):
  24. keys = self.db.keys('node_*')
  25. nodes = {}
  26. for key in keys:
  27. node_id = key[5:]
  28. node = self.get_node_data(node_id)
  29. nodes[node_id] = node
  30. return nodes
  31. def set_node_data(self, key, data):
  32. if data is None or len(data) == 0:
  33. self.db.delete('node_' + key)
  34. return
  35. thedata = {}
  36. for item in data:
  37. payload = data[item]
  38. if isinstance(payload, basestring):
  39. thedata[item] = data[item]
  40. thedata[item + self.FIELDSUFFIX_TYPE] = 'str'
  41. elif isinstance(payload, int):
  42. thedata[item] = str(data[item])
  43. thedata[item + self.FIELDSUFFIX_TYPE] = 'int'
  44. elif isinstance(payload, float):
  45. thedata[item] = str(data[item])
  46. thedata[item + self.FIELDSUFFIX_TYPE] = 'float'
  47. else:
  48. thedata[item] = json.dumps(data[item])
  49. thedata[item + self.FIELDSUFFIX_TYPE] = 'json'
  50. self.db.hmset('node_' + key, thedata)
  51. def get_node_data(self, key):
  52. node = {}
  53. thedata = self.db.hgetall('node_' + key)
  54. if thedata is None:
  55. return None
  56. for item in thedata:
  57. if item.endswith(self.FIELDSUFFIX_TYPE):
  58. continue
  59. fieldtype = thedata.get(item + self.FIELDSUFFIX_TYPE, 'str')
  60. payload = thedata[item]
  61. if fieldtype == 'json':
  62. node[item] = json.loads(payload)
  63. elif fieldtype == 'int':
  64. node[item] = int(payload)
  65. elif fieldtype == 'float':
  66. node[item] = float(payload)
  67. else:
  68. node[item] = payload
  69. return node
  70. def get_vpn_keys(self):
  71. keys = [key[4:] for key in self.db.keys('vpn_*')]
  72. return keys
  73. def get_vpn_item(self, key, create=False):
  74. self.check_vpn_key(key)
  75. rawdata = self.db.get('vpn_' + key)
  76. if rawdata is None:
  77. if not create:
  78. return None
  79. self.store_vpn_item(key, {'active': {}, 'last': {}})
  80. rawdata = self.db.get('vpn_' + key)
  81. data = json.loads(rawdata)
  82. return data
  83. def store_vpn_item(self, key, data):
  84. self.check_vpn_key(key)
  85. self.db.set('vpn_' + key, json.dumps(data))