1# (c) 2014, Brian Coca, Josh Drake, et al
2# (c) 2017 Ansible Project
3# GNU General Public License v3.0+ (see COPYING or https://www.gnu.org/licenses/gpl-3.0.txt)
4from __future__ import (absolute_import, division, print_function)
5__metaclass__ = type
6
7DOCUMENTATION = '''
8    cache: redis
9    short_description: Use Redis DB for cache
10    description:
11        - This cache uses JSON formatted, per host records saved in Redis.
12    version_added: "1.9"
13    requirements:
14      - redis>=2.4.5 (python lib)
15    options:
16      _uri:
17        description:
18          - A colon separated string of connection information for Redis.
19        required: True
20        env:
21          - name: ANSIBLE_CACHE_PLUGIN_CONNECTION
22        ini:
23          - key: fact_caching_connection
24            section: defaults
25      _prefix:
26        description: User defined prefix to use when creating the DB entries
27        env:
28          - name: ANSIBLE_CACHE_PLUGIN_PREFIX
29        ini:
30          - key: fact_caching_prefix
31            section: defaults
32      _timeout:
33        default: 86400
34        description: Expiration timeout for the cache plugin data
35        env:
36          - name: ANSIBLE_CACHE_PLUGIN_TIMEOUT
37        ini:
38          - key: fact_caching_timeout
39            section: defaults
40        type: integer
41'''
42
43import time
44import json
45
46from ansible import constants as C
47from ansible.errors import AnsibleError
48from ansible.plugins.cache import BaseCacheModule
49
50try:
51    from redis import StrictRedis, VERSION
52except ImportError:
53    raise AnsibleError("The 'redis' python module (version 2.4.5 or newer) is required for the redis fact cache, 'pip install redis'")
54
55
56class CacheModule(BaseCacheModule):
57    """
58    A caching module backed by redis.
59    Keys are maintained in a zset with their score being the timestamp
60    when they are inserted. This allows for the usage of 'zremrangebyscore'
61    to expire keys. This mechanism is used or a pattern matched 'scan' for
62    performance.
63    """
64    def __init__(self, *args, **kwargs):
65        if C.CACHE_PLUGIN_CONNECTION:
66            connection = C.CACHE_PLUGIN_CONNECTION.split(':')
67        else:
68            connection = []
69
70        self._timeout = float(C.CACHE_PLUGIN_TIMEOUT)
71        self._prefix = C.CACHE_PLUGIN_PREFIX
72        self._cache = {}
73        self._db = StrictRedis(*connection)
74        self._keys_set = 'ansible_cache_keys'
75
76    def _make_key(self, key):
77        return self._prefix + key
78
79    def get(self, key):
80
81        if key not in self._cache:
82            value = self._db.get(self._make_key(key))
83            # guard against the key not being removed from the zset;
84            # this could happen in cases where the timeout value is changed
85            # between invocations
86            if value is None:
87                self.delete(key)
88                raise KeyError
89            self._cache[key] = json.loads(value)
90
91        return self._cache.get(key)
92
93    def set(self, key, value):
94
95        value2 = json.dumps(value)
96        if self._timeout > 0:  # a timeout of 0 is handled as meaning 'never expire'
97            self._db.setex(self._make_key(key), int(self._timeout), value2)
98        else:
99            self._db.set(self._make_key(key), value2)
100
101        if VERSION[0] == 2:
102            self._db.zadd(self._keys_set, time.time(), key)
103        else:
104            self._db.zadd(self._keys_set, {key: time.time()})
105        self._cache[key] = value
106
107    def _expire_keys(self):
108        if self._timeout > 0:
109            expiry_age = time.time() - self._timeout
110            self._db.zremrangebyscore(self._keys_set, 0, expiry_age)
111
112    def keys(self):
113        self._expire_keys()
114        return self._db.zrange(self._keys_set, 0, -1)
115
116    def contains(self, key):
117        self._expire_keys()
118        return (self._db.zrank(self._keys_set, key) is not None)
119
120    def delete(self, key):
121        if key in self._cache:
122            del self._cache[key]
123        self._db.delete(self._make_key(key))
124        self._db.zrem(self._keys_set, key)
125
126    def flush(self):
127        for key in self.keys():
128            self.delete(key)
129
130    def copy(self):
131        # TODO: there is probably a better way to do this in redis
132        ret = dict()
133        for key in self.keys():
134            ret[key] = self.get(key)
135        return ret
136
137    def __getstate__(self):
138        return dict()
139
140    def __setstate__(self, data):
141        self.__init__()
142