|
1 | | -import itertools |
2 | 1 | from datetime import timedelta |
3 | 2 | from typing import ( |
4 | 3 | Dict, |
|
8 | 7 | Sequence, |
9 | 8 | Tuple, |
10 | 9 | Union, |
11 | | - cast, |
12 | 10 | overload, |
13 | 11 | ) |
14 | 12 |
|
@@ -47,206 +45,45 @@ def delete_all(self) -> None: |
47 | 45 |
|
48 | 46 |
|
49 | 47 | class RedisRepository: |
50 | | - # Deletes M expired keys from the hash. |
51 | | - # |
52 | | - # Complexity: O(log(N) + M) + O(M) + O(M * log(N)) where N is the size of hash |
53 | | - # and M is the number of elements deleted. For constant M (e.g., 3) the complexity |
54 | | - # is O(log(N)). |
55 | | - # |
56 | | - # KEYS: |
57 | | - # hash_name. |
58 | | - # zset_name. |
59 | | - # |
60 | | - # ARGV: |
61 | | - # delete_count. |
62 | | - # |
63 | | - # Returns: nil. |
64 | | - LUA_VACUUM_SCRIPT = """ |
65 | | - local hash_name = KEYS[1] |
66 | | - local zset_name = KEYS[2] |
67 | | - local delete_count = tonumber(ARGV[1]) |
68 | | - local time = redis.call('TIME') |
69 | | - local pnow = tonumber(time[1]) * 1000 + math.floor(tonumber(time[2]) / 1000) |
70 | | -
|
71 | | - local keys_to_delete = redis.call( |
72 | | - 'ZRANGEBYSCORE', zset_name, '-inf', pnow, 'LIMIT', 0, delete_count |
73 | | - ) |
74 | | -
|
75 | | - if not next(keys_to_delete) then |
76 | | - return |
77 | | - end |
78 | | -
|
79 | | - redis.call('HDEL', hash_name, unpack(keys_to_delete)) |
80 | | - redis.call('ZREM', zset_name, unpack(keys_to_delete)) |
81 | | - """ |
82 | | - |
83 | | - # Gets either all or only requested keys and values from hash. |
84 | | - # |
85 | | - # Gets all keys if no keys specified in ARGV. |
86 | | - # |
87 | | - # Complexity: O(N) where N is the number of requested keys. |
88 | | - # |
89 | | - # KEYS: |
90 | | - # hash_name. |
91 | | - # zset_name. |
92 | | - # |
93 | | - # Optional ARGV: |
94 | | - # key1. |
95 | | - # key2. |
96 | | - # ... |
97 | | - # |
98 | | - # Returns: list of names and values for non-expired keys. |
99 | | - LUA_GET_SCRIPT = """ |
100 | | - local hash_name = KEYS[1] |
101 | | - local zset_name = KEYS[2] |
102 | | - local time = redis.call('TIME') |
103 | | - local pnow = tonumber(time[1]) * 1000 + math.floor(tonumber(time[2]) / 1000) |
104 | | -
|
105 | | - local keys = ARGV |
106 | | - if next(ARGV) == nil then |
107 | | - keys = redis.call('HKEYS', hash_name) |
108 | | - end |
109 | | -
|
110 | | - if not next(keys) then |
111 | | - return {} |
112 | | - end |
113 | | -
|
114 | | - local hash = redis.call('HMGET', hash_name, unpack(keys)) |
115 | | -
|
116 | | - local result = {} |
117 | | -
|
118 | | - for i, key in ipairs(keys) do |
119 | | - local pexpireat = redis.call('ZSCORE', zset_name, key) |
120 | | -
|
121 | | - if not pexpireat or pnow < tonumber(pexpireat) then |
122 | | - table.insert(result, key) |
123 | | - table.insert(result, hash[i]) |
124 | | - end |
125 | | - end |
126 | | -
|
127 | | - return result |
128 | | - """ |
129 | | - |
130 | | - # Inserts list of keys-value-expiration tuples into the hash. |
131 | | - # |
132 | | - # Complexity: O(N * log(M)), where N is the number of inserted elements and |
133 | | - # M is the hash size. |
134 | | - # |
135 | | - # 1. If hash does not exist, it will automatically create one. |
136 | | - # 2. If the field already exists, its value and ttl will be overwritten. |
137 | | - # 3. Hash and zset ttl is always set to the biggest field's ttl. |
138 | | - # 4. When the field expires it may be deleted by: |
139 | | - # - Manually invoking `vacuum` script. |
140 | | - # - Redis automatically deleting expired hash (see note #3). |
141 | | - # |
142 | | - # KEYS: |
143 | | - # hash_name. |
144 | | - # zset_name. |
145 | | - # |
146 | | - # ARGV: |
147 | | - # key. |
148 | | - # value. |
149 | | - # ttl. |
150 | | - # ... |
151 | | - # |
152 | | - # Returns: nil. |
153 | | - LUA_SET_SCRIPT = """ |
154 | | - local hash_name = KEYS[1] |
155 | | - local zset_name = KEYS[2] |
156 | | - local time = redis.call('TIME') |
157 | | -
|
158 | | - for i, _ in ipairs(ARGV) do |
159 | | - if i % 3 == 1 then |
160 | | - local key = ARGV[i] |
161 | | - local value = ARGV[i + 1] |
162 | | - local ttl = ARGV[i + 2] |
163 | | - local pexpireat = ( |
164 | | - (tonumber(time[1]) + ttl) * 1000 + math.floor(tonumber(time[2]) / 1000) |
165 | | - ) |
166 | | -
|
167 | | - redis.call('HSET', hash_name, key, value) |
168 | | - redis.call('ZADD', zset_name, pexpireat, key) |
169 | | -
|
170 | | - end |
171 | | - end |
172 | | -
|
173 | | - local max_pexpireat = tonumber(redis.call( |
174 | | - 'ZREVRANGEBYSCORE', zset_name, '+inf', '-inf', 'WITHSCORES', 'LIMIT', 0, 1 |
175 | | - )[2]) |
176 | | -
|
177 | | - redis.call('PEXPIREAT', hash_name, max_pexpireat) |
178 | | - redis.call('PEXPIREAT', zset_name, max_pexpireat) |
179 | | - """ |
180 | | - |
181 | | - # Deletes all data from hash and zset. |
182 | | - # |
183 | | - # Complexity: O(N) where N is the hash size. |
184 | | - # |
185 | | - # KEYS: |
186 | | - # hash_name. |
187 | | - # zset_name. |
188 | | - # |
189 | | - # Returns: nil. |
190 | | - LUA_DELETE_ALL_SCRIPT = """ |
191 | | - local hash_name = KEYS[1] |
192 | | - local zset_name = KEYS[2] |
193 | | -
|
194 | | - redis.call('DEL', hash_name, zset_name) |
195 | | - """ |
196 | 48 |
|
197 | | - def __init__(self, hash_name: str, client: redis.Redis, use_lua_52: bool = False): |
| 49 | + def __init__(self, hash_name: str, client: redis.Redis): |
198 | 50 | self.hash_name = hash_name |
199 | | - self.zset_name = f'{hash_name}.EXPIREAT' |
200 | 51 | self.client = client |
201 | | - self.lua_set_many = self.client.register_script(self.LUA_SET_SCRIPT) |
202 | | - |
203 | | - lua_get_script = self.LUA_GET_SCRIPT |
204 | | - if use_lua_52: |
205 | | - # Hack for tests to work with fakeredis, as it uses Lua version > 5.1. |
206 | | - # In Lua 5.1, unpack was a global, but in 5.2 it's been moved to |
207 | | - # table.unpack. |
208 | | - lua_get_script = self.LUA_GET_SCRIPT.replace('unpack', 'table.unpack') |
209 | | - self.lua_get = self.client.register_script(lua_get_script) |
210 | | - |
211 | | - self.lua_vacuum = self.client.register_script(self.LUA_VACUUM_SCRIPT) |
212 | | - self.lua_delete_all = self.client.register_script(self.LUA_DELETE_ALL_SCRIPT) |
213 | 52 |
|
214 | 53 | def set(self, key: str, value: str, ttl: int) -> None: |
215 | 54 | self.set_many(data=[(key, value, ttl)]) |
216 | 55 |
|
217 | 56 | def set_many(self, data: Sequence[Tuple[str, str, int]]) -> None: |
218 | | - self.lua_set_many( |
219 | | - keys=[self.hash_name, self.zset_name], |
220 | | - args=list(itertools.chain.from_iterable(data)), |
221 | | - ) |
| 57 | + pipe = self.client.pipeline() |
| 58 | + for key, value, ttl in data: |
| 59 | + pipe.hset(self.hash_name, key, value) |
| 60 | + pipe.execute_command("HEXPIRE", self.hash_name, ttl, "FIELDS", 1, key) |
| 61 | + pipe.execute() |
222 | 62 |
|
223 | 63 | def get(self, key: str) -> Optional[str]: |
224 | | - return self.get_many(keys=[key]).get(key) |
| 64 | + val = self.client.hget(self.hash_name, key) |
| 65 | + return None if val is None else str(val) |
225 | 66 |
|
226 | 67 | def get_many(self, keys: Sequence[str]) -> Dict[str, Optional[str]]: |
227 | | - data = self.lua_get(keys=[self.hash_name, self.zset_name], args=list(keys)) |
228 | | - |
229 | | - data = dict(zip(data[::2], data[1::2])) |
230 | | - |
231 | | - for missing_key in set(keys) - set(data): |
232 | | - data[missing_key] = None |
233 | | - |
234 | | - return data |
| 68 | + if not keys: |
| 69 | + return {} |
| 70 | + values = self.client.hmget(self.hash_name, keys) |
| 71 | + # redis-py returns a list of values where non-existent/expired are None |
| 72 | + return {k: (None if v is None else str(v)) for k, v in zip(keys, values)} |
235 | 73 |
|
236 | 74 | def get_all(self) -> Dict[str, str]: |
237 | | - return cast(Dict[str, str], self.get_many(keys=[])) |
| 75 | + raw = self.client.hgetall(self.hash_name) |
| 76 | + return dict(raw) |
238 | 77 |
|
239 | 78 | def delete(self, key: str) -> None: |
240 | | - self.delete_many(keys=[key]) |
| 79 | + self.client.hdel(self.hash_name, key) |
241 | 80 |
|
242 | 81 | def delete_many(self, keys: Sequence[str]) -> None: |
243 | | - self.set_many(data=[(key, '', -1) for key in keys]) |
| 82 | + if keys: |
| 83 | + self.client.hdel(self.hash_name, *keys) |
244 | 84 |
|
245 | 85 | def delete_all(self) -> None: |
246 | | - self.lua_delete_all(keys=[self.hash_name, self.zset_name]) |
247 | | - |
248 | | - def vacuum(self, delete_count: int) -> None: |
249 | | - self.lua_vacuum(keys=[self.hash_name, self.zset_name], args=[delete_count]) |
| 86 | + self.client.delete(self.hash_name) |
250 | 87 |
|
251 | 88 |
|
252 | 89 | class DeprecatedRedisAdapter: |
|
0 commit comments