redisstorage.py 3.0 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495
  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_all_nodes_raw(self):
  22. keys = self.db.keys('node_*')
  23. nodes = {}
  24. for key in keys:
  25. node_id = key[5:]
  26. node = self.get_node_data(node_id)
  27. nodes[node_id] = node
  28. return nodes
  29. def set_node_data(self, key, data):
  30. thedata = {}
  31. for item in data:
  32. payload = data[item]
  33. if isinstance(payload, basestring):
  34. thedata[item] = data[item]
  35. thedata[item + self.FIELDSUFFIX_TYPE] = 'str'
  36. elif isinstance(payload, int):
  37. thedata[item] = str(data[item])
  38. thedata[item + self.FIELDSUFFIX_TYPE] = 'int'
  39. elif isinstance(payload, float):
  40. thedata[item] = str(data[item])
  41. thedata[item + self.FIELDSUFFIX_TYPE] = 'float'
  42. else:
  43. thedata[item] = json.dumps(data[item])
  44. thedata[item + self.FIELDSUFFIX_TYPE] = 'json'
  45. self.db.hmset('node_' + key, thedata)
  46. def get_node_data(self, key):
  47. node = {}
  48. thedata = self.db.hgetall('node_' + key)
  49. for item in thedata:
  50. if item.endswith(self.FIELDSUFFIX_TYPE):
  51. continue
  52. fieldtype = thedata.get(item + self.FIELDSUFFIX_TYPE, 'str')
  53. payload = thedata[item]
  54. if fieldtype == 'json':
  55. node[item] = json.loads(payload)
  56. elif fieldtype == 'int':
  57. node[item] = int(payload)
  58. elif fieldtype == 'float':
  59. node[item] = float(payload)
  60. else:
  61. node[item] = payload
  62. return node
  63. def get_vpn_keys(self):
  64. keys = [key[4:] for key in self.db.keys('vpn_*')]
  65. return keys
  66. def get_vpn_item(self, key, create=False):
  67. self.check_vpn_key(key)
  68. rawdata = self.db.get('vpn_' + key)
  69. if rawdata is None:
  70. if not create:
  71. return None
  72. self.store_vpn_item(key, {'active': {}, 'last': {}})
  73. rawdata = self.db.get('vpn_' + key)
  74. data = json.loads(rawdata)
  75. return data
  76. def store_vpn_item(self, key, data):
  77. self.check_vpn_key(key)
  78. self.db.set('vpn_' + key, json.dumps(data))